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
 

examples/detailed_balance.ipynb

@@ -0,0 +1,134 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "97050b3e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "%matplotlib widget\n",
+    "import numpy as np\n",
+    "\n",
+    "\n",
+    "from easydynamics.utils import detailed_balance_factor\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c1654720",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "temperatures=[1, 10, 100]\n",
+    "temperature_unit='K'\n",
+    "omega=np.linspace(-1,1,100)\n",
+    "# omega=1.0\n",
+    "omega_unit='meV'\n",
+    "\n",
+    "plt.figure()\n",
+    "for temperature in temperatures:\n",
+    "    DBF = detailed_balance_factor(omega, temperature, omega_unit,temperature_unit)\n",
+    "    plt.plot(omega, DBF, label=f'T={temperature} K')\n",
+    "plt.legend()\n",
+    "plt.xlabel('Energy transfer (meV)')\n",
+    "plt.ylabel('Detailed balance factor')\n",
+    "plt.title('Detailed balance factor for different temperatures, normalized to 1 at zero energy transfer')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a64fbe7c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "temperatures=[1, 10, 100]\n",
+    "temperature_unit='K'\n",
+    "omega=np.linspace(-1,1,100)\n",
+    "# omega=1.0\n",
+    "omega_unit='meV'\n",
+    "\n",
+    "plt.figure()\n",
+    "for temperature in temperatures:\n",
+    "    DBF = detailed_balance_factor(omega, temperature, omega_unit,temperature_unit,divide_by_temperature=False)\n",
+    "    plt.plot(omega, DBF, label=f'T={temperature} K')\n",
+    "plt.legend()\n",
+    "plt.xlabel('Energy transfer (meV)')\n",
+    "plt.ylabel('Detailed balance factor')\n",
+    "plt.title('Detailed balance factor for different temperatures, not normalized')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "896af4e3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "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": "e0a75d55",
+   "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
+}

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,179 @@
+import numpy as np
+
+from easyscience import Parameter
+
+import scipp as sc
+from scipp import UnitError
+from scipp.constants import Boltzmann as kB
+
+
+from typing import Union, Optional
+import warnings
+
+
+# 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: str = "meV",
+    temperature_unit: str = "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):
+        raise TypeError("energy_unit must be a string.")
+
+    if not isinstance(temperature_unit, str):
+        raise TypeError("temperature_unit must be a string.")
+
+    # Convert temperature and energy to sc variables to make units easy to handle
+    temperature = _convert_to_scipp_variable(
+        temperature, temperature_unit, "temperature"
+    )
+
+    if temperature.value < 0:
+        raise ValueError("Temperature must be non-negative.")
+
+    energy = _convert_to_scipp_variable(energy, energy_unit, "energy")
+
+    # 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}."
+        )
+        energy_unit = energy.unit
+
+    # Same for temperature
+    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}."
+        )
+        temperature_unit = 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 * sc.Unit(energy_unit), 0.0 * sc.Unit(energy_unit), energy
+        )
+
+        if DBF.sizes == {}:
+            DBF_values = np.array(DBF.value)
+        else:
+            DBF_values = DBF.values
+        return DBF_values
+
+    # Work with dimensionless x = energy/(kB*T). Here, we have divided by kB*T
+    x = energy / (kB * temperature)
+
+    try:
+        x = sc.to_unit(x, unit="1")  # Make sure the unit is 1 and not e.g. 1e3
+    except Exception as e:
+        raise UnitError(
+            f"Error converting energy/(kB*T) to dimensionless units: {e}. Check that energy and temperature have compatible units."
+        )
+
+    DBF = sc.zeros_like(x)
+
+    # Small x (small energy and/or high temperature): Taylor expansion
+    small = sc.abs(x) < SMALL_THRESHOLD
+
+    DBF = sc.where(small, 1 + x / 2 + x**2 / 12, DBF)
+
+    # Large x (large energy and/or low temperature): asymptotic form
+    large = x > LARGE_THRESHOLD
+    DBF = sc.where(large, x, DBF)
+
+    # General case: exact formula
+    mid = ~small & ~large
+    DBF = sc.where(mid, x / (1 - sc.exp(-x)), DBF)
+
+    #
+    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],
+    unit: Optional[str],
+    name: str = "value",
+) -> 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 sc.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 sc.UnitError as e:
+            raise UnitError(f"Invalid unit string '{unit}' for {name}: {e}")