performance improvements (#4089)

* performance improvements

* numpy changes

* fixes

* minor fixes

* final changes

* Update AUTHORS

* Update transformations.py

* Update topologyattrs.py

* Update bat.py

---------

Co-authored-by: g2707 <[email protected]>

package/AUTHORS

@@ -215,6 +215,7 @@ Chronological list of authors
   - Josh Vermaas
   - Xiaoxu Ruan
   - Egor Marin
+  - Shaivi Malik
 
 External code
 -------------

package/CHANGELOG

@@ -16,7 +16,8 @@ The rules for this file:
 
 ??/??/?? IAlibay, pgbarletta, mglagolev, hmacdope, manuel.nuno.melo, chrispfae, 
          ooprathamm, MeetB7, BFedder, v-parmar, MoSchaeffler, jbarnoud, jandom,
-         xhgchen, jaclark5, DrDomenicoMarson, AHMED-salah00, schlaicha, SophiaRuan
+         xhgchen, jaclark5, DrDomenicoMarson, AHMED-salah00, schlaicha, SophiaRuan,
+	 g2707
  * 2.5.0
 
 Fixes
@@ -49,6 +50,8 @@ Enhancements
     and SegmentGroup. (PR #3953)
 
 Changes
+  * einsum method for Einstein summation convention introduced to increase 
+    performance (Issue #2063, PR #4089)
   * Add progress bars to track the progress of _conclude() functions
     (_conclude_simple() and _conclude_fft()) in msd.py (Issue #4070, PR #4072)
   * As per NEP29 the minimum supported NumPy version has been raised to 1.21

package/MDAnalysis/analysis/bat.py

@@ -372,13 +372,13 @@ def _single_frame(self):
         v01 = p1 - p0
         v21 = p1 - p2
         # Internal coordinates
-        r01 = np.sqrt(np.sum(v01 *
-                             v01))  # Distance between first two root atoms
-        r12 = np.sqrt(np.sum(v21 *
-                             v21))  # Distance between second two root atoms
+        r01 = np.sqrt(np.einsum('i,i->',v01,v01))  
+        # Distance between first two root atoms
+        r12 = np.sqrt(np.einsum('i,i->',v21,v21))  
+        # Distance between second two root atoms
         # Angle between root atoms
-        a012 = np.arccos(max(-1.,min(1.,np.sum(v01*v21)/\
-                             np.sqrt(np.sum(v01*v01)*np.sum(v21*v21)))))
+        a012 = np.arccos(max(-1.,min(1.,np.einsum('i,i->',v01,v21)/\
+                              np.sqrt(np.einsum('i,i->',v01,v01)*np.einsum('i,i->',v21,v21)))))
         # External coordinates
         e = v01 / r01
         phi = np.arctan2(e[1], e[0])  # Polar angle
@@ -545,12 +545,12 @@ def Cartesian(self, bat_frame):
             cs_tor = np.cos(torsion)
 
             v21 = p1 - p2
-            v21 /= np.sqrt(np.sum(v21 * v21))
+            v21 /= np.sqrt(np.einsum('i,i->',v21,v21))
             v32 = p2 - p3
-            v32 /= np.sqrt(np.sum(v32 * v32))
+            v32 /= np.sqrt(np.einsum('i,i->',v32,v32))
 
             vp = np.cross(v32, v21)
-            cs = np.sum(v21 * v32)
+            cs = np.einsum('i,i->',v21,v32)
 
             sn = max(np.sqrt(1.0 - cs * cs), 0.0000000001)
             vp = vp / sn

package/MDAnalysis/core/groups.py

@@ -1071,7 +1071,7 @@ def center(self, weights, wrap=False, unwrap=False, compound='group'):
                 return coords.mean(axis=0)
             # promote weights to dtype if required:
             weights = weights.astype(dtype, copy=False)
-            return (coords * weights[:, None]).sum(axis=0) / weights.sum()
+            return np.einsum('ij,ij->j',coords,weights[:, None]) / weights.sum()
 
         # When compound split caching gets implemented it will be clever to
         # preempt at this point whether or not stable sorting will be needed
@@ -1100,7 +1100,7 @@ def center(self, weights, wrap=False, unwrap=False, compound='group'):
                 _centers = _coords.mean(axis=1)
             else:
                 _weights = weights[atom_mask]
-                _centers = (_coords * _weights[:, :, None]).sum(axis=1)
+                _centers = np.einsum('ijk,ijk->ik',_coords,_weights[:, :, None])
                 _centers /= _weights.sum(axis=1)[:, None]
             centers[compound_mask] = _centers
         if wrap:
@@ -1864,7 +1864,7 @@ def unwrap(self, compound='fragments', reference='com', inplace=True):
                         raise ValueError("Cannot perform unwrap with "
                                          "reference='com' because the total "
                                          "mass of the group is zero.")
-                    refpos = np.sum(positions * masses[:, None], axis=0)
+                    refpos = np.einsum('ij,ij->j',positions,masses[:, None])
                     refpos /= total_mass
                 else:  # reference == 'cog'
                     refpos = positions.mean(axis=0)
@@ -1893,8 +1893,8 @@ def unwrap(self, compound='fragments', reference='com', inplace=True):
                                              "reference='com' because the "
                                              "total mass of at least one of "
                                              "the {} is zero.".format(comp))
-                        refpos = np.sum(positions[atom_mask]
-                                        * masses[:, :, None], axis=1)
+                        refpos = np.einsum('ijk,ijk->ik',positions[atom_mask],
+                                           masses[:, :, None])
                         refpos /= total_mass[:, None]
                     else:  # reference == 'cog'
                         refpos = positions[atom_mask].mean(axis=1)

package/MDAnalysis/core/topologyattrs.py

@@ -1718,8 +1718,8 @@ def radius_of_gyration(group, wrap=False, **kwargs):
         else:
             recenteredpos = atomgroup.positions - com
 
-        rog_sq = np.sum(masses * np.sum(recenteredpos**2,
-                                        axis=1)) / atomgroup.total_mass()
+        rog_sq = np.einsum('i,i->',masses,np.einsum('ij,ij->i',
+                                     recenteredpos,recenteredpos))/atomgroup.total_mass()
 
         return np.sqrt(rog_sq)
 
@@ -2232,8 +2232,8 @@ def dipole_vector(group,
                 ) - ref)
             else:
                 recenteredpos = (atomgroup.positions - ref)
-            dipole_vector = np.sum(recenteredpos * charges[:, np.newaxis],
-                                   axis=0)
+            dipole_vector = np.einsum('ij,ij->j',recenteredpos, 
+                                       charges[:, np.newaxis])
         else:
             (atom_masks, compound_masks,
              n_compounds) = atomgroup._split_by_compound_indices(compound)
@@ -2248,10 +2248,10 @@ def dipole_vector(group,
 
             dipole_vector = np.empty((n_compounds, 3), dtype=np.float64)
             for compound_mask, atom_mask in zip(compound_masks, atom_masks):
-                dipole_vector[compound_mask] = np.sum(
-                    (coords[atom_mask] - ref[compound_mask][:, None, :]) *
-                    chgs[atom_mask][:, :, None],
-                    axis=1)
+                dipole_vector[compound_mask] = np.einsum('ijk,ijk->ik',
+                                                          (coords[atom_mask]-
+                                                           ref[compound_mask][:, None, :]),
+                                                          chgs[atom_mask][:, :, None])
 
         return dipole_vector
 
@@ -2315,9 +2315,9 @@ def dipole_moment(group, **kwargs):
         dipole_vector = atomgroup.dipole_vector(**kwargs)
 
         if len(dipole_vector.shape) > 1:
-            dipole_moment = np.sqrt(np.sum(dipole_vector**2, axis=1))
+            dipole_moment = np.sqrt(np.einsum('ij,ij->i',dipole_vector,dipole_vector))
         else:
-            dipole_moment = np.sqrt(np.sum(dipole_vector**2))
+            dipole_moment = np.sqrt(np.einsum('i,i->',dipole_vector,dipole_vector))
 
         return dipole_moment
 

package/MDAnalysis/lib/transformations.py

@@ -941,9 +941,9 @@ def superimposition_matrix(v0, v1, scaling=False, usesvd=True):
         M[:3, :3] = R
     else:
         # compute symmetric matrix N
-        xx, yy, zz = np.sum(v0 * v1, axis=1)
-        xy, yz, zx = np.sum(v0 * np.roll(v1, -1, axis=0), axis=1)
-        xz, yx, zy = np.sum(v0 * np.roll(v1, -2, axis=0), axis=1)
+        xx, yy, zz = np.einsum('ij,ij->i', v0 , v1)
+        xy, yz, zx = np.einsum('ij,ij->i', v0, np.roll(v1, -1, axis=0))
+        xz, yx, zy = np.einsum('ij,ij->i', v0, np.roll(v1, -2, axis=0))
         N = (
             (xx + yy + zz, 0.0, 0.0, 0.0),
             (yz - zy, xx - yy - zz, 0.0, 0.0),
@@ -958,9 +958,9 @@ def superimposition_matrix(v0, v1, scaling=False, usesvd=True):
 
     # scale: ratio of rms deviations from centroid
     if scaling:
-        v0 *= v0
-        v1 *= v1
-        M[:3, :3] *= math.sqrt(np.sum(v1) / np.sum(v0))
+        M[:3, :3] *= math.sqrt(np.einsum('ij,ij->',v1,v1) / 
+                                np.einsum('ij,ij->',v0,v0))
+
 
     # translation
     M[:3, 3] = t1