diff --git a/docs/docs/tutorials/diffusion_model.ipynb b/docs/docs/tutorials/diffusion_model.ipynb index fa264b75..dc19c92d 100644 --- a/docs/docs/tutorials/diffusion_model.ipynb +++ b/docs/docs/tutorials/diffusion_model.ipynb @@ -6,7 +6,9 @@ "metadata": {}, "source": [ "# Diffusion Model\n", - "We support several standard models of diffusion. Here we show examples of Brownian Translational Diffusion, Jump Translational Diffusion and Chudley-Elliott Jump Diffusion. In each model the scattering is a Lorentzian for every $Q$, and the models differ in how the width $\\Gamma(Q)$ depends on $Q$. For Brownian Translational Diffusion, the width is given by $\\Gamma = D Q^2$, where $D$ is the diffusion coefficient (in m$^2$/s) and $Q$ is the momentum transfer." + "We support several standard models of diffusion. Here we show examples of Brownian Translational Diffusion, Jump Translational Diffusion, Chudley-Elliott Jump Diffusion and Hall-Ross Jump Diffusion. In each model the scattering is a Lorentzian for every $Q$, and the models differ in how the width $\\Gamma(Q)$ depends on $Q$. For Brownian Translational Diffusion, the width is given by $\\Gamma = D Q^2$, where $D$ is the diffusion coefficient (in m$^2$/s) and $Q$ is the momentum transfer.\n", + "\n", + "Models with an elastic line and several Lorentzians have their own tutorials: [equivalent sites on a circle](equivalent_sites_circle.ipynb), [jump sites with a log-normal distribution](jump_sites_log_normal_distribution.ipynb) and [the 3D Gaussian model](gaussian_model_3d.ipynb)." ] }, { @@ -204,13 +206,7 @@ "metadata": {}, "source": [ "## Chudley-Elliott Jump Diffusion\n", - "The Chudley-Elliott model of\n", - "\n", - "> C. T. Chudley and R. J. Elliott, *Neutron Scattering from a Liquid on a Jump Diffusion Model*,\n", - "> Proc. Phys. Soc. **77**, 353 (1961),\n", - "> [doi:10.1088/0370-1328/77/2/319](https://doi.org/10.1088/0370-1328/77/2/319).\n", - "\n", - "describes a particle that sits on a site for a residence time $\\tau$ and then jumps a fixed distance $l$. The width is\n", + "The Chudley-Elliott model1 describes a particle that sits on a site for a residence time $\\tau$ and then jumps a fixed distance $l$. The width is\n", "\n", "$$\\Gamma(Q) = \\frac{\\hbar}{\\tau}\\left(1 - \\frac{\\sin(Ql)}{Ql}\\right),$$\n", "\n", @@ -295,6 +291,115 @@ "plt.title('HWHM vs Q for Chudley-Elliott Jump Diffusion')\n", "plt.show()" ] + }, + { + "cell_type": "markdown", + "id": "17", + "metadata": {}, + "source": [ + "## Hall-Ross Jump Diffusion\n", + "The Hall-Ross model2 describes a particle that sits on a site for a residence time $\\tau$ and then jumps a random distance drawn from a Gaussian distribution. The width is\n", + "\n", + "$$\\Gamma(Q) = \\frac{\\hbar}{\\tau}\\left(1 - \\exp\\left(-\\frac{Q^2 l^2}{2}\\right)\\right),$$\n", + "\n", + "where $\\tau$ is the residence time (in ps) and $l$ is the standard deviation of the jump-length distribution (in Å).3 The mean-square jump length is $\\langle r^2 \\rangle = 3 l^2$. Unlike Chudley-Elliott, the width does not oscillate: it rises monotonically from $\\hbar D Q^2$, with $D = l^2 / (2\\tau)$, to $\\hbar/\\tau$." + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "### Model for different Q values" + ] + }, + { + "cell_type": "code", + "id": "19", + "metadata": {}, + "source": [ + "# Create Hall-Ross Jump Diffusion model\n", + "# 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", + "residence_time = 1.0 # ps\n", + "jump_length = 1.5 / np.sqrt(3) # Å, the standard deviation l\n", + "\n", + "hall_ross_model = edyn.HallRossJumpDiffusion(\n", + " display_name='HallRossModel',\n", + " scale=scale,\n", + " residence_time=residence_time,\n", + " jump_length=jump_length,\n", + " Q=Q,\n", + ")\n", + "\n", + "component_collections = hall_ross_model.get_component_collections()\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.xlabel('Energy (meV)')\n", + "plt.ylabel('Intensity (arb. units)')\n", + "plt.title('Hall-Ross Jump Diffusion Model')\n", + "plt.show()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "### Width as a function of Q\n", + "The standard deviation above was chosen as $l = 1.5/\\sqrt{3}$ Å, so that the mean-square jump length, $3 l^2 = 2.25$ Å$^2$, equals that of the Chudley-Elliott model above with its fixed jump of 1.5 Å. Both models then have the same diffusion coefficient and agree at low $Q$; at high $Q$ the Hall-Ross width approaches $\\hbar/\\tau$ without oscillating." + ] + }, + { + "cell_type": "code", + "id": "21", + "metadata": {}, + "source": [ + "# Calculate and plot the half width at half maximum (HWHM) as function\n", + "# of Q, together with the Chudley-Elliott model with the same mean-square jump length\n", + "Q = np.linspace(0.1, 4, 201)\n", + "HWHM = hall_ross_model.calculate_width(Q)\n", + "plt.figure()\n", + "plt.plot(Q, HWHM, label='Hall-Ross')\n", + "plt.plot(Q, chudley_elliott_model.calculate_width(Q), linestyle='dashed', label='Chudley-Elliott')\n", + "plt.xlabel('Q (Å$^{-1}$)')\n", + "plt.ylabel('HWHM (meV)')\n", + "plt.xlim(0, 4.2)\n", + "plt.ylim(0, max(HWHM) * 1.2)\n", + "plt.title('HWHM vs Q for Hall-Ross Jump Diffusion')\n", + "plt.legend()\n", + "plt.show()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "1. C. T. Chudley and R. J. Elliott, *Neutron Scattering from a Liquid on a Jump Diffusion Model*, Proc. Phys. Soc. **77**, 353 (1961), [doi:10.1088/0370-1328/77/2/319](https://doi.org/10.1088/0370-1328/77/2/319).\n", + "\n", + "2. P. L. Hall and D. K. Ross, *Incoherent neutron scattering functions for random jump diffusion in bounded and infinite media*, Mol. Phys. **42**, 673 (1981), [doi:10.1080/00268978100100521](https://doi.org/10.1080/00268978100100521).\n", + "\n", + "3. The same convention as Mantid's [`HallRoss` fit function](https://github.com/mantidproject/mantid/blob/17c1cd96c1d46b39bee21dc75a60017b658d689c/Framework/PythonInterface/plugins/functions/HallRoss.py#L30-L39), whose `L` parameter is this $l$." + ] } ], "metadata": { diff --git a/docs/docs/tutorials/equivalent_sites_circle.ipynb b/docs/docs/tutorials/equivalent_sites_circle.ipynb new file mode 100644 index 00000000..ea9a6d66 --- /dev/null +++ b/docs/docs/tutorials/equivalent_sites_circle.ipynb @@ -0,0 +1,258 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Equivalent sites on a circle\n", + "The model1,2 describes jumps between $N$ equivalent sites on a circle of radius $R$, for example a methyl group ($N = 3$) rotating about its axis. The scattering is an elastic line plus $N - 1$ Lorentzians:\n", + "\n", + "$$ S(Q, E) = A_0(Q)\\,\\delta(E) + \\sum_{l=1}^{N-1} A_l(Q)\\, L(E, \\Gamma_l) $$\n", + "\n", + "with\n", + "\n", + "$$ A_l(Q) = \\frac{1}{N} \\sum_{k=1}^{N} j_0(Q r_k) \\cos\\left(\\frac{2\\pi l k}{N}\\right),\n", + "\\quad r_k = 2R \\sin\\left(\\frac{\\pi k}{N}\\right),\n", + "\\quad \\Gamma_l = \\frac{2\\hbar}{\\tau} \\sin^2\\left(\\frac{\\pi l}{N}\\right), $$\n", + "\n", + "where $\\tau$ is the residence time on a site and $j_0(x) = \\sin(x)/x$. $A_0$ is the Elastic Incoherent Structure Factor (EISF), $A_l$ is the Quasi-Elastic Incoherent Structure Factor (QISF) of order $l$, and they sum to one. The widths do not depend on $Q$; only the weights do.\n", + "\n", + "The width is also written as $4 \\sin^2(\\pi l / N) / \\tau$, with a $\\tau$ that is twice the residence time used here.3" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import easydynamics as edyn\n", + "\n", + "%matplotlib widget" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## The model at different Q values" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "# Q is in Angstrom^-1, the radius in Angstrom and the residence time in ps.\n", + "Q = np.linspace(0.5, 2, 7)\n", + "\n", + "circle_model = edyn.EquivalentSitesCircle(\n", + " scale=1.0,\n", + " radius=1.0,\n", + " residence_time=2.0,\n", + " number_of_sites=3,\n", + " Q=Q,\n", + ")\n", + "print(circle_model)" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "With $N = 3$ sites the model has $N - 1 = 2$ Lorentzians per $Q$. The radius is fixed by default, since it is normally known from the molecular geometry; free it with `circle_model.radius.fixed = False` to fit it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "# Plot the quasi-elastic part (all Lorentzians) at every Q. The delta function\n", + "# would be a spike at E = 0; its weight, the EISF, is plotted further down.\n", + "energy = np.linspace(-2, 2, 501)\n", + "component_collections = circle_model.get_component_collections()\n", + "\n", + "cmap = plt.cm.jet\n", + "nQ = len(component_collections)\n", + "plt.figure()\n", + "for Q_index, collection in enumerate(component_collections):\n", + " quasielastic = sum(\n", + " component.evaluate(energy)\n", + " for component in collection\n", + " if isinstance(component, edyn.Lorentzian)\n", + " )\n", + " color = cmap(Q_index / (nQ - 1))\n", + " plt.plot(energy, quasielastic, label=f'Q={Q[Q_index]:.2f} Å^-1', color=color)\n", + "\n", + "plt.legend()\n", + "plt.xlabel('Energy (meV)')\n", + "plt.ylabel('Intensity (arb. units)')\n", + "plt.title('Equivalent sites on a circle, N = 3')\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "## The widths of the Lorentzians" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "for order in range(1, circle_model.number_of_lorentzians + 1):\n", + " width = circle_model.calculate_width(order=order)[0]\n", + " print(f'Lorentzian of order {order}: HWHM = {width:.4f} meV')" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": {}, + "source": [ + "For $N = 3$ both orders have the same width, $(2\\hbar/\\tau)\\sin^2(\\pi/3) = (2\\hbar/\\tau)\\sin^2(2\\pi/3)$, so the quasi-elastic part is a single Lorentzian. For larger $N$ the orders have different widths." + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "## EISF and QISF as a function of Q" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "Q_fine = np.linspace(0.01, 4, 200)\n", + "\n", + "plt.figure()\n", + "plt.plot(Q_fine, circle_model.calculate_EISF(Q_fine), label='EISF, $A_0$')\n", + "for order in range(1, circle_model.number_of_lorentzians + 1):\n", + " plt.plot(\n", + " Q_fine,\n", + " circle_model.calculate_QISF_order(Q_fine, order=order),\n", + " label=f'QISF of order {order}, $A_{order}$',\n", + " )\n", + "plt.plot(\n", + " Q_fine,\n", + " circle_model.calculate_EISF(Q_fine) + circle_model.calculate_QISF(Q_fine),\n", + " color='k',\n", + " linestyle='dashed',\n", + " label='EISF + QISF',\n", + ")\n", + "plt.xlabel('Q (Å$^{-1}$)')\n", + "plt.ylabel('Weight')\n", + "plt.title('Weights for N = 3')\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "The EISF depends on the number of sites. For $N = 2$ this is the two-site jump model with jump distance $2R$; for large $N$ it approaches continuous rotational diffusion on the circle." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure()\n", + "for number_of_sites in [2, 3, 6]:\n", + " model = edyn.EquivalentSitesCircle(\n", + " radius=1.0, residence_time=2.0, number_of_sites=number_of_sites\n", + " )\n", + " plt.plot(Q_fine, model.calculate_EISF(Q_fine), label=f'N = {number_of_sites}')\n", + "plt.xlabel('Q (Å$^{-1}$)')\n", + "plt.ylabel('EISF')\n", + "plt.title('EISF for different numbers of sites')\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## Components and fit targets\n", + "`EquivalentSitesCircle` is built on `DeltaLorentzianDiffusionModelBase`, the same base class as `JumpSitesLogNormalDistribution` and `GaussianModel3D`. Every collection holds a delta function and the Lorentzians of order $l$, named `_l`. Their areas and widths are dependent Parameters, so they follow the model parameters during a fit." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "collection = circle_model.get_component_collections()[0]\n", + "print([component.name for component in collection])\n", + "print([fit_target.name for fit_target in circle_model.get_fit_targets()])" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "1. J. D. Barnes, J. Chem. Phys. **58**, 5193 (1973), [doi:10.1063/1.1679130](https://doi.org/10.1063/1.1679130).\n", + "\n", + "2. M. Bée, *Quasielastic Neutron Scattering*, Adam Hilger, Bristol (1988).\n", + "\n", + "3. Mantid, [`InelasticDiffRotDiscreteCircle`](https://github.com/mantidproject/mantid/blob/6ec66a6f524a7ec9d4258351fe3001c66c63dbe5/Framework/CurveFitting/src/Functions/InelasticDiffRotDiscreteCircle.cpp#L79-L108): its `Decay` parameter is this $2\\tau$." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "default", + "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.14.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/docs/tutorials/gaussian_model_3d.ipynb b/docs/docs/tutorials/gaussian_model_3d.ipynb new file mode 100644 index 00000000..3f6b82e9 --- /dev/null +++ b/docs/docs/tutorials/gaussian_model_3d.ipynb @@ -0,0 +1,250 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Diffusion in an isotropic 3D harmonic well\n", + "The model1 describes a particle that diffuses with diffusion coefficient $D$ in a harmonic potential that is the same in all three directions, so its position about the centre of the well has a Gaussian distribution. $\\langle u^2 \\rangle$ is the mean square distance from the centre, summed over the three directions. Because the well is isotropic, the scattering depends only on the magnitude of $Q$:\n", + "\n", + "$$ S(Q, E) = A_0(Q)\\,\\delta(E) + \\sum_{i=1}^{n} A_i(Q)\\, L(E, \\Gamma_i) $$\n", + "\n", + "with\n", + "\n", + "$$ A_i(Q) = e^{-y} \\frac{y^i}{i!}, \\quad y = \\frac{Q^2 \\langle u^2 \\rangle}{3},\n", + "\\quad \\Gamma_i = \\frac{3 i \\hbar D}{\\langle u^2 \\rangle}. $$\n", + "\n", + "$A_0 = e^{-y}$ is the Elastic Incoherent Structure Factor (EISF) and $A_i$ the Quasi-Elastic Incoherent Structure Factor (QISF) of order $i$. The paper1 uses the variance along one direction, $\\langle u_x^2 \\rangle = \\langle u^2 \\rangle / 3$.\n", + "\n", + "The QISF form a Poisson distribution in $i$ with mean $y$, and together with the EISF they sum to one. The infinite sum is therefore truncated at $n$ = `number_of_lorentzians`, which must be well above $y$ at the largest $Q$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import easydynamics as edyn\n", + "\n", + "%matplotlib widget" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## The model at different Q values" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "# Q is in Angstrom^-1, D in m^2/s and in Angstrom^2.\n", + "Q = np.linspace(0.5, 2, 7)\n", + "\n", + "gaussian_model = edyn.GaussianModel3D(\n", + " scale=1.0,\n", + " diffusion_coefficient=2e-9,\n", + " mean_u_squared=1.5,\n", + " number_of_lorentzians=20,\n", + " Q=Q,\n", + ")\n", + "print(gaussian_model)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "# Plot the quasi-elastic part (all Lorentzians) at every Q. The delta function\n", + "# would be a spike at E = 0; its weight, the EISF, is plotted further down.\n", + "energy = np.linspace(-2, 2, 501)\n", + "component_collections = gaussian_model.get_component_collections()\n", + "\n", + "cmap = plt.cm.jet\n", + "nQ = len(component_collections)\n", + "plt.figure()\n", + "for Q_index, collection in enumerate(component_collections):\n", + " quasielastic = sum(\n", + " component.evaluate(energy)\n", + " for component in collection\n", + " if isinstance(component, edyn.Lorentzian)\n", + " )\n", + " color = cmap(Q_index / (nQ - 1))\n", + " plt.plot(energy, quasielastic, label=f'Q={Q[Q_index]:.2f} Å^-1', color=color)\n", + "\n", + "plt.legend()\n", + "plt.xlabel('Energy (meV)')\n", + "plt.ylabel('Intensity (arb. units)')\n", + "plt.title('3D Gaussian model')\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## Widths and weights" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "The widths grow linearly with the order and do not depend on $Q$. The weights do: at small $Q$ almost all weight is in the elastic line and the first Lorentzians, while at large $Q$ it moves to higher orders." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "orders = np.arange(1, gaussian_model.number_of_lorentzians + 1)\n", + "widths = [gaussian_model.calculate_width(order=order)[0] for order in orders]\n", + "print(f'HWHM of order 1: {widths[0]:.4f} meV, of order 20: {widths[-1]:.4f} meV')\n", + "\n", + "plt.figure()\n", + "for Q_value in [0.5, 1.5, 3.0]:\n", + " weights = [gaussian_model.calculate_QISF_order(Q_value, order=order)[0] for order in orders]\n", + " plt.plot(orders, weights, marker='o', label=f'Q={Q_value} Å^-1')\n", + "plt.xlabel('Order i')\n", + "plt.ylabel('QISF of order i, $A_i$')\n", + "plt.title('QISF $A_i$ of each Lorentzian (a Poisson distribution in i)')\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": {}, + "source": [ + "## EISF and QISF as a function of Q" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "The EISF, $\\exp(-Q^2 \\langle u^2 \\rangle / 3)$, is the same Debye-Waller factor as in `DeltaLorentz`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "Q_fine = np.linspace(0.01, 4, 200)\n", + "delta_lorentz = edyn.DeltaLorentz(mean_u_squared=1.5)\n", + "\n", + "plt.figure()\n", + "plt.plot(Q_fine, gaussian_model.calculate_EISF(Q_fine), label='EISF')\n", + "plt.plot(Q_fine, gaussian_model.calculate_QISF(Q_fine), label='QISF')\n", + "plt.plot(\n", + " Q_fine,\n", + " delta_lorentz.calculate_EISF(Q_fine),\n", + " color='k',\n", + " linestyle='dotted',\n", + " label='DeltaLorentz EISF',\n", + ")\n", + "plt.xlabel('Q (Å$^{-1}$)')\n", + "plt.ylabel('Weight')\n", + "plt.title('EISF and QISF of the 3D Gaussian model')\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## Truncation" + ] + }, + { + "cell_type": "markdown", + "id": "12", + "metadata": {}, + "source": [ + "Truncating the sum leaves out the weight of the orders above $n$. A warning is raised when that missing weight exceeds $10^{-3}$ at any of the model's $Q$ values, and `calculate_missing_weight` returns it. As a rule, $n$ must be well above $y = Q^2 \\langle u^2 \\rangle / 3$ at the largest $Q$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "Q_wide = np.linspace(0.5, 3, 6)\n", + "\n", + "with warnings.catch_warnings(record=True) as caught:\n", + " warnings.simplefilter('always')\n", + " truncated_model = edyn.GaussianModel3D(\n", + " diffusion_coefficient=2e-9, mean_u_squared=1.5, number_of_lorentzians=5, Q=Q_wide\n", + " )\n", + "for warning in caught:\n", + " print(warning.message)\n", + "\n", + "print('Missing weight per Q:', np.round(truncated_model.calculate_missing_weight(), 4))\n", + "largest = gaussian_model.calculate_missing_weight(Q_wide).max()\n", + "print(f'With 20 Lorentzians the largest missing weight is {largest:.1e}')" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "1. F. Volino, J.-C. Perrin and S. Lyonnard, J. Phys. Chem. B **110**, 11217 (2006), [doi:10.1021/jp061103s](https://doi.org/10.1021/jp061103s)." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "dynamics-lib:py-314-env (3.14.8)", + "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.14.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/docs/tutorials/jump_sites_log_normal_distribution.ipynb b/docs/docs/tutorials/jump_sites_log_normal_distribution.ipynb new file mode 100644 index 00000000..c3051ddf --- /dev/null +++ b/docs/docs/tutorials/jump_sites_log_normal_distribution.ipynb @@ -0,0 +1,238 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Jump sites with a log-normal distribution of jump rates\n", + "The model1 extends the [equivalent sites on a circle](equivalent_sites_circle.ipynb) model to a disordered environment, for example methyl groups in a glassy polymer, where every group has its own barrier. The jump rates then follow a log-normal distribution of standard deviation $\\sigma$, centred on the residence time $\\tau$:\n", + "\n", + "$$ S(Q, E) = A_0(Q)\\,\\delta(E) + \\sum_{l=1}^{N-1} A_l(Q) \\sum_{s=0}^{20} g_s\\, L(E, \\Gamma_l r_s), $$\n", + "\n", + "with $A_l(Q)$ and $\\Gamma_l$ as in the circle model. The distribution is sampled at 21 log-spaced widths $\\Gamma_l r_s$, with $\\ln r_s = \\sigma c (s/10 - 1)$ and $c = \\sqrt{-2 \\ln 0.1}$, and weights $g_s \\propto \\exp(-\\ln^2 r_s / (2\\sigma^2))$ that sum to one.\n", + "\n", + "$A_0$ is the Elastic Incoherent Structure Factor (EISF) and $g_s A_l$ the Quasi-Elastic Incoherent Structure Factor (QISF) of one Lorentzian. The Lorentzian of order $j = (l - 1) \\cdot 21 + s + 1$ has circle order $l$ and sample $s$, so there are $21 (N - 1)$ Lorentzians per $Q$." + ] + }, + { + "cell_type": "code", + "id": "1", + "metadata": {}, + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import easydynamics as edyn\n", + "\n", + "%matplotlib widget" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## The model at different Q values" + ] + }, + { + "cell_type": "code", + "id": "3", + "metadata": {}, + "source": [ + "# Q is in Angstrom^-1, the radius in Angstrom and the residence time in ps.\n", + "Q = np.linspace(0.5, 2, 4)\n", + "\n", + "log_normal_model = edyn.JumpSitesLogNormalDistribution(\n", + " scale=1.0,\n", + " radius=1.0,\n", + " residence_time=2.0,\n", + " sigma=0.5,\n", + " number_of_sites=3,\n", + " Q=Q,\n", + ")\n", + "print(f'{log_normal_model.number_of_lorentzians} Lorentzians per Q')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "4", + "metadata": {}, + "source": [ + "# Plot the quasi-elastic part (all Lorentzians) at every Q. The delta function\n", + "# would be a spike at E = 0; its weight, the EISF, is plotted further down.\n", + "energy = np.linspace(-2, 2, 501)\n", + "component_collections = log_normal_model.get_component_collections()\n", + "\n", + "cmap = plt.cm.jet\n", + "nQ = len(component_collections)\n", + "plt.figure()\n", + "for Q_index, collection in enumerate(component_collections):\n", + " quasielastic = sum(\n", + " component.evaluate(energy)\n", + " for component in collection\n", + " if isinstance(component, edyn.Lorentzian)\n", + " )\n", + " color = cmap(Q_index / (nQ - 1))\n", + " plt.plot(energy, quasielastic, label=f'Q={Q[Q_index]:.2f} Å^-1', color=color)\n", + "\n", + "plt.legend()\n", + "plt.xlabel('Energy (meV)')\n", + "plt.ylabel('Intensity (arb. units)')\n", + "plt.title('Jump sites with a log-normal distribution, sigma = 0.5')\n", + "plt.show()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## The sampled distribution" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "The grid of widths scales with $\\sigma$, so the weights $g_s$ are the same for every $\\sigma$; only the widths change." + ] + }, + { + "cell_type": "code", + "id": "7", + "metadata": {}, + "source": [ + "plt.figure()\n", + "plt.bar(range(21), log_normal_model.sample_weights())\n", + "plt.xlabel('Sample s')\n", + "plt.ylabel('Weight $g_s$')\n", + "plt.title('Weights of the 21 sampled widths')\n", + "plt.show()\n", + "\n", + "for order in [1, 11, 21]:\n", + " width = log_normal_model.calculate_width(order=order)[0]\n", + " print(f'Lorentzian of order {order}: HWHM = {width:.4f} meV')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": {}, + "source": [ + "## The effect of sigma" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "A wider distribution gives a quasi-elastic line with a sharper peak and longer wings. For $\\sigma \\to 0$ the model becomes the equivalent-sites circle model." + ] + }, + { + "cell_type": "code", + "id": "10", + "metadata": {}, + "source": [ + "energy = np.linspace(-2, 2, 501)\n", + "Q_single = 1.0\n", + "\n", + "\n", + "def quasielastic(\n", + " model: edyn.EquivalentSitesCircle | edyn.JumpSitesLogNormalDistribution,\n", + ") -> np.ndarray:\n", + " collection = model.get_component_collections()[0]\n", + " return sum(\n", + " component.evaluate(energy)\n", + " for component in collection\n", + " if isinstance(component, edyn.Lorentzian)\n", + " )\n", + "\n", + "\n", + "circle_model = edyn.EquivalentSitesCircle(\n", + " radius=1.0, residence_time=2.0, number_of_sites=3, Q=Q_single\n", + ")\n", + "\n", + "plt.figure()\n", + "plt.plot(energy, quasielastic(circle_model), color='k', linestyle='dashed', label='circle model')\n", + "for sigma in [0.01, 0.5, 1.0]:\n", + " model = edyn.JumpSitesLogNormalDistribution(\n", + " radius=1.0, residence_time=2.0, sigma=sigma, number_of_sites=3, Q=Q_single\n", + " )\n", + " plt.plot(energy, quasielastic(model), label=f'sigma = {sigma}')\n", + "plt.yscale('log')\n", + "plt.xlabel('Energy (meV)')\n", + "plt.ylabel('Intensity (arb. units)')\n", + "plt.title('Quasi-elastic part at Q = 1 Å$^{-1}$')\n", + "plt.legend()\n", + "plt.show()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "The EISF does not depend on $\\sigma$: the distribution changes how fast the jumps are, not where the sites are." + ] + }, + { + "cell_type": "code", + "id": "12", + "metadata": {}, + "source": [ + "Q_fine = np.linspace(0.01, 4, 200)\n", + "eisf = log_normal_model.calculate_EISF(Q_fine)\n", + "print('Same EISF as the circle model:', np.allclose(eisf, circle_model.calculate_EISF(Q_fine)))\n", + "print('EISF + QISF = 1:', np.allclose(eisf + log_normal_model.calculate_QISF(Q_fine), 1))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "1. A. Chahid, A. Alegría and J. Colmenero, Macromolecules **27**, 3282 (1994), [doi:10.1021/ma00090a022](https://doi.org/10.1021/ma00090a022)." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "default", + "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.14.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index b8659848..b0b39d42 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -212,6 +212,9 @@ nav: - Detailed balance: tutorials/detailed_balance.ipynb - DeltaLorentz: tutorials/delta_lorentz.ipynb - Diffusion model: tutorials/diffusion_model.ipynb + - Equivalent sites on a circle: tutorials/equivalent_sites_circle.ipynb + - Jump sites, log-normal distribution: tutorials/jump_sites_log_normal_distribution.ipynb + - 3D Gaussian model: tutorials/gaussian_model_3d.ipynb - Mittag-Leffler relaxation: tutorials/mittag_leffler.ipynb - Sample model: tutorials/sample_model.ipynb - Instrument model: tutorials/instrument_model.ipynb diff --git a/src/easydynamics/__init__.py b/src/easydynamics/__init__.py index 8984f977..4237557a 100644 --- a/src/easydynamics/__init__.py +++ b/src/easydynamics/__init__.py @@ -31,10 +31,14 @@ from easydynamics.sample_model import DeltaFunction from easydynamics.sample_model import DeltaLorentz from easydynamics.sample_model import DiffusionDampedMittagLeffler +from easydynamics.sample_model import EquivalentSitesCircle from easydynamics.sample_model import Exponential from easydynamics.sample_model import ExpressionComponent from easydynamics.sample_model import Gaussian +from easydynamics.sample_model import GaussianModel3D +from easydynamics.sample_model import HallRossJumpDiffusion from easydynamics.sample_model import InstrumentModel +from easydynamics.sample_model import JumpSitesLogNormalDistribution from easydynamics.sample_model import JumpTranslationalDiffusion from easydynamics.sample_model import Lorentzian from easydynamics.sample_model import MittagLefflerDiffusion @@ -70,12 +74,16 @@ 'DiffusionDampedMittagLeffler', 'EasyDynamicsBase', 'EasyDynamicsModelBase', + 'EquivalentSitesCircle', 'Experiment', 'Exponential', 'ExpressionComponent', 'FitBinding', 'Gaussian', + 'GaussianModel3D', + 'HallRossJumpDiffusion', 'InstrumentModel', + 'JumpSitesLogNormalDistribution', 'JumpTranslationalDiffusion', 'Lorentzian', 'MittagLefflerDiffusion', diff --git a/src/easydynamics/sample_model/__init__.py b/src/easydynamics/sample_model/__init__.py index 384ba507..43f92e65 100644 --- a/src/easydynamics/sample_model/__init__.py +++ b/src/easydynamics/sample_model/__init__.py @@ -22,6 +22,14 @@ BrownianTranslationalDiffusion, ) from easydynamics.sample_model.diffusion_model.delta_lorentz import DeltaLorentz +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import EquivalentSitesCircle +from easydynamics.sample_model.diffusion_model.gaussian_model_3d import GaussianModel3D +from easydynamics.sample_model.diffusion_model.hall_ross_jump_diffusion import ( + HallRossJumpDiffusion, +) +from easydynamics.sample_model.diffusion_model.jump_sites_log_normal_distribution import ( + JumpSitesLogNormalDistribution, +) from easydynamics.sample_model.diffusion_model.jump_translational_diffusion import ( JumpTranslationalDiffusion, ) @@ -41,10 +49,14 @@ 'DeltaFunction', 'DeltaLorentz', 'DiffusionDampedMittagLeffler', + 'EquivalentSitesCircle', 'Exponential', 'ExpressionComponent', 'Gaussian', + 'GaussianModel3D', + 'HallRossJumpDiffusion', 'InstrumentModel', + 'JumpSitesLogNormalDistribution', 'JumpTranslationalDiffusion', 'Lorentzian', 'MittagLefflerDiffusion', diff --git a/src/easydynamics/sample_model/diffusion_model/__init__.py b/src/easydynamics/sample_model/diffusion_model/__init__.py index 2e92cfcd..73e301db 100644 --- a/src/easydynamics/sample_model/diffusion_model/__init__.py +++ b/src/easydynamics/sample_model/diffusion_model/__init__.py @@ -8,6 +8,14 @@ ChudleyElliottJumpDiffusion, ) from easydynamics.sample_model.diffusion_model.delta_lorentz import DeltaLorentz +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import EquivalentSitesCircle +from easydynamics.sample_model.diffusion_model.gaussian_model_3d import GaussianModel3D +from easydynamics.sample_model.diffusion_model.hall_ross_jump_diffusion import ( + HallRossJumpDiffusion, +) +from easydynamics.sample_model.diffusion_model.jump_sites_log_normal_distribution import ( + JumpSitesLogNormalDistribution, +) from easydynamics.sample_model.diffusion_model.jump_translational_diffusion import ( JumpTranslationalDiffusion, ) @@ -19,6 +27,10 @@ 'BrownianTranslationalDiffusion', 'ChudleyElliottJumpDiffusion', 'DeltaLorentz', + 'EquivalentSitesCircle', + 'GaussianModel3D', + 'HallRossJumpDiffusion', + 'JumpSitesLogNormalDistribution', 'JumpTranslationalDiffusion', 'MittagLefflerDiffusion', ] diff --git a/src/easydynamics/sample_model/diffusion_model/delta_lorentzian_diffusion_model_base.py b/src/easydynamics/sample_model/diffusion_model/delta_lorentzian_diffusion_model_base.py new file mode 100644 index 00000000..9ba160f4 --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/delta_lorentzian_diffusion_model_base.py @@ -0,0 +1,552 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +from abc import ABCMeta +from abc import abstractmethod + +import numpy as np +import scipp as sc +from easyscience.variable import DescriptorNumber +from easyscience.variable import Parameter + +from easydynamics.sample_model.component_collection import ComponentCollection +from easydynamics.sample_model.components import DeltaFunction +from easydynamics.sample_model.components import Lorentzian +from easydynamics.sample_model.diffusion_model.diffusion_model_base import DiffusionModelBase +from easydynamics.utils.fit_target import FitTarget +from easydynamics.utils.utils import CANONICAL_Q_UNIT +from easydynamics.utils.utils import Numeric +from easydynamics.utils.utils import Q_type + + +class DeltaLorentzianDiffusionModelBase(DiffusionModelBase, metaclass=ABCMeta): + r""" + Base class for diffusion models made of one delta function and n Lorentzians per Q-value. + + The model is + + $$ S(Q, E) = K \left[ A_0(Q)\,\delta(E) + \sum_{j=1}^{n} A_j(Q)\, L(E, \Gamma_j(Q)) \right] $$ + + where $K$ is the scale factor, $A_0(Q)$ is the Elastic Incoherent Structure Factor (EISF), + $A_j(Q)$ is the Quasi-Elastic Incoherent Structure Factor (QISF) of order $j$, and $L(E, + \Gamma)$ is a Lorentzian with half width at half maximum $\Gamma$. The total QISF is + $\sum_{j=1}^{n} A_j(Q)$, and EISF + QISF = 1 unless the model truncates an infinite sum. $Q$ is + assumed to have units of 1/angstrom. Creates ComponentCollections with one delta function and n + Lorentzian components for given Q-values. Their areas and widths are dependent Parameters, so + they follow the model parameters during a fit. + + The base class holds no physics. A subclass supplies $\Gamma_j$ through ``calculate_width``, + $A_0$ through ``calculate_EISF``, $A_j$ through ``calculate_QISF_order``, the dependent + Parameter of $A_j$ through ``_create_weight`` (order 0 is the EISF), and the dependent width + through ``_write_width_dependency_expression`` and ``_write_width_dependency_map_expression``. + It calls ``self.create_component_collections()`` at the end of its ``__init__``, after its + parameters exist. + + The Lorentzian of order $j$ is named ``'_'``, where the symbol is + the class attribute ``_order_symbol``, the conventional symbol of the subclass's order (``'l'`` + by default). ``get_fit_targets`` provides ``'elastic_area'`` and, for every order, + ``'area_'`` and ``'width_'``. + + See Also + -------- + EquivalentSitesCircle : Jumps between N equivalent sites on a circle. + JumpSitesLogNormalDistribution : The same with a log-normal distribution of jump rates. + GaussianModel3D : Diffusion in a three-dimensional harmonic confinement. + + Examples + -------- + **Using a model built on this base class** + + The base class is abstract. A subclass, here EquivalentSitesCircle with N = 3 sites, has N - 1 + = 2 Lorentzians per Q-value: + ```python + import numpy as np + import easydynamics as edyn + + model = edyn.EquivalentSitesCircle( + radius=1.0, + residence_time=2.0, + number_of_sites=3, + Q=np.linspace(0.5, 2, 7), + ) + model.number_of_lorentzians # 2 + model.lorentzian_component_name(1) # 'EquivalentSitesCircle_l1' + model.calculate_EISF() + model.calculate_QISF() # ones + ``` + """ + + _order_symbol: str = 'l' + + def __init__( + self, + number_of_lorentzians: int, + *, + name: str, + display_name: str | None, + scale: Numeric = 1.0, + Q: Q_type | None = None, + x_unit: str | sc.Unit = 'meV', + y_unit: str | sc.Unit = 'dimensionless', + lorentzian_name: str | None = None, + lorentzian_display_name: str | None = None, + delta_name: str = 'Delta function', + delta_display_name: str | None = None, + unique_name: str | None = None, + ) -> None: + """ + Initialize the shared part of a delta + n Lorentzians diffusion model. + + Builds no components; the subclass calls ``create_component_collections`` once its + parameters exist. + + Parameters + ---------- + number_of_lorentzians : int + Number n of Lorentzians per Q. Must be an integer of at least 1. + name : str + Name of the diffusion model. + display_name : str | None + Display name of the diffusion model. + scale : Numeric, default=1.0 + Scale factor for the diffusion model. Must be a non-negative number. + Q : Q_type | None, default=None + Q values for the model. If None, Q is not set. + x_unit : str | sc.Unit, default='meV' + Unit of the x-axis (energy/frequency). Must be convertible to meV. + y_unit : str | sc.Unit, default='dimensionless' + Unit of the model output (intensity). Determines scale.unit = x_unit * y_unit. + lorentzian_name : str | None, default=None + Base name of the Lorentzian components. If None, it will be set to the name of the + diffusion model. + lorentzian_display_name : str | None, default=None + Base display name of the Lorentzian components. If None, it will be set to + lorentzian_name. + delta_name : str, default='Delta function' + Name of the delta function component. + delta_display_name : str | None, default=None + Display name of the delta function component. If None, it will be set to delta_name. + unique_name : str | None, default=None + Unique name of the diffusion model. If None, a unique name will be generated. + + Raises + ------ + TypeError + If number_of_lorentzians is not an integer, or if delta_name or delta_display_name is + not a string (or None). + ValueError + If number_of_lorentzians is smaller than 1. + """ + super().__init__( + scale=scale, + Q=Q, + x_unit=x_unit, + y_unit=y_unit, + name=name, + display_name=display_name, + lorentzian_name=lorentzian_name, + lorentzian_display_name=lorentzian_display_name, + unique_name=unique_name, + ) + + if isinstance(number_of_lorentzians, bool) or not isinstance( + number_of_lorentzians, (int, np.integer) + ): + raise TypeError('number_of_lorentzians must be an integer.') + if number_of_lorentzians < 1: + raise ValueError('number_of_lorentzians must be at least 1.') + self._number_of_lorentzians = int(number_of_lorentzians) + + if not isinstance(delta_name, str): + raise TypeError('delta_name must be a string.') + if delta_display_name is None: + delta_display_name = delta_name + if not isinstance(delta_display_name, str): + raise TypeError('delta_display_name must be a string or None.') + self._delta_name = delta_name + self._delta_display_name = delta_display_name + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def number_of_lorentzians(self) -> int: + """ + Get the number of Lorentzians per Q. + + Returns + ------- + int + The number n of Lorentzians. Set at construction. + """ + return self._number_of_lorentzians + + @property + def delta_name(self) -> str: + """ + Get the name of the delta function component. + + Returns + ------- + str + Name of the delta function component. + """ + return self._delta_name + + @delta_name.setter + def delta_name(self, delta_name: str) -> None: + """ + Set the name of the delta function component. + + Parameters + ---------- + delta_name : str + The new name for the delta function component. + + Raises + ------ + TypeError + If delta_name is not a string. + """ + if not isinstance(delta_name, str): + raise TypeError('delta_name must be a string.') + self._delta_name = delta_name + + @property + def delta_display_name(self) -> str: + """ + Get the display name of the delta function component. + + Returns + ------- + str + Display name of the delta function component. + """ + return self._delta_display_name + + @delta_display_name.setter + def delta_display_name(self, delta_display_name: str | None) -> None: + """ + Set the display name of the delta function component. + + Parameters + ---------- + delta_display_name : str | None + The new display name for the delta function component. + + Raises + ------ + TypeError + If delta_display_name is not a string or None. + """ + if not isinstance(delta_display_name, (str, type(None))): + raise TypeError('delta_display_name must be a string or None.') + self._delta_display_name = delta_display_name + + # ------------------------------------------------------------------ + # Physics, supplied by the subclass + # ------------------------------------------------------------------ + + @abstractmethod + def calculate_width(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + r""" + Calculate the half-width at half-maximum (HWHM) $\Gamma_j(Q)$ of the Lorentzian of order + $j$. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + order : int, default=1 + The order j of the Lorentzian, between 1 and n. + + Returns + ------- + np.ndarray + HWHM values in the unit of the model (e.g., meV). + """ + + @abstractmethod + def calculate_EISF(self, Q: Q_type | None = None) -> np.ndarray: + """ + Calculate the Elastic Incoherent Structure Factor (EISF), $A_0(Q)$, the weight of the delta + function. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + + Returns + ------- + np.ndarray + EISF values (dimensionless). + """ + + @abstractmethod + def calculate_QISF_order(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + """ + Calculate the Quasi-Elastic Incoherent Structure Factor (QISF) of one order, $A_j(Q)$, the + weight of the Lorentzian of order $j$. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + order : int, default=1 + The order j of the Lorentzian, between 1 and n. + + Returns + ------- + np.ndarray + QISF values of that order (dimensionless). + """ + + @abstractmethod + def _create_weight(self, Q: float, order: int) -> Parameter: + """ + Create the EISF (order 0) or the QISF of one order as a dimensionless dependent Parameter. + + Parameters + ---------- + Q : float + Scattering vector in 1/angstrom. + order : int + The order j of the Lorentzian; 0 for the delta function. + + Returns + ------- + Parameter + The EISF or QISF, which becomes the weight of the component. + """ + + @abstractmethod + def _write_width_dependency_expression(self, order: int) -> str: + """ + Write the dependency expression for the width of the Lorentzian of a given order. + + Parameters + ---------- + order : int + The order of the Lorentzian. + + Returns + ------- + str + Dependency expression for the width. + """ + + @abstractmethod + def _write_width_dependency_map_expression(self) -> dict[str, DescriptorNumber]: + """ + Write the dependency map for the width. + + Returns + ------- + dict[str, DescriptorNumber] + Dependency map for the width. + """ + + # ------------------------------------------------------------------ + # Other methods + # ------------------------------------------------------------------ + + def lorentzian_component_name(self, order: int) -> str: + """ + Get the name of the Lorentzian component of a given order. + + Parameters + ---------- + order : int + The order of the Lorentzian, between 1 and n. + + Returns + ------- + str + The component name, '_'. + """ + self._verify_order(order) + return f'{self.lorentzian_name}_{self._order_symbol}{order}' + + def create_component_collections(self) -> list[ComponentCollection]: + r""" + Create ComponentCollections for the model at the current Q values. + + Each collection holds a delta function with area $K A_0(Q)$ (scale times EISF) and n + Lorentzians with widths $\Gamma_j(Q)$ from ``calculate_width`` and areas $K A_j(Q)$ (scale + times the QISF of order $j$). The created collections are installed on the model, so the + returned list is the live one. + + Returns + ------- + list[ComponentCollection] + One ComponentCollection per Q value. + """ + if self.Q is None: + self._component_collections = [] + return self._component_collections + + Q = self.Q.values + component_collection_list = [None] * len(Q) + + for i, Q_value in enumerate(Q): + component_collection_list[i] = ComponentCollection( + name=f'{self.name}_Q{Q_value:.2f}', + display_name=f'{self.display_name}_Q{Q_value:.2f}', + x_unit=self.x_unit, + y_unit=self.y_unit, + ) + + delta_component = DeltaFunction( + name=self.delta_name, + display_name=self.delta_display_name, + x_unit=self.x_unit, + y_unit=self.y_unit, + ) + with np.errstate(invalid='ignore'): + delta_component.area.make_dependent_on( + dependency_expression=self._write_area_dependency_expression(), + dependency_map=self._write_area_dependency_map_expression( + self._create_weight(float(Q_value), 0) + ), + ) + component_collection_list[i].append_component(delta_component) + + for order in range(1, self.number_of_lorentzians + 1): + lorentzian_base_display_name = self.lorentzian_display_name or self.lorentzian_name + lorentzian_component = Lorentzian( + name=self.lorentzian_component_name(order), + display_name=f'{lorentzian_base_display_name}_{self._order_symbol}{order}', + x_unit=self.x_unit, + y_unit=self.y_unit, + ) + + # easyscience propagates inf bounds through arithmetic, producing inf/inf=nan + # as a transient intermediate. Python's min/max ignore nan so the final bounds + # are correct; suppress the spurious numpy RuntimeWarning. + with np.errstate(invalid='ignore'): + lorentzian_component.width.make_dependent_on( + dependency_expression=self._write_width_dependency_expression(order), + dependency_map=self._write_width_dependency_map_expression(), + desired_unit=self.x_unit, + ) + lorentzian_component.area.make_dependent_on( + dependency_expression=self._write_area_dependency_expression(), + dependency_map=self._write_area_dependency_map_expression( + self._create_weight(float(Q_value), order) + ), + ) + component_collection_list[i].append_component(lorentzian_component) + + self._component_collections = component_collection_list + return self._component_collections + + def get_fit_targets(self) -> list[FitTarget]: + """ + Get the fittable predictions of the model as FitTargets. + + The model predicts ``'elastic_area'`` (``scale * EISF``) and, for every order from 1 to n, + ``'area_'`` (``scale`` times the QISF of that order) and ``'width_'`` (its + HWHM). The base class implementation of DiffusionModelBase is replaced because there is no + single Lorentzian. Select the orders to fit with the ``targets`` of a FitBinding. + + Returns + ------- + list[FitTarget] + The fittable predictions of this model. + """ + targets = [ + FitTarget( + name='elastic_area', + dataset_key=f'{self.delta_name} area', + function=lambda Q, model=self, **_: model.calculate_EISF(Q) * model.scale.value, + label=f'{self.display_name} elastic_area', + x_unit=CANONICAL_Q_UNIT, + y_unit=str(self.scale.unit), + ) + ] + for order in range(1, self.number_of_lorentzians + 1): + key = self.lorentzian_component_name(order) + targets.append( + FitTarget( + name=f'area_{order}', + dataset_key=f'{key} area', + function=lambda Q, model=self, order=order, **_: ( + model.calculate_QISF_order(Q, order) * model.scale.value + ), + label=f'{self.display_name} area_{order}', + x_unit=CANONICAL_Q_UNIT, + y_unit=str(self.scale.unit), + ) + ) + targets.append( + FitTarget( + name=f'width_{order}', + dataset_key=f'{key} width', + function=lambda Q, model=self, order=order, **_: model.calculate_width( + Q, order + ), + label=f'{self.display_name} width_{order}', + x_unit=CANONICAL_Q_UNIT, + y_unit=str(self.x_unit), + ) + ) + return targets + + # ------------------------------------------------------------------ + # Private methods + # ------------------------------------------------------------------ + + def _write_area_dependency_expression(self) -> str: + """ + Write the dependency expression for the area of a component. + + Returns + ------- + str + Dependency expression for the area, the scale times the EISF or QISF. + """ + return 'scale * weight' + + def _write_area_dependency_map_expression( + self, weight: Parameter + ) -> dict[str, DescriptorNumber]: + """ + Write the dependency map for the areas. + + Parameters + ---------- + weight : Parameter + The EISF or the QISF of one order, from ``_create_weight`` (dimensionless). + + Returns + ------- + dict[str, DescriptorNumber] + Dependency map for the area. + """ + return {'scale': self.scale, 'weight': weight} + + def _verify_order(self, order: int) -> None: + """ + Verify that order is a valid Lorentzian order. + + Parameters + ---------- + order : int + The order to verify. + + Raises + ------ + TypeError + If order is not an integer. + ValueError + If order is not between 1 and n. + """ + if isinstance(order, bool) or not isinstance(order, (int, np.integer)): + raise TypeError('order must be an integer.') + if not 1 <= order <= self.number_of_lorentzians: + raise ValueError(f'order must be between 1 and {self.number_of_lorentzians}.') + + def _on_Q_change(self) -> None: + """Rebuild the component collections when Q changes.""" + self.create_component_collections() diff --git a/src/easydynamics/sample_model/diffusion_model/equivalent_sites_circle.py b/src/easydynamics/sample_model/diffusion_model/equivalent_sites_circle.py new file mode 100644 index 00000000..fec7348e --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/equivalent_sites_circle.py @@ -0,0 +1,530 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import scipp as sc +from easyscience.variable import DescriptorNumber +from easyscience.variable import Parameter + +from easydynamics.sample_model.diffusion_model.delta_lorentzian_diffusion_model_base import ( + DeltaLorentzianDiffusionModelBase, +) +from easydynamics.utils.utils import Numeric +from easydynamics.utils.utils import Q_type +from easydynamics.utils.utils import angstrom +from easydynamics.utils.utils import hbar + +MINIMUM_RADIUS = 1e-10 # A zero radius is unphysical +MINIMUM_WIDTH = 1e-10 # To avoid division by zero +MAXIMUM_NUMBER_OF_SITES = 50 + + +class EquivalentSitesCircle(DeltaLorentzianDiffusionModelBase): + r""" + Model of jumps between N equivalent sites on a circle. + + $$ S(Q, E) = A_0(Q)\,\delta(E) + \sum_{l=1}^{N-1} A_l(Q)\, L(E, \Gamma_l) $$ + + with + + $$ A_l(Q) = \frac{1}{N} \sum_{k=1}^{N} j_0(Q r_k) \cos\left(\frac{2\pi l k}{N}\right), \quad + r_k = 2R \sin\left(\frac{\pi k}{N}\right), \quad \Gamma_l = \frac{2\hbar}{\tau} + \sin^2\left(\frac{\pi l}{N}\right) $$ + + where $R$ is the radius of the circle, $\tau$ the residence time on a site, $j_0(x) = \sin(x) / + x$, and $L(E, \Gamma)$ a Lorentzian with half width at half maximum $\Gamma$. $A_0$ is the EISF + and $A_l$ ($l \geq 1$) the QISF of order $l$; they sum to one. For $N = 2$ this is the two-site + jump model, with jump distance $2R$ and $\Gamma_1 = 2\hbar/\tau$. The width is also written as + $4 \sin^2(\pi l / N) / \tau$, with a $\tau$ that is twice the residence time used here + (reference 3). + + The radius is fixed by default, since it is normally known from the molecular geometry. + + References + ---------- + 1. J. D. Barnes, J. Chem. Phys. 58, 5193 (1973), doi:10.1063/1.1679130. + 2. M. Bée, Quasielastic Neutron Scattering, Adam Hilger, Bristol (1988). + 3. Mantid, ``InelasticDiffRotDiscreteCircle``, whose ``Decay`` parameter is this $2\tau$, + https://github.com/mantidproject/mantid/blob/6ec66a6f524a7ec9d4258351fe3001c66c63dbe5/Framework/CurveFitting/src/Functions/InelasticDiffRotDiscreteCircle.cpp#L79-L108. + + Examples + -------- + **Creating an EquivalentSitesCircle model** + + The model creates one delta function and N - 1 Lorentzians per Q-value: + ```python + import numpy as np + import easydynamics as edyn + + Q = np.linspace(0.5, 2, 7) + model = edyn.EquivalentSitesCircle( + scale=1.0, + radius=1.0, + residence_time=2.0, + number_of_sites=3, + Q=Q, + ) + component_collections = model.create_component_collections() + ``` + """ + + _order_symbol = 'l' # components are named '_l' + + def __init__( + self, + scale: Numeric = 1.0, + radius: Numeric = 1.0, + residence_time: Numeric = 1.0, + number_of_sites: int = 3, + Q: Q_type | None = None, + x_unit: str | sc.Unit = 'meV', + y_unit: str | sc.Unit = 'dimensionless', + name: str = 'EquivalentSitesCircle', + display_name: str | None = 'EquivalentSitesCircle', + lorentzian_name: str | None = None, + lorentzian_display_name: str | None = None, + delta_name: str = 'Delta function', + delta_display_name: str | None = None, + unique_name: str | None = None, + ) -> None: + """ + Initialize a new EquivalentSitesCircle model. + + Parameters + ---------- + scale : Numeric, default=1.0 + Scale factor for the diffusion model. Must be a non-negative number. + radius : Numeric, default=1.0 + Radius of the circle in angstrom. Must be positive. Fixed by default. + residence_time : Numeric, default=1.0 + Residence time tau on a site in ps. + number_of_sites : int, default=3 + Number of equivalent sites N on the circle, between 2 and MAXIMUM_NUMBER_OF_SITES. Set + at construction; not a fit parameter. + Q : Q_type | None, default=None + Q values for the model. If None, Q is not set. + x_unit : str | sc.Unit, default='meV' + Unit of the x-axis (energy/frequency). Must be convertible to meV. + y_unit : str | sc.Unit, default='dimensionless' + Unit of the model output (intensity). Determines scale.unit = x_unit * y_unit. + name : str, default='EquivalentSitesCircle' + Name of the diffusion model. + display_name : str | None, default='EquivalentSitesCircle' + Display name of the diffusion model. + lorentzian_name : str | None, default=None + Base name of the Lorentzian components. The Lorentzian of order l is named + '_l'. If None, it will be set to the name of the diffusion model. + lorentzian_display_name : str | None, default=None + Base display name of the Lorentzian components. If None, it will be set to + lorentzian_name. + delta_name : str, default='Delta function' + Name of the delta function component. + delta_display_name : str | None, default=None + Display name of the delta function component. If None, it will be set to delta_name. + unique_name : str | None, default=None + Unique name of the diffusion model. If None, a unique name will be generated. + + Raises + ------ + TypeError + If radius or residence_time is not a number, if number_of_sites is not an integer or if + delta_name or delta_display_name is not a string (or None). + ValueError + If radius or residence_time is smaller than its minimum, or if number_of_sites is not + between 2 and MAXIMUM_NUMBER_OF_SITES. + """ + self._verify_radius(radius) + self._verify_residence_time(residence_time) + + if isinstance(number_of_sites, bool) or not isinstance(number_of_sites, (int, np.integer)): + raise TypeError('number_of_sites must be an integer.') + if not 2 <= number_of_sites <= MAXIMUM_NUMBER_OF_SITES: + raise ValueError(f'number_of_sites must be between 2 and {MAXIMUM_NUMBER_OF_SITES}.') + self._number_of_sites = int(number_of_sites) + + # N sites give N - 1 Lorentzians, of orders l = 1 ... N - 1 + super().__init__( + self._number_of_sites - 1, + scale=scale, + Q=Q, + x_unit=x_unit, + y_unit=y_unit, + name=name, + display_name=display_name, + lorentzian_name=lorentzian_name, + lorentzian_display_name=lorentzian_display_name, + delta_name=delta_name, + delta_display_name=delta_display_name, + unique_name=unique_name, + ) + + self._radius = Parameter( + name='radius', + value=float(radius), + fixed=True, + min=MINIMUM_RADIUS, + unit='angstrom', + ) + self._residence_time = Parameter( + name='residence_time', + value=float(residence_time), + fixed=False, + min=MINIMUM_WIDTH, + unit='ps', + ) + + self._hbar = hbar + # Dimensionless one, which turns the plain numbers of the weight expressions into + # (dimensionless) descriptors, as dependency expressions must evaluate to one. + self._one = DescriptorNumber(name='one', value=1.0) + + # Dimensionless radius/angstrom for the weight dependency expressions. Deriving it through + # the dependency graph keeps the weights correct if the radius is converted to another + # unit (e.g. nm). + with np.errstate(invalid='ignore'): + self._radius_over_angstrom = Parameter.from_dependency( + name='radius_over_angstrom', + dependency_expression='radius / angstrom', + dependency_map={'radius': self._radius, 'angstrom': angstrom}, + ) + self._radius_over_angstrom.set_desired_unit('dimensionless') + + self.create_component_collections() + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def radius(self) -> Parameter: + """ + Get the radius of the circle. + + Returns + ------- + Parameter + Radius R. + """ + return self._radius + + @radius.setter + def radius(self, radius: Numeric) -> None: + """ + Set the radius of the circle. + + Parameters + ---------- + radius : Numeric + Radius R in the unit of the radius parameter (angstrom by default). + """ + self._verify_radius(radius) + self._radius.value = float(radius) + + @property + def residence_time(self) -> Parameter: + """ + Get the residence time on a site. + + Returns + ------- + Parameter + Residence time tau. + """ + return self._residence_time + + @residence_time.setter + def residence_time(self, residence_time: Numeric) -> None: + """ + Set the residence time on a site. + + Parameters + ---------- + residence_time : Numeric + Residence time tau in ps. + """ + self._verify_residence_time(residence_time) + self._residence_time.value = float(residence_time) + + @property + def number_of_sites(self) -> int: + """ + Get the number of equivalent sites on the circle. + + Returns + ------- + int + The number of sites N. Set at construction. + """ + return self._number_of_sites + + # ------------------------------------------------------------------ + # Other methods + # ------------------------------------------------------------------ + + def calculate_width(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + r""" + Calculate the half-width at half-maximum (HWHM) of the Lorentzian of a given order. + + The width is $(2\hbar/\tau) \sin^2(\pi l / N)$ and does not depend on Q. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + order : int, default=1 + The order l of the Lorentzian, between 1 and N - 1. + + Returns + ------- + np.ndarray + HWHM values in the unit of the model (e.g., meV). + """ + self._verify_order(order) + Q = self._ensure_Q(Q) + + rate = 2 * self._hbar / self.residence_time + rate.convert_unit(self.x_unit) + angle = sc.scalar(np.pi * order / self.number_of_sites, unit='rad') + return rate.value * sc.sin(angle).value ** 2 * np.ones_like(Q) + + def calculate_EISF(self, Q: Q_type | None = None) -> np.ndarray: + """ + Calculate the Elastic Incoherent Structure Factor (EISF), $A_0(Q)$. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + + Returns + ------- + np.ndarray + EISF values (dimensionless). + """ + Q = self._ensure_Q(Q) + return self._structure_factors(Q)['order', 0].values + + def calculate_QISF_order(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + """ + Calculate the Quasi-Elastic Incoherent Structure Factor (QISF) of one order, $A_l(Q)$, the + weight of the Lorentzian of order $l$. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + order : int, default=1 + The order l of the Lorentzian, between 1 and N - 1. + + Returns + ------- + np.ndarray + QISF values of that order (dimensionless). + """ + self._verify_order(order) + Q = self._ensure_Q(Q) + return self._structure_factors(Q)['order', int(order)].values + + def calculate_QISF(self, Q: Q_type | None = None) -> np.ndarray: + """ + Calculate the total Quasi-Elastic Incoherent Structure Factor (QISF), the sum of the QISF + of all N - 1 orders. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + + Returns + ------- + np.ndarray + QISF values (dimensionless). + """ + Q = self._ensure_Q(Q) + return sc.sum(self._structure_factors(Q)['order', 1:], 'order').values + + def get_global_variables(self) -> list[Parameter]: + """ + Get all global variables from the diffusion model. + + Returns + ------- + list[Parameter] + The scale, radius and residence time. + """ + return [self.scale, self.radius, self.residence_time] + + # ------------------------------------------------------------------ + # Private methods + # ------------------------------------------------------------------ + + def _structure_factors(self, Q: np.ndarray) -> sc.Variable: + r""" + Calculate the EISF and the QISF, $A_l(Q)$ for all orders $l = 0 \ldots N - 1$, with Scipp. + + $A_l(Q) = \frac{1}{N} \sum_{k=1}^{N} j_0(Q r_k) \cos(2\pi l k / N)$ with $r_k = 2R\sin(\pi + k / N)$. The variance of the radius is not propagated. + + Parameters + ---------- + Q : np.ndarray + Scattering vector in 1/angstrom. + + Returns + ------- + sc.Variable + Dimensionless $A_l(Q)$ with dims ``('Q', 'order')``; order 0 is the EISF. + """ + N = self.number_of_sites + site = sc.arange('k', 1.0, N + 1.0) + order = sc.arange('order', 0.0, float(N)) + half_turn = sc.scalar(np.pi, unit='rad') + + radius = sc.values(self.radius.full_value).to(unit='angstrom') + jump_distance = 2 * radius * sc.sin(half_turn * site / N) + Q_variable = sc.array(dims=['Q'], values=Q, unit='1/angstrom') + + weights = sc.sum( + sc.sinc(Q_variable * jump_distance) * sc.cos(2 * half_turn * order * site / N), 'k' + ) + return (weights / N).transpose(['Q', 'order']) + + def _create_weight(self, Q: float, order: int) -> Parameter: + r""" + Create the EISF ($l = 0$) or the QISF of order $l$, $A_l(Q)$, as a dimensionless dependent + Parameter that follows the radius. + + The dependency expressions cannot call Scipp, so the sum over sites is written out with + NumPy's sinc, sinc(x) = sin(pi x) / (pi x). Keeping the weight separate from the scale + makes updates of the scale cheap. + + Parameters + ---------- + Q : float + Scattering vector in 1/angstrom. + order : int + The order l; 0 for the delta function (the EISF). + + Raises + ------ + TypeError + If Q is not a float. + + Returns + ------- + Parameter + The weight. + """ + if not isinstance(Q, float): + raise TypeError('Q must be a float.') + + N = self.number_of_sites + terms = [] + for site in range(1, N + 1): + # float(): NumPy 2 writes np.float64 as 'np.float64(...)', which asteval cannot parse + coefficient = float(np.cos(2 * np.pi * order * site / N) / N) + # j0(Q r_k) = sinc(Q r_k / pi) with r_k = 2 R sin(pi k / N) + argument = float(Q * 2 * np.sin(np.pi * site / N) / np.pi) + terms.append(f'{coefficient!r} * sinc({argument!r} * radius_ratio.value)') + expression = f'one * ({" + ".join(terms)})' + with np.errstate(invalid='ignore'): + return Parameter.from_dependency( + name=f'weight_l{order}_Q{Q:.2f}', + dependency_expression=expression, + dependency_map={'one': self._one, 'radius_ratio': self._radius_over_angstrom}, + ) + + def _write_width_dependency_expression(self, order: int) -> str: + """ + Write the dependency expression for the width of the Lorentzian of a given order. + + Parameters + ---------- + order : int + The order l of the Lorentzian. + + Returns + ------- + str + Dependency expression for the width. + """ + factor = float(2 * np.sin(np.pi * order / self.number_of_sites) ** 2) + return f'{factor!r} * hbar / tau' + + def _write_width_dependency_map_expression(self) -> dict[str, DescriptorNumber]: + """ + Write the dependency map for the width. + + Returns + ------- + dict[str, DescriptorNumber] + Dependency map for the width. + """ + return {'hbar': self._hbar, 'tau': self.residence_time} + + @staticmethod + def _verify_radius(radius: Numeric) -> None: + """ + Verify a radius value. + + Parameters + ---------- + radius : Numeric + The radius to verify. + + Raises + ------ + TypeError + If radius is not a number. + ValueError + If radius is below MINIMUM_RADIUS. + """ + if not isinstance(radius, Numeric): + raise TypeError('radius must be a number.') + if float(radius) < MINIMUM_RADIUS: + raise ValueError(f'radius must be at least {MINIMUM_RADIUS}.') + + @staticmethod + def _verify_residence_time(residence_time: Numeric) -> None: + """ + Verify a residence time value. + + Parameters + ---------- + residence_time : Numeric + The residence time to verify. + + Raises + ------ + TypeError + If residence_time is not a number. + ValueError + If residence_time is below MINIMUM_WIDTH. + """ + if not isinstance(residence_time, Numeric): + raise TypeError('residence_time must be a number.') + if float(residence_time) < MINIMUM_WIDTH: + raise ValueError(f'residence_time must be at least {MINIMUM_WIDTH}.') + + # ------------------------------------------------------------------ + # dunder methods + # ------------------------------------------------------------------ + + def __repr__(self) -> str: + """ + String representation of the EquivalentSitesCircle model. + + Returns + ------- + str + String representation of the model. + """ + return ( + f'{self.__class__.__name__}(' + f'name={self.name!r}, display_name={self.display_name!r},\n' + f'x_unit={self.x_unit}, y_unit={self.y_unit},\n' + f' number_of_sites={self.number_of_sites},\n' + f' radius={self.radius},\n' + f' residence_time={self.residence_time},\n' + f' scale={self.scale})' + ) diff --git a/src/easydynamics/sample_model/diffusion_model/gaussian_model_3d.py b/src/easydynamics/sample_model/diffusion_model/gaussian_model_3d.py new file mode 100644 index 00000000..68e20b09 --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/gaussian_model_3d.py @@ -0,0 +1,614 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import math +import warnings + +import numpy as np +import scipp as sc +from easyscience.variable import DescriptorNumber +from easyscience.variable import Parameter + +from easydynamics.sample_model.component_collection import ComponentCollection +from easydynamics.sample_model.diffusion_model.delta_lorentzian_diffusion_model_base import ( + DeltaLorentzianDiffusionModelBase, +) +from easydynamics.utils.utils import Numeric +from easydynamics.utils.utils import Q_type +from easydynamics.utils.utils import angstrom +from easydynamics.utils.utils import hbar + +MINIMUM_DIFFUSION_COEFFICIENT = 1e-20 # In m^2/s, to keep the Lorentzian widths finite +MINIMUM_MEAN_U_SQUARED = 1e-10 # In angstrom^2; the widths scale with 1/ +MAXIMUM_NUMBER_OF_LORENTZIANS = 100 # As in QENSmodels +MISSING_WEIGHT_TOLERANCE = 1e-3 # Truncating the sum above this loss of weight raises a warning + + +class GaussianModel3D(DeltaLorentzianDiffusionModelBase): + r""" + Model of diffusion inside an isotropic three-dimensional harmonic well. + + A particle diffuses with diffusion coefficient $D$ in a harmonic potential that is the same in + all three directions, so its position about the centre of the well has a Gaussian distribution. + $\langle u^2 \rangle$ is the mean square distance from the centre, summed over the three + directions; along any one direction it is $\langle u^2 \rangle / 3$. Because the well is + isotropic, the scattering depends only on the magnitude of $Q$: + + $$ S(Q, E) = A_0(Q)\,\delta(E) + \sum_{i=1}^{n} A_i(Q)\, L(E, \Gamma_i) $$ + + with + + $$ A_i(Q) = e^{-y} \frac{y^i}{i!}, \quad y = \frac{Q^2 \langle u^2 \rangle}{3}, \quad \Gamma_i + = \frac{3 i \hbar D}{\langle u^2 \rangle} $$ + + where $L(E, \Gamma)$ is a Lorentzian with half width at half maximum $\Gamma$. $A_0 = e^{-y}$ + is the EISF, the same Debye-Waller factor as in ``DeltaLorentz``, and $A_i$ is the QISF of + order $i$. The widths $\Gamma_i$ grow linearly with $i$ and do not depend on $Q$; only the QISF + do. The paper writes the model with the variance along one direction, $\langle u_x^2 \rangle = + \langle u^2 \rangle / 3$. + + The QISF form a Poisson distribution in $i$ with mean $y$, and together with the EISF they sum + to one. The infinite sum is therefore truncated at $n$ = ``number_of_lorentzians``, which must + be well above $y$ at the largest $Q$. A warning is raised if the truncation leaves out more + than 1e-3 of the total weight at any of the model's Q values, and ``calculate_missing_weight`` + returns the weight left out. + + References + ---------- + F. Volino, J.-C. Perrin and S. Lyonnard, J. Phys. Chem. B 110, 11217 (2006), + doi:10.1021/jp061103s. + + Examples + -------- + **Creating a GaussianModel3D model** + + The model creates one delta function and ``number_of_lorentzians`` Lorentzians per Q-value: + ```python + import numpy as np + import easydynamics as edyn + + Q = np.linspace(0.5, 2, 7) + model = edyn.GaussianModel3D( + scale=1.0, + diffusion_coefficient=2e-9, # m^2/s + mean_u_squared=1.5, # angstrom^2 + number_of_lorentzians=20, + Q=Q, + ) + component_collections = model.create_component_collections() + ``` + """ + + _order_symbol = 'i' # components are named '_i' + + def __init__( + self, + scale: Numeric = 1.0, + diffusion_coefficient: Numeric = 1.0, + mean_u_squared: Numeric = 1.0, + number_of_lorentzians: int = 20, + Q: Q_type | None = None, + x_unit: str | sc.Unit = 'meV', + y_unit: str | sc.Unit = 'dimensionless', + name: str = 'GaussianModel3D', + display_name: str | None = 'GaussianModel3D', + lorentzian_name: str | None = None, + lorentzian_display_name: str | None = None, + delta_name: str = 'Delta function', + delta_display_name: str | None = None, + unique_name: str | None = None, + ) -> None: + """ + Initialize a new GaussianModel3D model. + + Parameters + ---------- + scale : Numeric, default=1.0 + Scale factor for the diffusion model. Must be a non-negative number. + diffusion_coefficient : Numeric, default=1.0 + Diffusion coefficient D in m^2/s. + mean_u_squared : Numeric, default=1.0 + Mean square displacement in angstrom^2, three times Volino's . + number_of_lorentzians : int, default=20 + Number of Lorentzians n per Q-value, i.e. the order at which the sum is truncated, + between 1 and MAXIMUM_NUMBER_OF_LORENTZIANS. Set at construction. + Q : Q_type | None, default=None + Q values for the model. If None, Q is not set. + x_unit : str | sc.Unit, default='meV' + Unit of the x-axis (energy/frequency). Must be convertible to meV. + y_unit : str | sc.Unit, default='dimensionless' + Unit of the model output (intensity). Determines scale.unit = x_unit * y_unit. + name : str, default='GaussianModel3D' + Name of the diffusion model. + display_name : str | None, default='GaussianModel3D' + Display name of the diffusion model. + lorentzian_name : str | None, default=None + Base name of the Lorentzian components. The Lorentzian of order i is named + '_i'. If None, it will be set to the name of the diffusion model. + lorentzian_display_name : str | None, default=None + Base display name of the Lorentzian components. If None, it will be set to + lorentzian_name. + delta_name : str, default='Delta function' + Name of the delta function component. + delta_display_name : str | None, default=None + Display name of the delta function component. If None, it will be set to delta_name. + unique_name : str | None, default=None + Unique name of the diffusion model. If None, a unique name will be generated. + + Raises + ------ + TypeError + If diffusion_coefficient or mean_u_squared is not a number, if number_of_lorentzians is + not an integer or if delta_name or delta_display_name is not a string (or None). + ValueError + If diffusion_coefficient or mean_u_squared is smaller than its minimum, or if + number_of_lorentzians is not between 1 and MAXIMUM_NUMBER_OF_LORENTZIANS. + """ + self._verify_diffusion_coefficient(diffusion_coefficient) + self._verify_mean_u_squared(mean_u_squared) + + if isinstance(number_of_lorentzians, bool) or not isinstance( + number_of_lorentzians, (int, np.integer) + ): + raise TypeError('number_of_lorentzians must be an integer.') + if not 1 <= number_of_lorentzians <= MAXIMUM_NUMBER_OF_LORENTZIANS: + raise ValueError( + f'number_of_lorentzians must be between 1 and {MAXIMUM_NUMBER_OF_LORENTZIANS}.' + ) + + super().__init__( + number_of_lorentzians, + scale=scale, + Q=Q, + x_unit=x_unit, + y_unit=y_unit, + name=name, + display_name=display_name, + lorentzian_name=lorentzian_name, + lorentzian_display_name=lorentzian_display_name, + delta_name=delta_name, + delta_display_name=delta_display_name, + unique_name=unique_name, + ) + + self._diffusion_coefficient = Parameter( + name='diffusion_coefficient', + value=float(diffusion_coefficient), + fixed=False, + min=MINIMUM_DIFFUSION_COEFFICIENT, + unit='m**2/s', + ) + self._mean_u_squared = Parameter( + name='mean_u_squared', + value=float(mean_u_squared), + fixed=False, + min=MINIMUM_MEAN_U_SQUARED, + unit='angstrom**2', + ) + + self._hbar = hbar + # Dimensionless one, which turns the plain numbers of the weight expressions into + # (dimensionless) descriptors, as dependency expressions must evaluate to one. + self._one = DescriptorNumber(name='one', value=1.0) + + # Dimensionless /angstrom^2 for the weight dependency expressions. Deriving it + # through the dependency graph keeps the weights correct if is converted to another + # unit (e.g. nm^2). + with np.errstate(invalid='ignore'): + self._mean_u_squared_over_angstrom_squared = Parameter.from_dependency( + name='mean_u_squared_over_angstrom_squared', + dependency_expression='mean_u_squared / angstrom**2', + dependency_map={'mean_u_squared': self._mean_u_squared, 'angstrom': angstrom}, + ) + self._mean_u_squared_over_angstrom_squared.set_desired_unit('dimensionless') + + self.create_component_collections() + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def diffusion_coefficient(self) -> Parameter: + """ + Get the diffusion coefficient. + + Returns + ------- + Parameter + Diffusion coefficient D. + """ + return self._diffusion_coefficient + + @diffusion_coefficient.setter + def diffusion_coefficient(self, diffusion_coefficient: Numeric) -> None: + """ + Set the diffusion coefficient. + + Parameters + ---------- + diffusion_coefficient : Numeric + Diffusion coefficient D in m^2/s. + """ + self._verify_diffusion_coefficient(diffusion_coefficient) + self._diffusion_coefficient.value = float(diffusion_coefficient) + + @property + def mean_u_squared(self) -> Parameter: + """ + Get the mean square displacement. + + Returns + ------- + Parameter + Mean square displacement . + """ + return self._mean_u_squared + + @mean_u_squared.setter + def mean_u_squared(self, mean_u_squared: Numeric) -> None: + """ + Set the mean square displacement. + + Parameters + ---------- + mean_u_squared : Numeric + Mean square displacement in the unit of the parameter (angstrom^2 by default). + """ + self._verify_mean_u_squared(mean_u_squared) + self._mean_u_squared.value = float(mean_u_squared) + + # ------------------------------------------------------------------ + # Other methods + # ------------------------------------------------------------------ + + def calculate_width(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + r""" + Calculate the half-width at half-maximum (HWHM) of the Lorentzian of a given order. + + The width is $3 i \hbar D / \langle u^2 \rangle$ and does not depend on Q. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + order : int, default=1 + The order i of the Lorentzian, between 1 and number_of_lorentzians. + + Returns + ------- + np.ndarray + HWHM values in the unit of the model (e.g., meV). + """ + self._verify_order(order) + Q = self._ensure_Q(Q) + + width = ( + 3 + * sc.values(self._hbar.full_value) + * sc.values(self.diffusion_coefficient.full_value) + / sc.values(self.mean_u_squared.full_value) + ).to(unit=self.x_unit) + return order * width.value * np.ones_like(Q) + + def calculate_EISF(self, Q: Q_type | None = None) -> np.ndarray: + r""" + Calculate the Elastic Incoherent Structure Factor (EISF), $\exp(-Q^2\langle u^2\rangle/3)$. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + + Returns + ------- + np.ndarray + EISF values (dimensionless). + """ + Q = self._ensure_Q(Q) + return sc.exp(-self._y(Q)).values + + def calculate_QISF_order(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + """ + Calculate the Quasi-Elastic Incoherent Structure Factor (QISF) of one order, $A_i(Q) = + e^{-y} y^i / i!$, the weight of the Lorentzian of order $i$. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + order : int, default=1 + The order i of the Lorentzian, between 1 and number_of_lorentzians. + + Returns + ------- + np.ndarray + QISF values of that order (dimensionless). + """ + self._verify_order(order) + Q = self._ensure_Q(Q) + return self._quasielastic_weights(Q)['order', int(order) - 1].values + + def calculate_QISF(self, Q: Q_type | None = None) -> np.ndarray: + """ + Calculate the total Quasi-Elastic Incoherent Structure Factor (QISF), the sum of the QISF + of all number_of_lorentzians orders. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + + Returns + ------- + np.ndarray + QISF values (dimensionless). + """ + Q = self._ensure_Q(Q) + return sc.sum(self._quasielastic_weights(Q), 'order').values + + def calculate_missing_weight(self, Q: Q_type | None = None) -> np.ndarray: + """ + Calculate the spectral weight lost by truncating the sum at number_of_lorentzians. + + The weights of all orders sum to one, so this is 1 - EISF - QISF. If it is not small at the + Q values of interest, increase number_of_lorentzians. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + + Returns + ------- + np.ndarray + The omitted weight (dimensionless). + """ + Q = self._ensure_Q(Q) + return 1.0 - self.calculate_EISF(Q) - self.calculate_QISF(Q) + + def create_component_collections(self) -> list[ComponentCollection]: + r""" + Create ComponentCollections for the model at the current Q values. + + Each collection holds a delta function with area $\mathrm{scale}\, A_0(Q)$ and + number_of_lorentzians Lorentzians with widths $\Gamma_i$ and areas $\mathrm{scale}\, + A_i(Q)$ (built by the base class). The sum over i is truncated, so a warning is raised if + it omits more than MISSING_WEIGHT_TOLERANCE of the weight. The created collections are + installed on the model, so the returned list is the live one. + + Returns + ------- + list[ComponentCollection] + One ComponentCollection per Q value. + """ + collections = super().create_component_collections() + self._warn_if_truncated() + return collections + + def get_global_variables(self) -> list[Parameter]: + """ + Get all global variables from the diffusion model. + + Returns + ------- + list[Parameter] + The scale, diffusion coefficient and mean square displacement. + """ + return [self.scale, self.diffusion_coefficient, self.mean_u_squared] + + # ------------------------------------------------------------------ + # Private methods + # ------------------------------------------------------------------ + + def _y(self, Q: np.ndarray) -> sc.Variable: + r""" + Calculate $y = Q^2 \langle u^2 \rangle / 3$ with Scipp. + + The variance of is not propagated. + + Parameters + ---------- + Q : np.ndarray + Scattering vector in 1/angstrom. + + Returns + ------- + sc.Variable + Dimensionless y with dims ``('Q',)``. + """ + Q_variable = sc.array(dims=['Q'], values=Q, unit='1/angstrom') + mean_u_squared = sc.values(self.mean_u_squared.full_value) + return (Q_variable**2 * mean_u_squared / 3).to(unit='dimensionless') + + def _quasielastic_weights(self, Q: np.ndarray) -> sc.Variable: + r""" + Calculate the QISF $A_i = e^{-y} y^i / i!$ of the orders $i = 1 \ldots n$ with Scipp. + + Computed in log space, $\exp(-y + i \ln y - \ln i!)$, so that neither $y^i$ nor $i!$ + overflows. Scipp has no factorial, so $\ln i!$ is the cumulative sum of $\ln k$. At $Q = 0$ + the logarithm is $-\infty$ and the weights are exactly 0. + + Parameters + ---------- + Q : np.ndarray + Scattering vector in 1/angstrom. + + Returns + ------- + sc.Variable + Dimensionless QISF with dims ``('Q', 'order')``. + """ + order = sc.arange('order', 1.0, self.number_of_lorentzians + 1.0) + log_factorial = sc.cumsum(sc.log(order), 'order') + y = self._y(Q) + weights = sc.exp(-y + order * sc.log(y) - log_factorial) + return weights.transpose(['Q', 'order']) + + def _create_weight(self, Q: float, order: int) -> Parameter: + r""" + Create the EISF ($i = 0$) or the QISF of order $i$, $A_i(Q)$, as a dimensionless dependent + Parameter that follows $\langle u^2 \rangle$. + + The dependency expressions cannot call Scipp, so the weight is written out in the same + log-space form as ``_quasielastic_weights``. Keeping the weight separate from the scale + makes updates of the scale cheap. + + Parameters + ---------- + Q : float + Scattering vector in 1/angstrom. + order : int + The order i; 0 for the delta function (the EISF). + + Raises + ------ + TypeError + If Q is not a float. + + Returns + ------- + Parameter + The weight. + """ + if not isinstance(Q, float): + raise TypeError('Q must be a float.') + + # float(): NumPy 2 writes np.float64 as 'np.float64(...)', which asteval cannot parse + y = f'{float(Q**2 / 3)!r} * u2.value' + if order == 0: + expression = f'one * exp(-{y})' + else: + log_factorial = math.lgamma(order + 1) + expression = f'y = {y}\none * exp(-y + {order} * log(y) - {log_factorial!r})' + with np.errstate(invalid='ignore'): + return Parameter.from_dependency( + name=f'weight_i{order}_Q{Q:.2f}', + dependency_expression=expression, + dependency_map={ + 'one': self._one, + 'u2': self._mean_u_squared_over_angstrom_squared, + }, + ) + + def _write_width_dependency_expression(self, order: int) -> str: + """ + Write the dependency expression for the width of the Lorentzian of a given order. + + Parameters + ---------- + order : int + The order i of the Lorentzian. + + Returns + ------- + str + Dependency expression for the width. + """ + # Dividing by the dimensionless value of /angstrom^2 instead of by the + # Parameter avoids inf/inf in Core's bound arithmetic (D and are both unbounded + # above), which warns on every update during a fit. + return f'{3 * order} * hbar * D / angstrom**2 / u2.value' + + def _write_width_dependency_map_expression(self) -> dict[str, DescriptorNumber]: + """ + Write the dependency map for the width. + + Returns + ------- + dict[str, DescriptorNumber] + Dependency map for the width. + """ + return { + 'hbar': self._hbar, + 'D': self.diffusion_coefficient, + 'angstrom': angstrom, + 'u2': self._mean_u_squared_over_angstrom_squared, + } + + def _warn_if_truncated(self) -> None: + """ + Warn if truncating the sum at number_of_lorentzians omits too much weight. + + The warning is raised if the missing weight exceeds MISSING_WEIGHT_TOLERANCE at any of the + model's Q values. Without Q values there is nothing to check. + """ + if self.Q is None: + return + if np.max(self.calculate_missing_weight(), initial=0.0) > MISSING_WEIGHT_TOLERANCE: + warnings.warn( + f'Truncating the sum at number_of_lorentzians={self.number_of_lorentzians} omits ' + f'more than {MISSING_WEIGHT_TOLERANCE} of the spectral weight at some Q values. ' + 'Increase number_of_lorentzians.', + UserWarning, + stacklevel=3, + ) + + @staticmethod + def _verify_diffusion_coefficient(diffusion_coefficient: Numeric) -> None: + """ + Verify a diffusion coefficient value. + + Parameters + ---------- + diffusion_coefficient : Numeric + The diffusion coefficient to verify. + + Raises + ------ + TypeError + If diffusion_coefficient is not a number. + ValueError + If diffusion_coefficient is below MINIMUM_DIFFUSION_COEFFICIENT. + """ + if not isinstance(diffusion_coefficient, Numeric): + raise TypeError('diffusion_coefficient must be a number.') + if float(diffusion_coefficient) < MINIMUM_DIFFUSION_COEFFICIENT: + raise ValueError( + f'diffusion_coefficient must be at least {MINIMUM_DIFFUSION_COEFFICIENT}.' + ) + + @staticmethod + def _verify_mean_u_squared(mean_u_squared: Numeric) -> None: + """ + Verify a mean square displacement value. + + Parameters + ---------- + mean_u_squared : Numeric + The mean square displacement to verify. + + Raises + ------ + TypeError + If mean_u_squared is not a number. + ValueError + If mean_u_squared is below MINIMUM_MEAN_U_SQUARED. + """ + if not isinstance(mean_u_squared, Numeric): + raise TypeError('mean_u_squared must be a number.') + if float(mean_u_squared) < MINIMUM_MEAN_U_SQUARED: + raise ValueError(f'mean_u_squared must be at least {MINIMUM_MEAN_U_SQUARED}.') + + # ------------------------------------------------------------------ + # dunder methods + # ------------------------------------------------------------------ + + def __repr__(self) -> str: + """ + String representation of the GaussianModel3D model. + + Returns + ------- + str + String representation of the model. + """ + return ( + f'{self.__class__.__name__}(' + f'name={self.name!r}, display_name={self.display_name!r},\n' + f'x_unit={self.x_unit}, y_unit={self.y_unit},\n' + f' number_of_lorentzians={self.number_of_lorentzians},\n' + f' diffusion_coefficient={self.diffusion_coefficient},\n' + f' mean_u_squared={self.mean_u_squared},\n' + f' scale={self.scale})' + ) diff --git a/src/easydynamics/sample_model/diffusion_model/hall_ross_jump_diffusion.py b/src/easydynamics/sample_model/diffusion_model/hall_ross_jump_diffusion.py new file mode 100644 index 00000000..a9282fda --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/hall_ross_jump_diffusion.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + + +import numpy as np +import scipp as sc +from easyscience.variable import DescriptorNumber +from easyscience.variable import Parameter + +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 angstrom +from easydynamics.utils.utils import hbar + +MINIMUM_WIDTH = 1e-10 # To avoid division by zero + + +class HallRossJumpDiffusion(DiffusionModelBase): + r""" + Model of Hall-Ross jump diffusion, with a Gaussian distribution of jump lengths. + + The model consists of a Lorentzian function for each Q-value, where the width is given by + + $$ \Gamma(Q) = \frac{\hbar}{\tau} \left( 1 - \exp\left(-\frac{Q^2 l^2}{2}\right) \right) $$ + + where $\tau$ is the residence time and $l$ is the standard deviation of the jump-length + distribution, the convention of reference 2. The mean-square jump length is $\langle r^2 + \rangle = 3 l^2$. At small $Q$ the width tends to $\hbar D Q^2$ with $D = l^2 / (2 \tau)$, and + at large $Q$ to $\hbar / \tau$. $Q$ is assumed to have units of 1/angstrom. Creates + ComponentCollections with Lorentzian components for given Q-values. + + References + ---------- + 1. P. L. Hall and D. K. Ross, *Incoherent neutron scattering functions for random jump + diffusion in bounded and infinite media*, Mol. Phys. **42**, 673 (1981), + [doi:10.1080/00268978100100521](https://doi.org/10.1080/00268978100100521). + 2. Mantid, ``HallRoss`` fit function, whose ``L`` parameter is this $l$, + https://github.com/mantidproject/mantid/blob/17c1cd96c1d46b39bee21dc75a60017b658d689c/Framework/PythonInterface/plugins/functions/HallRoss.py#L30-L39. + + Examples + -------- + **Creating a HallRossJumpDiffusion model** + + Pass the residence time (in ps) and jump length (in angstroms) along with Q values: + ```python + import numpy as np + import easydynamics as edyn + + Q = np.linspace(0.5, 2, 7) + diffusion_model = edyn.HallRossJumpDiffusion( + scale=1.0, + residence_time=1.0, + jump_length=1.5, + Q=Q, + ) + component_collections = diffusion_model.create_component_collections() + ``` + + See also the tutorials. + """ + + def __init__( + self, + scale: Numeric = 1.0, + residence_time: Numeric = 1.0, + jump_length: Numeric = 1.0, + Q: Q_type | None = None, + x_unit: str | sc.Unit = 'meV', + y_unit: str | sc.Unit = 'dimensionless', + name: str = 'HallRossJumpDiffusion', + display_name: str | None = 'HallRossJumpDiffusion', + lorentzian_name: str | None = None, + lorentzian_display_name: str | None = None, + unique_name: str | None = None, + ) -> None: + """ + Initialize a new HallRossJumpDiffusion model. + + Parameters + ---------- + scale : Numeric, default=1.0 + Scale factor for the diffusion model. Must be a non-negative number. + residence_time : Numeric, default=1.0 + Residence time parameter tau in ps. + jump_length : Numeric, default=1.0 + Standard deviation l of the jump-length distribution in angstrom. + Q : Q_type | None, default=None + Q values for the model. If None, Q is not set. + x_unit : str | sc.Unit, default='meV' + Unit of the x-axis (energy/frequency). Must be convertible to meV. + y_unit : str | sc.Unit, default='dimensionless' + Unit of the model output (intensity). Determines scale.unit = x_unit * y_unit. + name : str, default='HallRossJumpDiffusion' + Name of the diffusion model. + display_name : str | None, default='HallRossJumpDiffusion' + Display name of the diffusion model. + lorentzian_name : str | None, default=None + Name of the Lorentzian component. If None, it will be set to the name of the diffusion + model. By default, None. + lorentzian_display_name : str | None, default=None + Display name of the Lorentzian component. If None, it will be set to the + lorentzian_name. By default, None + unique_name : str | None, default=None + Unique name of the diffusion model. If None, a unique name will be generated. By + default, None. + + Raises + ------ + TypeError + If ``residence_time`` or ``jump_length`` is not a number. + ValueError + If ``jump_length`` is negative or ``residence_time`` is below ``MINIMUM_WIDTH``. + """ + super().__init__( + Q=Q, + x_unit=x_unit, + y_unit=y_unit, + scale=scale, + name=name, + display_name=display_name, + lorentzian_name=lorentzian_name, + lorentzian_display_name=lorentzian_display_name, + unique_name=unique_name, + ) + + if not isinstance(residence_time, Numeric): + raise TypeError('residence_time must be a number.') + + if float(residence_time) < MINIMUM_WIDTH: + raise ValueError(f'residence_time must be at least {MINIMUM_WIDTH}.') + + if not isinstance(jump_length, Numeric): + raise TypeError('jump_length must be a number.') + + if float(jump_length) < 0: + raise ValueError('jump_length must be non-negative.') + + residence_time = Parameter( + name='residence_time', + value=float(residence_time), + fixed=False, + unit='ps', + min=MINIMUM_WIDTH, + ) + + jump_length = Parameter( + name='jump_length', + value=float(jump_length), + fixed=False, + unit='angstrom', + min=0.0, + ) + + self._hbar = hbar + self._angstrom = angstrom + self._residence_time = residence_time + self._jump_length = jump_length + + self._component_collections = self.create_component_collections() + + ################################ + # Properties + ################################ + + @property + def residence_time(self) -> Parameter: + """ + Get the residence time parameter tau. + + Returns + ------- + Parameter + Residence time tau. + """ + return self._residence_time + + @residence_time.setter + def residence_time(self, residence_time: Numeric) -> None: + """ + Set the residence time parameter tau. + + Parameters + ---------- + residence_time : Numeric + Residence time tau in ps. + + Raises + ------ + TypeError + If residence_time is not a number. + ValueError + If residence_time is below ``MINIMUM_WIDTH``. + """ + if not isinstance(residence_time, Numeric): + raise TypeError('residence_time must be a number.') + if float(residence_time) < MINIMUM_WIDTH: + raise ValueError(f'residence_time must be at least {MINIMUM_WIDTH}.') + + self._residence_time.value = float(residence_time) + + @property + def jump_length(self) -> Parameter: + """ + Get the jump length parameter l. + + Returns + ------- + Parameter + Standard deviation l of the jump-length distribution in angstrom. + """ + return self._jump_length + + @jump_length.setter + def jump_length(self, jump_length: Numeric) -> None: + """ + Set the jump length parameter l. + + Parameters + ---------- + jump_length : Numeric + Standard deviation l of the jump-length distribution in angstrom. + + Raises + ------ + TypeError + If jump_length is not a number. + ValueError + If jump_length is negative. + """ + if not isinstance(jump_length, Numeric): + raise TypeError('jump_length must be a number.') + + if float(jump_length) < 0: + raise ValueError('jump_length must be non-negative.') + self._jump_length.value = float(jump_length) + + ################################ + # Other methods + ################################ + + def calculate_width(self, Q: Q_type | None = None) -> np.ndarray: + r""" + Calculate the half-width at half-maximum (HWHM) for the diffusion model. + + $\Gamma(Q) = (\hbar / \tau) (1 - \exp(-Q^2 l^2 / 2))$, where $\tau$ is the residence time + and $l$ is the standard deviation of the jump-length distribution. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. Can be a single value or an array of values. If None, + Q values stored in the model are used. + + Returns + ------- + np.ndarray + HWHM values in the unit of the model (e.g., meV). + """ + Q = self._ensure_Q(Q) + + conversion_factor = self._jump_length / self._angstrom + conversion_factor.convert_unit('dimensionless') + + prefactor = self._hbar / self._residence_time + prefactor.convert_unit(self.x_unit) + + argument = Q * conversion_factor.value + + # -expm1(-x) is 1 - exp(-x) without the cancellation at small Q + return prefactor.value * -np.expm1(-(argument**2) / 2) + + def calculate_EISF(self, Q: Q_type) -> np.ndarray: + """ + Calculate the Elastic Incoherent Structure Factor (EISF), which is 0 for this model. + + Parameters + ---------- + Q : Q_type + Scattering vector in 1/angstrom. Can be a single value or an array of values. + + Returns + ------- + np.ndarray + EISF values (dimensionless). + """ + Q = self._ensure_Q(Q) + + return np.zeros_like(Q) + + def calculate_QISF(self, Q: Q_type) -> np.ndarray: + """ + Calculate the Quasi-Elastic Incoherent Structure Factor (QISF), which is 1 for this model. + + Parameters + ---------- + Q : Q_type + Scattering vector in 1/angstrom. Can be a single value or an array of values. + + Returns + ------- + np.ndarray + QISF values (dimensionless). + """ + Q = self._ensure_Q(Q) + + return np.ones_like(Q) + + def create_component_collections( + self, + ) -> list[ComponentCollection]: + """ + Create ComponentCollection components for the diffusion model at given Q values. + + The created collections are installed on the model (they become the collections returned by + ``get_component_collections``), so the returned list is the live one. + + Returns + ------- + list[ComponentCollection] + List of ComponentCollections with Hall-Ross Lorentzian components. + """ + if self.Q is None: + self._component_collections = [] + return self._component_collections + + Q = self.Q.values + component_collection_list = [None] * len(Q) + QISF = self.calculate_QISF(Q) + + # One Lorentzian per Q-value, with the Hall-Ross width 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( + name=f'{self.name}_Q{Q_value:.2f}', + display_name=f'{self.display_name}_Q{Q_value:.2f}', + x_unit=self.x_unit, + y_unit=self.y_unit, + ) + + lorentzian_component = Lorentzian( + name=self.lorentzian_name, + display_name=self.lorentzian_display_name, + x_unit=self.x_unit, + y_unit=self.y_unit, + ) + + dependency_expression = self._write_width_dependency_expression(Q[i]) + dependency_map = self._write_width_dependency_map_expression() + + # easyscience propagates inf bounds through arithmetic, producing inf/inf=nan + # as a transient intermediate. Python's min/max ignore nan so the final bounds + # are correct; suppress the spurious numpy RuntimeWarning. + with np.errstate(invalid='ignore'): + lorentzian_component.width.make_dependent_on( + dependency_expression=dependency_expression, + dependency_map=dependency_map, + desired_unit=self.x_unit, + ) + + area_dependency_map = self._write_area_dependency_map_expression() + lorentzian_component.area.make_dependent_on( + dependency_expression=self._write_area_dependency_expression(QISF[i]), + dependency_map=area_dependency_map, + ) + + component_collection_list[i].append_component(lorentzian_component) + + self._component_collections = component_collection_list + return self._component_collections + + ################################ + # Private methods + ################################ + def _on_Q_change(self) -> None: + """ + Update the component collections when Q changes. This is called automatically when the Q + property is set. It regenerates the component collections based on the new Q values. + """ + self._component_collections = self.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. + + Raises + ------ + TypeError + If Q is not a float. + + Returns + ------- + str + Dependency expression for the width. + """ + if not isinstance(Q, (float)): + raise TypeError('Q must be a float.') + + # Q is in 1/angstrom and l in angstrom; -expm1(-x) is 1 - exp(-x) without cancellation + return f'(hbar / tau) * -expm1(-(({Q} * l.value) ** 2) / 2)' + + def _write_width_dependency_map_expression(self) -> dict[str, DescriptorNumber]: + """ + Write the dependency map expression to make dependent Parameters. + + Returns + ------- + dict[str, DescriptorNumber] + Dependency map for the width. + """ + return { + 'tau': self.residence_time, + 'l': self.jump_length, + 'hbar': self._hbar, + } + + ################################ + # dunder methods + ################################ + + def __repr__(self) -> str: + """ + String representation of the HallRossJumpDiffusion model. + + Returns + ------- + str + String representation of the HallRossJumpDiffusion model. + """ + return ( + f'{self.__class__.__name__}(' + f'name={self.name!r}, display_name={self.display_name!r},\n' + f'x_unit={self.x_unit}, y_unit={self.y_unit}, \n' + f' residence_time={self.residence_time},\n' + f' jump_length={self.jump_length}),\n' + f' scale={self.scale})' + ) diff --git a/src/easydynamics/sample_model/diffusion_model/jump_sites_log_normal_distribution.py b/src/easydynamics/sample_model/diffusion_model/jump_sites_log_normal_distribution.py new file mode 100644 index 00000000..61cbb887 --- /dev/null +++ b/src/easydynamics/sample_model/diffusion_model/jump_sites_log_normal_distribution.py @@ -0,0 +1,651 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import scipp as sc +from easyscience.variable import DescriptorNumber +from easyscience.variable import Parameter + +from easydynamics.sample_model.diffusion_model.delta_lorentzian_diffusion_model_base import ( + DeltaLorentzianDiffusionModelBase, +) +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import ( + MAXIMUM_NUMBER_OF_SITES, +) +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import MINIMUM_RADIUS +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import MINIMUM_WIDTH +from easydynamics.utils.utils import Numeric +from easydynamics.utils.utils import Q_type +from easydynamics.utils.utils import angstrom +from easydynamics.utils.utils import hbar + +# The sampling of the log-normal distribution, as in QENSmodels jump_sites_log_norm_dist.py: +# https://github.com/QENSlibrary/QENSmodels/blob/main/QENSmodels/jump_sites_log_norm_dist.py +NUMBER_OF_SAMPLES = 21 # 2 n_max + 1 widths per circle order, with n_max = 10 +LOWER_LIMIT = 0.1 # Smallest sampled weight relative to the largest +MINIMUM_SIGMA = 1e-10 # A zero width of the distribution gives a 0/0 in the sampling grid + +# Half the number of sampling steps, and log(ratio) / sigma at the edge of the grid +_HALF_STEPS = (NUMBER_OF_SAMPLES - 1) // 2 +_EDGE = float(np.sqrt(-2.0 * np.log(LOWER_LIMIT))) + + +class JumpSitesLogNormalDistribution(DeltaLorentzianDiffusionModelBase): + r""" + Model of jumps between N equivalent sites on a circle with a log-normal distribution of jump + rates (Chahid, Alegría and Colmenero). + + $$ S(Q, E) = A_0(Q)\,\delta(E) + \sum_{l=1}^{N-1} A_l(Q) \sum_{s=0}^{20} g_s\, L(E, \Gamma_l + r_s) $$ + + with the weights and widths of the equivalent-sites circle model, + + $$ A_l(Q) = \frac{1}{N} \sum_{k=1}^{N} j_0(Q r_k) \cos\left(\frac{2\pi l k}{N}\right), \quad + r_k = 2R \sin\left(\frac{\pi k}{N}\right), \quad \Gamma_l = \frac{2\hbar}{\tau} + \sin^2\left(\frac{\pi l}{N}\right), $$ + + spread over 21 log-spaced widths $\Gamma_l r_s$ with $\ln r_s = \sigma c (s/10 - 1)$, $c = + \sqrt{-2 \ln 0.1}$, and weights $g_s \propto \exp(-\ln^2 r_s / (2\sigma^2))$ normalised to + $\sum_s g_s = 1$. Because the grid scales with $\sigma$, the $g_s$ do not depend on $\sigma$; + only the widths do. For $\sigma \to 0$ this is the equivalent-sites circle model. + + $A_0$ is the EISF and $g_s A_l$ the QISF of order $j$, where the Lorentzian of order $j = (l - + 1) \cdot 21 + s + 1$, from 1 to $21 (N - 1)$, has circle order $l$ and sample $s$. The 21 + widths are cut where the weight falls to 0.1 of its maximum. + + The radius is fixed by default, since it is normally known from the molecular geometry. + + References + ---------- + 1. A. Chahid, A. Alegría and J. Colmenero, Macromolecules 27, 3282 (1994), + doi:10.1021/ma00090a022. + 2. QENSmodels, jump_sites_log_norm_dist, + https://github.com/QENSlibrary/QENSmodels/blob/main/QENSmodels/jump_sites_log_norm_dist.py. + + Examples + -------- + **Creating a JumpSitesLogNormalDistribution model** + + The model creates one delta function and 21 (N - 1) Lorentzians per Q-value: + ```python + import numpy as np + import easydynamics as edyn + + model = edyn.JumpSitesLogNormalDistribution( + radius=1.0, + residence_time=2.0, + sigma=0.5, + number_of_sites=3, + Q=np.linspace(0.5, 2, 7), + ) + ``` + """ + + _order_symbol = 'j' # components are named '_j' + + def __init__( + self, + scale: Numeric = 1.0, + radius: Numeric = 1.0, + residence_time: Numeric = 1.0, + sigma: Numeric = 1.0, + number_of_sites: int = 3, + Q: Q_type | None = None, + x_unit: str | sc.Unit = 'meV', + y_unit: str | sc.Unit = 'dimensionless', + name: str = 'JumpSitesLogNormalDistribution', + display_name: str | None = 'JumpSitesLogNormalDistribution', + lorentzian_name: str | None = None, + lorentzian_display_name: str | None = None, + delta_name: str = 'Delta function', + delta_display_name: str | None = None, + unique_name: str | None = None, + ) -> None: + """ + Initialize a new JumpSitesLogNormalDistribution model. + + Parameters + ---------- + scale : Numeric, default=1.0 + Scale factor for the diffusion model. Must be a non-negative number. + radius : Numeric, default=1.0 + Radius of the circle in angstrom. Must be positive. Fixed by default. + residence_time : Numeric, default=1.0 + Residence time tau on a site in ps, the centre of the distribution. + sigma : Numeric, default=1.0 + Standard deviation of the log-normal distribution of jump rates (dimensionless). Must + be at least MINIMUM_SIGMA. + number_of_sites : int, default=3 + Number of equivalent sites N on the circle, between 2 and MAXIMUM_NUMBER_OF_SITES. Set + at construction; not a fit parameter. + Q : Q_type | None, default=None + Q values for the model. If None, Q is not set. + x_unit : str | sc.Unit, default='meV' + Unit of the x-axis (energy/frequency). Must be convertible to meV. + y_unit : str | sc.Unit, default='dimensionless' + Unit of the model output (intensity). Determines scale.unit = x_unit * y_unit. + name : str, default='JumpSitesLogNormalDistribution' + Name of the diffusion model. + display_name : str | None, default='JumpSitesLogNormalDistribution' + Display name of the diffusion model. + lorentzian_name : str | None, default=None + Base name of the Lorentzian components. The Lorentzian of order j is named + '_j'. If None, it will be set to the name of the diffusion model. + lorentzian_display_name : str | None, default=None + Base display name of the Lorentzian components. If None, it will be set to + lorentzian_name. + delta_name : str, default='Delta function' + Name of the delta function component. + delta_display_name : str | None, default=None + Display name of the delta function component. If None, it will be set to delta_name. + unique_name : str | None, default=None + Unique name of the diffusion model. If None, a unique name will be generated. + + Raises + ------ + TypeError + If radius, residence_time or sigma is not a number, if number_of_sites is not an + integer or if delta_name or delta_display_name is not a string (or None). + ValueError + If radius, residence_time or sigma is smaller than its minimum, or if number_of_sites + is not between 2 and MAXIMUM_NUMBER_OF_SITES. + """ + self._verify_radius(radius) + self._verify_residence_time(residence_time) + self._verify_sigma(sigma) + + if isinstance(number_of_sites, bool) or not isinstance(number_of_sites, (int, np.integer)): + raise TypeError('number_of_sites must be an integer.') + if not 2 <= number_of_sites <= MAXIMUM_NUMBER_OF_SITES: + raise ValueError(f'number_of_sites must be between 2 and {MAXIMUM_NUMBER_OF_SITES}.') + self._number_of_sites = int(number_of_sites) + + super().__init__( + NUMBER_OF_SAMPLES * (self._number_of_sites - 1), + scale=scale, + Q=Q, + x_unit=x_unit, + y_unit=y_unit, + name=name, + display_name=display_name, + lorentzian_name=lorentzian_name, + lorentzian_display_name=lorentzian_display_name, + delta_name=delta_name, + delta_display_name=delta_display_name, + unique_name=unique_name, + ) + + self._radius = Parameter( + name='radius', + value=float(radius), + fixed=True, + min=MINIMUM_RADIUS, + unit='angstrom', + ) + self._residence_time = Parameter( + name='residence_time', + value=float(residence_time), + fixed=False, + min=MINIMUM_WIDTH, + unit='ps', + ) + self._sigma = Parameter( + name='sigma', + value=float(sigma), + fixed=False, + min=MINIMUM_SIGMA, + unit='dimensionless', + ) + + self._hbar = hbar + # Dimensionless one, which turns the plain numbers of the weight expressions into + # (dimensionless) descriptors, as dependency expressions must evaluate to one. + self._one = DescriptorNumber(name='one', value=1.0) + + # Dimensionless radius/angstrom for the weight dependency expressions. Deriving it through + # the dependency graph keeps the weights correct if the radius is converted to another + # unit (e.g. nm). + with np.errstate(invalid='ignore'): + self._radius_over_angstrom = Parameter.from_dependency( + name='radius_over_angstrom', + dependency_expression='radius / angstrom', + dependency_map={'radius': self._radius, 'angstrom': angstrom}, + ) + self._radius_over_angstrom.set_desired_unit('dimensionless') + + self.create_component_collections() + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def radius(self) -> Parameter: + """ + Get the radius of the circle. + + Returns + ------- + Parameter + Radius R. + """ + return self._radius + + @radius.setter + def radius(self, radius: Numeric) -> None: + """ + Set the radius of the circle. + + Parameters + ---------- + radius : Numeric + Radius R in the unit of the radius parameter (angstrom by default). + """ + self._verify_radius(radius) + self._radius.value = float(radius) + + @property + def residence_time(self) -> Parameter: + """ + Get the residence time on a site, the centre of the distribution. + + Returns + ------- + Parameter + Residence time tau. + """ + return self._residence_time + + @residence_time.setter + def residence_time(self, residence_time: Numeric) -> None: + """ + Set the residence time on a site. + + Parameters + ---------- + residence_time : Numeric + Residence time tau in ps. + """ + self._verify_residence_time(residence_time) + self._residence_time.value = float(residence_time) + + @property + def sigma(self) -> Parameter: + """ + Get the standard deviation of the log-normal distribution of jump rates. + + Returns + ------- + Parameter + Sigma (dimensionless). + """ + return self._sigma + + @sigma.setter + def sigma(self, sigma: Numeric) -> None: + """ + Set the standard deviation of the log-normal distribution of jump rates. + + Parameters + ---------- + sigma : Numeric + Sigma (dimensionless). + """ + self._verify_sigma(sigma) + self._sigma.value = float(sigma) + + @property + def number_of_sites(self) -> int: + """ + Get the number of equivalent sites on the circle. + + Returns + ------- + int + The number of sites N. Set at construction. + """ + return self._number_of_sites + + # ------------------------------------------------------------------ + # Other methods + # ------------------------------------------------------------------ + + def sample_weights(self) -> np.ndarray: + r""" + Get the weights $g_s$ of the 21 sampled widths, which sum to one. + + $g_s \propto \exp(-c^2 (s/10 - 1)^2 / 2)$: the log-normal weight $\exp(-\ln^2 r_s / + (2\sigma^2))$ with $\ln r_s = \sigma c (s/10 - 1)$, so it does not depend on $\sigma$. + + Returns + ------- + np.ndarray + The 21 weights. + """ + log_ratio_per_sigma = _EDGE * (np.arange(NUMBER_OF_SAMPLES) / _HALF_STEPS - 1.0) + weights = np.exp(-0.5 * log_ratio_per_sigma**2) + return weights / np.sum(weights) + + def calculate_width(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + r""" + Calculate the half-width at half-maximum (HWHM) of the Lorentzian of a given order. + + The width is $(2\hbar/\tau) \sin^2(\pi l / N) \exp(\sigma c (s/10 - 1))$ and does not + depend on Q. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + order : int, default=1 + The order j of the Lorentzian, between 1 and 21 (N - 1). + + Returns + ------- + np.ndarray + HWHM values in the unit of the model (e.g., meV). + """ + self._verify_order(order) + Q = self._ensure_Q(Q) + circle_order, s = self._split_order(order) + + rate = 2 * self._hbar / self.residence_time + rate.convert_unit(self.x_unit) + angle = sc.scalar(np.pi * circle_order / self.number_of_sites, unit='rad') + ratio = np.exp(self.sigma.value * _EDGE * (s / _HALF_STEPS - 1.0)) + return rate.value * sc.sin(angle).value ** 2 * ratio * np.ones_like(Q) + + def calculate_EISF(self, Q: Q_type | None = None) -> np.ndarray: + """ + Calculate the Elastic Incoherent Structure Factor (EISF), $A_0(Q)$. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + + Returns + ------- + np.ndarray + EISF values (dimensionless). + """ + Q = self._ensure_Q(Q) + return self._circle_structure_factors(Q)['l', 0].values + + def calculate_QISF_order(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + """ + Calculate the Quasi-Elastic Incoherent Structure Factor (QISF) of one order, $g_s A_l(Q)$, + the weight of the Lorentzian of order $j$ (circle order $l$, sample $s$). + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + order : int, default=1 + The order j of the Lorentzian, between 1 and 21 (N - 1). + + Returns + ------- + np.ndarray + QISF values of that order (dimensionless). + """ + self._verify_order(order) + Q = self._ensure_Q(Q) + circle_order, s = self._split_order(order) + return ( + self.sample_weights()[s] * self._circle_structure_factors(Q)['l', circle_order].values + ) + + def calculate_QISF(self, Q: Q_type | None = None) -> np.ndarray: + """ + Calculate the total Quasi-Elastic Incoherent Structure Factor (QISF), the sum of the QISF + of all 21 (N - 1) orders. + + Parameters + ---------- + Q : Q_type | None, default=None + Scattering vector in 1/angstrom. + + Returns + ------- + np.ndarray + QISF values (dimensionless). + """ + Q = self._ensure_Q(Q) + circle_QISF = sc.sum(self._circle_structure_factors(Q)['l', 1:], 'l').values + return np.sum(self.sample_weights()) * circle_QISF + + def get_global_variables(self) -> list[Parameter]: + """ + Get all global variables from the diffusion model. + + Returns + ------- + list[Parameter] + The scale, radius, residence time and sigma. + """ + return [self.scale, self.radius, self.residence_time, self.sigma] + + # ------------------------------------------------------------------ + # Private methods + # ------------------------------------------------------------------ + + def _split_order(self, order: int) -> tuple[int, int]: + """ + Split a Lorentzian order j into its circle order l and sample s. + + Parameters + ---------- + order : int + The order j, between 1 and 21 (N - 1). + + Returns + ------- + tuple[int, int] + The circle order l (1 to N - 1) and the sample s (0 to 20). + """ + return (int(order) - 1) // NUMBER_OF_SAMPLES + 1, (int(order) - 1) % NUMBER_OF_SAMPLES + + def _circle_structure_factors(self, Q: np.ndarray) -> sc.Variable: + r""" + Calculate the circle model's EISF and QISF, $A_l(Q)$ for all circle orders $l = 0 \ldots N + - 1$, with Scipp. + + $A_l(Q) = \frac{1}{N} \sum_{k=1}^{N} j_0(Q r_k) \cos(2\pi l k / N)$ with $r_k = 2R\sin(\pi + k / N)$, as in EquivalentSitesCircle. The variance of the radius is not propagated. + + Parameters + ---------- + Q : np.ndarray + Scattering vector in 1/angstrom. + + Returns + ------- + sc.Variable + Dimensionless $A_l(Q)$ with dims ``('Q', 'l')``; l = 0 is the EISF. + """ + N = self.number_of_sites + site = sc.arange('k', 1.0, N + 1.0) + circle_order = sc.arange('l', 0.0, float(N)) + half_turn = sc.scalar(np.pi, unit='rad') + + radius = sc.values(self.radius.full_value).to(unit='angstrom') + jump_distance = 2 * radius * sc.sin(half_turn * site / N) + Q_variable = sc.array(dims=['Q'], values=Q, unit='1/angstrom') + + weights = sc.sum( + sc.sinc(Q_variable * jump_distance) * sc.cos(2 * half_turn * circle_order * site / N), + 'k', + ) + return (weights / N).transpose(['Q', 'l']) + + def _create_weight(self, Q: float, order: int) -> Parameter: + """ + Create the EISF ($j = 0$) or the QISF of order $j$, $g_s A_l(Q)$, as a dimensionless + dependent Parameter that follows the radius. + + The dependency expressions cannot call Scipp, so the circle's sum over sites is written out + with NumPy's sinc, sinc(x) = sin(pi x) / (pi x), and multiplied by the constant g_s. + Keeping the weight separate from the scale makes updates of the scale cheap. + + Parameters + ---------- + Q : float + Scattering vector in 1/angstrom. + order : int + The order j; 0 for the delta function (the EISF). + + Raises + ------ + TypeError + If Q is not a float. + + Returns + ------- + Parameter + The weight. + """ + if not isinstance(Q, float): + raise TypeError('Q must be a float.') + + if order == 0: + circle_order, sample_weight = 0, 1.0 + else: + circle_order, s = self._split_order(order) + sample_weight = float(self.sample_weights()[s]) + + N = self.number_of_sites + terms = [] + for site in range(1, N + 1): + # float(): NumPy 2 writes np.float64 as 'np.float64(...)', which asteval cannot parse + coefficient = float(sample_weight * np.cos(2 * np.pi * circle_order * site / N) / N) + # j0(Q r_k) = sinc(Q r_k / pi) with r_k = 2 R sin(pi k / N) + argument = float(Q * 2 * np.sin(np.pi * site / N) / np.pi) + terms.append(f'{coefficient!r} * sinc({argument!r} * radius_ratio.value)') + expression = f'one * ({" + ".join(terms)})' + with np.errstate(invalid='ignore'): + return Parameter.from_dependency( + name=f'weight_j{order}_Q{Q:.2f}', + dependency_expression=expression, + dependency_map={'one': self._one, 'radius_ratio': self._radius_over_angstrom}, + ) + + def _write_width_dependency_expression(self, order: int) -> str: + """ + Write the dependency expression for the width of the Lorentzian of a given order. + + Parameters + ---------- + order : int + The order j of the Lorentzian. + + Returns + ------- + str + Dependency expression for the width. + """ + circle_order, s = self._split_order(order) + factor = float(2 * np.sin(np.pi * circle_order / self.number_of_sites) ** 2) + log_ratio_per_sigma = float(_EDGE * (s / _HALF_STEPS - 1.0)) + return f'{factor!r} * exp({log_ratio_per_sigma!r} * sigma.value) * hbar / tau' + + def _write_width_dependency_map_expression(self) -> dict[str, DescriptorNumber]: + """ + Write the dependency map for the width. + + Returns + ------- + dict[str, DescriptorNumber] + Dependency map for the width. + """ + return {'hbar': self._hbar, 'tau': self.residence_time, 'sigma': self.sigma} + + @staticmethod + def _verify_radius(radius: Numeric) -> None: + """ + Verify a radius value. + + Parameters + ---------- + radius : Numeric + The radius to verify. + + Raises + ------ + TypeError + If radius is not a number. + ValueError + If radius is below MINIMUM_RADIUS. + """ + if not isinstance(radius, Numeric): + raise TypeError('radius must be a number.') + if float(radius) < MINIMUM_RADIUS: + raise ValueError(f'radius must be at least {MINIMUM_RADIUS}.') + + @staticmethod + def _verify_residence_time(residence_time: Numeric) -> None: + """ + Verify a residence time value. + + Parameters + ---------- + residence_time : Numeric + The residence time to verify. + + Raises + ------ + TypeError + If residence_time is not a number. + ValueError + If residence_time is below MINIMUM_WIDTH. + """ + if not isinstance(residence_time, Numeric): + raise TypeError('residence_time must be a number.') + if float(residence_time) < MINIMUM_WIDTH: + raise ValueError(f'residence_time must be at least {MINIMUM_WIDTH}.') + + @staticmethod + def _verify_sigma(sigma: Numeric) -> None: + """ + Verify a sigma value. + + Parameters + ---------- + sigma : Numeric + The standard deviation of the distribution to verify. + + Raises + ------ + TypeError + If sigma is not a number. + ValueError + If sigma is below MINIMUM_SIGMA. + """ + if not isinstance(sigma, Numeric): + raise TypeError('sigma must be a number.') + if float(sigma) < MINIMUM_SIGMA: + raise ValueError(f'sigma must be at least {MINIMUM_SIGMA}.') + + # ------------------------------------------------------------------ + # dunder methods + # ------------------------------------------------------------------ + + def __repr__(self) -> str: + """ + String representation of the JumpSitesLogNormalDistribution model. + + Returns + ------- + str + String representation of the model. + """ + return ( + f'{self.__class__.__name__}(' + f'name={self.name!r}, display_name={self.display_name!r},\n' + f'x_unit={self.x_unit}, y_unit={self.y_unit},\n' + f' number_of_sites={self.number_of_sites},\n' + f' radius={self.radius},\n' + f' residence_time={self.residence_time},\n' + f' sigma={self.sigma},\n' + f' scale={self.scale})' + ) diff --git a/tests/integration/fitting/test_fitting_with_multi_lorentzian_models.py b/tests/integration/fitting/test_fitting_with_multi_lorentzian_models.py new file mode 100644 index 00000000..0c34b442 --- /dev/null +++ b/tests/integration/fitting/test_fitting_with_multi_lorentzian_models.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +# Round trip for the delta-plus-Lorentzians diffusion models: synthetic data from the model, +# convolved with a Gaussian resolution, then a simultaneous fit over all Q values. The unit tests +# fit one Q value at a time without resolution; this exercises the shared Parameters across Q and +# the convolution through Experiment, Analysis and the EasyScience fitter. + +import numpy as np +import pytest +import scipp as sc + +import easydynamics as edyn + +Q = sc.array(dims=['Q'], values=[0.5, 1.0, 1.5], unit='1/angstrom') +ENERGY = sc.linspace('energy', -2.0, 2.0, 401, unit='meV') +SHAPE = (Q.shape[0], ENERGY.shape[0]) +NOISE = 0.003 + +# (model class, fixed arguments, true values of the fitted parameters, start values) +CASES = { + 'HallRossJumpDiffusion': ( + edyn.HallRossJumpDiffusion, + {}, + {'scale': 1.0, 'residence_time': 2.0, 'jump_length': 1.5}, + {'scale': 0.8, 'residence_time': 1.5, 'jump_length': 1.2}, + ), + 'EquivalentSitesCircle': ( + edyn.EquivalentSitesCircle, + {'radius': 1.2, 'number_of_sites': 3}, + {'scale': 1.0, 'residence_time': 2.0}, + {'scale': 0.8, 'residence_time': 1.5}, + ), + 'GaussianModel3D': ( + edyn.GaussianModel3D, + {}, + {'scale': 1.0, 'diffusion_coefficient': 2e-9, 'mean_u_squared': 1.5}, + {'scale': 0.8, 'diffusion_coefficient': 1.5e-9, 'mean_u_squared': 1.2}, + ), + 'JumpSitesLogNormalDistribution': ( + edyn.JumpSitesLogNormalDistribution, + {'radius': 1.2, 'number_of_sites': 3}, + {'scale': 1.0, 'residence_time': 2.0, 'sigma': 0.5}, + {'scale': 0.8, 'residence_time': 1.5, 'sigma': 0.7}, + ), +} + + +def as_data(values: np.ndarray, variances: np.ndarray) -> sc.DataArray: + return sc.DataArray( + sc.array(dims=['Q', 'energy'], values=values, variances=variances), + coords={'Q': Q, 'energy': ENERGY}, + ) + + +def make_analysis(experiment: edyn.Experiment, model: object) -> edyn.Analysis: + resolution = edyn.ResolutionModel(components=edyn.Gaussian(width=0.05, area=1.0)) + resolution.fix_all_parameters() + analysis = edyn.Analysis( + experiment=experiment, + sample_model=edyn.SampleModel(diffusion_models=model), + instrument_model=edyn.InstrumentModel(resolution_model=resolution), + ) + analysis.fix_energy_offset() + return analysis + + +@pytest.mark.parametrize('case', CASES.values(), ids=CASES.keys()) +def test_simultaneous_fit_recovers_parameters(case): + # WHEN + model_class, fixed, true_values, start_values = case + true_model = model_class(**fixed, **true_values) + experiment = edyn.Experiment('synthetic', data=as_data(np.ones(SHAPE), np.ones(SHAPE))) + clean = np.array(make_analysis(experiment, true_model).calculate()) + rng = np.random.default_rng(seed=1) + experiment.data = as_data(clean + rng.normal(0.0, NOISE, SHAPE), np.full(SHAPE, NOISE**2)) + fit_model = model_class(**fixed, **start_values) + analysis = make_analysis(experiment, fit_model) + + # THEN + results = analysis.fit(fit_method='simultaneous') + + # EXPECT + assert all(result.success for result in results) + residuals = (experiment.data.values - np.array(analysis.calculate())) / NOISE + reduced_chi2 = float(np.sum(residuals**2) / residuals.size) + assert 0.85 < reduced_chi2 < 1.15 + for name in true_values: + fitted = getattr(fit_model, name) + true_value = getattr(true_model, name).value + assert fitted.variance > 0 + assert abs(fitted.value - true_value) < 5 * fitted.error, name diff --git a/tests/unit/easydynamics/sample_model/diffusion_model/test_delta_lorentzian_diffusion_model_base.py b/tests/unit/easydynamics/sample_model/diffusion_model/test_delta_lorentzian_diffusion_model_base.py new file mode 100644 index 00000000..c5ec5d93 --- /dev/null +++ b/tests/unit/easydynamics/sample_model/diffusion_model/test_delta_lorentzian_diffusion_model_base.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import pytest +from easyscience.variable import DescriptorNumber +from easyscience.variable import Parameter + +from easydynamics.sample_model.diffusion_model.delta_lorentzian_diffusion_model_base import ( + DeltaLorentzianDiffusionModelBase, +) +from easydynamics.utils.utils import Q_type + +# Weights of orders 0 (delta function) to 3; they sum to 1. +WEIGHTS = [0.5, 0.3, 0.15, 0.05] + + +class DummyDiffusionModel(DeltaLorentzianDiffusionModelBase): + """Test-only model with constant weights and widths order * w.""" + + _order_symbol = 'k' + + def __init__(self, number_of_lorentzians: int = 3, **kwargs) -> None: + kwargs.setdefault('name', 'Dummy') + kwargs.setdefault('display_name', 'Dummy') + super().__init__(number_of_lorentzians, **kwargs) + self.w = Parameter(name='w', value=0.1, unit='meV') + self.create_component_collections() + + def calculate_width(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + self._verify_order(order) + return order * self.w.value * np.ones_like(self._ensure_Q(Q)) + + def calculate_EISF(self, Q: Q_type | None = None) -> np.ndarray: + return WEIGHTS[0] * np.ones_like(self._ensure_Q(Q)) + + def calculate_QISF_order(self, Q: Q_type | None = None, order: int = 1) -> np.ndarray: + self._verify_order(order) + return WEIGHTS[order] * np.ones_like(self._ensure_Q(Q)) + + def _create_weight(self, Q: float, order: int) -> Parameter: + return Parameter(name=f'weight_k{order}_Q{Q:.2f}', value=WEIGHTS[order], fixed=True) + + def _write_width_dependency_expression(self, order: int) -> str: + return f'{order} * w' + + def _write_width_dependency_map_expression(self) -> dict[str, DescriptorNumber]: + return {'w': self.w} + + +@pytest.fixture +def dummy(): + return DummyDiffusionModel(scale=2.0, Q=np.array([0.5, 1.0])) + + +# ---------------------------------------------------------------------- +# Components +# ---------------------------------------------------------------------- + + +def test_component_names_use_symbol(dummy): + # WHEN + components = dummy.get_component_collections()[0][:] + + # THEN EXPECT + assert dummy.lorentzian_component_name(2) == 'Dummy_k2' + assert [c.name for c in components] == ['Delta function', 'Dummy_k1', 'Dummy_k2', 'Dummy_k3'] + assert components[1].display_name == 'Dummy_k1' + + +def test_components_have_weights_and_widths(dummy): + # WHEN THEN EXPECT + for collection in dummy.get_component_collections(): + delta, *lorentzians = collection[:] + assert delta.area.value == pytest.approx(1.0) + assert [c.area.value for c in lorentzians] == pytest.approx([0.6, 0.3, 0.1]) + assert [c.width.value for c in lorentzians] == pytest.approx([0.1, 0.2, 0.3]) + + +def test_areas_follow_scale(dummy): + # WHEN + dummy.scale.value = 4.0 + + # THEN EXPECT + assert dummy.get_component_collections()[0][1].area.value == pytest.approx(1.2) + + +def test_setting_Q_builds_collections(): + # WHEN + dummy = DummyDiffusionModel() + + # THEN EXPECT + assert dummy.get_component_collections() == [] + dummy.Q = [0.5] + assert len(dummy.get_component_collections()[0][:]) == 4 + + +# ---------------------------------------------------------------------- +# Orders and validation +# ---------------------------------------------------------------------- + + +@pytest.mark.parametrize('order', [0, 4, 1.5, True]) +def test_order_out_of_range_raises(dummy, order): + # WHEN THEN EXPECT + with pytest.raises((TypeError, ValueError), match='order must be'): + dummy.lorentzian_component_name(order) + + +@pytest.mark.parametrize('n, error', [(0, ValueError), (1.0, TypeError), (True, TypeError)]) +def test_number_of_lorentzians_validated(n, error): + # WHEN THEN EXPECT + with pytest.raises(error, match='number_of_lorentzians must be'): + DummyDiffusionModel(number_of_lorentzians=n) + + +def test_number_of_lorentzians_read_only(dummy): + # WHEN THEN EXPECT + assert dummy.number_of_lorentzians == 3 + with pytest.raises(AttributeError): + dummy.number_of_lorentzians = 2 + + +@pytest.mark.parametrize('kwargs', [{'delta_name': 1}, {'delta_display_name': 1}]) +def test_delta_names_validated(kwargs): + # WHEN THEN EXPECT + with pytest.raises(TypeError): + DummyDiffusionModel(**kwargs) + + +@pytest.mark.parametrize('display_name', ['Elastic line', None]) +def test_delta_name_setters(dummy, display_name): + # WHEN + dummy.delta_name = 'Elastic' + dummy.delta_display_name = display_name + + # THEN EXPECT + assert dummy.delta_name == 'Elastic' + assert dummy.delta_display_name == display_name + + +@pytest.mark.parametrize( + 'attribute, message', + [ + ('delta_name', 'delta_name must be a string'), + ('delta_display_name', 'delta_display_name must be a string or None'), + ], +) +def test_delta_name_setters_reject_non_strings(dummy, attribute, message): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match=message): + setattr(dummy, attribute, 1) + + +def test_incomplete_subclass_cannot_be_created(): + # WHEN + class NoEISF(DummyDiffusionModel): + calculate_EISF = DeltaLorentzianDiffusionModelBase.calculate_EISF + + # THEN EXPECT + with pytest.raises(TypeError): + NoEISF() + + +# ---------------------------------------------------------------------- +# Fit targets +# ---------------------------------------------------------------------- + + +def test_fit_target_names(dummy): + # WHEN THEN EXPECT + assert [t.name for t in dummy.get_fit_targets()] == [ + 'elastic_area', + 'area_1', + 'width_1', + 'area_2', + 'width_2', + 'area_3', + 'width_3', + ] + + +def test_fit_target_values(dummy): + # WHEN + targets = {t.name: t for t in dummy.get_fit_targets()} + Q = np.array([0.5]) + + # THEN EXPECT + np.testing.assert_allclose(targets['elastic_area'].function(Q), [1.0]) + np.testing.assert_allclose(targets['area_2'].function(Q), [0.3]) + np.testing.assert_allclose(targets['width_3'].function(Q), [0.3]) + + +def test_fit_target_keys_follow_names(dummy): + # WHEN + dummy.delta_name = 'Elastic' + targets = dummy.get_fit_targets() + + # THEN EXPECT + assert targets[0].dataset_key == 'Elastic area' + assert targets[1].dataset_key == 'Dummy_k1 area' diff --git a/tests/unit/easydynamics/sample_model/diffusion_model/test_equivalent_sites_circle.py b/tests/unit/easydynamics/sample_model/diffusion_model/test_equivalent_sites_circle.py new file mode 100644 index 00000000..00e2a2d6 --- /dev/null +++ b/tests/unit/easydynamics/sample_model/diffusion_model/test_equivalent_sites_circle.py @@ -0,0 +1,404 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import pytest +import scipp as sc +from scipp.constants import hbar as scipp_hbar + +from easydynamics.sample_model.components import DeltaFunction +from easydynamics.sample_model.components import Lorentzian +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import ( + MAXIMUM_NUMBER_OF_SITES, +) +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import MINIMUM_RADIUS +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import MINIMUM_WIDTH +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import EquivalentSitesCircle + +# A weight is a sum of N terms of size up to 1/N that can nearly cancel, so its rounding error is +# absolute, a few N * machine epsilon, not relative to the (possibly tiny) result. +WEIGHT_ATOL = 1e-14 + + +def hbar_over(residence_time_ps: float) -> float: + """hbar / tau in meV, computed with Scipp's constants.""" + return (scipp_hbar / sc.scalar(residence_time_ps, unit='ps')).to(unit='meV').value + + +class TestEquivalentSitesCircle: + @pytest.fixture + def model(self): + return EquivalentSitesCircle() + + def test_create_weight_requires_float_Q(self, model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match='Q must be a float'): + model._create_weight(1, 1) + + # ------------------------------------------------------------------ + # Construction and parameters + # ------------------------------------------------------------------ + + def test_init_default(self, model): + # WHEN THEN EXPECT + assert model.display_name == 'EquivalentSitesCircle' + assert model.x_unit == 'meV' + assert model.y_unit == 'dimensionless' + assert model.scale.value == pytest.approx(1.0) + assert model.number_of_sites == 3 + assert model.radius.value == pytest.approx(1.0) + assert model.radius.unit == 'Å' + assert model.radius.fixed is True + assert model.radius.min == MINIMUM_RADIUS + assert model.residence_time.value == pytest.approx(1.0) + assert model.residence_time.unit == 'ps' + assert model.residence_time.fixed is False + assert model.residence_time.min == MINIMUM_WIDTH + + @pytest.mark.parametrize( + 'kwargs, expected_exception, expected_message', + [ + ({'radius': 'invalid'}, TypeError, 'radius must be a number'), + ({'radius': 0.0}, ValueError, 'radius must be at least'), + ({'residence_time': 'invalid'}, TypeError, 'residence_time must be a number'), + ({'residence_time': 0.0}, ValueError, 'residence_time must be at least'), + ({'number_of_sites': 2.5}, TypeError, 'number_of_sites must be an integer'), + ({'number_of_sites': True}, TypeError, 'number_of_sites must be an integer'), + ({'number_of_sites': 1}, ValueError, 'number_of_sites must be between 2 and'), + ( + {'number_of_sites': MAXIMUM_NUMBER_OF_SITES + 1}, + ValueError, + 'number_of_sites must be between 2 and', + ), + ({'delta_name': 1}, TypeError, 'delta_name must be a string'), + ], + ids=[ + 'radius_type', + 'radius_zero', + 'residence_time_type', + 'residence_time_zero', + 'sites_float', + 'sites_bool', + 'sites_too_few', + 'sites_too_many', + 'delta_name_type', + ], + ) + def test_init_rejects_invalid_input(self, kwargs, expected_exception, expected_message): + with pytest.raises(expected_exception, match=expected_message): + EquivalentSitesCircle(**kwargs) + + def test_number_of_sites_accepts_numpy_integer(self): + # WHEN THEN EXPECT + assert EquivalentSitesCircle(number_of_sites=np.int64(4)).number_of_sites == 4 + + def test_number_of_sites_is_read_only(self, model): + with pytest.raises(AttributeError): + model.number_of_sites = 4 + + def test_radius_setter(self, model): + # WHEN + model.radius = 2.5 + + # THEN EXPECT + assert model.radius.value == pytest.approx(2.5) + + @pytest.mark.parametrize( + 'value, expected_exception', + [('invalid', TypeError), (0.0, ValueError)], + ids=['type', 'zero'], + ) + def test_radius_setter_raises(self, model, value, expected_exception): + with pytest.raises(expected_exception, match='radius must be'): + model.radius = value + + def test_residence_time_setter(self, model): + # WHEN + model.residence_time = 3.0 + + # THEN EXPECT + assert model.residence_time.value == pytest.approx(3.0) + + @pytest.mark.parametrize( + 'value, expected_exception', + [('invalid', TypeError), (0.0, ValueError), (MINIMUM_WIDTH / 10, ValueError)], + ids=['type', 'zero', 'below_minimum'], + ) + def test_residence_time_setter_raises(self, model, value, expected_exception): + with pytest.raises(expected_exception, match='residence_time must be'): + model.residence_time = value + + # ------------------------------------------------------------------ + # Physics of the structure factors and widths + # ------------------------------------------------------------------ + + @pytest.mark.parametrize('number_of_sites', [2, 3, 6, 11]) + def test_weights_sum_to_one(self, number_of_sites): + """EISF + sum of all quasi-elastic weights is exactly 1 (no truncation).""" + # WHEN + model = EquivalentSitesCircle(number_of_sites=number_of_sites, radius=1.3) + Q = np.array([1e-3, 0.3, 1.0, 2.5, 7.0]) + + # THEN + total = model.calculate_EISF(Q) + model.calculate_QISF(Q) + + # EXPECT + np.testing.assert_allclose(total, 1.0, rtol=0, atol=1e-12) + + def test_QISF_is_sum_of_orders(self): + # WHEN + model = EquivalentSitesCircle(number_of_sites=5, radius=1.3) + Q = np.array([0.5, 1.5]) + + # THEN + orders = [model.calculate_QISF_order(Q, order) for order in range(1, 5)] + + # EXPECT + np.testing.assert_allclose( + model.calculate_QISF(Q), np.sum(orders, axis=0), rtol=0, atol=WEIGHT_ATOL + ) + + def test_two_sites_is_the_two_site_jump_model(self): + """N = 2: EISF = (1 + j0(2QR)) / 2 and Gamma = 2 hbar / tau (Bee).""" + # WHEN + radius = 1.2 + model = EquivalentSitesCircle(number_of_sites=2, radius=radius, residence_time=1.5) + Q = np.array([0.3, 1.0, 2.0]) + + # THEN + EISF = model.calculate_EISF(Q) + width = model.calculate_width(Q, order=1) + + # EXPECT + jump_distance = 2 * radius + expected_EISF = (1 + np.sin(Q * jump_distance) / (Q * jump_distance)) / 2 + np.testing.assert_allclose(EISF, expected_EISF, rtol=0, atol=WEIGHT_ATOL) + np.testing.assert_allclose( + model.calculate_QISF_order(Q, 1), 1 - expected_EISF, rtol=0, atol=WEIGHT_ATOL + ) + np.testing.assert_allclose(width, 2 * hbar_over(1.5), rtol=1e-12) + + def test_EISF_tends_to_one_at_small_Q(self): + # WHEN + model = EquivalentSitesCircle(number_of_sites=4, radius=2.0) + + # THEN EXPECT + assert model.calculate_EISF(np.array([1e-8]))[0] == pytest.approx(1.0, rel=1e-12) + + def test_orders_l_and_N_minus_l_are_equal(self): + # WHEN + model = EquivalentSitesCircle(number_of_sites=6, radius=1.1, residence_time=2.0) + Q = np.array([0.4, 1.7]) + + # THEN EXPECT + for order in range(1, 6): + np.testing.assert_allclose( + model.calculate_QISF_order(Q, order), + model.calculate_QISF_order(Q, 6 - order), + rtol=0, + atol=WEIGHT_ATOL, + ) + np.testing.assert_allclose( + model.calculate_width(Q, order), model.calculate_width(Q, 6 - order), rtol=1e-12 + ) + + def test_width_follows_sine_squared(self): + """Gamma_l = (2 hbar / tau) sin^2(pi l / N), the residence-time convention.""" + # WHEN + model = EquivalentSitesCircle(number_of_sites=5, residence_time=2.0) + Q = np.array([0.5, 1.5]) + + # THEN EXPECT + for order in range(1, 5): + expected = 2 * hbar_over(2.0) * np.sin(np.pi * order / 5) ** 2 + np.testing.assert_allclose(model.calculate_width(Q, order), expected, rtol=1e-12) + + def test_matches_QENSmodels_reference_values(self): + """ + Doctest values of QENSmodels hwhmEquivalentSitesCircle([1., 2.], 6, 0.5, 1.5), rounded + there to the digits given. Its widths are in 1/ps, so they are multiplied by hbar here. + """ + # WHEN + model = EquivalentSitesCircle(number_of_sites=6, radius=0.5, residence_time=1.5) + Q = np.array([1.0, 2.0]) + + # THEN + EISF = model.calculate_EISF(Q) + + # EXPECT + np.testing.assert_allclose(EISF, [0.92, 0.713], atol=5e-4) + assert model.calculate_QISF_order(Q, 2)[0] == pytest.approx(0.000503, abs=5e-7) + assert model.calculate_QISF_order(Q, 5)[1] == pytest.approx(0.13616, abs=5e-6) + hbar_per_ps = hbar_over(1.0) + assert model.calculate_width(Q, 1)[0] / hbar_per_ps == pytest.approx(0.333, abs=5e-4) + assert model.calculate_width(Q, 3)[0] / hbar_per_ps == pytest.approx(1.333, abs=5e-4) + + def test_weights_follow_radius_unit(self): + # WHEN + model = EquivalentSitesCircle(number_of_sites=4, radius=1.5) + Q = np.array([0.8, 1.6]) + EISF_angstrom = model.calculate_EISF(Q) + + # THEN + model.radius.convert_unit('nm') + + # EXPECT the same physical radius gives the same weights + np.testing.assert_allclose( + model.calculate_EISF(Q), EISF_angstrom, rtol=0, atol=WEIGHT_ATOL + ) + + @pytest.mark.parametrize('order', [0, 3, 1.0, True], ids=['zero', 'N', 'float', 'bool']) + def test_order_validation(self, model, order): + with pytest.raises((TypeError, ValueError), match='order must be'): + model.calculate_QISF_order(np.array([1.0]), order) + with pytest.raises((TypeError, ValueError), match='order must be'): + model.calculate_width(np.array([1.0]), order) + + def test_calculate_EISF_type_error(self, model): + with pytest.raises(TypeError, match='Q must be '): + model.calculate_EISF(Q='invalid') + + # ------------------------------------------------------------------ + # Components + # ------------------------------------------------------------------ + + def test_no_Q_gives_no_collections(self, model): + # WHEN THEN EXPECT + assert model.get_component_collections() == [] + + def test_component_names_follow_the_base_convention(self): + # WHEN + model = EquivalentSitesCircle(number_of_sites=3, Q=np.array([1.0])) + + # THEN EXPECT + assert model.number_of_lorentzians == 2 + assert model.lorentzian_component_name(2) == 'EquivalentSitesCircle_l2' + + def test_components_per_Q(self): + # WHEN + model = EquivalentSitesCircle(number_of_sites=4, Q=np.array([0.5, 1.0])) + + # THEN + collections = model.get_component_collections() + + # EXPECT one delta function and N - 1 Lorentzians per Q + assert len(collections) == 2 + for collection in collections: + assert len(collection) == 4 + assert isinstance(collection[0], DeltaFunction) + assert all(isinstance(component, Lorentzian) for component in collection[1:]) + assert [component.name for component in collection[1:]] == [ + model.lorentzian_component_name(order) for order in range(1, 4) + ] + + def test_component_areas_and_widths_match_calculations(self): + """The dependency strings must give what the Scipp calculations give.""" + # WHEN + Q = np.array([0.4, 1.1, 2.3]) + model = EquivalentSitesCircle( + scale=2.5, radius=1.4, residence_time=0.8, number_of_sites=5, Q=Q + ) + + # THEN + collections = model.get_component_collections() + + # EXPECT + EISF = model.calculate_EISF(Q) + for i, collection in enumerate(collections): + assert collection[0].area.value == pytest.approx(2.5 * EISF[i], rel=0, abs=WEIGHT_ATOL) + for order in range(1, 5): + lorentzian = collection[order] + weight = model.calculate_QISF_order(Q, order)[i] + width = model.calculate_width(Q, order)[i] + assert lorentzian.area.value == pytest.approx(2.5 * weight, rel=0, abs=WEIGHT_ATOL) + assert lorentzian.width.value == pytest.approx(width, rel=1e-12) + assert lorentzian.width.unit == 'meV' + assert lorentzian.area.independent is False + + def test_components_follow_parameter_changes(self): + # WHEN + Q = np.array([0.6, 1.8]) + model = EquivalentSitesCircle(number_of_sites=3, Q=Q) + + # THEN + model.radius = 2.2 + model.residence_time = 3.5 + model.scale = 4.0 + + # EXPECT + collections = model.get_component_collections() + EISF = model.calculate_EISF(Q) + for i, collection in enumerate(collections): + assert collection[0].area.value == pytest.approx(4.0 * EISF[i], rel=0, abs=WEIGHT_ATOL) + assert collection[1].area.value == pytest.approx( + 4.0 * model.calculate_QISF_order(Q, 1)[i], rel=0, abs=WEIGHT_ATOL + ) + assert collection[1].width.value == pytest.approx( + model.calculate_width(Q, 1)[i], rel=1e-12 + ) + + def test_component_weights_follow_radius_unit(self): + # WHEN + Q = np.array([0.8, 1.6]) + model = EquivalentSitesCircle(number_of_sites=4, radius=1.5, Q=Q) + areas_angstrom = [c[1].area.value for c in model.get_component_collections()] + + # THEN + model.radius.convert_unit('nm') + + # EXPECT + areas_nm = [c[1].area.value for c in model.get_component_collections()] + np.testing.assert_allclose(areas_nm, areas_angstrom, rtol=0, atol=WEIGHT_ATOL) + + def test_create_component_collections_installs_collections(self): + # WHEN + model = EquivalentSitesCircle(Q=np.array([1.0, 2.0])) + + # THEN + collections = model.create_component_collections() + + # EXPECT + assert collections is model.get_component_collections() + + # ------------------------------------------------------------------ + # Fit targets, variables and representation + # ------------------------------------------------------------------ + + def test_fit_targets(self): + # WHEN + Q = np.array([0.5, 1.5]) + model = EquivalentSitesCircle(scale=2.0, number_of_sites=3, Q=Q) + + # THEN + targets = {target.name: target for target in model.get_fit_targets()} + + # EXPECT + assert set(targets) == {'elastic_area', 'area_1', 'width_1', 'area_2', 'width_2'} + np.testing.assert_allclose( + targets['elastic_area'].function(Q), 2.0 * model.calculate_EISF(Q), rtol=1e-12 + ) + np.testing.assert_allclose( + targets['area_2'].function(Q), 2.0 * model.calculate_QISF_order(Q, 2), rtol=1e-12 + ) + np.testing.assert_allclose( + targets['width_1'].function(Q), model.calculate_width(Q, 1), rtol=1e-12 + ) + + def test_get_global_variables(self, model): + # WHEN THEN EXPECT + assert model.get_global_variables() == [ + model.scale, + model.radius, + model.residence_time, + ] + + def test_repr(self, model): + # WHEN THEN + repr_str = repr(model) + + # EXPECT + assert 'EquivalentSitesCircle' in repr_str + assert 'number_of_sites=3' in repr_str + assert 'radius' in repr_str + assert 'residence_time' in repr_str + assert 'x_unit=meV, y_unit=dimensionless' in repr_str diff --git a/tests/unit/easydynamics/sample_model/diffusion_model/test_gaussian_model_3d.py b/tests/unit/easydynamics/sample_model/diffusion_model/test_gaussian_model_3d.py new file mode 100644 index 00000000..88e86743 --- /dev/null +++ b/tests/unit/easydynamics/sample_model/diffusion_model/test_gaussian_model_3d.py @@ -0,0 +1,452 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import warnings + +import numpy as np +import pytest +import scipp as sc +from scipp.constants import hbar as scipp_hbar +from scipy.stats import poisson + +from easydynamics.sample_model.components import DeltaFunction +from easydynamics.sample_model.components import Lorentzian +from easydynamics.sample_model.diffusion_model.delta_lorentz import DeltaLorentz +from easydynamics.sample_model.diffusion_model.gaussian_model_3d import ( + MAXIMUM_NUMBER_OF_LORENTZIANS, +) +from easydynamics.sample_model.diffusion_model.gaussian_model_3d import ( + MINIMUM_DIFFUSION_COEFFICIENT, +) +from easydynamics.sample_model.diffusion_model.gaussian_model_3d import MINIMUM_MEAN_U_SQUARED +from easydynamics.sample_model.diffusion_model.gaussian_model_3d import GaussianModel3D + +# Volino's example below in this model's units: D = 0.5 angstrom^2/ps = 5e-9 m^2/s, and +# = 3 = 4.5 angstrom^2. +VOLINO_D = 5e-9 +VOLINO_MEAN_U_SQUARED = 4.5 + + +def hbar_rate(rate_per_ps: float) -> float: + """hbar * rate in meV, for a rate in 1/ps, computed with Scipp's constants.""" + return (scipp_hbar * sc.scalar(rate_per_ps, unit='1/ps')).to(unit='meV').value + + +def model_without_warning(**kwargs) -> GaussianModel3D: + """Create a model and fail if it warns about truncation.""" + with warnings.catch_warnings(): + warnings.simplefilter('error') + return GaussianModel3D(**kwargs) + + +class TestGaussianModel3D: + @pytest.fixture + def model(self): + return GaussianModel3D() + + def test_create_weight_requires_float_Q(self, model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match='Q must be a float'): + model._create_weight(1, 1) + + # ------------------------------------------------------------------ + # Construction and parameters + # ------------------------------------------------------------------ + + def test_init_default(self, model): + # WHEN THEN EXPECT + assert model.display_name == 'GaussianModel3D' + assert model.x_unit == 'meV' + assert model.y_unit == 'dimensionless' + assert model.scale.value == pytest.approx(1.0) + assert model.number_of_lorentzians == 20 + assert model.diffusion_coefficient.value == pytest.approx(1.0) + assert model.diffusion_coefficient.unit == 'm^2/s' + assert model.diffusion_coefficient.min == MINIMUM_DIFFUSION_COEFFICIENT + assert model.mean_u_squared.value == pytest.approx(1.0) + assert model.mean_u_squared.unit == 'Å^2' + assert model.mean_u_squared.min == MINIMUM_MEAN_U_SQUARED + + @pytest.mark.parametrize( + 'kwargs, expected_exception, expected_message', + [ + ({'diffusion_coefficient': 'x'}, TypeError, 'diffusion_coefficient must be a number'), + ({'diffusion_coefficient': 0.0}, ValueError, 'diffusion_coefficient must be at least'), + ({'mean_u_squared': 'x'}, TypeError, 'mean_u_squared must be a number'), + ({'mean_u_squared': 0.0}, ValueError, 'mean_u_squared must be at least'), + ({'number_of_lorentzians': 2.0}, TypeError, 'number_of_lorentzians must be an'), + ({'number_of_lorentzians': True}, TypeError, 'number_of_lorentzians must be an'), + ({'number_of_lorentzians': 0}, ValueError, 'number_of_lorentzians must be between'), + ( + {'number_of_lorentzians': MAXIMUM_NUMBER_OF_LORENTZIANS + 1}, + ValueError, + 'number_of_lorentzians must be between', + ), + ({'delta_name': 1}, TypeError, 'delta_name must be a string'), + ], + ids=[ + 'D_type', + 'D_zero', + 'u2_type', + 'u2_zero', + 'n_float', + 'n_bool', + 'n_zero', + 'n_too_many', + 'delta_name_type', + ], + ) + def test_init_rejects_invalid_input(self, kwargs, expected_exception, expected_message): + with pytest.raises(expected_exception, match=expected_message): + GaussianModel3D(**kwargs) + + def test_number_of_lorentzians_is_read_only(self, model): + with pytest.raises(AttributeError): + model.number_of_lorentzians = 5 + + def test_setters(self, model): + # WHEN + model.diffusion_coefficient = 2e-9 + model.mean_u_squared = 3.0 + + # THEN EXPECT + assert model.diffusion_coefficient.value == pytest.approx(2e-9) + assert model.mean_u_squared.value == pytest.approx(3.0) + + @pytest.mark.parametrize( + 'name, value, expected_exception', + [ + ('diffusion_coefficient', 'x', TypeError), + ('diffusion_coefficient', 0.0, ValueError), + ('mean_u_squared', 'x', TypeError), + ('mean_u_squared', 0.0, ValueError), + ], + ids=['D_type', 'D_zero', 'u2_type', 'u2_zero'], + ) + def test_setters_raise(self, model, name, value, expected_exception): + with pytest.raises(expected_exception, match=f'{name} must be'): + setattr(model, name, value) + + # ------------------------------------------------------------------ + # Physics of the structure factors and widths + # ------------------------------------------------------------------ + + def test_weights_are_poisson(self): + """A_i(Q) = exp(-y) y^i / i! with y = Q^2 / 3.""" + # WHEN + model = GaussianModel3D(mean_u_squared=2.4, number_of_lorentzians=30) + Q = np.array([0.3, 1.0, 2.2]) + y = Q**2 * 2.4 / 3 + + # THEN EXPECT + np.testing.assert_allclose(model.calculate_EISF(Q), poisson.pmf(0, y), rtol=1e-12) + for order in (1, 2, 7, 30): + np.testing.assert_allclose( + model.calculate_QISF_order(Q, order), poisson.pmf(order, y), rtol=1e-12 + ) + + def test_EISF_equals_DeltaLorentz_debye_waller_factor(self): + """The same mean_u_squared gives the same exp(- Q^2 / 3) as DeltaLorentz.""" + # WHEN + Q = np.array([0.5, 1.0, 1.5]) + model = GaussianModel3D(mean_u_squared=1.7) + delta_lorentz = DeltaLorentz(mean_u_squared=1.7, A_0=1.0) + + # THEN EXPECT + np.testing.assert_allclose( + model.calculate_EISF(Q), delta_lorentz.calculate_EISF(Q), rtol=1e-12 + ) + + def test_weights_and_missing_weight_sum_to_one(self): + # WHEN + model = GaussianModel3D(mean_u_squared=3.0, number_of_lorentzians=5) + Q = np.array([0.5, 1.5, 3.0]) + + # THEN + total = ( + model.calculate_EISF(Q) + model.calculate_QISF(Q) + model.calculate_missing_weight(Q) + ) + + # EXPECT + np.testing.assert_allclose(total, 1.0, rtol=0, atol=1e-14) + np.testing.assert_allclose( + model.calculate_missing_weight(Q), poisson.sf(5, Q**2), rtol=1e-9 + ) + + def test_width_is_order_times_3_hbar_D_over_u2(self): + # WHEN + model = GaussianModel3D(diffusion_coefficient=2e-9, mean_u_squared=1.2) + Q = np.array([0.5, 2.0]) + + # THEN EXPECT Gamma_i = 3 i hbar D / , with D / = 2e-9 m^2/s / 1.2 angstrom^2 + rate_per_ps = 3 * 2e-9 * 1e20 / 1e12 / 1.2 + for order in (1, 4, 20): + np.testing.assert_allclose( + model.calculate_width(Q, order), order * hbar_rate(rate_per_ps), rtol=1e-12 + ) + + def test_matches_QENSmodels_reference_values(self): + """ + Doctest values of QENSmodels hwhmGaussianModel3D([1., 2.], 0.5, 1.5), with D in + angstrom^2/ps and in angstrom^2, rounded there to the digits given. + """ + # WHEN + model = GaussianModel3D( + diffusion_coefficient=VOLINO_D, + mean_u_squared=VOLINO_MEAN_U_SQUARED, + number_of_lorentzians=99, + ) + Q = np.array([1.0, 2.0]) + + # THEN EXPECT + np.testing.assert_allclose(model.calculate_EISF(Q), [0.223, 0.002], atol=5e-4) + np.testing.assert_allclose(model.calculate_QISF_order(Q, 1), [0.3347, 0.0149], atol=5e-5) + width_10 = model.calculate_width(Q, 10) / hbar_rate(1.0) + width_99 = model.calculate_width(Q, 99) / hbar_rate(1.0) + np.testing.assert_allclose(width_10, 3.333, atol=5e-4) + np.testing.assert_allclose(width_99, 33.0, atol=5e-2) + + def test_Q_zero_is_purely_elastic(self): + # WHEN + model = GaussianModel3D() + + # THEN + Q = np.array([0.0]) + + # EXPECT exactly: exp(0) = 1 and exp(-inf) = 0, with no NaN from 0 * log(0) + np.testing.assert_array_equal(model.calculate_EISF(Q), [1.0]) + np.testing.assert_array_equal(model.calculate_QISF(Q), [0.0]) + np.testing.assert_array_equal(model.calculate_QISF_order(Q, 1), [0.0]) + + def test_large_Q_gives_no_nan(self): + """y = Q^2 / 3 up to about 1000: the weights underflow to 0, never NaN.""" + # WHEN + model = GaussianModel3D(mean_u_squared=30.0, number_of_lorentzians=100) + Q = np.array([5.0, 10.0]) + + # THEN + weights = [model.calculate_QISF_order(Q, order) for order in (1, 50, 100)] + + # EXPECT + assert np.all(np.isfinite(weights)) + assert np.all(np.isfinite(model.calculate_EISF(Q))) + + def test_weights_follow_mean_u_squared_unit(self): + # WHEN + model = GaussianModel3D(mean_u_squared=1.5) + Q = np.array([0.8, 1.6]) + EISF_angstrom = model.calculate_EISF(Q) + width_angstrom = model.calculate_width(Q, 1) + + # THEN + model.mean_u_squared.convert_unit('nm**2') + + # EXPECT + np.testing.assert_allclose(model.calculate_EISF(Q), EISF_angstrom, rtol=1e-12) + np.testing.assert_allclose(model.calculate_width(Q, 1), width_angstrom, rtol=1e-12) + + @pytest.mark.parametrize( + 'order', [0, 21, 1.0, True], ids=['zero', 'n_plus_1', 'float', 'bool'] + ) + def test_order_validation(self, model, order): + with pytest.raises((TypeError, ValueError), match='order must be'): + model.calculate_QISF_order(np.array([1.0]), order) + with pytest.raises((TypeError, ValueError), match='order must be'): + model.calculate_width(np.array([1.0]), order) + + def test_calculate_EISF_type_error(self, model): + with pytest.raises(TypeError, match='Q must be '): + model.calculate_EISF(Q='invalid') + + # ------------------------------------------------------------------ + # Truncation warning + # ------------------------------------------------------------------ + + def test_warns_when_truncation_drops_weight(self): + # WHEN THEN EXPECT y = 3 at Q = 3 needs more than 3 Lorentzians + with pytest.warns(UserWarning, match='number_of_lorentzians'): + GaussianModel3D(mean_u_squared=1.0, number_of_lorentzians=3, Q=np.array([3.0])) + + def test_no_warning_when_truncation_is_negligible(self): + # WHEN THEN EXPECT y = 0.03 at Q = 0.3: 3 Lorentzians leave about 3e-8 + model_without_warning(mean_u_squared=1.0, number_of_lorentzians=3, Q=np.array([0.3])) + + def test_warns_again_when_Q_changes(self): + """Guards the refactor: new Q values rebuild the components and check the truncation.""" + # WHEN + model = model_without_warning( + mean_u_squared=1.0, number_of_lorentzians=3, Q=np.array([0.3]) + ) + model.clear_Q(confirm=True) + + # THEN EXPECT y = 25/3 at Q = 5: 3 Lorentzians leave about 0.97 of the weight + with pytest.warns(UserWarning, match='number_of_lorentzians'): + model.Q = np.array([5.0]) + + def test_component_names_follow_the_base_convention(self): + # WHEN + model = model_without_warning(number_of_lorentzians=3, Q=np.array([0.5])) + + # THEN EXPECT + assert model.lorentzian_component_name(3) == 'GaussianModel3D_i3' + + # ------------------------------------------------------------------ + # Components + # ------------------------------------------------------------------ + + def test_no_Q_gives_no_collections(self, model): + # WHEN THEN EXPECT + assert model.get_component_collections() == [] + + def test_components_per_Q(self): + # WHEN + model = model_without_warning(number_of_lorentzians=4, Q=np.array([0.5, 1.0])) + + # THEN + collections = model.get_component_collections() + + # EXPECT one delta function and n Lorentzians per Q + assert len(collections) == 2 + for collection in collections: + assert len(collection) == 5 + assert isinstance(collection[0], DeltaFunction) + assert all(isinstance(component, Lorentzian) for component in collection[1:]) + assert [component.name for component in collection[1:]] == [ + model.lorentzian_component_name(order) for order in range(1, 5) + ] + + def test_component_areas_and_widths_match_calculations(self): + """The dependency strings must give what the Scipp calculations give.""" + # WHEN + Q = np.array([0.4, 1.1, 2.3]) + model = model_without_warning( + scale=2.5, + diffusion_coefficient=3e-9, + mean_u_squared=0.9, + number_of_lorentzians=12, + Q=Q, + ) + + # THEN + collections = model.get_component_collections() + + # EXPECT + EISF = model.calculate_EISF(Q) + for i, collection in enumerate(collections): + assert collection[0].area.value == pytest.approx(2.5 * EISF[i], rel=1e-12) + for order in range(1, 13): + lorentzian = collection[order] + weight = model.calculate_QISF_order(Q, order)[i] + width = model.calculate_width(Q, order)[i] + assert lorentzian.area.value == pytest.approx(2.5 * weight, rel=1e-12, abs=0) + assert lorentzian.width.value == pytest.approx(width, rel=1e-12) + assert lorentzian.width.unit == 'meV' + assert lorentzian.area.independent is False + + def test_components_follow_parameter_changes(self): + # WHEN + Q = np.array([0.6, 1.8]) + model = model_without_warning(number_of_lorentzians=15, mean_u_squared=0.5, Q=Q) + + # THEN + model.mean_u_squared = 0.8 + model.diffusion_coefficient = 4e-9 + model.scale = 4.0 + + # EXPECT + collections = model.get_component_collections() + EISF = model.calculate_EISF(Q) + for i, collection in enumerate(collections): + assert collection[0].area.value == pytest.approx(4.0 * EISF[i], rel=1e-12) + assert collection[3].area.value == pytest.approx( + 4.0 * model.calculate_QISF_order(Q, 3)[i], rel=1e-12 + ) + assert collection[3].width.value == pytest.approx( + model.calculate_width(Q, 3)[i], rel=1e-12 + ) + + def test_parameter_updates_raise_no_runtime_warning(self): + """A fit changes D and on every step; the dependent widths must not warn.""" + # WHEN + model = model_without_warning(number_of_lorentzians=3, Q=np.array([0.5, 1.0])) + + # THEN EXPECT + with warnings.catch_warnings(): + warnings.simplefilter('error', RuntimeWarning) + model.diffusion_coefficient = 3e-9 + model.mean_u_squared = 0.7 + + def test_component_weights_follow_mean_u_squared_unit(self): + # WHEN + Q = np.array([0.8, 1.6]) + model = model_without_warning(mean_u_squared=1.5, Q=Q) + areas_angstrom = [c[2].area.value for c in model.get_component_collections()] + widths_angstrom = [c[2].width.value for c in model.get_component_collections()] + + # THEN + model.mean_u_squared.convert_unit('nm**2') + + # EXPECT + collections = model.get_component_collections() + np.testing.assert_allclose( + [c[2].area.value for c in collections], areas_angstrom, rtol=1e-12 + ) + np.testing.assert_allclose( + [c[2].width.value for c in collections], widths_angstrom, rtol=1e-12 + ) + + def test_create_component_collections_installs_collections(self): + # WHEN + model = model_without_warning(Q=np.array([1.0, 2.0])) + + # THEN + collections = model.create_component_collections() + + # EXPECT + assert collections is model.get_component_collections() + + # ------------------------------------------------------------------ + # Fit targets, variables and representation + # ------------------------------------------------------------------ + + def test_fit_targets(self): + # WHEN + Q = np.array([0.5, 1.5]) + model = model_without_warning(scale=2.0, number_of_lorentzians=6, Q=Q) + + # THEN + targets = {target.name: target for target in model.get_fit_targets()} + + # EXPECT + expected_names = {'elastic_area'} | { + f'{kind}_{order}' for kind in ('area', 'width') for order in range(1, 7) + } + assert set(targets) == expected_names + np.testing.assert_allclose( + targets['elastic_area'].function(Q), 2.0 * model.calculate_EISF(Q), rtol=1e-12 + ) + np.testing.assert_allclose( + targets['area_2'].function(Q), 2.0 * model.calculate_QISF_order(Q, 2), rtol=1e-12 + ) + np.testing.assert_allclose( + targets['width_3'].function(Q), model.calculate_width(Q, 3), rtol=1e-12 + ) + + def test_get_global_variables(self, model): + # WHEN THEN EXPECT + assert model.get_global_variables() == [ + model.scale, + model.diffusion_coefficient, + model.mean_u_squared, + ] + + def test_repr(self, model): + # WHEN THEN + repr_str = repr(model) + + # EXPECT + assert 'GaussianModel3D' in repr_str + assert 'number_of_lorentzians=20' in repr_str + assert 'diffusion_coefficient' in repr_str + assert 'mean_u_squared' in repr_str + assert 'x_unit=meV, y_unit=dimensionless' in repr_str diff --git a/tests/unit/easydynamics/sample_model/diffusion_model/test_hall_ross_jump_diffusion.py b/tests/unit/easydynamics/sample_model/diffusion_model/test_hall_ross_jump_diffusion.py new file mode 100644 index 00000000..d22e7b80 --- /dev/null +++ b/tests/unit/easydynamics/sample_model/diffusion_model/test_hall_ross_jump_diffusion.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import pytest +import scipp as sc +from scipp import UnitError +from scipp.constants import hbar as scipp_hbar + +from easydynamics.sample_model.diffusion_model.hall_ross_jump_diffusion import MINIMUM_WIDTH +from easydynamics.sample_model.diffusion_model.hall_ross_jump_diffusion import ( + HallRossJumpDiffusion, +) + + +class TestHallRossJumpDiffusion: + @pytest.fixture + def hall_ross_model(self): + return HallRossJumpDiffusion() + + def test_init_default(self, hall_ross_model): + # WHEN THEN EXPECT + assert hall_ross_model.display_name == 'HallRossJumpDiffusion' + assert hall_ross_model.x_unit == 'meV' + assert hall_ross_model.y_unit == 'dimensionless' + assert hall_ross_model.scale.value == pytest.approx(1.0) + assert hall_ross_model.scale.unit == 'meV' + assert hall_ross_model.residence_time.value == pytest.approx(1.0) + assert hall_ross_model.jump_length.value == pytest.approx(1.0) + + @pytest.mark.parametrize( + 'kwargs,expected_exception, expected_message', + [ + ( + { + 'x_unit': 123, + 'scale': 1.0, + 'residence_time': 1.0, + 'jump_length': 1.0, + }, + UnitError, + 'Invalid unit', + ), + ( + { + 'y_unit': 123, + 'scale': 1.0, + 'residence_time': 1.0, + 'jump_length': 1.0, + }, + UnitError, + 'unit must be None, a string', + ), + ( + { + 'x_unit': 'meV', + 'scale': 'invalid', + 'residence_time': 1.0, + 'jump_length': 1.0, + }, + TypeError, + 'scale must be a number', + ), + ( + { + 'x_unit': 'meV', + 'scale': 1.0, + 'residence_time': 'invalid', + 'jump_length': 1.0, + }, + TypeError, + 'residence_time must be a number', + ), + ( + { + 'x_unit': 'meV', + 'scale': 1.0, + 'residence_time': -1.0, + 'jump_length': 1.0, + }, + ValueError, + 'residence_time must be at least', + ), + ( + { + 'x_unit': 'meV', + 'scale': 1.0, + 'residence_time': 0.0, + 'jump_length': 1.0, + }, + ValueError, + 'residence_time must be at least', + ), + ( + { + 'x_unit': 'meV', + 'scale': 1.0, + 'residence_time': 1.0, + 'jump_length': 'invalid', + }, + TypeError, + 'jump_length must be a number', + ), + ( + { + 'x_unit': 'meV', + 'scale': 1.0, + 'residence_time': 1.0, + 'jump_length': -1.0, + }, + ValueError, + 'jump_length must be non-negative', + ), + ], + ids=[ + 'invalid_x_unit', + 'invalid_y_unit', + 'invalid_scale_type', + 'invalid_residence_time_type', + 'invalid_residence_time_negative', + 'invalid_residence_time_zero', + 'invalid_jump_length_type', + 'invalid_jump_length_negative', + ], + ) + def test_input_type_validation_raises(self, kwargs, expected_exception, expected_message): + with pytest.raises(expected_exception, match=expected_message): + HallRossJumpDiffusion(display_name='HallRossJumpDiffusion', **kwargs) + + def test_residence_time_setter(self, hall_ross_model): + # WHEN THEN + hall_ross_model.residence_time = 3.0 + + # EXPECT + assert hall_ross_model.residence_time.value == pytest.approx(3.0) + + def test_residence_time_setter_raises(self, hall_ross_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match=r'residence_time must be a number.'): + hall_ross_model.residence_time = 'invalid' + + @pytest.mark.parametrize('invalid_value', [-1.0, 0.0, MINIMUM_WIDTH / 10]) + def test_residence_time_setter_below_minimum_raises(self, hall_ross_model, invalid_value): + # WHEN THEN EXPECT + with pytest.raises(ValueError, match=r'residence_time must be at least'): + hall_ross_model.residence_time = invalid_value + + def test_residence_time_has_minimum_width(self, hall_ross_model): + # WHEN THEN EXPECT + assert hall_ross_model.residence_time.min == MINIMUM_WIDTH + + def test_jump_length_setter(self, hall_ross_model): + # WHEN THEN + hall_ross_model.jump_length = 2.5 + + # EXPECT + assert hall_ross_model.jump_length.value == pytest.approx(2.5) + + def test_jump_length_setter_raises(self, hall_ross_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match=r'jump_length must be a number.'): + hall_ross_model.jump_length = 'invalid' + + def test_jump_length_setter_negative_raises(self, hall_ross_model): + # WHEN THEN EXPECT + with pytest.raises(ValueError, match=r'jump_length must be non-negative.'): + hall_ross_model.jump_length = -1.0 + + def test_calculate_width_type_error(self, hall_ross_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match='Q must be '): + hall_ross_model.calculate_width(Q='invalid') + + def test_calculate_width(self, hall_ross_model): + "Test the calculation against a scipp implementation instead of our Parameters" + # WHEN + hall_ross_model.residence_time = 2.0 + hall_ross_model.jump_length = 1.5 + Q_values = sc.linspace('Q', 0.5, 1.5, num=6, unit='1/angstrom') + # sc.values drops the (zero) variance, which scipp cannot broadcast against Q_values + residence_time_sc = sc.values(hall_ross_model.residence_time.full_value) + jump_length_sc = sc.values(hall_ross_model.jump_length.full_value) + + # THEN + widths = hall_ross_model.calculate_width(Q_values) + + # EXPECT + argument = (Q_values * jump_length_sc).to(unit='dimensionless') + expected_widths = scipp_hbar / residence_time_sc * (1 - sc.exp(-(argument**2) / 2)) + expected_widths = expected_widths.to(unit=hall_ross_model.x_unit) + np.testing.assert_allclose(widths, expected_widths.values, rtol=1e-5) + + def test_calculate_width_small_Q_limit_is_brownian(self, hall_ross_model): + """ + At small Q the width is hbar * D * Q**2 with D = l**2 / (2 tau). Q = 1e-6 also checks + that the width keeps its precision where 1 - exp(-x) would cancel. + """ + # WHEN + hall_ross_model.residence_time = 2.0 + hall_ross_model.jump_length = 1.5 + Q_values = np.array([1e-6, 1e-4]) + + # THEN + widths = hall_ross_model.calculate_width(Q_values) + + # EXPECT + diffusion_coefficient = sc.scalar(1.5**2 / (2 * 2.0), unit='angstrom**2/ps') + Q_sc = sc.array(dims=['Q'], values=Q_values, unit='1/angstrom') + expected_widths = (scipp_hbar * diffusion_coefficient * Q_sc**2).to(unit='meV') + np.testing.assert_allclose(widths, expected_widths.values, rtol=1e-8) + + def test_calculate_width_large_Q_limit(self, hall_ross_model): + # WHEN + hall_ross_model.residence_time = 2.0 + Q_values = np.array([50.0, 100.0]) + + # THEN + widths = hall_ross_model.calculate_width(Q_values) + + # EXPECT width tends to hbar / tau + expected_width = (scipp_hbar / sc.scalar(2.0, unit='ps')).to(unit='meV') + np.testing.assert_allclose(widths, expected_width.value, rtol=1e-12) + + def test_calculate_EISF(self, hall_ross_model): + # WHEN + Q_values = np.array([0.1, 0.2, 0.3]) + + # THEN + EISF = hall_ross_model.calculate_EISF(Q_values) + + # EXPECT + np.testing.assert_array_equal(EISF, np.zeros_like(Q_values)) + + def test_calculate_EISF_type_error(self, hall_ross_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match='Q must be '): + hall_ross_model.calculate_EISF(Q='invalid') + + def test_calculate_QISF(self, hall_ross_model): + # WHEN + Q_values = np.array([0.1, 0.2, 0.3]) + + # THEN + QISF = hall_ross_model.calculate_QISF(Q_values) + + # EXPECT + np.testing.assert_array_equal(QISF, np.ones_like(Q_values)) + + def test_calculate_QISF_type_error(self, hall_ross_model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match='Q must be '): + hall_ross_model.calculate_QISF(Q='invalid') + + @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, hall_ross_model, Q): + # WHEN + hall_ross_model.Q = Q + + # THEN + component_collections = hall_ross_model.create_component_collections() + + # EXPECT + expected_widths = hall_ross_model.calculate_width(Q) + for model_index in range(len(component_collections)): + model = component_collections[model_index] + assert len(model) == 1 + component = model[0] + assert component.width.unit == hall_ross_model.x_unit + assert np.isclose(component.width.value, expected_widths[model_index]) + assert component.width.independent is False + # area.unit = area_unit = x_unit * y_unit + assert component.area.unit == 'meV' + + def test_component_widths_follow_parameters(self): + # WHEN + model = HallRossJumpDiffusion(Q=np.array([0.5, 1.0])) + + # THEN + model.residence_time = 4.0 + model.jump_length = 2.0 + + # EXPECT the dependent widths update with the parameters + widths = [collection[0].width.value for collection in model.get_component_collections()] + np.testing.assert_allclose(widths, model.calculate_width(np.array([0.5, 1.0]))) + + def test_component_width_keeps_precision_at_small_Q(self): + """The dependent widths a fit uses must not lose digits where 1 - exp(-x) cancels.""" + # WHEN + model = HallRossJumpDiffusion(Q=np.array([1e-6]), residence_time=2.0, jump_length=1.5) + + # THEN + width = model.get_component_collections()[0][0].width.value + + # EXPECT hbar * D * Q**2 with D = l**2 / (2 tau) + diffusion_coefficient = sc.scalar(1.5**2 / (2 * 2.0), unit='angstrom**2/ps') + Q_sc = sc.scalar(1e-6, unit='1/angstrom') + expected_width = (scipp_hbar * diffusion_coefficient * Q_sc**2).to(unit='meV') + # abs=0: pytest.approx's default absolute tolerance of 1e-12 exceeds this width (~4e-13) + assert width == pytest.approx(expected_width.value, rel=1e-8, abs=0) + + def test_create_component_collections_installs_collections(self): + # WHEN + model = HallRossJumpDiffusion(Q=np.array([1.0, 2.0])) + + # THEN + collections = model.create_component_collections() + + # EXPECT the returned collections are the installed (live) ones + assert collections is model.get_component_collections() + + def test_write_width_dependency_expression(self, hall_ross_model): + # WHEN THEN + expression = hall_ross_model._write_width_dependency_expression(0.5) + + # EXPECT + assert expression == '(hbar / tau) * -expm1(-((0.5 * l.value) ** 2) / 2)' + + def test_write_width_dependency_map_expression(self, hall_ross_model): + # WHEN THEN + expression_map = hall_ross_model._write_width_dependency_map_expression() + + # EXPECT + expected_map = { + 'tau': hall_ross_model.residence_time, + 'l': hall_ross_model.jump_length, + 'hbar': hall_ross_model._hbar, + } + assert expression_map == expected_map + + def test_write_width_dependency_expression_raises(self, hall_ross_model): + with pytest.raises(TypeError, match='Q must be a float'): + hall_ross_model._write_width_dependency_expression('invalid') + + def test_write_area_dependency_expression_raises(self, hall_ross_model): + with pytest.raises(TypeError, match='QISF must be a float'): + hall_ross_model._write_area_dependency_expression('invalid') + + def test_y_unit_setter_raises(self, hall_ross_model): + # WHEN THEN EXPECT + with pytest.raises(AttributeError, match=r'read-only'): + hall_ross_model.y_unit = '1/meV' + + def test_repr(self, hall_ross_model): + # WHEN THEN + repr_str = repr(hall_ross_model) + + # EXPECT + assert 'HallRossJumpDiffusion' in repr_str + assert 'residence_time' in repr_str + assert 'jump_length' in repr_str + assert 'scale=' in repr_str + assert 'x_unit=meV, y_unit=dimensionless' in repr_str diff --git a/tests/unit/easydynamics/sample_model/diffusion_model/test_jump_sites_log_normal_distribution.py b/tests/unit/easydynamics/sample_model/diffusion_model/test_jump_sites_log_normal_distribution.py new file mode 100644 index 00000000..b7548064 --- /dev/null +++ b/tests/unit/easydynamics/sample_model/diffusion_model/test_jump_sites_log_normal_distribution.py @@ -0,0 +1,398 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import warnings + +import numpy as np +import pytest +import scipp as sc +from scipp.constants import hbar as scipp_hbar + +import easydynamics as edyn +from easydynamics.sample_model.components import DeltaFunction +from easydynamics.sample_model.components import Lorentzian +from easydynamics.sample_model.diffusion_model.equivalent_sites_circle import EquivalentSitesCircle +from easydynamics.sample_model.diffusion_model.jump_sites_log_normal_distribution import ( + MINIMUM_SIGMA, +) +from easydynamics.sample_model.diffusion_model.jump_sites_log_normal_distribution import ( + NUMBER_OF_SAMPLES, +) +from easydynamics.sample_model.diffusion_model.jump_sites_log_normal_distribution import ( + JumpSitesLogNormalDistribution, +) + +# A weight is a sum of N terms of size up to 1/N that can nearly cancel, so its rounding error is +# absolute, a few N * machine epsilon, not relative to the (possibly tiny) result. +WEIGHT_ATOL = 1e-14 + +# Rows of tests/manual/jump_sites_log_norm_dist_ref_data.dat (QENSmodels tests/reference_data): +# energy transfer w in 1/ps and S(q, w), for Nsites = 7, radius = 5, resTime = 2, sigma = 0.6, +# scale = 2, center = 0.8 and q = float32(0.7). +REFERENCE_ROWS = [ + (-1.0, 0.08158287388105), + (-0.5, 0.1336353341603), + (0.5, 0.6730922522463), + (1.0, 0.8950268953231), + (1.5, 0.2983075456912), +] + + +def hbar_over(residence_time_ps: float) -> float: + """hbar / tau in meV, computed with Scipp's constants.""" + return (scipp_hbar / sc.scalar(residence_time_ps, unit='ps')).to(unit='meV').value + + +def order_of(circle_order: int, s: int) -> int: + """The Lorentzian order j of circle order l and sample s.""" + return (circle_order - 1) * NUMBER_OF_SAMPLES + s + 1 + + +class TestJumpSitesLogNormalDistribution: + @pytest.fixture + def model(self): + return JumpSitesLogNormalDistribution() + + def test_create_weight_requires_float_Q(self, model): + # WHEN THEN EXPECT + with pytest.raises(TypeError, match='Q must be a float'): + model._create_weight(1, 1) + + # ------------------------------------------------------------------ + # Construction and parameters + # ------------------------------------------------------------------ + + def test_init_default(self, model): + # WHEN THEN EXPECT + assert model.display_name == 'JumpSitesLogNormalDistribution' + assert model.number_of_sites == 3 + assert model.number_of_lorentzians == 2 * NUMBER_OF_SAMPLES + assert model.radius.value == pytest.approx(1.0) + assert model.radius.unit == 'Å' + assert model.radius.fixed is True + assert model.residence_time.unit == 'ps' + assert model.residence_time.fixed is False + assert model.sigma.value == pytest.approx(1.0) + assert model.sigma.unit == 'dimensionless' + assert model.sigma.fixed is False + assert model.sigma.min == MINIMUM_SIGMA + + @pytest.mark.parametrize( + 'kwargs, expected_exception, expected_message', + [ + ({'radius': 'invalid'}, TypeError, 'radius must be a number'), + ({'radius': 0.0}, ValueError, 'radius must be at least'), + ({'residence_time': 'invalid'}, TypeError, 'residence_time must be a number'), + ({'residence_time': 0.0}, ValueError, 'residence_time must be at least'), + ({'sigma': 'invalid'}, TypeError, 'sigma must be a number'), + ({'sigma': 0.0}, ValueError, 'sigma must be at least'), + ({'number_of_sites': 2.5}, TypeError, 'number_of_sites must be an integer'), + ({'number_of_sites': 1}, ValueError, 'number_of_sites must be between 2 and'), + ({'delta_name': 1}, TypeError, 'delta_name must be a string'), + ], + ) + def test_init_rejects_invalid_input(self, kwargs, expected_exception, expected_message): + # WHEN THEN EXPECT + with pytest.raises(expected_exception, match=expected_message): + JumpSitesLogNormalDistribution(**kwargs) + + def test_sigma_setter(self, model): + # WHEN + model.sigma = 0.4 + + # THEN EXPECT + assert model.sigma.value == pytest.approx(0.4) + with pytest.raises(ValueError, match='sigma must be at least'): + model.sigma = 0.0 + + # ------------------------------------------------------------------ + # Physics + # ------------------------------------------------------------------ + + def test_matches_QENSmodels_doctest_values(self): + """ + Doctest values of QENSmodels hwhmJumpSitesLogNormDist([1., 2.], 4, 0.5, 1.5, 1.0), rounded + there to the digits given. Its widths are in 1/ps, so they are divided by hbar here. + """ + # WHEN + model = JumpSitesLogNormalDistribution( + number_of_sites=4, radius=0.5, residence_time=1.5, sigma=1.0 + ) + Q = np.array([1.0, 2.0]) + hbar_per_ps = hbar_over(1.0) + + # THEN EXPECT + np.testing.assert_allclose(model.calculate_EISF(Q), [0.92, 0.713], atol=5e-4) + assert model.calculate_width(Q, order_of(3, 20))[1] / hbar_per_ps == pytest.approx( + 5.7, abs=5e-2 + ) + assert model.calculate_width(Q, order_of(1, 5))[1] / hbar_per_ps == pytest.approx( + 0.228, abs=5e-4 + ) + assert model.calculate_QISF_order(Q, order_of(3, 19))[0] == pytest.approx( + 0.000538, abs=5e-7 + ) + assert model.calculate_QISF_order(Q, order_of(2, 5))[1] == pytest.approx( + 0.000712, abs=5e-7 + ) + + def test_matches_QENSmodels_reference_spectrum(self): + """The QENSmodels reference spectrum, with w in 1/ps and hbar = 1 there.""" + # WHEN + model = JumpSitesLogNormalDistribution( + scale=2.0, + radius=5.0, + residence_time=2.0, + sigma=0.6, + number_of_sites=7, + Q=np.array([float(np.float32(0.7))]), + ) + hbar_meV_ps = scipp_hbar.to(unit='meV*ps').value + w = np.array([row[0] for row in REFERENCE_ROWS]) + + # THEN + # Lorentzians only: on these few points the delta function would land in the point + # nearest its centre (w = 1.0), while the reference has it at w = 0.8, not listed here. + x = hbar_meV_ps * (w - 0.8) + lorentzians = model.get_component_collections()[0][1:] + spectrum = hbar_meV_ps * sum(component.evaluate(x) for component in lorentzians) + + # EXPECT + np.testing.assert_allclose(spectrum, [row[1] for row in REFERENCE_ROWS], rtol=1e-9) + + @pytest.mark.parametrize('number_of_sites', [2, 3, 7]) + def test_weights_sum_to_one(self, number_of_sites): + # WHEN + model = JumpSitesLogNormalDistribution(number_of_sites=number_of_sites, radius=1.3) + Q = np.array([0.1, 0.9, 2.5]) + + # THEN EXPECT + np.testing.assert_allclose( + model.calculate_EISF(Q) + model.calculate_QISF(Q), 1.0, rtol=0, atol=WEIGHT_ATOL + ) + + def test_sample_weights_do_not_depend_on_sigma(self): + # WHEN + narrow = JumpSitesLogNormalDistribution(sigma=0.2).sample_weights() + wide = JumpSitesLogNormalDistribution(sigma=1.5).sample_weights() + + # THEN EXPECT + assert narrow.shape == (NUMBER_OF_SAMPLES,) + np.testing.assert_allclose(narrow, wide, rtol=1e-15, atol=0) + assert narrow.sum() == pytest.approx(1.0, rel=1e-15) + # gaussian in log(ratio) with the half: g_0 / g_10 = exp(-0.5 * 2 ln 10) = 1/10 + assert narrow[0] / narrow[10] == pytest.approx(0.1, rel=1e-12) + + def test_minimum_sigma_is_the_circle(self): + # WHEN + Q = np.array([0.5, 1.7]) + model = JumpSitesLogNormalDistribution( + number_of_sites=5, radius=1.2, residence_time=0.7, sigma=MINIMUM_SIGMA + ) + circle = EquivalentSitesCircle(number_of_sites=5, radius=1.2, residence_time=0.7) + + # THEN EXPECT + for circle_order in range(1, 5): + for s in range(NUMBER_OF_SAMPLES): + np.testing.assert_allclose( + model.calculate_width(Q, order_of(circle_order, s)), + circle.calculate_width(Q, circle_order), + rtol=1e-9, + ) + summed = sum( + model.calculate_QISF_order(Q, order_of(circle_order, s)) + for s in range(NUMBER_OF_SAMPLES) + ) + np.testing.assert_allclose( + summed, circle.calculate_QISF_order(Q, circle_order), rtol=0, atol=WEIGHT_ATOL + ) + np.testing.assert_allclose( + model.calculate_EISF(Q), circle.calculate_EISF(Q), rtol=0, atol=WEIGHT_ATOL + ) + + def test_weights_follow_radius_unit(self): + # WHEN + Q = np.array([0.8, 1.6]) + model = JumpSitesLogNormalDistribution(number_of_sites=4, radius=1.5, Q=Q) + EISF_angstrom = model.calculate_EISF(Q) + areas_angstrom = [c[3].area.value for c in model.get_component_collections()] + + # THEN + model.radius.convert_unit('nm') + + # EXPECT the same physical radius gives the same weights + np.testing.assert_allclose(model.calculate_EISF(Q), EISF_angstrom, rtol=1e-12) + areas_nm = [c[3].area.value for c in model.get_component_collections()] + np.testing.assert_allclose(areas_nm, areas_angstrom, rtol=0, atol=WEIGHT_ATOL) + + @pytest.mark.parametrize('order', [0, 2 * NUMBER_OF_SAMPLES + 1, 1.5]) + def test_order_validation(self, model, order): + # WHEN THEN EXPECT + with pytest.raises((TypeError, ValueError), match='order must be'): + model.calculate_QISF_order(np.array([1.0]), order) + + # ------------------------------------------------------------------ + # Components + # ------------------------------------------------------------------ + + @pytest.mark.parametrize('number_of_sites', [3, 7]) + def test_components_per_Q(self, number_of_sites): + # WHEN + model = JumpSitesLogNormalDistribution( + number_of_sites=number_of_sites, Q=np.array([0.5, 1.0]) + ) + n = NUMBER_OF_SAMPLES * (number_of_sites - 1) + + # THEN + collections = model.get_component_collections() + + # EXPECT + assert len(collections) == 2 + for collection in collections: + assert len(collection) == n + 1 + assert isinstance(collection[0], DeltaFunction) + assert all(isinstance(component, Lorentzian) for component in collection[1:]) + assert collection[n].name == model.lorentzian_component_name(n) + assert model.lorentzian_component_name(5) == 'JumpSitesLogNormalDistribution_j5' + + def test_component_areas_and_widths_match_calculations(self): + """The dependency strings must give what the Scipp calculations give.""" + # WHEN + Q = np.array([0.4, 2.3]) + model = JumpSitesLogNormalDistribution( + scale=2.5, radius=1.4, residence_time=0.8, sigma=0.7, number_of_sites=4, Q=Q + ) + + # THEN + collections = model.get_component_collections() + + # EXPECT + EISF = model.calculate_EISF(Q) + for i, collection in enumerate(collections): + assert collection[0].area.value == pytest.approx(2.5 * EISF[i], rel=0, abs=WEIGHT_ATOL) + for order in range(1, model.number_of_lorentzians + 1): + lorentzian = collection[order] + weight = model.calculate_QISF_order(Q, order)[i] + assert lorentzian.area.value == pytest.approx(2.5 * weight, rel=0, abs=WEIGHT_ATOL) + assert lorentzian.width.value == pytest.approx( + model.calculate_width(Q, order)[i], rel=1e-12 + ) + assert lorentzian.width.unit == 'meV' + + def test_components_follow_parameter_changes(self): + # WHEN + Q = np.array([0.6, 1.8]) + model = JumpSitesLogNormalDistribution(number_of_sites=3, Q=Q) + + # THEN + model.radius = 2.2 + model.residence_time = 3.5 + model.sigma = 0.3 + model.scale = 4.0 + + # EXPECT + collections = model.get_component_collections() + for order in (1, 7, 21, 30, 42): + for i, collection in enumerate(collections): + assert collection[order].area.value == pytest.approx( + 4.0 * model.calculate_QISF_order(Q, order)[i], rel=0, abs=WEIGHT_ATOL + ) + assert collection[order].width.value == pytest.approx( + model.calculate_width(Q, order)[i], rel=1e-12 + ) + + def test_component_parameters_are_dependent(self): + # WHEN + model = JumpSitesLogNormalDistribution(Q=np.array([1.0])) + collection = model.get_component_collections()[0] + + # THEN + fit_parameters = model.get_fit_parameters() + + # EXPECT + for component in collection: + assert component.area.independent is False + assert all(component.area is not p for p in fit_parameters) + for component in collection[1:]: + assert component.width.independent is False + assert all(component.width is not p for p in fit_parameters) + + def test_parameter_updates_raise_no_runtime_warning(self): + """A fit changes tau and sigma on every step; the dependent values must not warn.""" + # WHEN + model = JumpSitesLogNormalDistribution(Q=np.array([0.5, 1.0])) + + # THEN EXPECT + with warnings.catch_warnings(): + warnings.simplefilter('error', RuntimeWarning) + model.residence_time = 3.0 + model.sigma = 0.4 + + def test_fit_targets(self, model): + # WHEN THEN + names = [target.name for target in model.get_fit_targets()] + + # EXPECT + assert names[0] == 'elastic_area' + assert len(names) == 1 + 2 * model.number_of_lorentzians + + def test_get_global_variables(self, model): + # WHEN THEN EXPECT + assert model.get_global_variables() == [ + model.scale, + model.radius, + model.residence_time, + model.sigma, + ] + + def test_repr(self, model): + # WHEN THEN EXPECT + assert 'sigma=' in repr(model) + assert 'number_of_sites=3' in repr(model) + + # ------------------------------------------------------------------ + # Fit loop + # ------------------------------------------------------------------ + + def test_fit_recovers_parameters(self): + """One Q, synthetic data from the model; tau and sigma are fitted, as in a real fit.""" + # WHEN + Q = sc.array(dims=['Q'], values=[1.2], unit='1/angstrom') + energy = sc.linspace('energy', -2.0, 2.0, 201, unit='meV') + shape = (1, energy.shape[0]) + + def as_data(values, variances): + return sc.DataArray( + sc.array(dims=['Q', 'energy'], values=values, variances=variances), + coords={'Q': Q, 'energy': energy}, + ) + + true_model = JumpSitesLogNormalDistribution( + radius=1.5, residence_time=2.0, sigma=0.5, number_of_sites=2 + ) + experiment = edyn.Experiment('synthetic', data=as_data(np.ones(shape), np.ones(shape))) + clean = np.array( + edyn.Analysis( + experiment=experiment, sample_model=edyn.SampleModel(diffusion_models=true_model) + ).calculate() + ) + noise = 0.003 * clean.max() + rng = np.random.default_rng(seed=1) + experiment.data = as_data(clean + rng.normal(0.0, noise, shape), np.full(shape, noise**2)) + fit_model = JumpSitesLogNormalDistribution( + radius=1.5, residence_time=1.5, sigma=0.8, number_of_sites=2 + ) + analysis = edyn.Analysis( + experiment=experiment, sample_model=edyn.SampleModel(diffusion_models=fit_model) + ) + analysis.fix_energy_offset() + + # THEN + analysis.fit(fit_method='independent') + + # EXPECT + for name, true_value in {'residence_time': 2.0, 'sigma': 0.5}.items(): + parameter = getattr(fit_model, name) + assert parameter.variance > 0 + assert abs(parameter.value - true_value) < 5 * parameter.error + assert fit_model.radius.value == pytest.approx(1.5) # fixed, so unchanged