Use easy parallelization

dev/proposed_interface_to_chemcoord.py

@@ -3,7 +3,6 @@
 import os
 
 import ase.io
-import numpy as np
 from ase.calculators.emt import EMT
 from ase.constraints import FixAtoms
 from ase.lattice.surface import add_adsorbate, fcc111
@@ -69,7 +68,7 @@
 # Now I can see what the projection vector is for the displacement
 infinitesimal_mode_vector = InternalCoordinates.get_mode_infitisimal(0)
 # project forces on the mode vector
-forces_projected = np.dot(infinitesimal_mode_vector, f)
+forces_projected = infinitesimal_mode_vector @ f
 
 # Reset the displacements so that I can do a displacement along a different
 # internal mode

src/chemcoord/_cartesian_coordinates/_cart_transformation.py

@@ -1,6 +1,6 @@
 import numba as nb
 import numpy as np
-from numba import jit
+from numba import njit, prange
 from numba.extending import overload
 from numpy import arccos, arctan2, sqrt
 
@@ -44,12 +44,12 @@ def f(X, indices):
         raise AssertionError("Should not be here")
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def get_ref_pos(X, indices):
     return _stub_get_ref_pos(X, indices)
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def get_B(X, c_table, j):
     B = np.empty((3, 3))
     ref_pos = get_ref_pos(X, c_table[:, j])
@@ -66,7 +66,7 @@ def get_B(X, c_table, j):
     return (ERR_CODE_OK, B)
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def get_grad_B(X, c_table, j):
     grad_B = np.empty((3, 3, 3, 3))
     ref_pos = get_ref_pos(X, c_table[:, j])
@@ -1088,7 +1088,7 @@ def get_grad_B(X, c_table, j):
     return grad_B
 
 
-@jit(nb.f8[:](nb.f8[:]), nopython=True)
+@njit(nb.f8[:](nb.f8[:]))
 def get_S_inv(v):
     x, y, z = v
     r = np.linalg.norm(v)
@@ -1099,7 +1099,7 @@ def get_S_inv(v):
     return np.array([r, alpha, delta])
 
 
-@jit(nb.f8[:, :](nb.f8[:]), nopython=True)
+@njit(nb.f8[:, :](nb.f8[:]))
 def get_grad_S_inv(v):
     x, y, z = v
     grad_S_inv = np.zeros((3, 3))
@@ -1130,22 +1130,22 @@ def get_grad_S_inv(v):
     return grad_S_inv
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def get_T(X, c_table, j):
     err, B = get_B(X, c_table, j)
     if err == ERR_CODE_OK:
         v_b = get_ref_pos(X, c_table[0, j])
-        result = np.dot(B.T, X[:, j] - v_b)
+        result = B.T @ (X[:, j] - v_b)
     else:
         result = np.empty(3)
     return err, result
 
 
-@jit(nopython=True, cache=True)
+@njit(parallel=True, cache=True)
 def get_C(X, c_table):
     C = np.empty((3, c_table.shape[1]))
 
-    for j in range(C.shape[1]):
+    for j in prange(C.shape[1]):
         err, v = get_T(X, c_table, j)
         if err == ERR_CODE_OK:
             C[:, j] = get_S_inv(v)
@@ -1154,12 +1154,12 @@ def get_C(X, c_table):
     return (ERR_CODE_OK, C)
 
 
-@jit(nopython=True, cache=True)
+@njit(parallel=True, cache=True)
 def get_grad_C(X, c_table):
     n_atoms = X.shape[1]
     grad_C = np.zeros((3, n_atoms, n_atoms, 3))
 
-    for j in range(X.shape[1]):
+    for j in prange(X.shape[1]):
         IB = (X[:, j] - get_ref_pos(X, c_table[0, j])).reshape((3, 1, 1))
         grad_S_inv = get_grad_S_inv(get_T(X, c_table, j)[1])
         err, B = get_B(X, c_table, j)
@@ -1168,26 +1168,26 @@ def get_grad_C(X, c_table):
         grad_B = get_grad_B(X, c_table, j)
 
         # Derive for j
-        grad_C[:, j, j, :] = np.dot(grad_S_inv, B.T)
+        grad_C[:, j, j, :] = grad_S_inv @ B.T
 
         # Derive for b(j)
         if c_table[0, j] > constants.keys_below_are_abs_refs:
             A = np.sum(grad_B[:, :, 0, :] * IB, axis=0)
-            grad_C[:, j, c_table[0, j], :] = np.dot(grad_S_inv, A - B.T)
+            grad_C[:, j, c_table[0, j], :] = grad_S_inv @ (A - B.T)
         else:
             grad_C[:, j, c_table[0, j], :] = 0.0
 
         # Derive for a(j)
         if c_table[1, j] > constants.keys_below_are_abs_refs:
             A = np.sum(grad_B[:, :, 1, :] * IB, axis=0)
-            grad_C[:, j, c_table[1, j], :] = np.dot(grad_S_inv, A)
+            grad_C[:, j, c_table[1, j], :] = grad_S_inv @ A
         else:
             grad_C[:, j, c_table[1, j], :] = 0.0
 
         # Derive for d(j)
         if c_table[2, j] > constants.keys_below_are_abs_refs:
             A = np.sum(grad_B[:, :, 2, :] * IB, axis=0)
-            grad_C[:, j, c_table[2, j], :] = np.dot(grad_S_inv, A)
+            grad_C[:, j, c_table[2, j], :] = grad_S_inv @ A
         else:
             grad_C[:, j, c_table[2, j], :] = 0.0
     return (ERR_CODE_OK, j, grad_C)  # pylint:disable=undefined-loop-variable

src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py

@@ -6,7 +6,7 @@
 import numba as nb
 import numpy as np
 import pandas as pd
-from numba import jit
+from numba import njit
 from sortedcontainers import SortedSet
 
 import chemcoord._cartesian_coordinates.xyz_functions as xyz_functions
@@ -236,7 +236,7 @@ def __matmul__(self, other):
     def __rmatmul__(self, other):
         coords = ["x", "y", "z"]
         new = self.copy()
-        new.loc[:, coords] = (np.dot(other, new.loc[:, coords].T)).T
+        new.loc[:, coords] = (other @ new.loc[:, coords].T).T
         return new
 
     def __eq__(self, other):
@@ -298,7 +298,7 @@ def subs_function(x):
         return out
 
     @staticmethod
-    @jit(nopython=True, cache=True)
+    @njit(cache=True)
     def _jit_give_bond_array(pos, bond_radii, self_bonding_allowed=False):
         """Calculate a boolean array where ``A[i,j] is True`` indicates a
         bond between the i-th and j-th atom.
@@ -979,7 +979,7 @@ def get_without(self, fragments, use_lookup=None):
         return sorted(missing_part, key=len, reverse=True)
 
     @staticmethod
-    @jit(nopython=True, cache=True)
+    @njit(cache=True)
     def _jit_pairwise_distances(pos1, pos2):
         """Optimized function for calculating the distance between each pair
         of points in positions1 and positions2.
@@ -1120,9 +1120,9 @@ def basistransform(self, new_basis, old_basis=None, orthonormalize=True):
             is_rotation_matrix = True
 
         if is_rotation_matrix:
-            return np.dot(new_basis.T, old_basis) @ self
+            return new_basis.T @ old_basis @ self
         else:
-            return np.dot(np.linalg.inv(new_basis), old_basis) @ self
+            return np.linalg.inv(new_basis) @ old_basis @ self
 
     def _get_positions(self, indices):
         old_index = self.index

src/chemcoord/_cartesian_coordinates/asymmetric_unit_cartesian_class.py

@@ -1,4 +1,3 @@
-import numpy as np
 import pandas as pd
 
 from chemcoord._cartesian_coordinates.cartesian_class_main import Cartesian
@@ -36,5 +35,5 @@ def get_cartesian(self):
         frame.loc[self.index, coords] = self.loc[:, coords]
         for i in eq_sets:
             for j in eq_sets[i]:
-                frame.loc[j, coords] = np.dot(sym_ops[i][j], frame.loc[i, coords])
+                frame.loc[j, coords] = sym_ops[i][j] @ frame.loc[i, coords]
         return Cartesian(frame)

src/chemcoord/_cartesian_coordinates/xyz_functions.py

@@ -8,7 +8,7 @@
 import numpy as np
 import pandas as pd
 import sympy
-from numba import jit
+from numba import njit
 
 from chemcoord.configuration import settings
 
@@ -299,7 +299,7 @@ def normalize(vector):
     return normed_vector
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def _jit_normalize(vector):
     """Normalizes a vector"""
     normed_vector = vector / np.linalg.norm(vector)
@@ -327,7 +327,7 @@ def get_rotation_matrix(axis, angle):
     return _jit_get_rotation_matrix(axis, angle)
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def _jit_get_rotation_matrix(axis, angle):
     """Returns the rotation matrix.
 

src/chemcoord/_internal_coordinates/_zmat_transformation.py

@@ -8,7 +8,7 @@
 
 import numba as nb
 import numpy as np
-from numba import jit
+from numba import njit, prange
 from numpy import cos, cross, sin
 from numpy.linalg import inv
 
@@ -21,7 +21,7 @@
 from chemcoord.exceptions import ERR_CODE_OK, ERR_CODE_InvalidReference
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def get_S(C, j):
     S = np.zeros(3)
     r, alpha, delta = C[:, j]
@@ -36,7 +36,7 @@ def get_S(C, j):
     return S
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def get_grad_S(C, j):
     grad_S = np.empty((3, 3), dtype=nb.f8)
     r, alpha, delta = C[:, j]
@@ -58,7 +58,7 @@ def get_grad_S(C, j):
     return grad_S
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def get_X(C, c_table):
     X = np.empty_like(C)
     n_atoms = X.shape[1]
@@ -70,7 +70,7 @@ def get_X(C, c_table):
     return (ERR_CODE_OK, j, X)  # pylint:disable=undefined-loop-variable
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def chain_grad(X, grad_X, C, c_table, j, l):
     """Chain the gradients.
 
@@ -122,7 +122,7 @@ def chain_grad(X, grad_X, C, c_table, j, l):
     return new_grad_X
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def get_grad_X(C, c_table, chain=True):
     n_atoms = C.shape[1]
     grad_X = np.zeros((3, n_atoms, n_atoms, 3))
@@ -135,24 +135,24 @@ def get_grad_X(C, c_table, chain=True):
     return grad_X
 
 
-@jit(nopython=True, cache=True)
+@njit(nogil=True, cache=True)
 def to_barycenter(X, masses):
     M = masses.sum()
     v = (X * masses).sum(axis=1).reshape((3, 1)) / M
     return X - v
 
 
-@jit(nopython=True, cache=True)
+@njit(paralell=True, cache=True)
 def remove_translation(grad_X, masses):
     clean_grad_X = np.empty_like(grad_X)
     n_atoms = grad_X.shape[1]
-    for j in range(3):
-        for i in range(n_atoms):
+    for j in prange(3):
+        for i in prange(n_atoms):
             clean_grad_X[:, :, i, j] = to_barycenter(grad_X[:, :, i, j], masses)
     return clean_grad_X
 
 
-@jit(nopython=True, cache=True)
+@njit(parallel=True, cache=True)
 def pure_internal_grad(X, grad_X, masses, theta):
     """Return a gradient for the transformation to X
     that only contains internal degrees of freedom
@@ -172,8 +172,8 @@ def pure_internal_grad(X, grad_X, masses, theta):
     X = to_barycenter(X, masses)
     grad_X = remove_translation(grad_X, masses)
     inv_theta = inv(theta)
-    for j in range(3):
-        for i in range(n_atoms):
+    for j in prange(3):
+        for i in prange(n_atoms):
             L = (cross(X.T, grad_X[:, :, i, j].T).T * masses).sum(axis=1)
 
             grad_X[:, :, i, j] = grad_X[:, :, i, j] + cross(-inv_theta @ L, X.T).T

src/chemcoord/constants.py

@@ -15,7 +15,7 @@
 
 import numpy as np
 import pandas as pd
-from numba import jit
+from numba import njit
 
 keys_below_are_abs_refs = -sys.maxsize + 100
 int_label = {
@@ -39,7 +39,7 @@
 }
 
 
-@jit(nopython=True, cache=True)
+@njit(cache=True)
 def _jit_absolute_refs(j):
     # Because dicts are not supported in numba :(
     if j == -sys.maxsize - 1: