diff --git a/examples/diffusion_model.ipynb b/examples/diffusion_model.ipynb new file mode 100644 index 0000000..6ec8051 --- /dev/null +++ b/examples/diffusion_model.ipynb @@ -0,0 +1,128 @@ +{ + "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", + "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", + "\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') " + ] + }, + { + "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": { + "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..5efdc9e --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/brownian_translational_diffusion.py @@ -0,0 +1,289 @@ +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 +from easydynamics.sample_model.components import Lorentzian +from easydynamics.sample_model.diffusion_model.diffusion_model_base import ( + DiffusionModelBase, +) + +Numeric = Union[float, int] + +Q_type = np.ndarray | Numeric | list | ArrayLike + + +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: str | None = "BrownianTranslationalDiffusion", + unique_name: str | None = None, + 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. + + Parameters + ---------- + 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 + 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 m**2/s. Options are 'meV*Å**2' or 'm**2/s' + + """ + if not isinstance(scale, (Parameter, Numeric)): + raise TypeError("scale must be a number.") + + if not isinstance(diffusion_coefficient, (Parameter, 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, + ) + super().__init__( + display_name=display_name, + unique_name=unique_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: + """ + Set the scale parameter of the diffusion model. + """ + 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: + """ + 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 + + def calculate_width(self, Q: Q_type) -> np.ndarray: + """ + Calculate the half-width at half-maximum (HWHM) for the diffusion model. + + Parameters + ---------- + Q : np.ndarray | Numeric | list | ArrayLike + Scattering vector in 1/angstrom + + Returns + ------- + np.ndarray + HWHM values in the unit of the model (e.g., meV). + """ + + Q = self._validate_and_convert_Q(Q) + + 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: Q_type) -> np.ndarray: + """ + Calculate the Elastic Incoherent Structure Factor (EISF) for the Brownian translational diffusion model. + + Parameters + ---------- + Q : np.ndarray | Numeric | list | ArrayLike + Scattering vector in 1/angstrom + + Returns + ------- + np.ndarray + EISF values (dimensionless). + """ + Q = self._validate_and_convert_Q(Q) + EISF = np.zeros_like(Q) + return EISF + + def calculate_QISF(self, Q: Q_type) -> np.ndarray: + """ + Calculate the Quasi-Elastic Incoherent Structure Factor (QISF). + + Parameters + ---------- + Q : np.ndarray | Numeric | list | ArrayLike + Scattering vector in 1/angstrom + + Returns + ------- + np.ndarray + QISF values (dimensionless). + """ + + Q = self._validate_and_convert_Q(Q) + QISF = np.ones_like(Q) + return QISF + + def create_component_collections( + self, + Q: Q_type, + component_display_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_display_name : str + Name of the Lorentzian component. + Returns + ------- + List[ComponentCollection] + List of ComponentCollections with Lorentzian components. + """ + Q = self._validate_and_convert_Q(Q) + + if not isinstance(component_display_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, Q_value in enumerate(Q): + component_collection_list[i] = ComponentCollection( + display_name=f"{self.display_name}_Q{Q_value:.2f}", unit=self.unit + ) + + lorentzian_component = Lorentzian( + display_name=component_display_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. + 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.") + + # 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, DescriptorNumber]: + """ + 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..f7bcc55 --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/diffusion_model_base.py @@ -0,0 +1,96 @@ +from typing import Union + +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] + +Q_type = np.ndarray | Numeric | list | ArrayLike + + +class DiffusionModelBase(ModelBase): + """ + Base class for constructing diffusion models. + """ + + def __init__( + self, + display_name="MyDiffusionModel", + unique_name: str | None = None, + 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". + """ + + 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 + + @property + def unit(self) -> str: + """ + Get the unit of the DiffusionModel. + + Returns + ------- + str or sc.Unit or None + """ + return str(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 + + 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 + + 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 new file mode 100644 index 0000000..0de1d6f --- /dev/null +++ b/tests/unit_tests/sample_model/diffusion_model/test_brownian_translational_diffusion.py @@ -0,0 +1,277 @@ +import numpy as np +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 ( + 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_exception, expected_message", + [ + ( + { + "unit": 123, + "scale": 1.0, + "diffusion_coefficient": 1.0, + "diffusion_unit": "m**2/s", + }, + UnitError, + "Invalid unit", + ), + ( + { + "unit": 123, + "scale": "invalid", + "diffusion_coefficient": 1.0, + "diffusion_unit": "m**2/s", + }, + TypeError, + "scale must be a number", + ), + ( + { + "unit": 123, + "scale": 1.0, + "diffusion_coefficient": "invalid", + "diffusion_unit": "m**2/s", + }, + TypeError, + "diffusion_coefficient must be a number", + ), + ( + { + "unit": 123, + "scale": 1.0, + "diffusion_coefficient": 1.0, + "diffusion_unit": 123, + }, + TypeError, + "diffusion_unit must be ", + ), + ], + ) + 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 + ) + + 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_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 "): + 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 "): + 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 "): + 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 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( + 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_display_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/unit_tests/sample_model/diffusion_model/test_diffusion_model.py b/tests/unit_tests/sample_model/diffusion_model/test_diffusion_model.py new file mode 100644 index 0000000..1c50266 --- /dev/null +++ b/tests/unit_tests/sample_model/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"