enthalpy 0.1.0__py3-none-any.whl

enthalpy/__init__.py

@@ -0,0 +1,55 @@
+"""EnthalPy: global analysis of molecular binding equilibria from titration experiments."""
+
+from enthalpy import io
+from enthalpy.chemistry import ChemicalSystem, GibbsEnthalpy
+from enthalpy.experiments import (
+    Conditions,
+    Experiment,
+    ITCHeat,
+    ITCInjectionProtocol,
+    NMRShift,
+    Stock,
+    TitrationProtocol,
+    UVVisAbsorbance,
+    synthesize,
+)
+from enthalpy.inference import FitResult, ScipyLeastSquares
+from enthalpy.project import GlobalFit
+from enthalpy.statistics import (
+    bootstrap,
+    compare_models,
+    information_criteria,
+    likelihood_ratio_test,
+    linear_standard_errors,
+    profile_likelihood,
+)
+from enthalpy.units import Q_, ureg
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "ChemicalSystem",
+    "Conditions",
+    "Experiment",
+    "FitResult",
+    "GibbsEnthalpy",
+    "GlobalFit",
+    "ITCHeat",
+    "ITCInjectionProtocol",
+    "NMRShift",
+    "Q_",
+    "ScipyLeastSquares",
+    "Stock",
+    "TitrationProtocol",
+    "UVVisAbsorbance",
+    "__version__",
+    "bootstrap",
+    "io",
+    "compare_models",
+    "information_criteria",
+    "likelihood_ratio_test",
+    "linear_standard_errors",
+    "profile_likelihood",
+    "synthesize",
+    "ureg",
+]

enthalpy/backend/__init__.py

@@ -0,0 +1,78 @@
+"""Array backends for the numerical kernels (decision D3).
+
+Examples
+--------
+>>> from enthalpy.backend import get_backend
+>>> backend = get_backend("numpy")
+>>> backend.xp.exp(backend.asarray([0.0]))
+array([1.])
+"""
+
+from __future__ import annotations
+
+import importlib.util
+
+from enthalpy.backend._base import Array, Backend
+from enthalpy.backend._jax import jax_unavailable_reason
+from enthalpy.backend._numpy import NumpyBackend
+from enthalpy.exceptions import BackendError
+
+__all__ = ["Array", "Backend", "available_backends", "get_backend"]
+
+_BACKEND_NAMES = ("numpy", "jax")
+_instances: dict[str, Backend] = {}
+
+
+def available_backends() -> list[str]:
+    """Return the names of backends that can be used on this machine.
+
+    JAX counts as available only if it is installed *and* imports successfully (jaxlib can fail at
+    import time, e.g. on CPUs without AVX). Checking this imports JAX, which takes about a second the
+    first time.
+    """
+    names = ["numpy"]
+    if importlib.util.find_spec("jax") is not None and jax_unavailable_reason() is None:
+        names.append("jax")
+    return names
+
+
+def get_backend(backend: str | Backend = "numpy", *, enable_x64: bool = False) -> Backend:
+    """Return a backend instance.
+
+    Parameters
+    ----------
+    backend
+        Backend name (``"numpy"`` or ``"jax"``) or an existing :class:`Backend`, which is
+        validated and returned unchanged.
+    enable_x64
+        JAX only: enable 64-bit floats globally if they are not enabled yet.
+
+    Raises
+    ------
+    BackendError
+        If the backend is unknown, its dependencies are missing, or it is misconfigured.
+    """
+    if isinstance(backend, Backend):
+        backend.validate()
+        return backend
+
+    name = backend.lower()
+    if name not in _BACKEND_NAMES:
+        raise BackendError(f"Unknown backend {backend!r}. Available backends: {', '.join(_BACKEND_NAMES)}.")
+
+    instance = _instances.get(name)
+    if instance is None:
+        if name == "numpy":
+            instance = NumpyBackend()
+        else:
+            from enthalpy.backend._jax import JaxBackend
+
+            instance = JaxBackend(enable_x64=enable_x64)
+        _instances[name] = instance
+    elif name == "jax" and enable_x64:
+        import jax
+
+        jax.config.update("jax_enable_x64", True)
+
+    instance.validate()
+    return instance

enthalpy/backend/_base.py

@@ -0,0 +1,144 @@
+"""Abstract array backend (decision D3).
+
+Numerical kernels (θ → thermodynamics → equilibrium → observables → log-likelihood) are written
+against a :class:`Backend`. It provides
+
+- an Array-API namespace (:attr:`Backend.xp`) for elementwise operations and linear algebra,
+- functional control flow (:meth:`Backend.while_loop`, :meth:`Backend.fori_loop`) that stays
+  traceable under ``jax.jit``,
+- functional updates (:meth:`Backend.at_set`, :meth:`Backend.at_add`) because JAX arrays are
+  immutable,
+- root finding with implicit differentiation (:meth:`Backend.custom_root`).
+
+Kernel rules: no in-place assignment, no Python ``if``/``while`` on array values, and float64
+throughout.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from types import ModuleType
+from typing import Any, ClassVar, TypeAlias, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+__all__ = ["Array", "Backend"]
+
+Array: TypeAlias = Any
+"""Backend array type (``numpy.ndarray`` or ``jax.Array``)."""
+
+T = TypeVar("T")
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class Backend(ABC):
+    """Interface that every array backend implements."""
+
+    name: ClassVar[str]
+
+    @property
+    @abstractmethod
+    def xp(self) -> ModuleType:
+        """Array-API compatible namespace (use like ``numpy``)."""
+
+    def validate(self) -> None:  # noqa: B027 - optional hook, deliberately not abstract
+        """Check that the backend is usable in its current configuration.
+
+        Raises
+        ------
+        BackendError
+            If the backend is misconfigured.
+        """
+
+    def asarray(self, x: Any) -> Array:
+        """Convert ``x`` to a ``float64`` array of this backend."""
+        return self.xp.asarray(x, dtype=self.xp.float64)
+
+    @abstractmethod
+    def to_numpy(self, x: Array) -> NDArray[np.float64]:
+        """Convert a backend array to a NumPy ``float64`` array."""
+
+    @abstractmethod
+    def while_loop(self, cond_fun: Callable[[T], Any], body_fun: Callable[[T], T], init_val: T) -> T:
+        """Run ``body_fun`` while ``cond_fun`` is true (``jax.lax.while_loop`` semantics).
+
+        Reverse-mode differentiation through this loop is not supported. Differentiate through
+        :meth:`custom_root` instead.
+        """
+
+    @abstractmethod
+    def fori_loop(self, lower: int, upper: int, body_fun: Callable[[int, T], T], init_val: T) -> T:
+        """Apply ``val = body_fun(i, val)`` for ``i`` in ``range(lower, upper)``."""
+
+    @abstractmethod
+    def at_set(self, x: Array, index: Any, values: Any) -> Array:
+        """Return a copy of ``x`` with ``x[index] = values`` (``x`` is not modified)."""
+
+    @abstractmethod
+    def at_add(self, x: Array, index: Any, values: Any) -> Array:
+        """Return a copy of ``x`` with ``values`` added at ``index``; repeated indices accumulate."""
+
+    @abstractmethod
+    def custom_root(
+        self,
+        f: Callable[[Array], Array],
+        initial_guess: Array,
+        solve: Callable[[Callable[[Array], Array], Array], Any],
+        tangent_solve: Callable[[Callable[[Array], Array], Array], Array],
+        *,
+        has_aux: bool = False,
+    ) -> Any:
+        """Find a root of ``f`` and make it differentiable by the implicit-function theorem.
+
+        Parameters
+        ----------
+        f
+            Residual function. Differentiable parameters are the variables ``f`` closes over.
+        initial_guess
+            Starting point passed to ``solve``.
+        solve
+            ``solve(f, initial_guess) -> x`` with ``f(x) ≈ 0``.
+        tangent_solve
+            ``tangent_solve(g, y) -> x`` solving the linear system ``g(x) = y``, where ``g`` is the
+            Jacobian-vector product of ``f`` at the root. :meth:`dense_tangent_solve` is a default.
+        has_aux
+            If true, ``solve`` returns ``(x, aux)``. The auxiliary output (e.g. convergence flags)
+            is passed through without being differentiated.
+
+        Notes
+        -----
+        On backends without automatic differentiation this simply returns
+        ``solve(f, initial_guess)``. Derivatives are then computed explicitly by the caller with
+        the same implicit-function formula.
+        """
+
+    @abstractmethod
+    def dense_tangent_solve(self, g: Callable[[Array], Array], y: Array) -> Array:
+        """Solve the linear system ``g(x) = y`` by materialising the matrix of the linear map ``g``."""
+
+    def block_tangent_solve(self, g: Callable[[Array], Array], y: Array) -> Array:
+        """Solve ``g(x) = y`` for a linear map ``g`` that is block-diagonal over the leading axes.
+
+        ``y`` has shape ``(..., n)``, and ``g`` must act independently on every leading index (e.g.
+        independent titration points). The ``n × n`` blocks are obtained from ``n`` evaluations
+        of ``g``, so the cost is linear in the number of blocks.
+        """
+        xp = self.xp
+        n = y.shape[-1]
+        identity = xp.eye(n, dtype=xp.float64)
+        columns = [g(xp.zeros(y.shape, dtype=xp.float64) + identity[k]) for k in range(n)]
+        blocks = xp.stack(columns, axis=-1)  # blocks[..., i, k] = g(e_k)[..., i]
+        return xp.linalg.solve(blocks, y[..., None])[..., 0]
+
+    @abstractmethod
+    def stop_gradient(self, x: T) -> T:
+        """Treat ``x`` as a constant for differentiation."""
+
+    @abstractmethod
+    def jit(self, fun: F) -> F:
+        """Compile ``fun`` if the backend supports it; otherwise return it unchanged."""
+
+    def __repr__(self) -> str:
+        return f"<{type(self).__name__} {self.name!r}>"

enthalpy/backend/_jax.py

@@ -0,0 +1,149 @@
+"""JAX backend (automatic differentiation, JIT compilation, accelerators).
+
+JAX is an optional dependency (``pip install enthalpy[jax]``). Equilibrium concentrations span
+many orders of magnitude, so float64 is mandatory. The backend refuses to run while JAX is in its
+default float32 mode.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from types import ModuleType
+from typing import Any, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+from enthalpy.backend._base import Array, Backend
+from enthalpy.exceptions import BackendError
+
+__all__ = ["JaxBackend", "import_jax", "jax_unavailable_reason"]
+
+T = TypeVar("T")
+F = TypeVar("F", bound=Callable[..., Any])
+
+_X64_HELP = (
+    "The JAX backend requires 64-bit floats. Enable them before creating any arrays with "
+    "`jax.config.update('jax_enable_x64', True)`, by setting the environment variable "
+    "JAX_ENABLE_X64=1, or by calling `enthalpy.backend.get_backend('jax', enable_x64=True)`."
+)
+
+
+def import_jax() -> tuple[Any, Any]:
+    """Import ``jax`` and ``jax.numpy``.
+
+    Raises
+    ------
+    BackendError
+        If JAX is not installed, or if it is installed but cannot be imported. The latter happens,
+        for example, when jaxlib requires CPU instructions (AVX) that the machine does not provide.
+    """
+    try:
+        import jax
+        import jax.numpy as jnp
+    except ImportError as exc:
+        raise BackendError(f"The JAX backend requires the optional dependency 'jax' (enthalpy[jax]): {exc}") from exc
+    except Exception as exc:  # e.g. RuntimeError from jaxlib's CPU feature check
+        raise BackendError(f"JAX is installed but cannot be imported: {exc}") from exc
+    return jax, jnp
+
+
+def jax_unavailable_reason() -> str | None:
+    """Return why the JAX backend cannot be used, or ``None`` if JAX imports successfully."""
+    try:
+        import_jax()
+    except BackendError as exc:
+        return str(exc)
+    return None
+
+
+class JaxBackend(Backend):
+    """Backend based on ``jax.numpy`` and ``jax.lax``.
+
+    Parameters
+    ----------
+    enable_x64
+        If true, switch JAX to 64-bit mode globally. Otherwise 64-bit mode must already be enabled.
+    """
+
+    name = "jax"
+
+    def __init__(self, *, enable_x64: bool = False) -> None:
+        jax, jnp = import_jax()
+        self._jax: Any = jax
+        self._jnp: Any = jnp
+        if enable_x64:
+            jax.config.update("jax_enable_x64", True)
+        self.validate()
+
+    def validate(self) -> None:
+        if not bool(self._jax.config.jax_enable_x64):
+            raise BackendError(_X64_HELP)
+
+    @property
+    def xp(self) -> ModuleType:
+        return self._jnp
+
+    def to_numpy(self, x: Array) -> NDArray[np.float64]:
+        return np.asarray(x, dtype=np.float64)
+
+    def while_loop(self, cond_fun: Callable[[T], Any], body_fun: Callable[[T], T], init_val: T) -> T:
+        return self._jax.lax.while_loop(cond_fun, body_fun, init_val)
+
+    def fori_loop(self, lower: int, upper: int, body_fun: Callable[[int, T], T], init_val: T) -> T:
+        return self._jax.lax.fori_loop(lower, upper, body_fun, init_val)
+
+    def at_set(self, x: Array, index: Any, values: Any) -> Array:
+        return self._jnp.asarray(x).at[index].set(values)
+
+    def at_add(self, x: Array, index: Any, values: Any) -> Array:
+        return self._jnp.asarray(x).at[index].add(values)
+
+    def custom_root(
+        self,
+        f: Callable[[Array], Array],
+        initial_guess: Array,
+        solve: Callable[[Callable[[Array], Array], Array], Any],
+        tangent_solve: Callable[[Callable[[Array], Array], Array], Array],
+        *,
+        has_aux: bool = False,
+    ) -> Any:
+        lax = self._jax.lax
+        if not has_aux:
+            return lax.custom_root(f, initial_guess, solve, tangent_solve)
+
+        # Some JAX versions (e.g. 0.10.x) produce tangents of the wrong dtype for integer or boolean
+        # auxiliary outputs of custom_root. Carry all auxiliary leaves as float64 and restore their
+        # dtypes afterwards; the auxiliary output is never differentiated.
+        tree_util = self._jax.tree_util
+        jnp = self._jnp
+        meta: dict[str, Any] = {}
+
+        def solve_with_float_aux(fun: Callable[[Array], Array], guess: Array) -> Any:
+            solution, aux = solve(fun, guess)
+            leaves, treedef = tree_util.tree_flatten(aux)
+            meta["treedef"], meta["dtypes"] = treedef, [jnp.asarray(leaf).dtype for leaf in leaves]
+            return solution, tree_util.tree_unflatten(
+                treedef, [jnp.asarray(leaf, dtype=jnp.float64) for leaf in leaves]
+            )
+
+        solution, aux = lax.custom_root(f, initial_guess, solve_with_float_aux, tangent_solve, has_aux=True)
+        leaves = [
+            lax.stop_gradient(leaf).astype(dtype)
+            for leaf, dtype in zip(tree_util.tree_leaves(aux), meta["dtypes"], strict=True)
+        ]
+        return solution, tree_util.tree_unflatten(meta["treedef"], leaves)
+
+    def dense_tangent_solve(self, g: Callable[[Array], Array], y: Array) -> Array:
+        jnp = self._jnp
+        y = jnp.asarray(y)
+        if y.ndim == 0:
+            return y / g(jnp.ones_like(y))
+        matrix = self._jax.jacfwd(g)(y).reshape(y.size, y.size)
+        return jnp.linalg.solve(matrix, y.ravel()).reshape(y.shape)
+
+    def stop_gradient(self, x: T) -> T:
+        return self._jax.lax.stop_gradient(x)
+
+    def jit(self, fun: F) -> F:
+        return self._jax.jit(fun)

enthalpy/backend/_numpy.py

@@ -0,0 +1,83 @@
+"""NumPy backend (reference implementation, no automatic differentiation)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from types import ModuleType
+from typing import Any, TypeVar
+
+import array_api_compat.numpy as xp_numpy
+import numpy as np
+from numpy.typing import NDArray
+
+from enthalpy.backend._base import Array, Backend
+
+__all__ = ["NumpyBackend"]
+
+T = TypeVar("T")
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class NumpyBackend(Backend):
+    """Backend based on NumPy, using the Array-API compatible ``array_api_compat.numpy`` namespace."""
+
+    name = "numpy"
+
+    @property
+    def xp(self) -> ModuleType:
+        return xp_numpy
+
+    def to_numpy(self, x: Array) -> NDArray[np.float64]:
+        return np.asarray(x, dtype=np.float64)
+
+    def while_loop(self, cond_fun: Callable[[T], Any], body_fun: Callable[[T], T], init_val: T) -> T:
+        val = init_val
+        while bool(cond_fun(val)):
+            val = body_fun(val)
+        return val
+
+    def fori_loop(self, lower: int, upper: int, body_fun: Callable[[int, T], T], init_val: T) -> T:
+        val = init_val
+        for i in range(lower, upper):
+            val = body_fun(i, val)
+        return val
+
+    def at_set(self, x: Array, index: Any, values: Any) -> Array:
+        out = np.array(x, copy=True)
+        out[index] = values
+        return out
+
+    def at_add(self, x: Array, index: Any, values: Any) -> Array:
+        out = np.array(x, copy=True)
+        np.add.at(out, index, values)  # pyright: ignore[reportUnknownMemberType] (incomplete NumPy stub)
+        return out
+
+    def custom_root(
+        self,
+        f: Callable[[Array], Array],
+        initial_guess: Array,
+        solve: Callable[[Callable[[Array], Array], Array], Any],
+        tangent_solve: Callable[[Callable[[Array], Array], Array], Array],
+        *,
+        has_aux: bool = False,
+    ) -> Any:
+        return solve(f, initial_guess)
+
+    def dense_tangent_solve(self, g: Callable[[Array], Array], y: Array) -> Array:
+        rhs: NDArray[np.float64] = np.asarray(y, dtype=np.float64)
+        if rhs.ndim == 0:
+            return rhs / np.asarray(g(np.ones_like(rhs)), dtype=np.float64)
+        # g is linear, so the columns of its matrix are the images of the unit vectors.
+        matrix = np.empty((rhs.size, rhs.size), dtype=np.float64)
+        for i in range(rhs.size):
+            unit = np.zeros(rhs.size, dtype=np.float64)
+            unit[i] = 1.0
+            matrix[:, i] = np.asarray(g(unit.reshape(rhs.shape)), dtype=np.float64).ravel()
+        solution = np.asarray(np.linalg.solve(matrix, rhs.ravel()), dtype=np.float64)
+        return solution.reshape(rhs.shape)
+
+    def stop_gradient(self, x: T) -> T:
+        return x
+
+    def jit(self, fun: F) -> F:
+        return fun

enthalpy/chemistry/__init__.py

@@ -0,0 +1,34 @@
+"""Chemical model: species, reactions, thermodynamics and equilibrium (independent of techniques)."""
+
+from enthalpy.chemistry.equilibrium import (
+    EquilibriumSensitivities,
+    EquilibriumSolver,
+    EquilibriumState,
+    NewtonSolver,
+    OneToOneSolver,
+    default_solver,
+    equilibrium_sensitivities,
+)
+from enthalpy.chemistry.reactions import Reaction, format_reaction, parse_reaction
+from enthalpy.chemistry.species import Species
+from enthalpy.chemistry.system import ChemicalSystem, FormationModel
+from enthalpy.chemistry.thermodynamics import GibbsEnthalpy, ThermodynamicModel, ln_k
+
+__all__ = [
+    "ChemicalSystem",
+    "EquilibriumSensitivities",
+    "EquilibriumSolver",
+    "EquilibriumState",
+    "FormationModel",
+    "GibbsEnthalpy",
+    "NewtonSolver",
+    "OneToOneSolver",
+    "Reaction",
+    "Species",
+    "ThermodynamicModel",
+    "default_solver",
+    "equilibrium_sensitivities",
+    "format_reaction",
+    "ln_k",
+    "parse_reaction",
+]