Update numerical backend and add example

Author: Sampreet

docs/pages/api.rst

@@ -178,7 +178,6 @@ class :class:`oqupy.pt_tebd.PtTebd`
     dictionary.
 
 
-
 Results
 -------
 
@@ -207,3 +206,6 @@ module :mod:`oqupy.operators`
 function :func:`oqupy.helpers.plot_correlations_with_parameters`
   A helper function to plot an auto-correlation function and the sampling
   points given by a set of parameters for a TEMPO computation.
+
+function :func:`oqupy.backends.enable_jax_features`
+  Option to use JAX to support multiple device backends (CPUs/GPUs/TPUs).

docs/pages/gpu_features.rst

@@ -22,8 +22,8 @@ To enable experimental features, switch to the ``dev/jax`` branch and use
 
 .. code:: python
 
-   from oqupy.backends import enable_gpu_features
-   enable_gpu_features()
+   from oqupy.backends import enable_jax_features
+   enable_jax_features()
 
 Contributing Guidelines
 ~~~~~~~~~~~~~~~~~~~~~~~

examples/simple_dynamics_jax.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+
+import sys
+sys.path.insert(0, '.')
+# set the 'OQUPY_BACKEND' environment variable
+# to 'jax' to initialize JAX backend by default
+# or switch to JAX backend using oqupy.backends
+import oqupy
+from oqupy.backends import enable_jax_features
+# import NumPy from numerical_backend
+from oqupy.backends.numerical_backend import np
+enable_jax_features()
+
+import matplotlib.pyplot as plt
+sigma_x = oqupy.operators.sigma("x")
+sigma_z = oqupy.operators.sigma("z")
+up_density_matrix = oqupy.operators.spin_dm("z+")
+Omega = 1.0
+omega_cutoff = 5.0
+alpha = 0.3
+
+system = oqupy.System(0.5 * Omega * sigma_x)
+correlations = oqupy.PowerLawSD(alpha=alpha,
+                                zeta=1,
+                                cutoff=omega_cutoff,
+                                cutoff_type='exponential')
+bath = oqupy.Bath(0.5 * sigma_z, correlations)
+tempo_parameters = oqupy.TempoParameters(dt=0.1, tcut=3.0, epsrel=10**(-4))
+
+dynamics = oqupy.tempo_compute(system=system,
+                               bath=bath,
+                               initial_state=up_density_matrix,
+                               start_time=0.0,
+                               end_time=2.0,
+                               parameters=tempo_parameters,
+                               unique=True)
+t, s_z = dynamics.expectations(0.5*sigma_z, real=True)
+print(s_z)
+plt.plot(t, s_z, label=r'$\alpha=0.3$')
+plt.xlabel(r'$t\,\Omega$')
+plt.ylabel(r'$\langle\sigma_z\rangle$')
+#plt.savefig('simple_dynamics.png')
+plt.show()

oqupy/backends/__init__.py

@@ -0,0 +1,9 @@
+"""Module to initialize OQuPy's backends."""
+
+from oqupy.backends.numerical_backend import set_numerical_backends
+
+def enable_jax_features():
+    """Function to enable experimental features."""
+
+    # set numerical backend to JAX
+    set_numerical_backends('jax')

oqupy/backends/numerical_backend.py

@@ -13,20 +13,81 @@
 Module containing NumPy-like and SciPy-like numerical backends.
 """
 
+import os
+
 import numpy as default_np
 import scipy.linalg as default_la
 
+from tensornetwork.backend_contextmanager import \
+    set_default_backend
+
 import oqupy.config as oc
 
+# store instances of the initialized backends
+# this way, `oqupy.config` remains unchanged
+# and `ocupy.config.DEFAULT_BACKEND` is used
+# when NumPy and LinAlg are initialized
+NUMERICAL_BACKEND_INSTANCES = {}
+
+def get_numerical_backends(
+        backend_name: str,
+    ):
+    """Function to get numerical backend.
+
+    Parameters
+    ----------
+    backend_name: str
+        Name of the backend. Options are `'jax'` and `'numpy'`.
+
+    Returns
+    -------
+    backends: list
+        NumPy and LinAlg backends.
+    """
+
+    _bn = backend_name.lower()
+    if _bn in NUMERICAL_BACKEND_INSTANCES:
+        set_default_backend(_bn)
+        return NUMERICAL_BACKEND_INSTANCES[_bn]
+    assert _bn in ['jax', 'numpy'], \
+        "currently supported backends are `'jax'` and `'numpy'`"
+
+    if 'jax' in _bn:
+        try:
+            # explicitly import and configure jax
+            import jax
+            import jax.numpy as jnp
+            import jax.scipy.linalg as jla
+            jax.config.update('jax_enable_x64', True)
+
+            # # TODO: GPU memory allocation (default is 0.75)
+            # os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'
+            # os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '.5'
+
+            # set TensorNetwork backend
+            set_default_backend('jax')
+
+            NUMERICAL_BACKEND_INSTANCES['jax'] = [jnp, jla]
+            return NUMERICAL_BACKEND_INSTANCES['jax']
+        except ImportError:
+            print("JAX not installed, defaulting to NumPy")
+
+    # set TensorNetwork backend
+    set_default_backend('numpy')
+
+    NUMERICAL_BACKEND_INSTANCES['numpy'] = [default_np, default_la]
+    return NUMERICAL_BACKEND_INSTANCES['numpy']
+
 class NumPy:
     """
     The NumPy backend employing
     dynamic switching through `oqupy.config`.
     """
-    @property
-    def backend(self) -> default_np:
+    def __init__(self,
+                 backend_name=oc.DEFAULT_BACKEND,
+                 ):
         """Getter for the backend."""
-        return oc.NUMERICAL_BACKEND_NUMPY
+        self.backend = get_numerical_backends(backend_name)[0]
 
     @property
     def dtype_complex(self) -> default_np.dtype:
@@ -42,12 +103,11 @@ def __getattr__(self,
                     name: str,
                     ):
         """Return the backend's default attribute."""
-        backend = object.__getattribute__(self, 'backend')
-        return getattr(backend, name)
+        return getattr(self.backend, name)
 
     def update(self,
                array,
-               indices:tuple,
+               indices: tuple,
                values,
                ) -> default_np.ndarray:
         """Option to update select indices of an array with given values."""
@@ -61,26 +121,46 @@ def get_random_floats(self,
                           shape,
                           ):
         """Method to obtain random floats with a given seed and shape."""
-        backend = object.__getattribute__(self, 'backend')
         random_floats = default_np.random.default_rng(seed).random(shape, \
                         dtype=default_np.float64)
-        return backend.array(random_floats, dtype=self.dtype_float)
+        return self.backend.array(random_floats, dtype=self.dtype_float)
 
 class LinAlg:
     """
     The Linear Algebra backend employing
     dynamic switching through `oqupy.config`.
     """
-    @property
-    def backend(self) -> default_la:
+    def __init__(self,
+                 backend_name=oc.DEFAULT_BACKEND,
+                 ):
         """Getter for the backend."""
-        return oc.NUMERICAL_BACKEND_LINALG
+        self.backend = get_numerical_backends(backend_name)[1]
 
-    def __getattr__(self, name: str):
+    def __getattr__(self,
+                    name: str,
+                    ):
         """Return the backend's default attribute."""
-        backend = object.__getattribute__(self, 'backend')
-        return getattr(backend, name)
+        return getattr(self.backend, name)
+
+# setup libraries using environment variable
+# fall back to oqupy.config.DEFAULT_BACKEND
+try:
+    BACKEND_NAME = os.environ[oc.BACKEND_ENV_VAR]
+except KeyError:
+    BACKEND_NAME = oc.DEFAULT_BACKEND
+np = NumPy(backend_name=BACKEND_NAME)
+la = LinAlg(backend_name=BACKEND_NAME)
 
-# initialize for import
-np = NumPy()
-la = LinAlg()
+def set_numerical_backends(
+        backend_name: str
+    ):
+    """Function to set numerical backend.
+
+    Parameters
+    ----------
+    backend_name: str
+        Name of the backend. Options are `'jax'` and `'numpy'`.
+    """
+    backends = get_numerical_backends(backend_name)
+    np.backend = backends[0]
+    la.backend = backends[1]

oqupy/config.py

@@ -15,9 +15,8 @@
 
 # numerical backend
 import numpy as default_np
-import scipy.linalg as default_la
-NUMERICAL_BACKEND_NUMPY = default_np
-NUMERICAL_BACKEND_LINALG = default_la
+BACKEND_ENV_VAR = 'OQUPY_BACKEND'
+DEFAULT_BACKEND = 'numpy'
 NumPyDtypeComplex = default_np.complex128   # earlier NpDtype
 NumPyDtypeFloat = default_np.float64        # earlier NpDtypeReal
 

pylintrc

@@ -79,6 +79,7 @@ disable=raw-checker-failed,
         consider-using-f-string,          # GEFux: added for development
         too-many-arguments,               # piperfw: added for development
         too-many-positional-arguments,    # piperfw: added 2024-09-27
+        import-outside-toplevel,          # sampreet: added 2024-10-24 for JAX
         possibly-used-before-assignment,  # GEFux: added by hand
         unnecessary-lambda-assignment     # GEFux: added by hand
 

tests/coverage/__init__.py

@@ -1,20 +1,2 @@
-from importlib.util import find_spec
-
-if find_spec('jax') is not None:
-    # JAX configuration
-    import jax
-    import jax.numpy as jnp
-    import jax.scipy.linalg as jla
-    import oqupy.config as oc
-    import tensornetwork as tn
-    jax.config.update('jax_enable_x64', True)
-    oc.NUMERICAL_BACKEND_NUMPY = jnp
-    oc.NumPyDtypeComplex = jnp.complex128
-    oc.NumPyDtypeFloat = jnp.float64
-    oc.NUMERICAL_BACKEND_LINALG = jla
-    tn.set_default_backend('jax')
-
-    # # TODO: GPU memory allocation (default is 0.75)
-    # import os
-    # os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'
-    # os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '.5'
+# from oqupy.backends import enable_jax_features
+# enable_jax_features()

tests/physics/__init__.py

@@ -1,20 +1,2 @@
-from importlib.util import find_spec
-
-if find_spec('jax') is not None:
-    # JAX configuration
-    import jax
-    import jax.numpy as jnp
-    import jax.scipy.linalg as jla
-    import oqupy.config as oc
-    import tensornetwork as tn
-    jax.config.update('jax_enable_x64', True)
-    oc.NUMERICAL_BACKEND_NUMPY = jnp
-    oc.NumPyDtypeComplex = jnp.complex128
-    oc.NumPyDtypeFloat = jnp.float64
-    oc.NUMERICAL_BACKEND_LINALG = jla
-    tn.set_default_backend('jax')
-
-    # # TODO: GPU memory allocation (default is 0.75)
-    # import os
-    # os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'
-    # os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '.5'
+# from oqupy.backends import enable_jax_features
+# enable_jax_features()