JAX support for refactor

.github/workflows/ci.yaml

@@ -31,7 +31,7 @@ jobs:
       - name: Install dyson
         run: |
           python -m pip install wheel
-          python -m pip install .[dev]
+          python -m pip install .[dev,jax]
       - name: Linting
         run: |
           ruff check

dyson/__init__.py

@@ -122,9 +122,7 @@
 
 __version__ = "1.0.0"
 
-import numpy
-import scipy
-
+from dyson._backend import set_backend, numpy, scipy
 from dyson.printing import console, quiet
 from dyson.representations import Lehmann, Spectral, Dynamic
 from dyson.solvers import (

dyson/_backend.py

@@ -0,0 +1,94 @@
+"""Backend management for :mod:`dyson`."""
+
+from __future__ import annotations
+
+import functools
+import importlib
+import os
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, Callable
+
+try:
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+    _HAVE_JAX = True
+except ImportError:
+    _HAVE_JAX = False
+
+_BACKEND = os.environ.get("DYSON_BACKEND", "numpy")
+_BACKEND_WARNINGS = os.environ.get("DYSON_BACKEND_WARNINGS", "0") == "1"
+
+_MODULE_CACHE: dict[tuple[str, str], ModuleType] = {}
+_BACKENDS: dict[str, dict[str, str]] = {}
+_BACKENDS["numpy"] = {
+    "numpy": "numpy",
+    "scipy": "scipy",
+}
+if _HAVE_JAX:
+    _BACKENDS["jax"] = {
+        "numpy": "jax.numpy",
+        "scipy": "jax.scipy",
+    }
+
+
+def set_backend(backend: str) -> None:
+    """Set the backend for :mod:`dyson`."""
+    global _BACKEND  # noqa: PLW0603
+    if backend not in _BACKENDS:
+        raise ValueError(
+            f"Invalid backend: {backend}. Available backends are: {list(_BACKENDS.keys())}"
+        )
+    _BACKEND = backend
+
+
+def cast_returned_array(func: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorate a function to coerce its returned array to the backend type."""
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        result = func(*args, **kwargs)
+        if isinstance(result, tuple):
+            return tuple(numpy.asarray(r) for r in result)
+        return numpy.asarray(result)
+
+    return wrapper
+
+
+class ProxyModule(ModuleType):
+    """Dynamic proxy module for backend-specific imports."""
+
+    def __init__(self, key: str) -> None:
+        """Initialise the object."""
+        super().__init__(f"{__name__}.{key}")
+        self._key = key
+
+    def __getattr__(self, attr: str) -> ModuleType:
+        """Get the attribute from the backend module."""
+        mod = self._load()
+        return getattr(mod, attr)
+
+    def _load(self) -> ModuleType:
+        """Load the backend module."""
+        # Check the cache
+        key = (self._key, _BACKEND)
+        if key in _MODULE_CACHE:
+            return _MODULE_CACHE[key]
+
+        # Load the module
+        keys = self._key.split(".")
+        module = _BACKENDS[_BACKEND][keys[0]]
+        if len(keys) > 1:
+            module += "." + ".".join(keys[1:])
+        _MODULE_CACHE[key] = importlib.import_module(module)
+
+        return _MODULE_CACHE[key]
+
+
+if TYPE_CHECKING:
+    import numpy
+    import scipy
+else:
+    numpy = ProxyModule("numpy")  # type: ignore[assignment]
+    scipy = ProxyModule("scipy")  # type: ignore[assignment]
+    scipy.optimize = ProxyModule("scipy.optimize")  # type: ignore[assignment]

dyson/expressions/adc.py

@@ -19,6 +19,7 @@
 
 from dyson import numpy as np
 from dyson import util
+from dyson._backend import cast_returned_array
 from dyson.expressions.expression import BaseExpression, ExpressionCollection
 from dyson.representations.enums import Reduction
 
@@ -89,6 +90,7 @@ def from_mf(cls, mf: RHF) -> BaseADC:
         adc_obj.kernel_gs()
         return cls.from_adc(adc_obj)
 
+    @cast_returned_array
     def apply_hamiltonian(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector.
 
@@ -115,6 +117,7 @@ def apply_hamiltonian_left(self, vector: Array) -> Array:
         """
         raise NotImplementedError("Left application of Hamiltonian is not implemented for ADC.")
 
+    @cast_returned_array
     def diagonal(self) -> Array:
         """Get the diagonal of the Hamiltonian.
 
@@ -234,6 +237,12 @@ def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) ->
         ooov = ooov.reshape(eo.size, eo.size, eo.size, ev.size)
         left = ooov * 2 - ooov.swapaxes(1, 2)
 
+        # Cast arrays
+        eo = np.asarray(eo)
+        ev = np.asarray(ev)
+        ooov = np.asarray(ooov)
+        left = np.asarray(left)
+
         # Get the subscript based on the reduction
         if Reduction(reduction) == Reduction.NONE:
             subscript = "ikla,jkla->ij"
@@ -302,6 +311,12 @@ def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) ->
         vvvo = vvvo.reshape(ev.size, ev.size, ev.size, eo.size)
         left = vvvo * 2 - vvvo.swapaxes(1, 2)
 
+        # Cast arrays
+        eo = np.asarray(eo)
+        ev = np.asarray(ev)
+        vvvo = np.asarray(vvvo)
+        left = np.asarray(left)
+
         # Get the subscript based on the reduction
         if Reduction(reduction) == Reduction.NONE:
             subscript = "acdi,bcdi->ab"

dyson/expressions/ccsd.py

@@ -21,6 +21,7 @@
 
 from dyson import numpy as np
 from dyson import util
+from dyson._backend import cast_returned_array
 from dyson.expressions.expression import BaseExpression, ExpressionCollection
 from dyson.representations.enums import Reduction
 
@@ -63,10 +64,10 @@ def __init__(
             imds: Intermediate integrals.
         """
         self._mol = mol
-        self._t1 = t1
-        self._t2 = t2
-        self._l1 = l1
-        self._l2 = l2
+        self._t1 = np.asarray(t1)
+        self._t2 = np.asarray(t2)
+        self._l1 = np.asarray(l1)
+        self._l2 = np.asarray(l2)
         self._imds = imds
         self._precompute_imds()
 
@@ -199,6 +200,7 @@ def _precompute_imds(self) -> None:
         """Precompute intermediate integrals."""
         self._imds.make_ip()
 
+    @cast_returned_array
     def vector_to_amplitudes(self, vector: Array, *args: Any) -> tuple[Array, Array]:
         """Convert a vector to amplitudes.
 
@@ -211,6 +213,7 @@ def vector_to_amplitudes(self, vector: Array, *args: Any) -> tuple[Array, Array]
         """
         return self.PYSCF_EOM.vector_to_amplitudes_ip(vector, self.nphys, self.nocc)
 
+    @cast_returned_array
     def amplitudes_to_vector(self, t1: Array, t2: Array) -> Array:
         """Convert amplitudes to a vector.
 
@@ -223,6 +226,7 @@ def amplitudes_to_vector(self, t1: Array, t2: Array) -> Array:
         """
         return self.PYSCF_EOM.amplitudes_to_vector_ip(t1, t2)
 
+    @cast_returned_array
     def apply_hamiltonian_right(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector on the right.
 
@@ -240,6 +244,7 @@ def apply_hamiltonian_right(self, vector: Array) -> Array:
         """
         return -self.PYSCF_EOM.lipccsd_matvec(self, vector, imds=self._imds)
 
+    @cast_returned_array
     def apply_hamiltonian_left(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector on the left.
 
@@ -260,6 +265,7 @@ def apply_hamiltonian_left(self, vector: Array) -> Array:
     apply_hamiltonian = apply_hamiltonian_right
     apply_hamiltonian.__doc__ = BaseCCSD.apply_hamiltonian.__doc__
 
+    @cast_returned_array
     def diagonal(self) -> Array:
         """Get the diagonal of the Hamiltonian.
 
@@ -367,6 +373,7 @@ def _precompute_imds(self) -> None:
         """Precompute intermediate integrals."""
         self._imds.make_ea()
 
+    @cast_returned_array
     def vector_to_amplitudes(self, vector: Array, *args: Any) -> tuple[Array, Array]:
         """Convert a vector to amplitudes.
 
@@ -379,6 +386,7 @@ def vector_to_amplitudes(self, vector: Array, *args: Any) -> tuple[Array, Array]
         """
         return self.PYSCF_EOM.vector_to_amplitudes_ea(vector, self.nphys, self.nocc)
 
+    @cast_returned_array
     def amplitudes_to_vector(self, t1: Array, t2: Array) -> Array:
         """Convert amplitudes to a vector.
 
@@ -391,6 +399,7 @@ def amplitudes_to_vector(self, t1: Array, t2: Array) -> Array:
         """
         return self.PYSCF_EOM.amplitudes_to_vector_ea(t1, t2)
 
+    @cast_returned_array
     def apply_hamiltonian_right(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector on the right.
 
@@ -402,6 +411,7 @@ def apply_hamiltonian_right(self, vector: Array) -> Array:
         """
         return self.PYSCF_EOM.eaccsd_matvec(self, vector, imds=self._imds)
 
+    @cast_returned_array
     def apply_hamiltonian_left(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector on the left.
 
@@ -416,6 +426,7 @@ def apply_hamiltonian_left(self, vector: Array) -> Array:
     apply_hamiltonian = apply_hamiltonian_right
     apply_hamiltonian.__doc__ = BaseCCSD.apply_hamiltonian.__doc__
 
+    @cast_returned_array
     def diagonal(self) -> Array:
         """Get the diagonal of the Hamiltonian.
 

dyson/expressions/fci.py

@@ -12,6 +12,7 @@
 
 from pyscf import ao2mo, fci
 
+from dyson._backend import cast_returned_array
 from dyson.expressions.expression import BaseExpression, ExpressionCollection
 from dyson.representations.enums import Reduction
 
@@ -103,6 +104,7 @@ def from_mf(cls, mf: RHF) -> BaseFCI:
         ci.kernel(h1e, h2e, mf.mol.nao, mf.mol.nelec)
         return cls.from_fci(ci, h1e, h2e)
 
+    @cast_returned_array
     def apply_hamiltonian(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector.
 
@@ -123,6 +125,7 @@ def apply_hamiltonian(self, vector: Array) -> Array:
         result -= (self.e_fci + self.chempot) * vector
         return self.SIGN * result
 
+    @cast_returned_array
     def diagonal(self) -> Array:
         """Get the diagonal of the Hamiltonian.
 
@@ -131,6 +134,7 @@ def diagonal(self) -> Array:
         """
         return self.SIGN * (self._diagonal - (self.e_fci + self.chempot))
 
+    @cast_returned_array
     def get_excitation_vector(self, orbital: int) -> Array:
         r"""Obtain the vector corresponding to a fermionic operator acting on the ground state.
 

dyson/expressions/gw.py

@@ -21,6 +21,7 @@
 
 from dyson import numpy as np
 from dyson import util
+from dyson._backend import cast_returned_array
 from dyson.expressions.expression import BaseExpression, ExpressionCollection
 from dyson.representations.enums import Reduction
 
@@ -151,6 +152,7 @@ def non_dyson(self) -> bool:
 class TDAGW_Dyson(BaseGW_Dyson):
     """GW expressions with Tamm--Dancoff (TDA) approximation for the Dyson Green's function."""
 
+    @cast_returned_array
     def apply_hamiltonian(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector.
 
@@ -215,6 +217,7 @@ def apply_hamiltonian_left(self, vector: Array) -> Array:
         """
         raise NotImplementedError("Left application of Hamiltonian is not implemented for TDA-GW.")
 
+    @cast_returned_array
     def diagonal(self) -> Array:
         """Get the diagonal of the Hamiltonian.
 

dyson/expressions/hamiltonian.py

@@ -7,6 +7,7 @@
 
 from dyson import numpy as np
 from dyson import util
+from dyson._backend import cast_returned_array
 from dyson.expressions.expression import BaseExpression
 from dyson.representations.enums import Reduction
 
@@ -37,8 +38,8 @@ def __init__(
                 :meth:`~dyson.expressions.expression.BaseExpression.get_excitation_ket`.
         """
         self._hamiltonian = hamiltonian
-        self._bra = bra
-        self._ket = ket
+        self._bra = np.asarray(bra) if bra is not None else None
+        self._ket = np.asarray(ket) if bra is not None else None
 
         if isinstance(hamiltonian, np.ndarray):
             self.hermitian_upfolded = np.allclose(hamiltonian, hamiltonian.T.conj())
@@ -58,6 +59,7 @@ def from_mf(cls, mf: RHF) -> Hamiltonian:
         """
         raise NotImplementedError("Cannot create Hamiltonian expression from mean-field object.")
 
+    @cast_returned_array
     def apply_hamiltonian(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector.
 
@@ -69,6 +71,7 @@ def apply_hamiltonian(self, vector: Array) -> Array:
         """
         return self._hamiltonian @ vector
 
+    @cast_returned_array
     def apply_hamiltonian_left(self, vector: Array) -> Array:
         """Apply the Hamiltonian to a vector on the left.
 
@@ -80,6 +83,7 @@ def apply_hamiltonian_left(self, vector: Array) -> Array:
         """
         return vector @ self._hamiltonian
 
+    @cast_returned_array
     def diagonal(self) -> Array:
         """Get the diagonal of the Hamiltonian.
 
@@ -88,6 +92,7 @@ def diagonal(self) -> Array:
         """
         return self._hamiltonian.diagonal()
 
+    @cast_returned_array
     def build_matrix(self) -> Array:
         """Build the Hamiltonian matrix.