Update functions and tests for JAX backend

Author: Sampreet

oqupy/backends/numerical_backend.py

@@ -18,7 +18,7 @@
 
 import oqupy.config as oc
 
-class Numpy:
+class NumPy:
     """
     The NumPy backend employing
     dynamic switching through `oqupy.config`.
@@ -38,19 +38,28 @@ def dtype_float(self) -> default_np.dtype:
         """Getter for the float datatype."""
         return oc.NumPyDtypeFloat
 
-    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)
 
-    def update(self, array, indices:tuple, values) -> default_np.ndarray:
+    def update(self,
+               array,
+               indices:tuple,
+               values,
+               ) -> default_np.ndarray:
         """Option to update select indices of an array with given values."""
         if not isinstance(array, default_np.ndarray):
             return array.at[indices].set(values)
         array[indices] = values
         return array
 
-    def get_random_floats(self, seed, shape):
+    def get_random_floats(self,
+                          seed,
+                          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, \
@@ -73,5 +82,5 @@ def __getattr__(self, name: str):
         return getattr(backend, name)
 
 # initialize for import
-np = Numpy()
+np = NumPy()
 la = LinAlg()

oqupy/backends/pt_tempo_backend.py

@@ -150,10 +150,6 @@ def initialize(self) -> None:
         # the inner `for` loop can be optimized for parallelization
         if self._num_infl > 2:
             indices = list(range(1, self._num_infl - 1))
-            # influences_mpo += create_deltas(self._influence, indices,
-            #                                 [0, 1, 1, 0])
-            # influences_mps += create_deltas(self._influence, indices,
-            #                                 [0, 1, 0], scale)
             for index in indices:
                 infl = self._influence(index)
                 influences_mpo.append(create_delta(infl, [0, 1, 1, 0]))

oqupy/backends/tempo_backend.py

@@ -434,8 +434,6 @@ def initialize_mps_mpo(self):
         # the inner `for` loop can be optimized for parallelization
         if dkmax_pre_compute > 1:
             indices = list(range(1, dkmax_pre_compute))
-            # influences += create_deltas(self._influence, indices,
-            #                             [1, 0, 0, 1])
             for index in indices:
                 infl = self._influence(index)
                 influences.append(create_delta(infl, [1, 0, 0, 1]))

oqupy/bath_dynamics.py

@@ -195,7 +195,7 @@ def occupation(
             np.sum(_sys_correlations.real*re_kernel \
                           + 1j*_sys_correlations.imag*im_kernel, axis = 0)
                 ).real * coup
-        bath_occupation = np.append([0], bath_occupation)
+        bath_occupation = np.append(np.array([0]), bath_occupation)
         if not change_only and self._temp > 0:
             bath_occupation += np.exp(-freq/self._temp) \
                 / (1 - np.exp(-freq/self._temp))
@@ -396,54 +396,121 @@ def phase(region, swap_ts = False):
                 ph += np.exp(a * tk*dt + b * tkp*dt) / (a * b)
             return ph
 
-
         if dagg == (0, 1):
-            re_kernel[regions['a']] = phase('a') + phase('a', 1)
-
-            re_kernel[regions['b']] = phase('b')
-
-            im_kernel[regions['a']] = ((2*n_1 + 1) * phase('a') -
-                                       (2*n_2 + 1) * phase('a', 1))
-
-            im_kernel[regions['b']] = (2*n_1 + 1) * phase('b')
-
-            im_kernel[regions['c']] = -2 * (n_1 + 1) * phase('c')
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['a'],
+                values=phase('a') + phase('a', 1)
+            )
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['b'],
+                values=phase('b')
+            )
+
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['a'],
+                values=((2 * n_1 + 1) * phase('a') -
+                        (2 * n_2 + 1) * phase('a', 1))
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['b'],
+                values=(2 * n_1 + 1) * phase('b')
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['c'],
+                values=- 2 * (n_1 + 1) * phase('c')
+            )
 
         elif dagg == (1, 0):
-            re_kernel[regions['a']] = phase('a') + phase('a', 1)
-
-            re_kernel[regions['b']] = phase('b')
-
-            im_kernel[regions['a']] = ((2*n_1 + 1) * phase('a') -
-                                       (2*n_2 + 1) * phase('a', 1))
-
-            im_kernel[regions['b']] = (2*n_1 + 1) * phase('b')
-
-            im_kernel[regions['c']] = 2 * n_1 * phase('c')
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['a'],
+                values=phase('a') + phase('a', 1)
+            )
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['b'],
+                values=phase('b')
+            )
+
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['a'],
+                values=((2*n_1 + 1) * phase('a') -
+                        (2*n_2 + 1) * phase('a', 1))
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['b'],
+                values=(2*n_1 + 1) * phase('b')
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['c'],
+                values=2 * n_1 * phase('c')
+            )
 
         elif dagg == (1, 1):
-            re_kernel[regions['a']] = -(phase('a') + phase('a', 1))
-
-            re_kernel[regions['b']] = -phase('b')
-
-            im_kernel[regions['a']] = ((2*n_1 + 1) * phase('a') +
-                                       (2*n_2 + 1) * phase('a', 1))
-
-            im_kernel[regions['b']] = (2*n_1 + 1) * phase('b')
-
-            im_kernel[regions['c']] = 2 * (n_1 + 1) * phase('c')
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['a'],
+                values=- (phase('a') + phase('a', 1))
+            )
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['b'],
+                values=- phase('b')
+            )
+
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['a'],
+                values=((2*n_1 + 1) * phase('a') +
+                        (2*n_2 + 1) * phase('a', 1))
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['b'],
+                values=(2*n_1 + 1) * phase('b')
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['c'],
+                values=2 * (n_1 + 1) * phase('c')
+            )
 
         elif dagg == (0, 0):
-            re_kernel[regions['a']] = -(phase('a') + phase('a', 1))
-
-            re_kernel[regions['b']] = -phase('b')
-
-            im_kernel[regions['a']] = -((2*n_2 + 1) * phase('a', 1) +
-                                        (2*n_1 + 1) * phase('a'))
-
-            im_kernel[regions['b']] = -(2*n_1 + 1) * phase('b')
-
-            im_kernel[regions['c']] = -2 * n_1 * phase('c')
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['a'],
+                values=- (phase('a') + phase('a', 1))
+            )
+            re_kernel = np.update(
+                array=re_kernel,
+                indices=regions['b'],
+                values=- phase('b')
+            )
+
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['a'],
+                values=- ((2*n_1 + 1) * phase('a') +
+                        (2*n_2 + 1) * phase('a', 1))
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['b'],
+                values=- (2*n_1 + 1) * phase('b')
+            )
+            im_kernel = np.update(
+                array=im_kernel,
+                indices=regions['c'],
+                values=- 2 * n_1 * phase('c')
+            )
 
         re_kernel = np.triu(re_kernel) #only keep triangular region
         im_kernel = np.triu(im_kernel)

oqupy/config.py

@@ -18,8 +18,8 @@
 import scipy.linalg as default_la
 NUMERICAL_BACKEND_NUMPY = default_np
 NUMERICAL_BACKEND_LINALG = default_la
-NumPyDtypeComplex = default_np.complex128
-NumPyDtypeFloat = default_np.float64
+NumPyDtypeComplex = default_np.complex128   # earlier NpDtype
+NumPyDtypeFloat = default_np.float64        # earlier NpDtypeReal
 
 # Separator string for __str__ functions
 SEPERATOR = "----------------------------------------------\n"

oqupy/dynamics.py

@@ -52,9 +52,9 @@ def __str__(self) -> Text:
         ret.append("  length        = {} timesteps \n".format(len(self)))
         if len(self) > 0:
             ret.append("  min time      = {} \n".format(
-                np.min(self._times)))
+                np.min(self.times)))
             ret.append("  max time      = {} \n".format(
-                np.max(self._times)))
+                np.max(self.times)))
         return "".join(ret)
 
     def __len__(self) -> int:
@@ -181,7 +181,7 @@ def add(
     @property
     def shape(self):
         """The shape of the states. """
-        return np.copy(self._shape)
+        return copy(self._shape)
 
 class MeanFieldDynamics(BaseAPIClass):
     """

oqupy/gradient.py

@@ -76,7 +76,7 @@ def state_gradient(
         'gradient' : derivatives of Z with respect to the parameters
         'dynamics' : the dynamics of the system
     """
-    check_isinstance(parameters, ndarray, 'parameters')
+    check_isinstance(parameters, np.ndarray, 'parameters')
 
     num_steps = len(process_tensors[0])
     dt = process_tensors[0].dt
@@ -154,14 +154,20 @@ def combine_derivs(
         prog_bar.update(i)
 
         for j in range(0,num_parameters):
-            total_derivs[2*i][j] = combine_derivs(
-                            adjoint_tensor[i],
-                            first_half_prop_derivs[j].T,
-                            second_half_prop.T)
-            total_derivs[2*i+1][j] = combine_derivs(
-                adjoint_tensor[i],
-                first_half_prop.T,
-                second_half_prop_derivs[j].T)
+            total_derivs = np.update(
+                array=total_derivs,
+                indices=(2 * i, j),
+                values=combine_derivs(adjoint_tensor[i],
+                                      first_half_prop_derivs[j].T,
+                                      second_half_prop.T)
+            )
+            total_derivs = np.update(
+                array=total_derivs,
+                indices=(2 * i + 1, j),
+                values=combine_derivs(adjoint_tensor[i],
+                                      first_half_prop.T,
+                                      second_half_prop_derivs[j].T)
+            )
 
     prog_bar.update(num_steps)
     prog_bar.exit()
@@ -265,7 +271,7 @@ def controls(step: int):
     #    edges 0, 1, .., num_envs-1    are the bond legs of the environments
     #    edge  -1                      is the state leg
     initial_ndarray = initial_state.reshape(hs_dim**2)
-    initial_ndarray.shape = tuple([1]*num_envs+[hs_dim**2])
+    initial_ndarray = initial_ndarray.reshape(tuple([1]*num_envs+[hs_dim**2]))
     current_node = tn.Node(initial_ndarray)
     current_edges = current_node[:]
 
@@ -356,7 +362,7 @@ def controls(step: int):
 
     target_ndarray = target_derivative
     target_ndarray = target_ndarray.reshape(hs_dim**2)
-    target_ndarray.shape = tuple([1]*num_envs+[hs_dim**2])
+    target_ndarray = target_ndarray.reshape(tuple([1]*num_envs+[hs_dim**2]))
     current_node = tn.Node(target_ndarray)
     current_edges = current_node[:]
 

oqupy/mps_mpo.py

@@ -199,10 +199,10 @@ def compute_nn_gate(
     # exponentiate the liouvillian to become a propagator
     propagator = la.expm(dt*liouvillian)
     # split leg 0 and leg 1 each into left and right.
-    propagator.shape = [hs_dim_l**2, # left output
-                        hs_dim_r**2, # right output
-                        hs_dim_l**2, # left input
-                        hs_dim_r**2] # right input
+    propagator = propagator.reshape([hs_dim_l**2,   # left output
+                                     hs_dim_r**2,   # right output
+                                     hs_dim_l**2,   # left input
+                                     hs_dim_r**2])  # right input
     temp = np.swapaxes(propagator, 1, 2)
     temp = temp.reshape([hs_dim_l**2 * hs_dim_l**2,
                             hs_dim_r**2 * hs_dim_r**2])
@@ -396,11 +396,14 @@ def __init__(
             if rank == 4:
                 tmp_gamma = np.swapaxes(tmp_gamma,2,3)
             elif rank == 3:
-                tmp_gamma.shape = (shape[0], shape[1], 1, shape[2])
+                tmp_gamma = tmp_gamma.reshape((shape[0], shape[1], \
+                                               1, shape[2]))
             elif rank == 2:
-                tmp_gamma.shape = (1, shape[0]*shape[1], 1, 1)
+                tmp_gamma = tmp_gamma.reshape((1, shape[0] * shape[1], \
+                                               1, 1))
             elif rank == 1:
-                tmp_gamma.shape = (1, shape[0], 1, 1)
+                tmp_gamma = tmp_gamma.reshape((1, shape[0], \
+                                               1, 1))
             else:
                 raise ValueError()
             tmp_gammas.append(tmp_gamma)