Detailed balance

.github/workflows/coverage.yaml

@@ -48,7 +48,7 @@ jobs:
       - name: Run unit tests with coverage
         run: >
           pytest tests/unit_tests 
-          --cov=src/easydynamics 
+          --cov=easydynamics 
           --cov-report=term-missing 
           --cov-report=xml:coverage-unit.xml
 

.gitignore

@@ -2,3 +2,5 @@ examples/QENS_example/*
 examples/INS_example/*
 examples/Anesthetics
 src/easydynamics/__pycache__
+.vscode/*
+**/__pycache__/*

src/easydynamics/utils/__init__.py

@@ -0,0 +1,3 @@
+from .detailed_balance import _detailed_balance_factor
+
+__all__ = ["_detailed_balance_factor"]

src/easydynamics/utils/detailed_balance.py

@@ -0,0 +1,184 @@
+import warnings
+from typing import Optional, Union
+
+import numpy as np
+import scipp as sc
+from easyscience import Parameter
+from scipp import UnitError
+from scipp.constants import Boltzmann as kB
+
+# Small and large values of x need special treatment.
+SMALL_THRESHOLD = 0.001  # For small values of x, the denominator is close to zero, which can give numerical issues. The issues don't start until x<~1e-6, but we use a larger threshold to be safe.
+LARGE_THRESHOLD = 100  # For large values of x, the exponential term becomes negligible. This happens around x>~10, but we use a larger threshold to be safe. At very large x, exp(-x) can be rounded to 0, which can give numerical issues.
+
+
+def _detailed_balance_factor(
+    energy: Union[int, float, list, np.ndarray, sc.Variable],
+    temperature: Union[int, float, sc.Variable, Parameter],
+    energy_unit: Union[str, sc.Unit] = "meV",
+    temperature_unit: Union[str, sc.Unit] = "K",
+    divide_by_temperature: bool = True,
+) -> np.ndarray:
+    """
+    Compute the detailed balance factor (DBF):
+    DBF(energy, T) = energy*(n(energy)+1)=energy / (1 - exp(-energy / (kB*T))), where n(energy) is the Bose-Einstein distribution.
+    If divide_by_temperature is True, the result is normalized by kB*T to have value 1 at energy=0.
+
+    Args:
+        energy : number, list, np.ndarray, or scipp Variable. If number, assumed to be in meV unless energy_unit is set.
+            Energy transfer
+        T : number, scipp Variable, or Parameter. If number, assumed to be in K unless temperature_unit is set.
+            Temperature
+        energy_unit : str, optional
+            Unit for energy if energy is given as a number or list. Default is 'meV'
+        temperature_unit : str, optional
+            Unit for temperature if temperature is given as a number. Default is 'K'
+        divide_by_temperature : True or False, optional
+            If True, divide the result by kB*T to make it dimensionless and have value 1 at energy=0. Default is True.
+
+    Returns:
+        DBF : np.ndarray (may be changed to scipp Variable in the future)
+            Detailed balance factor
+
+    Examples
+    --------
+    >>> detailed_balance_factor(1.0, 300)  # 1 meV at 300 K
+    >>> detailed_balance_factor(energy=[1.0, 2.0], temperature=300, energy_unit='microeV', temperature_unit='K', divide_by_temperature=False)
+    """
+
+    # Input validation
+    if not isinstance(divide_by_temperature, bool):
+        raise TypeError("divide_by_temperature must be True or False.")
+
+    if not isinstance(energy_unit, (str, sc.Unit)):
+        raise TypeError("energy_unit must be a string or scipp.Unit.")
+
+    if not isinstance(temperature_unit, (str, sc.Unit)):
+        raise TypeError("temperature_unit must be a string or scipp.Unit.")
+
+    # Convert temperature and energy to sc variables to make units easy to handle
+    temperature = _convert_to_scipp_variable(
+        value=temperature, unit=temperature_unit, name="temperature"
+    )
+
+    if temperature.value < 0:
+        raise ValueError("Temperature must be non-negative.")
+
+    energy = _convert_to_scipp_variable(value=energy, unit=energy_unit, name="energy")
+
+    # What if people give units that don't make sense?
+    try:
+        sc.to_unit(energy, unit="meV")
+    except Exception as e:
+        raise UnitError(
+            f"The unit of energy is wrong: {energy.unit}: {e} Check that energy has a valid unit."
+        )
+    # We give users the option to specify the unit of the energy, but if the input has a unit, they might clash
+    if energy.unit != energy_unit:
+        warnings.warn(
+            f"Input energy has unit {energy.unit}, but energy_unit was set to {energy_unit}. Using {energy.unit}."
+        )
+
+    # Same for temperature
+    try:
+        sc.to_unit(temperature, unit="K")
+    except Exception as e:
+        raise UnitError(
+            f"The unit of temperature is wrong: {temperature.unit}: {e} Check that temperature has a valid unit."
+        )
+
+    if temperature.unit != temperature_unit:
+        warnings.warn(
+            f"Input temperature has unit {temperature.unit}, but temperature_unit was set to {temperature_unit}. Using {temperature.unit}."
+        )
+
+    # Zero temperature deserves special treatment. Here, DBF is 0 for energy<0 and energy for energy>0
+    if temperature.value == 0:
+        if divide_by_temperature:
+            raise ZeroDivisionError("Cannot divide by T when T = 0.")
+        DBF = sc.where(energy < 0.0 * energy.unit, 0.0 * energy.unit, energy)
+
+        if DBF.sizes == {}:
+            DBF_values = np.array([DBF.value])
+        else:
+            DBF_values = DBF.values
+        return DBF_values
+
+    # Now work with finite temperatures. Here, it helps to work with dimensionless x = energy/(kB*T), where we have divided by kB*T
+    # We first check if the units are OK.
+
+    x = energy / (kB * temperature)
+
+    x = sc.to_unit(x, unit="1")  # Make sure the unit is 1 and not e.g. 1e3
+
+    # Now compute DBF. First handle small and large x, then the general case.
+
+    # Small x (small energy and/or high temperature): Taylor expansion. Works and is needed for both positive and negative energies
+    small = sc.abs(x) < SMALL_THRESHOLD
+
+    DBF = sc.where(small, 1 + x / 2 + x**2 / 12, sc.zeros_like(x))
+
+    # Large x (large positive energy and/or low temperature): asymptotic form. Only needed for positive energies.
+    large = x > LARGE_THRESHOLD
+    DBF = sc.where(large, x, DBF)
+
+    # General case: exact formula
+    mid = sc.logical_not(small) & sc.logical_not(large)
+    DBF = sc.where(
+        mid, x / (1 - sc.exp(-x)), DBF
+    )  # zeros in x are handled by SMALL_THRESHOLD
+
+    #
+    if not divide_by_temperature:
+        DBF = DBF * (kB * temperature)
+        DBF = sc.to_unit(DBF, unit=energy.unit)
+
+    if DBF.sizes == {}:
+        DBF_values = np.array([DBF.value])
+    else:
+        DBF_values = DBF.values
+    return DBF_values
+
+
+def _convert_to_scipp_variable(
+    value: Union[int, float, list, np.ndarray, Parameter, sc.Variable],
+    name: str,
+    unit: Optional[str] = None,
+) -> sc.Variable:
+    """Convert various input types to a scipp Variable with proper units."""
+    if isinstance(value, sc.Variable):
+        return value
+
+    # Convert to numpy array first for consistent handling
+    if isinstance(value, (int, float)):
+        array_value = np.array(value)
+    elif isinstance(value, (list)):
+        array_value = np.array(value)
+    elif isinstance(value, np.ndarray):
+        array_value = value
+    elif isinstance(value, Parameter):
+        array_value = np.array(value.value)
+        unit = value.unit
+    else:
+        if name == "energy":
+            raise TypeError(
+                f"{name} must be a number, list, numpy array or scipp Variable"
+            )
+        else:
+            raise TypeError(
+                f"{name} must be a number, list, numpy array, Parameter or scipp Variable"
+            )
+
+    # Create appropriate scipp variable based on shape
+    if array_value.shape == () or (array_value.shape == (1,)):
+        # Scalar or single-element array
+        try:
+            return sc.scalar(value=float(array_value.flat[0]), unit=unit)
+        except UnitError as e:
+            raise UnitError(f"Invalid unit string '{unit}' for {name}: {e}")
+    else:
+        # Multi-element array
+        try:
+            return sc.array(dims=["x"], values=array_value, unit=unit)
+        except UnitError as e:
+            raise UnitError(f"Invalid unit string '{unit}' for {name}: {e}")

tests/performance_tests/utils/detailed_balance_approximations.ipynb

@@ -0,0 +1,83 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3ca3ab48",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "%matplotlib widget\n",
+    "import numpy as np\n",
+    " "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7ce48af1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x= np.linspace(0.1,50,1000)\n",
+    "\n",
+    "y = x / (1 - np.exp(-x))\n",
+    "\n",
+    "y_approx = x\n",
+    "\n",
+    "plt.figure()\n",
+    "plt.plot(x,y,marker='o')\n",
+    "plt.plot(x,y_approx,marker='x')\n",
+    "plt.xlabel('x')\n",
+    "plt.ylabel('x/(1-exp(-x))')\n",
+    "plt.legend(['Exact','Approximation'])\n",
+    "plt.title('Comparison of exact and approximate expressions for large x')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e731f810",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "x= np.linspace(1e-10,1e-5,1000)\n",
+    "\n",
+    "y = x / (1 - np.exp(-x))\n",
+    "\n",
+    "y_approx = 1 + x/2 + x**2/12\n",
+    "\n",
+    "plt.figure()\n",
+    "plt.plot(x,y,marker='o')\n",
+    "plt.plot(x,y_approx,marker='x')\n",
+    "plt.xlabel('x')\n",
+    "plt.ylabel('x/(1-exp(-x))')\n",
+    "plt.legend(['Exact','Approximation'])\n",
+    "plt.title('Comparison of exact and approximate expressions for small x')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "newdynamics",
+   "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.11.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}