diff --git a/dev/proposed_interface_to_chemcoord.py b/dev/proposed_interface_to_chemcoord.py index 5db1be32..0f23254e 100644 --- a/dev/proposed_interface_to_chemcoord.py +++ b/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 diff --git a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py index ebca3c61..985047a5 100644 --- a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py +++ b/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 diff --git a/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py b/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py index 2e97c08d..a7854bf5 100644 --- a/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py +++ b/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 diff --git a/src/chemcoord/_cartesian_coordinates/asymmetric_unit_cartesian_class.py b/src/chemcoord/_cartesian_coordinates/asymmetric_unit_cartesian_class.py index 5494e0fa..b1d51112 100644 --- a/src/chemcoord/_cartesian_coordinates/asymmetric_unit_cartesian_class.py +++ b/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) diff --git a/src/chemcoord/_cartesian_coordinates/xyz_functions.py b/src/chemcoord/_cartesian_coordinates/xyz_functions.py index 3b6a18f9..353a1dc9 100644 --- a/src/chemcoord/_cartesian_coordinates/xyz_functions.py +++ b/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. diff --git a/src/chemcoord/_internal_coordinates/_zmat_transformation.py b/src/chemcoord/_internal_coordinates/_zmat_transformation.py index 37c51812..b708a06f 100644 --- a/src/chemcoord/_internal_coordinates/_zmat_transformation.py +++ b/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 diff --git a/src/chemcoord/constants.py b/src/chemcoord/constants.py index 881695ab..58f16b27 100644 --- a/src/chemcoord/constants.py +++ b/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: