From 774c0a84f41ee40b1b2bd125d4d626b3b64bd3be Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Tue, 6 Jan 2026 14:51:21 +0100 Subject: [PATCH 1/6] Brownian diffusion with example and test --- examples/diffusion_model.ipynb | 77 +++++ src/easydynamics/sample_model/__init__.py | 4 + .../sample_model/diffusion_model/__init__.py | 7 + .../brownian_translational_diffusion.py | 298 ++++++++++++++++++ .../diffusion_model/diffusion_model_base.py | 50 +++ .../test_brownian_translational_diffusion.py | 292 +++++++++++++++++ tests/diffusion_model/test_diffusion_model.py | 24 ++ 7 files changed, 752 insertions(+) create mode 100644 examples/diffusion_model.ipynb create mode 100644 src/easydynamics/sample_model/diffusion_model/__init__.py create mode 100644 src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py create mode 100644 src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py create mode 100644 tests/diffusion_model/test_brownian_translational_diffusion.py create mode 100644 tests/diffusion_model/test_diffusion_model.py diff --git a/examples/diffusion_model.ipynb b/examples/diffusion_model.ipynb new file mode 100644 index 0000000..f2fb212 --- /dev/null +++ b/examples/diffusion_model.ipynb @@ -0,0 +1,77 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "64deaa41", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "from easydynamics.sample_model import BrownianTranslationalDiffusion\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "%matplotlib widget" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "784d9e82", + "metadata": {}, + "outputs": [], + "source": [ + "# Create Brownian Translational Diffusion model and plot the model for different Q values.\n", + "# Q is in Angstrom^-1 and energy in meV.\n", + "\n", + "Q=np.linspace(0.5,2,7)\n", + "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(display_name=\"DiffusionModel\", scale=scale, diffusion_coefficient= diffusion_coefficient, diffusion_unit=diffusion_unit)\n", + "\n", + "component_collections=diffusion_model.create_component_collections(Q)\n", + "\n", + "\n", + "cmap = plt.cm.jet\n", + "nQ = len(component_collections)\n", + "plt.figure()\n", + "for Q_index in range(len(component_collections)):\n", + " color = cmap(Q_index / (nQ - 1))\n", + " y=component_collections[Q_index].evaluate(energy)\n", + " plt.plot(energy, y, label=f'Q={Q[Q_index]} Å^-1',color=color)\n", + " \n", + "plt.legend()\n", + "plt.show()\n", + "plt.xlabel('Energy (meV)')\n", + "plt.ylabel('Intensity (arb. units)')\n", + "plt.title('Brownian Translational Diffusion Model') " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "easydynamics_newbase", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/easydynamics/sample_model/__init__.py b/src/easydynamics/sample_model/__init__.py index 875020f..f9e2395 100644 --- a/src/easydynamics/sample_model/__init__.py +++ b/src/easydynamics/sample_model/__init__.py @@ -7,6 +7,9 @@ Polynomial, Voigt, ) +from .diffusion_model.brownian_translational_diffusion import ( + BrownianTranslationalDiffusion, +) __all__ = [ "ComponentCollection", @@ -16,4 +19,5 @@ "DeltaFunction", "DampedHarmonicOscillator", "Polynomial", + "BrownianTranslationalDiffusion", ] diff --git a/src/easydynamics/sample_model/diffusion_model/__init__.py b/src/easydynamics/sample_model/diffusion_model/__init__.py new file mode 100644 index 0000000..be41cbc --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/__init__.py @@ -0,0 +1,7 @@ +from .brownian_translational_diffusion import BrownianTranslationalDiffusion +from .diffusion_model_base import DiffusionModelBase + +__all__ = [ + "DiffusionModelBase", + "BrownianTranslationalDiffusion", +] diff --git a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py new file mode 100644 index 0000000..36c7cdf --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py @@ -0,0 +1,298 @@ +from numbers import Number +from typing import Dict, List, Optional, Union + +import numpy as np +import scipp as sc +from easyscience.variable import DescriptorNumber, Parameter +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, +) + +Numeric = Union[float, int] + + +class BrownianTranslationalDiffusion(DiffusionModelBase): + """ + Model of Brownian translational diffusion, consisting of a Lorentzian + function for each Q-value, where the width is given by :math:`DQ^2`. + Q is assumed to have units of 1/angstrom. + Creates ComponentCollections with Lorentzian components for given Q-values. + + Example usage: + Q=np.linspace(0.5,2,7) + energy=np.linspace(-2, 2, 501) + scale=1.0 + diffusion_coefficient = 2.4e-9 # m^2/s + diffusion_model=BrownianTranslationalDiffusion(display_name="DiffusionModel", scale=scale, diffusion_coefficient= diffusion_coefficient) + component_collections=diffusion_model.create_component_collections(Q) + See also the examples. + """ + + def __init__( + self, + display_name: Optional[str] = "BrownianTranslationalDiffusion", + unit: Optional[Union[str, sc.Unit]] = "meV", + scale: Optional[Union[Parameter, Numeric]] = 1.0, + diffusion_coefficient: Optional[Union[Parameter, Numeric]] = 1.0, + diffusion_unit: Optional[str] = "m**2/s", + ): + """ + Initialize a new BrownianTranslationalDiffusion model. + + Parameters + ---------- + display_name : str + Display name of the diffusion model. + unit : str or sc.Unit, optional + Energy unit for the underlying Lorentzian components. 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. Defaults to 1.0. + diffusion_unit : str, optional + Unit for the diffusion coefficient D. Default is "meV*Å**2". Options are 'meV*Å**2' or 'm**2/s' + + """ + if not isinstance(scale, (Parameter, Numeric)): + raise TypeError("scale must be a Parameter or a number.") + + if not isinstance(diffusion_coefficient, (Parameter, Numeric)): + raise TypeError("diffusion_coefficient must be a Parameter or 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'.") + + if not isinstance(scale, Parameter): + scale = Parameter(name="scale", value=float(scale), fixed=False, min=0.0) + + if not isinstance(diffusion_coefficient, Parameter): + diffusion_coefficient = Parameter( + name="diffusion_coefficient", + value=float(diffusion_coefficient), + fixed=False, + unit=diffusion_unit, + ) + super().__init__( + display_name=display_name, + unit=unit, + ) + 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: + if not isinstance(scale, (Numeric)): + raise TypeError("scale must be a number.") + self._scale.value = scale + + @property + def diffusion_coefficient(self) -> Parameter: + """ + Get the diffusion coefficient parameter D. + + Returns + ------- + Parameter + Diffusion coefficient D. + """ + return self._diffusion_coefficient + + @diffusion_coefficient.setter + def diffusion_coefficient(self, diffusion_coefficient: Numeric) -> None: + if not isinstance(diffusion_coefficient, (Numeric)): + raise TypeError("diffusion_coefficient must be a number.") + self._diffusion_coefficient.value = diffusion_coefficient + + def calculate_width(self, Q: np.ndarray) -> np.ndarray: + """ + Calculate the half-width at half-maximum (HWHM) for the diffusion model. + + Parameters + ---------- + Q : np.ndarray + Scattering vector in 1/angstrom + + Returns + ------- + np.ndarray + HWHM values in the unit of the model (e.g., meV). + """ + + if isinstance(Q, Numeric): + Q = np.array([Q]) + + if isinstance(Q, list): + Q = np.array(Q) + + if not isinstance(Q, np.ndarray): + raise TypeError("Q must be a numpy array.") + + width_list = [] + for Q_value in Q: + # Q is given as a float, so we need to divide by angstrom**2 to get the right units + width = ( + self._hbar + * self.diffusion_coefficient + * Q_value**2 + / (self._angstrom**2) + ) + width.convert_unit(self.unit) + width_list.append(width.value) + width = np.array(width_list) + + return width + + def calculate_EISF(self, Q: np.ndarray) -> np.ndarray: + """ + Calculate the Elastic Incoherent Structure Factor (EISF) for the Brownian translational diffusion model. + + Parameters + ---------- + Q : np.ndarray + Scattering vector in 1/angstrom + + Returns + ------- + np.ndarray + EISF values (dimensionless). + """ + if not isinstance(Q, np.ndarray): + raise TypeError("Q must be a numpy array.") + EISF = np.zeros_like(Q) + return EISF + + def calculate_QISF(self, Q: np.ndarray) -> np.ndarray: + """ + Calculate the Quasi-Elastic Incoherent Structure Factor (QISF). + + Parameters + ---------- + Q : np.ndarray + Scattering vector in 1/angstrom + + Returns + ------- + np.ndarray + QISF values (dimensionless). + """ + + if not isinstance(Q, np.ndarray): + raise TypeError("Q must be a numpy array.") + QISF = np.ones_like(Q) + return QISF + + def create_component_collections( + self, + Q: Union[Number, list, np.ndarray], + component_name: str = "Lorentzian", + ) -> List[ComponentCollection]: + """ + Create ComponentCollection components for the Brownian translational diffusion model at given Q values. + Args: + ---------- + Q : Number, list, or np.ndarray + Scattering vector values. + component_name : str + Name of the Lorentzian component. + width_name : str + Name of the width parameter. + Returns + ------- + List[ComponentCollection] + List of ComponentCollections with Lorentzian components. + """ + + if isinstance(Q, Numeric): + Q = np.array([Q]) + + if isinstance(Q, list): + Q = np.array(Q) + + if not isinstance(Q, np.ndarray): + raise TypeError("Q must be a number, list, or numpy array.") + + if Q.ndim > 1: + raise ValueError("Q must be a 1-dimensional array.") + + if not isinstance(component_name, str): + raise TypeError("component_name must be a string.") + + component_collection_list = [None] * len(Q) + # In more complex models, this is used to scale the area of the Lorentzians and the delta function. + QISF = self.calculate_QISF(Q) + + # Create a Lorentzian component for each Q-value, with width D*Q^2 and area equal to scale. No delta function, as the EISF is 0. + for i in range(len(Q)): + component_collection_list[i] = ComponentCollection( + display_name=f"{self.display_name}_Q{Q[i]:.2f}", unit=self.unit + ) + + lorentzian_component = Lorentzian( + display_name=component_name, area=self.scale * QISF[i], unit=self.unit + ) + + # Make the width dependent on Q + dependency_expression = self._write_width_dependency_expression(Q[i]) + dependency_map = self._write_width_dependency_map_expression() + + lorentzian_component.width.make_dependent_on( + dependency_expression=dependency_expression, + dependency_map=dependency_map, + ) + + # Resolving the dependency can do weird things to the units, so we make sure it's correct. + lorentzian_component.width.convert_unit(self.unit) + component_collection_list[i].add_component(lorentzian_component) + + return component_collection_list + + 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. + """ + if not isinstance(Q, (float)): + raise TypeError("Q must be a float.") + + # Q is given as a float, so we need to add the units + return f"hbar * D* {Q} **2*1/(angstrom**2)" + + def _write_width_dependency_map_expression(self) -> Dict[str, str]: + """ + Write the dependency map expression to make dependent Parameters. + """ + return { + "D": self.diffusion_coefficient, + "hbar": self._hbar, + "angstrom": self._angstrom, + } + + def __repr__(self): + """ + String representation of the BrownianTranslationalDiffusion model. + """ + return f"BrownianTranslationalDiffusion(display_name={self.display_name}, diffusion_coefficient={self.diffusion_coefficient}, scale={self.scale})" diff --git a/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py new file mode 100644 index 0000000..5a75159 --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py @@ -0,0 +1,50 @@ +import scipp as sc +from easyscience.base_classes.model_base import ModelBase + + +class DiffusionModelBase(ModelBase): + """ + Base class for constructing diffusion models. + """ + + def __init__( + self, + display_name="MyDiffusionModel", + unit: str | sc.Unit = "meV", + ): + """ + Initialize a new DiffusionModel. + + Parameters + ---------- + display_name : str + Display name of the diffusion model. + unit : str or sc.Unit, optional + Unit of the diffusion model. Defaults to "meV". + """ + + if not (unit is None or isinstance(unit, (str, sc.Unit))): + raise TypeError("unit must be None, a string, or a scipp Unit") + + super().__init__(display_name=display_name) + self._unit = unit + + @property + def unit(self) -> str | sc.Unit: + """ + Get the unit of the DiffusionModel. + + Returns + ------- + str or sc.Unit or None + """ + return self._unit + + @unit.setter + def unit(self, unit_str: str) -> None: + raise AttributeError( + ( + f"Unit is read-only. Use convert_unit to change the unit between allowed types " + f"or create a new {self.__class__.__name__} with the desired unit." + ) + ) # noqa: E501 diff --git a/tests/diffusion_model/test_brownian_translational_diffusion.py b/tests/diffusion_model/test_brownian_translational_diffusion.py new file mode 100644 index 0000000..908a2ed --- /dev/null +++ b/tests/diffusion_model/test_brownian_translational_diffusion.py @@ -0,0 +1,292 @@ +import numpy as np +import pytest +import scipp as sc +from easyscience.variable import DescriptorNumber, Parameter +from scipp.constants import hbar as scipp_hbar + +from easydynamics.sample_model.diffusion_model.brownian_translational_diffusion import ( + BrownianTranslationalDiffusion, +) + +hbar_1 = DescriptorNumber("hbar", 1.0) +hbar = DescriptorNumber.from_scipp("hbar", scipp_hbar) +angstrom = DescriptorNumber("angstrom", 1e-10, unit="m") + + +class TestBrownianTranslationalDiffusion: + @pytest.fixture + def brownian_diffusion_model(self): + return BrownianTranslationalDiffusion() + + def test_init_default(self, brownian_diffusion_model): + # WHEN THEN EXPECT + assert brownian_diffusion_model.display_name == "BrownianTranslationalDiffusion" + assert brownian_diffusion_model.unit == "meV" + assert brownian_diffusion_model.scale.value == 1.0 + assert brownian_diffusion_model.diffusion_coefficient.value == 1.0 + + @pytest.mark.parametrize( + "kwargs, expected_message", + [ + ( + { + "unit": 123, + "scale": 1.0, + "diffusion_coefficient": 1.0, + "diffusion_unit": "m**2/s", + }, + "unit must be None, a string, or a scipp Unit", + ), + ( + { + "unit": 123, + "scale": "invalid", + "diffusion_coefficient": 1.0, + "diffusion_unit": "m**2/s", + }, + "scale must be a Parameter or a number.", + ), + ( + { + "unit": 123, + "scale": 1.0, + "diffusion_coefficient": "invalid", + "diffusion_unit": "m**2/s", + }, + "diffusion_coefficient must be a Parameter or a number.", + ), + ( + { + "unit": 123, + "scale": 1.0, + "diffusion_coefficient": 1.0, + "diffusion_unit": 123, + }, + "diffusion_unit must be ", + ), + ], + ) + def test_input_type_validation_raises(self, kwargs, expected_message): + with pytest.raises(TypeError, match=expected_message): + BrownianTranslationalDiffusion( + display_name="BrownianTranslationalDiffusion", **kwargs + ) + + def test_diffusion_unit_value_error(self): + # WHEN THEN EXPECT + with pytest.raises(ValueError, match="diffusion_unit must be ."): + BrownianTranslationalDiffusion( + display_name="BrownianTranslationalDiffusion", + unit="meV", + scale=1.0, + diffusion_coefficient=1.0, + diffusion_unit="invalid_unit", + ) + + def test_init_with_parameters(self): + # WHEN + + scale = Parameter(name="scale_param", value=2.0) + diffusion_coefficient = Parameter( + name="diffusion_coefficient", value=3.0, unit="m**2/s" + ) + + # THEN + brownian_diffusion_model = BrownianTranslationalDiffusion( + display_name="CustomBrownianDiffusion", + unit="meV", + scale=scale, + diffusion_coefficient=diffusion_coefficient, + ) + + # EXPECT + assert brownian_diffusion_model.display_name == "CustomBrownianDiffusion" + assert brownian_diffusion_model.unit == "meV" + assert brownian_diffusion_model.scale is scale + assert brownian_diffusion_model.diffusion_coefficient is diffusion_coefficient + + def test_scale_setter(self, brownian_diffusion_model): + # WHEN + brownian_diffusion_model.scale = 2.0 + + # THEN EXPECT + assert brownian_diffusion_model.scale.value == 2.0 + + def test_scale_setter_raises(self, brownian_diffusion_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match="scale must be a number."): + brownian_diffusion_model.scale = "invalid" # Invalid type + + def test_diffusion_coefficient_setter(self, brownian_diffusion_model): + # WHEN + brownian_diffusion_model.diffusion_coefficient = 3.0 + + # THEN EXPECT + assert brownian_diffusion_model.diffusion_coefficient.value == 3.0 + + def test_diffusion_coefficient_setter_raises(self, brownian_diffusion_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match="diffusion_coefficient must be a number."): + brownian_diffusion_model.diffusion_coefficient = "invalid" # Invalid type + + def test_calculate_width_type_error(self, brownian_diffusion_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match="Q must be a numpy array."): + brownian_diffusion_model.calculate_width(Q="invalid") # Invalid type + + def test_calculate_width(self, brownian_diffusion_model): + # WHEN + Q_values = np.array([0.1, 0.2, 0.3]) # Example Q values in Å^-1 + + # WHEN + widths = brownian_diffusion_model.calculate_width(Q_values) + + # THEN EXPECT + unit_conversion_factor = sc.to_unit( + 1 + * sc.Unit(brownian_diffusion_model.diffusion_coefficient.unit) + * scipp_hbar + / (1 * sc.Unit("Å") ** 2), + "meV", + ) + expected_widths = 1.0 * unit_conversion_factor.value * (Q_values**2) + np.testing.assert_allclose(widths, expected_widths, rtol=1e-5) + + def test_calculate_width_diffusion_unit_mev_angstrom2(self): + # WHEN + diffusion_model = BrownianTranslationalDiffusion( + diffusion_coefficient=2.0, diffusion_unit="meV*Å**2" + ) + Q_values = np.array([0.1, 0.2, 0.3]) # Example Q values in Å^-1 + + # WHEN + widths = diffusion_model.calculate_width(Q_values) + + # THEN EXPECT + expected_widths = 2.0 * (Q_values**2) + np.testing.assert_allclose(widths, expected_widths, rtol=1e-5) + + def test_calculate_EISF(self, brownian_diffusion_model): + # WHEN + Q_values = np.array([0.1, 0.2, 0.3]) # Example Q values in Å^-1 + + # THEN + EISF = brownian_diffusion_model.calculate_EISF(Q_values) + + # EXPECT + expected_EISHF = np.zeros_like(Q_values) + np.testing.assert_array_equal(EISF, expected_EISHF) + + def test_calculate_EISF_type_error(self, brownian_diffusion_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match="Q must be a numpy array."): + brownian_diffusion_model.calculate_EISF(Q="invalid") # Invalid type + + def test_calculate_QISF(self, brownian_diffusion_model): + # WHEN + Q_values = np.array([0.1, 0.2, 0.3]) # Example Q values in Å^-1 + + # THEN + QISF = brownian_diffusion_model.calculate_QISF(Q_values) + + # EXPECT + expected_QISF = np.ones_like(Q_values) + np.testing.assert_array_equal(QISF, expected_QISF) + + def test_calculate_QISF_type_error(self, brownian_diffusion_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match="Q must be a numpy array."): + brownian_diffusion_model.calculate_QISF(Q="invalid") # Invalid type + + @pytest.mark.parametrize( + "Q", + [ + (0.5), + ([1.0, 2.0, 3.0]), + (np.array([1.0, 2.0, 3.0])), + ], + ids=[ + "python_scalar", + "python_list", + "numpy_array", + ], + ) + def test_create_component_collections(self, brownian_diffusion_model, Q): + # WHEN + + # THEN + component_collections = brownian_diffusion_model.create_component_collections( + Q=Q + ) + + # EXPECT + expected_widths = brownian_diffusion_model.calculate_width(Q) + for model_index in range(len(component_collections)): + model = component_collections[model_index] + assert len(model.components) == 1 + component = model.components[0] + assert component.display_name == "Lorentzian" + assert component.width.unit == brownian_diffusion_model.unit + assert component.width.value == expected_widths[model_index] + assert component.width.independent is False + + def test_create_component_collections_component_name_must_be_string( + self, brownian_diffusion_model + ): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match="component_name must be a string."): + brownian_diffusion_model.create_component_collections( + Q=np.array([0.1, 0.2, 0.3]), component_name=123 + ) + + def test_create_component_collections_Q_type_error(self, brownian_diffusion_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match="Q must be a "): + brownian_diffusion_model.create_component_collections( + Q="invalid" + ) # Invalid type + + def test_create_component_collections_Q_1dimensional_error( + self, brownian_diffusion_model + ): + # WHEN THEN EXPECT + with pytest.raises(ValueError, match="Q must be a 1-dimensional array."): + brownian_diffusion_model.create_component_collections( + Q=np.array([[0.1, 0.2], [0.3, 0.4]]) + ) # Invalid shape + + def test_write_width_dependency_expression(self, brownian_diffusion_model): + # WHEN THEN + expression = brownian_diffusion_model._write_width_dependency_expression(0.5) + + # EXPECT + expected_expression = "hbar * D* 0.5 **2*1/(angstrom**2)" + assert expression == expected_expression + + def test_write_width_dependency_map_expression(self, brownian_diffusion_model): + # WHEN THEN + expression_map = ( + brownian_diffusion_model._write_width_dependency_map_expression() + ) + + # EXPECT + expected_map = { + "D": brownian_diffusion_model.diffusion_coefficient, + "hbar": brownian_diffusion_model._hbar, + "angstrom": brownian_diffusion_model._angstrom, + } + + assert expression_map == expected_map + + def test_write_width_dependency_expression_raises(self, brownian_diffusion_model): + with pytest.raises(TypeError, match="Q must be a float"): + brownian_diffusion_model._write_width_dependency_expression("invalid") + + def test_repr(self, brownian_diffusion_model): + # WHEN THEN + repr_str = repr(brownian_diffusion_model) + + # EXPECT + assert "BrownianTranslationalDiffusion" in repr_str + assert "diffusion_coefficient" in repr_str + assert "scale=" in repr_str diff --git a/tests/diffusion_model/test_diffusion_model.py b/tests/diffusion_model/test_diffusion_model.py new file mode 100644 index 0000000..1c50266 --- /dev/null +++ b/tests/diffusion_model/test_diffusion_model.py @@ -0,0 +1,24 @@ +import pytest + +from easydynamics.sample_model.diffusion_model.diffusion_model_base import ( + DiffusionModelBase, +) + + +class TestDiffusionModel: + @pytest.fixture + def diffusion_model(self): + return DiffusionModelBase(display_name="TestDiffusionModel", unit="meV") + + def test_init_default(self, diffusion_model): + # WHEN THEN EXPECT + assert diffusion_model.display_name == "TestDiffusionModel" + assert diffusion_model.unit == "meV" + + def test_unit_setter_raises(self, diffusion_model): + # WHEN THEN EXPECT + with pytest.raises( + AttributeError, + match="Unit is read-only. Use convert_unit to change the unit between allowed types", + ): + diffusion_model.unit = "eV" From fab4aee03b4f744446aa8190c47b81ff213badde Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Tue, 6 Jan 2026 14:58:09 +0100 Subject: [PATCH 2/6] polishing --- .../brownian_translational_diffusion.py | 22 +++++++++++-------- .../diffusion_model/diffusion_model_base.py | 3 ++- .../test_brownian_translational_diffusion.py | 2 +- 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py index 36c7cdf..35d6237 100644 --- a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py +++ b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py @@ -1,5 +1,5 @@ from numbers import Number -from typing import Dict, List, Optional, Union +from typing import Dict, List, Union import numpy as np import scipp as sc @@ -34,11 +34,12 @@ class BrownianTranslationalDiffusion(DiffusionModelBase): def __init__( self, - display_name: Optional[str] = "BrownianTranslationalDiffusion", - unit: Optional[Union[str, sc.Unit]] = "meV", - scale: Optional[Union[Parameter, Numeric]] = 1.0, - diffusion_coefficient: Optional[Union[Parameter, Numeric]] = 1.0, - diffusion_unit: Optional[str] = "m**2/s", + display_name: str | None = "BrownianTranslationalDiffusion", + unique_name: str | None = None, + unit: str | sc.Unit = "meV", + scale: Parameter | Numeric = 1.0, + diffusion_coefficient: Parameter | Numeric = 1.0, + diffusion_unit: str = "m**2/s", ): """ Initialize a new BrownianTranslationalDiffusion model. @@ -86,6 +87,7 @@ def __init__( ) super().__init__( display_name=display_name, + unique_name=unique_name, unit=unit, ) self._angstrom = DescriptorNumber("angstrom", 1e-10, unit="m") @@ -209,7 +211,7 @@ def calculate_QISF(self, Q: np.ndarray) -> np.ndarray: def create_component_collections( self, Q: Union[Number, list, np.ndarray], - component_name: str = "Lorentzian", + component_display_name: str = "Lorentzian", ) -> List[ComponentCollection]: """ Create ComponentCollection components for the Brownian translational diffusion model at given Q values. @@ -239,7 +241,7 @@ def create_component_collections( if Q.ndim > 1: raise ValueError("Q must be a 1-dimensional array.") - if not isinstance(component_name, str): + if not isinstance(component_display_name, str): raise TypeError("component_name must be a string.") component_collection_list = [None] * len(Q) @@ -253,7 +255,9 @@ def create_component_collections( ) lorentzian_component = Lorentzian( - display_name=component_name, area=self.scale * QISF[i], unit=self.unit + display_name=component_display_name, + area=self.scale * QISF[i], + unit=self.unit, ) # Make the width dependent on Q diff --git a/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py index 5a75159..91a937c 100644 --- a/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py +++ b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py @@ -10,6 +10,7 @@ class DiffusionModelBase(ModelBase): def __init__( self, display_name="MyDiffusionModel", + unique_name: str | None = None, unit: str | sc.Unit = "meV", ): """ @@ -26,7 +27,7 @@ def __init__( if not (unit is None or isinstance(unit, (str, sc.Unit))): raise TypeError("unit must be None, a string, or a scipp Unit") - super().__init__(display_name=display_name) + super().__init__(display_name=display_name, unique_name=unique_name) self._unit = unit @property diff --git a/tests/diffusion_model/test_brownian_translational_diffusion.py b/tests/diffusion_model/test_brownian_translational_diffusion.py index 908a2ed..4904be2 100644 --- a/tests/diffusion_model/test_brownian_translational_diffusion.py +++ b/tests/diffusion_model/test_brownian_translational_diffusion.py @@ -236,7 +236,7 @@ def test_create_component_collections_component_name_must_be_string( # WHEN THEN EXPECT with pytest.raises(TypeError, match="component_name must be a string."): brownian_diffusion_model.create_component_collections( - Q=np.array([0.1, 0.2, 0.3]), component_name=123 + Q=np.array([0.1, 0.2, 0.3]), component_display_name=123 ) def test_create_component_collections_Q_type_error(self, brownian_diffusion_model): From a118a31f2aa4163ead8c5eb50eb71527e019da9b Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Tue, 6 Jan 2026 15:03:41 +0100 Subject: [PATCH 3/6] Move tests to where they belong --- .../diffusion_model/test_brownian_translational_diffusion.py | 0 .../sample_model}/diffusion_model/test_diffusion_model.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename tests/{ => unit_tests/sample_model}/diffusion_model/test_brownian_translational_diffusion.py (100%) rename tests/{ => unit_tests/sample_model}/diffusion_model/test_diffusion_model.py (100%) diff --git a/tests/diffusion_model/test_brownian_translational_diffusion.py b/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py similarity index 100% rename from tests/diffusion_model/test_brownian_translational_diffusion.py rename to tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py diff --git a/tests/diffusion_model/test_diffusion_model.py b/tests/unit_tests/sample_model/diffusion_model/test_diffusion_model.py similarity index 100% rename from tests/diffusion_model/test_diffusion_model.py rename to tests/unit_tests/sample_model/diffusion_model/test_diffusion_model.py From c5fafafd9934dcb33b3efe67557a4852b3cbb16f Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 8 Jan 2026 14:44:30 +0100 Subject: [PATCH 4/6] Updates based on PR comments --- .../brownian_translational_diffusion.py | 61 ++++++------------- .../diffusion_model/diffusion_model_base.py | 36 ++++++++++- .../test_brownian_translational_diffusion.py | 8 +-- 3 files changed, 56 insertions(+), 49 deletions(-) diff --git a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py index 35d6237..01b2922 100644 --- a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py +++ b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py @@ -1,9 +1,9 @@ -from numbers import Number from typing import Dict, List, Union import numpy as np import scipp as sc from easyscience.variable import DescriptorNumber, Parameter +from numpy.typing import ArrayLike from scipp.constants import hbar as scipp_hbar from easydynamics.sample_model.component_collection import ComponentCollection @@ -14,6 +14,8 @@ Numeric = Union[float, int] +Q_type = np.ndarray | Numeric | list | ArrayLike + class BrownianTranslationalDiffusion(DiffusionModelBase): """ @@ -55,7 +57,7 @@ def __init__( 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. Defaults to 1.0. diffusion_unit : str, optional - Unit for the diffusion coefficient D. Default is "meV*Å**2". Options are 'meV*Å**2' or 'm**2/s' + 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)): @@ -130,7 +132,7 @@ def diffusion_coefficient(self, diffusion_coefficient: Numeric) -> None: raise TypeError("diffusion_coefficient must be a number.") self._diffusion_coefficient.value = diffusion_coefficient - def calculate_width(self, Q: np.ndarray) -> np.ndarray: + def calculate_width(self, Q: Q_type) -> np.ndarray: """ Calculate the half-width at half-maximum (HWHM) for the diffusion model. @@ -145,31 +147,17 @@ def calculate_width(self, Q: np.ndarray) -> np.ndarray: HWHM values in the unit of the model (e.g., meV). """ - if isinstance(Q, Numeric): - Q = np.array([Q]) - - if isinstance(Q, list): - Q = np.array(Q) - - if not isinstance(Q, np.ndarray): - raise TypeError("Q must be a numpy array.") + Q = self._validate_and_convert_Q(Q) - width_list = [] - for Q_value in Q: - # Q is given as a float, so we need to divide by angstrom**2 to get the right units - width = ( - self._hbar - * self.diffusion_coefficient - * Q_value**2 - / (self._angstrom**2) - ) - width.convert_unit(self.unit) - width_list.append(width.value) - width = np.array(width_list) + unit_conversion_factor = ( + self._hbar * self.diffusion_coefficient / (self._angstrom**2) + ) + unit_conversion_factor.convert_unit(self.unit) + width = Q**2 * unit_conversion_factor.value return width - def calculate_EISF(self, Q: np.ndarray) -> np.ndarray: + def calculate_EISF(self, Q: Q_type) -> np.ndarray: """ Calculate the Elastic Incoherent Structure Factor (EISF) for the Brownian translational diffusion model. @@ -183,12 +171,11 @@ def calculate_EISF(self, Q: np.ndarray) -> np.ndarray: np.ndarray EISF values (dimensionless). """ - if not isinstance(Q, np.ndarray): - raise TypeError("Q must be a numpy array.") + Q = self._validate_and_convert_Q(Q) EISF = np.zeros_like(Q) return EISF - def calculate_QISF(self, Q: np.ndarray) -> np.ndarray: + def calculate_QISF(self, Q: Q_type) -> np.ndarray: """ Calculate the Quasi-Elastic Incoherent Structure Factor (QISF). @@ -203,14 +190,13 @@ def calculate_QISF(self, Q: np.ndarray) -> np.ndarray: QISF values (dimensionless). """ - if not isinstance(Q, np.ndarray): - raise TypeError("Q must be a numpy array.") + Q = self._validate_and_convert_Q(Q) QISF = np.ones_like(Q) return QISF def create_component_collections( self, - Q: Union[Number, list, np.ndarray], + Q: Q_type, component_display_name: str = "Lorentzian", ) -> List[ComponentCollection]: """ @@ -228,18 +214,7 @@ def create_component_collections( List[ComponentCollection] List of ComponentCollections with Lorentzian components. """ - - if isinstance(Q, Numeric): - Q = np.array([Q]) - - if isinstance(Q, list): - Q = np.array(Q) - - if not isinstance(Q, np.ndarray): - raise TypeError("Q must be a number, list, or numpy array.") - - if Q.ndim > 1: - raise ValueError("Q must be a 1-dimensional array.") + Q = self._validate_and_convert_Q(Q) if not isinstance(component_display_name, str): raise TypeError("component_name must be a string.") @@ -285,7 +260,7 @@ def _write_width_dependency_expression(self, Q: float) -> str: # Q is given as a float, so we need to add the units return f"hbar * D* {Q} **2*1/(angstrom**2)" - def _write_width_dependency_map_expression(self) -> Dict[str, str]: + def _write_width_dependency_map_expression(self) -> Dict[str, DescriptorNumber]: """ Write the dependency map expression to make dependent Parameters. """ diff --git a/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py index 91a937c..e90f5eb 100644 --- a/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py +++ b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py @@ -1,5 +1,13 @@ +from typing import Union + +import numpy as np import scipp as sc from easyscience.base_classes.model_base import ModelBase +from numpy.typing import ArrayLike + +Numeric = Union[float, int] + +Q_type = np.ndarray | Numeric | list | ArrayLike class DiffusionModelBase(ModelBase): @@ -31,7 +39,7 @@ def __init__( self._unit = unit @property - def unit(self) -> str | sc.Unit: + def unit(self) -> str: """ Get the unit of the DiffusionModel. @@ -39,7 +47,7 @@ def unit(self) -> str | sc.Unit: ------- str or sc.Unit or None """ - return self._unit + return str(self._unit) @unit.setter def unit(self, unit_str: str) -> None: @@ -49,3 +57,27 @@ def unit(self, unit_str: str) -> None: f"or create a new {self.__class__.__name__} with the desired unit." ) ) # noqa: E501 + + def _validate_and_convert_Q(self, Q: Q_type) -> np.ndarray: + """ + Validate and convert Q to a numpy array. + Parameters + ---------- + Q : Number, list, or np.ndarray + Scattering vector values in 1/angstrom. + Returns + ------- + np.ndarray + Q as a numpy array. + """ + if isinstance(Q, Numeric): + Q = np.array([Q]) + if isinstance(Q, list): + Q = np.array(Q) + if not isinstance(Q, np.ndarray): + raise TypeError("Q must be a number, list, or numpy array.") + + if Q.ndim > 1: + raise ValueError("Q must be a 1-dimensional array.") + + return Q diff --git a/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py b/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py index 4904be2..755761e 100644 --- a/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py +++ b/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py @@ -131,7 +131,7 @@ def test_diffusion_coefficient_setter_raises(self, brownian_diffusion_model): def test_calculate_width_type_error(self, brownian_diffusion_model): # WHEN THEN EXPECT - with pytest.raises(TypeError, match="Q must be a numpy array."): + with pytest.raises(TypeError, match="Q must be "): brownian_diffusion_model.calculate_width(Q="invalid") # Invalid type def test_calculate_width(self, brownian_diffusion_model): @@ -179,7 +179,7 @@ def test_calculate_EISF(self, brownian_diffusion_model): def test_calculate_EISF_type_error(self, brownian_diffusion_model): # WHEN THEN EXPECT - with pytest.raises(TypeError, match="Q must be a numpy array."): + with pytest.raises(TypeError, match="Q must be "): brownian_diffusion_model.calculate_EISF(Q="invalid") # Invalid type def test_calculate_QISF(self, brownian_diffusion_model): @@ -195,7 +195,7 @@ def test_calculate_QISF(self, brownian_diffusion_model): def test_calculate_QISF_type_error(self, brownian_diffusion_model): # WHEN THEN EXPECT - with pytest.raises(TypeError, match="Q must be a numpy array."): + with pytest.raises(TypeError, match="Q must be "): brownian_diffusion_model.calculate_QISF(Q="invalid") # Invalid type @pytest.mark.parametrize( @@ -227,7 +227,7 @@ def test_create_component_collections(self, brownian_diffusion_model, Q): component = model.components[0] assert component.display_name == "Lorentzian" assert component.width.unit == brownian_diffusion_model.unit - assert component.width.value == expected_widths[model_index] + assert np.isclose(component.width.value, expected_widths[model_index]) assert component.width.independent is False def test_create_component_collections_component_name_must_be_string( From ff2ddca372392ed4ddac599d416c5e98d3eb51c3 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Fri, 9 Jan 2026 11:28:13 +0100 Subject: [PATCH 5/6] Remove parameters as input --- examples/diffusion_model.ipynb | 53 ++++++++++++++++++- .../brownian_translational_diffusion.py | 24 ++++----- .../test_brownian_translational_diffusion.py | 28 ++-------- 3 files changed, 66 insertions(+), 39 deletions(-) diff --git a/examples/diffusion_model.ipynb b/examples/diffusion_model.ipynb index f2fb212..6ec8051 100644 --- a/examples/diffusion_model.ipynb +++ b/examples/diffusion_model.ipynb @@ -27,7 +27,8 @@ "# Q is in Angstrom^-1 and energy in meV.\n", "\n", "Q=np.linspace(0.5,2,7)\n", - "energy=np.linspace(-2, 2, 501)\n", + "# energy=np.linspace(-2, 2, 501)\n", + "energy=sc.linspace(start=-2,stop=2,num=501,unit=\"meV\",dim='energy')\n", "scale=1.0\n", "diffusion_coefficient = 2.4e-9 # m^2/s\n", "diffusion_unit= \"m**2/s\"\n", @@ -51,6 +52,56 @@ "plt.ylabel('Intensity (arb. units)')\n", "plt.title('Brownian Translational Diffusion Model') " ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dce619d8", + "metadata": {}, + "outputs": [], + "source": [ + "energy=np.linspace(-2, 2, 501)\n", + "energy=sc.linspace(start=-2,stop=2,num=501,unit=\"meV\",dim='energy')\n", + "energy" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "555cb19a", + "metadata": {}, + "outputs": [], + "source": [ + "component_collections[0].get_all_parameters()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7247d79d", + "metadata": {}, + "outputs": [], + "source": [ + "component_collections[0].get_all_parameters()[2].dependency_expression" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "04e99b93", + "metadata": {}, + "outputs": [], + "source": [ + "diffusion_model.diffusion_coefficient=5.0e-9" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "486cb003", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py index 01b2922..afec564 100644 --- a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py +++ b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py @@ -39,8 +39,8 @@ def __init__( display_name: str | None = "BrownianTranslationalDiffusion", unique_name: str | None = None, unit: str | sc.Unit = "meV", - scale: Parameter | Numeric = 1.0, - diffusion_coefficient: Parameter | Numeric = 1.0, + scale: Numeric = 1.0, + diffusion_coefficient: Numeric = 1.0, diffusion_unit: str = "m**2/s", ): """ @@ -61,10 +61,10 @@ def __init__( """ if not isinstance(scale, (Parameter, Numeric)): - raise TypeError("scale must be a Parameter or a number.") + raise TypeError("scale must be a number.") if not isinstance(diffusion_coefficient, (Parameter, Numeric)): - raise TypeError("diffusion_coefficient must be a Parameter or a number.") + 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'.") @@ -77,16 +77,14 @@ def __init__( else: raise ValueError("diffusion_unit must be 'meV*Å**2' or 'm**2/s'.") - if not isinstance(scale, Parameter): - scale = Parameter(name="scale", value=float(scale), fixed=False, min=0.0) + scale = Parameter(name="scale", value=float(scale), fixed=False, min=0.0) - if not isinstance(diffusion_coefficient, Parameter): - diffusion_coefficient = Parameter( - name="diffusion_coefficient", - value=float(diffusion_coefficient), - fixed=False, - unit=diffusion_unit, - ) + diffusion_coefficient = Parameter( + name="diffusion_coefficient", + value=float(diffusion_coefficient), + fixed=False, + unit=diffusion_unit, + ) super().__init__( display_name=display_name, unique_name=unique_name, diff --git a/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py b/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py index 755761e..cfa9ba3 100644 --- a/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py +++ b/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py @@ -1,7 +1,7 @@ import numpy as np import pytest import scipp as sc -from easyscience.variable import DescriptorNumber, Parameter +from easyscience.variable import DescriptorNumber from scipp.constants import hbar as scipp_hbar from easydynamics.sample_model.diffusion_model.brownian_translational_diffusion import ( @@ -44,7 +44,7 @@ def test_init_default(self, brownian_diffusion_model): "diffusion_coefficient": 1.0, "diffusion_unit": "m**2/s", }, - "scale must be a Parameter or a number.", + "scale must be a number", ), ( { @@ -53,7 +53,7 @@ def test_init_default(self, brownian_diffusion_model): "diffusion_coefficient": "invalid", "diffusion_unit": "m**2/s", }, - "diffusion_coefficient must be a Parameter or a number.", + "diffusion_coefficient must be a number", ), ( { @@ -83,28 +83,6 @@ def test_diffusion_unit_value_error(self): diffusion_unit="invalid_unit", ) - def test_init_with_parameters(self): - # WHEN - - scale = Parameter(name="scale_param", value=2.0) - diffusion_coefficient = Parameter( - name="diffusion_coefficient", value=3.0, unit="m**2/s" - ) - - # THEN - brownian_diffusion_model = BrownianTranslationalDiffusion( - display_name="CustomBrownianDiffusion", - unit="meV", - scale=scale, - diffusion_coefficient=diffusion_coefficient, - ) - - # EXPECT - assert brownian_diffusion_model.display_name == "CustomBrownianDiffusion" - assert brownian_diffusion_model.unit == "meV" - assert brownian_diffusion_model.scale is scale - assert brownian_diffusion_model.diffusion_coefficient is diffusion_coefficient - def test_scale_setter(self, brownian_diffusion_model): # WHEN brownian_diffusion_model.scale = 2.0 From 5df823dcb4e5d3a6bf6e21ba88723ee549b6fb0e Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Fri, 9 Jan 2026 11:45:07 +0100 Subject: [PATCH 6/6] Updates based on PR comments --- .../brownian_translational_diffusion.py | 30 ++++++++++++++----- .../diffusion_model/diffusion_model_base.py | 17 +++++++++-- .../test_brownian_translational_diffusion.py | 15 +++++++--- 3 files changed, 48 insertions(+), 14 deletions(-) diff --git a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py index afec564..5efdc9e 100644 --- a/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py +++ b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py @@ -50,6 +50,8 @@ def __init__( ---------- display_name : str Display name of the diffusion model. + unique_name : str or None + Unique name of the diffusion model. If None, a unique name is automatically generated. unit : str or sc.Unit, optional Energy unit for the underlying Lorentzian components. Defaults to "meV". scale : float or Parameter, optional @@ -108,6 +110,9 @@ def scale(self) -> Parameter: @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 @@ -126,6 +131,9 @@ 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)): raise TypeError("diffusion_coefficient must be a number.") self._diffusion_coefficient.value = diffusion_coefficient @@ -136,7 +144,7 @@ def calculate_width(self, Q: Q_type) -> np.ndarray: Parameters ---------- - Q : np.ndarray + Q : np.ndarray | Numeric | list | ArrayLike Scattering vector in 1/angstrom Returns @@ -161,7 +169,7 @@ def calculate_EISF(self, Q: Q_type) -> np.ndarray: Parameters ---------- - Q : np.ndarray + Q : np.ndarray | Numeric | list | ArrayLike Scattering vector in 1/angstrom Returns @@ -179,7 +187,7 @@ def calculate_QISF(self, Q: Q_type) -> np.ndarray: Parameters ---------- - Q : np.ndarray + Q : np.ndarray | Numeric | list | ArrayLike Scattering vector in 1/angstrom Returns @@ -203,10 +211,8 @@ def create_component_collections( ---------- Q : Number, list, or np.ndarray Scattering vector values. - component_name : str + component_display_name : str Name of the Lorentzian component. - width_name : str - Name of the width parameter. Returns ------- List[ComponentCollection] @@ -222,9 +228,9 @@ def create_component_collections( QISF = self.calculate_QISF(Q) # Create a Lorentzian component for each Q-value, with width D*Q^2 and area equal to scale. No delta function, as the EISF is 0. - for i in range(len(Q)): + for i, Q_value in enumerate(Q): component_collection_list[i] = ComponentCollection( - display_name=f"{self.display_name}_Q{Q[i]:.2f}", unit=self.unit + display_name=f"{self.display_name}_Q{Q_value:.2f}", unit=self.unit ) lorentzian_component = Lorentzian( @@ -251,6 +257,14 @@ def create_component_collections( 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. + Parameters + ---------- + Q : float + Scattering vector in 1/angstrom + Returns + ------- + str + Dependency expression for the width. """ if not isinstance(Q, (float)): raise TypeError("Q must be a float.") diff --git a/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py index e90f5eb..f7bcc55 100644 --- a/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py +++ b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py @@ -3,7 +3,9 @@ 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 scipp import UnitError Numeric = Union[float, int] @@ -32,8 +34,13 @@ def __init__( Unit of the diffusion model. Defaults to "meV". """ - if not (unit is None or isinstance(unit, (str, sc.Unit))): - raise TypeError("unit must be None, a string, or a scipp Unit") + try: + test = DescriptorNumber(name="test", value=1, unit=unit) + test.convert_unit("meV") + except Exception as e: + raise UnitError( + f"Invalid unit: {unit}. Unit must be a string or scipp Unit and convertible to meV." + ) from e super().__init__(display_name=display_name, unique_name=unique_name) self._unit = unit @@ -81,3 +88,9 @@ def _validate_and_convert_Q(self, Q: Q_type) -> np.ndarray: raise ValueError("Q must be a 1-dimensional array.") return Q + + def __repr__(self): + """ + String representation of the Diffusion model. + """ + return f"{self.__class__.__name__}(display_name={self.display_name}, unit={self.unit})" diff --git a/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py b/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py index cfa9ba3..0de1d6f 100644 --- a/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py +++ b/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py @@ -2,6 +2,7 @@ import pytest import scipp as sc from easyscience.variable import DescriptorNumber +from scipp import UnitError from scipp.constants import hbar as scipp_hbar from easydynamics.sample_model.diffusion_model.brownian_translational_diffusion import ( @@ -26,7 +27,7 @@ def test_init_default(self, brownian_diffusion_model): assert brownian_diffusion_model.diffusion_coefficient.value == 1.0 @pytest.mark.parametrize( - "kwargs, expected_message", + "kwargs,expected_exception, expected_message", [ ( { @@ -35,7 +36,8 @@ def test_init_default(self, brownian_diffusion_model): "diffusion_coefficient": 1.0, "diffusion_unit": "m**2/s", }, - "unit must be None, a string, or a scipp Unit", + UnitError, + "Invalid unit", ), ( { @@ -44,6 +46,7 @@ def test_init_default(self, brownian_diffusion_model): "diffusion_coefficient": 1.0, "diffusion_unit": "m**2/s", }, + TypeError, "scale must be a number", ), ( @@ -53,6 +56,7 @@ def test_init_default(self, brownian_diffusion_model): "diffusion_coefficient": "invalid", "diffusion_unit": "m**2/s", }, + TypeError, "diffusion_coefficient must be a number", ), ( @@ -62,12 +66,15 @@ def test_init_default(self, brownian_diffusion_model): "diffusion_coefficient": 1.0, "diffusion_unit": 123, }, + TypeError, "diffusion_unit must be ", ), ], ) - def test_input_type_validation_raises(self, kwargs, expected_message): - with pytest.raises(TypeError, match=expected_message): + def test_input_type_validation_raises( + self, kwargs, expected_exception, expected_message + ): + with pytest.raises(expected_exception, match=expected_message): BrownianTranslationalDiffusion( display_name="BrownianTranslationalDiffusion", **kwargs )