Jump diffusion model

docs/docs/tutorials/diffusion_model.ipynb

@@ -39,13 +39,11 @@
     "energy = np.linspace(-2, 2, 501)\n",
     "scale = 1.0\n",
     "diffusion_coefficient = 2.4e-9  # m^2/s\n",
-    "diffusion_unit = 'm**2/s'\n",
     "\n",
     "diffusion_model = BrownianTranslationalDiffusion(\n",
     "    display_name='DiffusionModel',\n",
     "    scale=scale,\n",
     "    diffusion_coefficient=diffusion_coefficient,\n",
-    "    diffusion_unit=diffusion_unit,\n",
     ")\n",
     "\n",
     "component_collections = diffusion_model.create_component_collections(Q)\n",

docs/docs/tutorials/sample_model.ipynb

@@ -48,12 +48,10 @@
     "\n",
     "scale = 1.0\n",
     "diffusion_coefficient = 2.4e-9  # m^2/s\n",
-    "diffusion_unit = 'm**2/s'\n",
     "diffusion_model = BrownianTranslationalDiffusion(\n",
     "    display_name='DiffusionModel',\n",
     "    scale=scale,\n",
     "    diffusion_coefficient=diffusion_coefficient,\n",
-    "    diffusion_unit=diffusion_unit,\n",
     ")\n",
     "\n",
     "\n",

src/easydynamics/sample_model/diffusion_model/__init__.py

@@ -2,9 +2,9 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from .brownian_translational_diffusion import BrownianTranslationalDiffusion
-from .diffusion_model_base import DiffusionModelBase
+from .jump_translational_diffusion import JumpTranslationalDiffusion
 
 __all__ = [
-    'DiffusionModelBase',
     'BrownianTranslationalDiffusion',
+    'JumpTranslationalDiffusion',
 ]

src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py

@@ -3,24 +3,20 @@
 
 from typing import Dict
 from typing import List
-from typing import Union
 
 import numpy as np
 import scipp as sc
 from easyscience.variable import DescriptorNumber
 from easyscience.variable import Parameter
-from numpy.typing import ArrayLike
 from scipp.constants import hbar as scipp_hbar
 
 from easydynamics.sample_model.component_collection import ComponentCollection
 from easydynamics.sample_model.components import Lorentzian
 from easydynamics.sample_model.diffusion_model.diffusion_model_base import DiffusionModelBase
+from easydynamics.utils.utils import Numeric
+from easydynamics.utils.utils import Q_type
 from easydynamics.utils.utils import _validate_and_convert_Q
 
-Numeric = Union[float, int]
-
-Q_type = np.ndarray | Numeric | list | ArrayLike
-
 
 class BrownianTranslationalDiffusion(DiffusionModelBase):
     """Model of Brownian translational diffusion, consisting of a
@@ -46,7 +42,6 @@ def __init__(
         unit: str | sc.Unit = 'meV',
         scale: Numeric = 1.0,
         diffusion_coefficient: Numeric = 1.0,
-        diffusion_unit: str = 'm**2/s',
     ):
         """Initialize a new BrownianTranslationalDiffusion model.
 
@@ -62,65 +57,35 @@ def __init__(
             Defaults to "meV".
         scale : float or Parameter, optional
             Scale factor for the diffusion model.
-        diffusion_coefficient : float or Parameter, optional
-            Diffusion coefficient D. If a number is provided,
-            it is assumed to be in the unit given by diffusion_unit.
+        diffusion_coefficient : Number, optional
+            Diffusion coefficient D in m^2/s.
             Defaults to 1.0.
-        diffusion_unit : str, optional
-            Unit for the diffusion coefficient D. Default is m**2/s.
-            Options are 'meV*Å**2' or 'm**2/s'
         """
-        if not isinstance(scale, (Parameter, Numeric)):
+        if not isinstance(scale, Numeric):
             raise TypeError('scale must be a number.')
 
-        if not isinstance(diffusion_coefficient, (Parameter, Numeric)):
+        if not isinstance(diffusion_coefficient, Numeric):
             raise TypeError('diffusion_coefficient must be a number.')
 
-        if not isinstance(diffusion_unit, str):
-            raise TypeError("diffusion_unit must be 'meV*Å**2' or 'm**2/s'.")
-
-        if diffusion_unit == 'meV*Å**2' or diffusion_unit == 'meV*angstrom**2':
-            # In this case, hbar is absorbed in the unit of D
-            self._hbar = DescriptorNumber('hbar', 1.0)
-        elif diffusion_unit == 'm**2/s' or diffusion_unit == 'm^2/s':
-            self._hbar = DescriptorNumber.from_scipp('hbar', scipp_hbar)
-        else:
-            raise ValueError("diffusion_unit must be 'meV*Å**2' or 'm**2/s'.")
-
-        scale = Parameter(name='scale', value=float(scale), fixed=False, min=0.0)
-
         diffusion_coefficient = Parameter(
             name='diffusion_coefficient',
             value=float(diffusion_coefficient),
             fixed=False,
-            unit=diffusion_unit,
+            unit='m**2/s',
         )
         super().__init__(
             display_name=display_name,
             unique_name=unique_name,
             unit=unit,
+            scale=scale,
         )
+        self._hbar = DescriptorNumber.from_scipp('hbar', scipp_hbar)
         self._angstrom = DescriptorNumber('angstrom', 1e-10, unit='m')
-        self._scale = scale
         self._diffusion_coefficient = diffusion_coefficient
 
-    @property
-    def scale(self) -> Parameter:
-        """Get the scale parameter of the diffusion model.
-
-        Returns
-        -------
-        Parameter
-            Scale parameter.
-        """
-        return self._scale
-
-    @scale.setter
-    def scale(self, scale: Numeric) -> None:
-        """Set the scale parameter of the diffusion model."""
-        if not isinstance(scale, (Numeric)):
-            raise TypeError('scale must be a number.')
-        self._scale.value = scale
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
 
     @property
     def diffusion_coefficient(self) -> Parameter:
@@ -136,10 +101,14 @@ def diffusion_coefficient(self) -> Parameter:
     @diffusion_coefficient.setter
     def diffusion_coefficient(self, diffusion_coefficient: Numeric) -> None:
         """Set the diffusion coefficient parameter D."""
-        if not isinstance(diffusion_coefficient, (Numeric)):
+        if not isinstance(diffusion_coefficient, Numeric):
             raise TypeError('diffusion_coefficient must be a number.')
         self._diffusion_coefficient.value = diffusion_coefficient
 
+    # ------------------------------------------------------------------
+    # Other methods
+    # ------------------------------------------------------------------
+
     def calculate_width(self, Q: Q_type) -> np.ndarray:
         """Calculate the half-width at half-maximum (HWHM) for the
         diffusion model.
@@ -265,6 +234,10 @@ def create_component_collections(
 
         return component_collection_list
 
+    # ------------------------------------------------------------------
+    # Private methods
+    # ------------------------------------------------------------------
+
     def _write_width_dependency_expression(self, Q: float) -> str:
         """Write the dependency expression for the width as a function
         of Q to make dependent Parameters.
@@ -316,6 +289,10 @@ def _write_area_dependency_map_expression(self) -> Dict[str, DescriptorNumber]:
             'scale': self.scale,
         }
 
+    # ------------------------------------------------------------------
+    # dunder methods
+    # ------------------------------------------------------------------
+
     def __repr__(self):
         """String representation of the BrownianTranslationalDiffusion
         model.

src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py

@@ -1,17 +1,14 @@
 # SPDX-FileCopyrightText: 2025-2026 EasyDynamics contributors <https://github.com/easyscience>
 # SPDX-License-Identifier: BSD-3-Clause
 
-import numpy as np
 import scipp as sc
 from easyscience.base_classes.model_base import ModelBase
 from easyscience.variable import DescriptorNumber
-from numpy.typing import ArrayLike
+from easyscience.variable import Parameter
 from scipp import UnitError
 
 from easydynamics.utils.utils import Numeric
 
-Q_type = np.ndarray | Numeric | list | ArrayLike
-
 
 class DiffusionModelBase(ModelBase):
     """Base class for constructing diffusion models."""
@@ -20,6 +17,7 @@ def __init__(
         self,
         display_name='MyDiffusionModel',
         unique_name: str | None = None,
+        scale: Numeric = 1.0,
         unit: str | sc.Unit = 'meV',
     ):
         """Initialize a new DiffusionModel.
@@ -31,6 +29,10 @@ def __init__(
         unit : str or sc.Unit, optional
             Unit of the diffusion model. Defaults to "meV".
         """
+        if not isinstance(scale, Numeric):
+            raise TypeError('scale must be a number.')
+
+        scale = Parameter(name='scale', value=float(scale), fixed=False, min=0.0)
 
         try:
             test = DescriptorNumber(name='test', value=1, unit=unit)
@@ -42,6 +44,11 @@ def __init__(
 
         super().__init__(display_name=display_name, unique_name=unique_name)
         self._unit = unit
+        self._scale = scale
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
 
     @property
     def unit(self) -> str:
@@ -62,6 +69,28 @@ def unit(self, unit_str: str) -> None:
             )
         )  # noqa: E501
 
+    @property
+    def scale(self) -> Parameter:
+        """Get the scale parameter of the diffusion model.
+
+        Returns
+        -------
+        Parameter
+            Scale parameter.
+        """
+        return self._scale
+
+    @scale.setter
+    def scale(self, scale: Numeric) -> None:
+        """Set the scale parameter of the diffusion model."""
+        if not isinstance(scale, Numeric):
+            raise TypeError('scale must be a number.')
+        self._scale.value = scale
+
+    # ------------------------------------------------------------------
+    # dunder methods
+    # ------------------------------------------------------------------
+
     def __repr__(self):
         """String representation of the Diffusion model."""
         return f'{self.__class__.__name__}(display_name={self.display_name}, unit={self.unit})'