Merge pull request #54 from jd15489/einsum

Faster implementation of the get_disp function
bjmorgan · Jan 5, 2024 · ff0c536 · ff0c536
2 parents a8997a9 + 56a1d52
commit ff0c536
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 15 deletions.
diff --git a/kinisi/diffusion.py b/kinisi/diffusion.py
@@ -112,7 +112,7 @@ def to_dict(self) -> dict:
             'intercept': None,
             'gradient': None,
             'start_dt': self._start_dt,
-            'model' : self._model
+            'model': self._model
         }
         if len(self._distributions) != 0:
             my_dict['distributions'] = [d.to_dict() for d in self._distributions]
@@ -409,7 +409,7 @@ def _model_variance(dt: np.ndarray, a: float) -> np.ndarray:
             self._popt, _ = curve_fit(_model_variance, self.dt[diff_regime:], self._v[diff_regime:])
             self._model_v = _model_variance(self.dt[diff_regime:], *self._popt)
         else:
-            self._model_v = self._v[diff_regime:] 
+            self._model_v = self._v[diff_regime:]
         self._covariance_matrix = _populate_covariance_matrix(self._model_v, self._n_o[diff_regime:])
         self._npd_covariance_matrix = self._covariance_matrix
         return cov_nearest(self._covariance_matrix)
@@ -557,7 +557,7 @@ def __init__(self,
         if block:
             import pyblock
             print('You are using the blocking method to estimate variances, please cite '
-                'doi:10.1063/1.457480 and the pyblock pacakge.')
+                  'doi:10.1063/1.457480 and the pyblock pacakge.')
         for i in self._iterator:
             disp_slice = self._displacements[i][:, :, self._slice].reshape(self._displacements[i].shape[0],
                                                                            self._displacements[i].shape[1], self.dims)
@@ -576,10 +576,10 @@ def __init__(self,
                 opt_block = pyblock.blocking.find_optimal_block(d_squared.flatten().size, reblock)
                 try:
                     mean = reblock[opt_block[0]].mean
-                    var = reblock[opt_block[0]].std_err ** 2
+                    var = reblock[opt_block[0]].std_err**2
                 except TypeError:
                     mean = reblock[-1].mean
-                    var = reblock[-1].std_err ** 2
+                    var = reblock[-1].std_err**2
                 self._n = np.append(self._n, mean)
                 self._v = np.append(self._v, var)
                 self._s = np.append(self._s, np.sqrt(self._v[i]))
@@ -641,7 +641,7 @@ def __init__(self,
         if block:
             import pyblock
             print('You are using the blocking method to estimate variances, please cite '
-                'doi:10.1063/1.457480 and the pyblock pacakge.')
+                  'doi:10.1063/1.457480 and the pyblock pacakge.')
         for i in self._iterator:
             disp_slice = self._displacements[i][:, :, self._slice].reshape(self._displacements[i].shape[0],
                                                                            self._displacements[i].shape[1], self.dims)
@@ -662,10 +662,10 @@ def __init__(self,
                 opt_block = pyblock.blocking.find_optimal_block(coll_motion.flatten().size, reblock)
                 try:
                     mean = reblock[opt_block[0]].mean
-                    var = reblock[opt_block[0]].std_err ** 2
+                    var = reblock[opt_block[0]].std_err**2
                 except TypeError:
                     mean = reblock[-1].mean
-                    var = reblock[-1].std_err ** 2
+                    var = reblock[-1].std_err**2
                 self._n = np.append(self._n, mean)
                 self._v = np.append(self._v, var)
                 self._s = np.append(self._s, np.sqrt(self._v[i]))
@@ -733,7 +733,7 @@ def __init__(self,
         if block:
             import pyblock
             print('You are using the blocking method to estimate variances, please cite '
-                'doi:10.1063/1.457480 and the pyblock pacakge.')
+                  'doi:10.1063/1.457480 and the pyblock pacakge.')
         for i in self._iterator:
             disp_slice = self._displacements[i][:, :, self._slice].reshape(self._displacements[i].shape[0],
                                                                            self._displacements[i].shape[1], self.dims)
@@ -754,10 +754,10 @@ def __init__(self,
                 opt_block = pyblock.blocking.find_optimal_block(sq_chg_motion.flatten().size, reblock)
                 try:
                     mean = reblock[opt_block[0]].mean
-                    var = reblock[opt_block[0]].std_err ** 2
+                    var = reblock[opt_block[0]].std_err**2
                 except TypeError:
                     mean = reblock[-1].mean
-                    var = reblock[-1].std_err ** 2
+                    var = reblock[-1].std_err**2
                 self._n = np.append(self._n, mean)
                 self._v = np.append(self._v, var)
                 self._s = np.append(self._s, np.sqrt(self._v[i]))

diff --git a/kinisi/parser.py b/kinisi/parser.py
@@ -109,10 +109,9 @@ def get_disp(coords: List[np.ndarray], latt: List[np.ndarray]) -> np.ndarray:
         d_coords = coords[:, 1:] - coords[:, :-1]
         d_coords = d_coords - np.round(d_coords)
         f_disp = np.cumsum(d_coords, axis=1)
-        c_disp = []
-        for i in f_disp:
-            c_disp.append([np.dot(d, m) for d, m in zip(i, latt[1:])])
-        disp = np.array(c_disp)
+        latt = np.array(latt)
+        disp = np.einsum('ijk,jkl->jik', f_disp, latt[1:])
+        disp = np.transpose(disp, (1, 0, 2))
         return disp
 
     @staticmethod