From 9d26a8ab724491337bb43ce6cffb089a72e71868 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 14:27:29 -0500 Subject: [PATCH 1/9] use nb.cross instead of custom jit function --- .../_cartesian_coordinates/_cart_transformation.py | 7 +++---- src/chemcoord/_cartesian_coordinates/xyz_functions.py | 10 ---------- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py index 5960976..68423e6 100644 --- a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py +++ b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py @@ -6,7 +6,6 @@ import chemcoord.constants as constants from chemcoord._cartesian_coordinates.xyz_functions import ( - _jit_cross, _jit_isclose, _jit_normalize, ) @@ -60,11 +59,11 @@ def get_B(X, c_table, j): return (ERR_CODE_InvalidReference, B) AD = ref_pos[:, 2] - ref_pos[:, 1] B[:, 2] = -_jit_normalize(BA) - N = _jit_cross(AD, BA) + N = np.cross(AD, BA) if _jit_isclose(N, 0.0).all(): return (ERR_CODE_InvalidReference, B) B[:, 1] = _jit_normalize(N) - B[:, 0] = _jit_cross(B[:, 1], B[:, 2]) + B[:, 0] = np.cross(B[:, 1], B[:, 2]) return (ERR_CODE_OK, B) @@ -77,7 +76,7 @@ def get_grad_B(X, c_table, j): x_a, y_a, z_a = v_a x_d, y_d, z_d = v_d BA, AD = v_a - v_b, v_d - v_a - norm_AD_cross_BA = np.linalg.norm(_jit_cross(AD, BA)) + norm_AD_cross_BA = np.linalg.norm(np.cross(AD, BA)) norm_BA = np.linalg.norm(BA) grad_B[0, 0, 0, 0] = ( (x_a - x_b) diff --git a/src/chemcoord/_cartesian_coordinates/xyz_functions.py b/src/chemcoord/_cartesian_coordinates/xyz_functions.py index 127fad7..0f3ab44 100644 --- a/src/chemcoord/_cartesian_coordinates/xyz_functions.py +++ b/src/chemcoord/_cartesian_coordinates/xyz_functions.py @@ -5,7 +5,6 @@ import warnings from threading import Thread -import numba as nb import numpy as np import pandas as pd import sympy @@ -309,15 +308,6 @@ def _jit_allclose(a, b, atol=1e-5, rtol=1e-8): return True -@jit(nb.f8[:](nb.f8[:], nb.f8[:]), nopython=True) -def _jit_cross(A, B): - C = np.empty_like(A) - C[0] = A[1] * B[2] - A[2] * B[1] - C[1] = A[2] * B[0] - A[0] * B[2] - C[2] = A[0] * B[1] - A[1] * B[0] - return C - - def normalize(vector): """Normalizes a vector""" normed_vector = vector / np.linalg.norm(vector) From 3fff076878fbb50add3048e519c101e55fb2d11c Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 14:31:38 -0500 Subject: [PATCH 2/9] use np.isclose instead of custom one --- .../_cartesian_coordinates/_cart_transformation.py | 9 ++++----- src/chemcoord/_cartesian_coordinates/xyz_functions.py | 5 ----- .../_internal_coordinates/_zmat_transformation.py | 5 ++--- 3 files changed, 6 insertions(+), 13 deletions(-) diff --git a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py index 68423e6..ebca3c6 100644 --- a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py +++ b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py @@ -6,7 +6,6 @@ import chemcoord.constants as constants from chemcoord._cartesian_coordinates.xyz_functions import ( - _jit_isclose, _jit_normalize, ) from chemcoord.exceptions import ERR_CODE_OK, ERR_CODE_InvalidReference @@ -55,12 +54,12 @@ def get_B(X, c_table, j): B = np.empty((3, 3)) ref_pos = get_ref_pos(X, c_table[:, j]) BA = ref_pos[:, 1] - ref_pos[:, 0] - if _jit_isclose(BA, 0.0).all(): + if np.allclose(BA, 0.0): return (ERR_CODE_InvalidReference, B) AD = ref_pos[:, 2] - ref_pos[:, 1] B[:, 2] = -_jit_normalize(BA) N = np.cross(AD, BA) - if _jit_isclose(N, 0.0).all(): + if np.allclose(N, 0.0): return (ERR_CODE_InvalidReference, B) B[:, 1] = _jit_normalize(N) B[:, 0] = np.cross(B[:, 1], B[:, 2]) @@ -1106,9 +1105,9 @@ def get_grad_S_inv(v): grad_S_inv = np.zeros((3, 3)) r = np.linalg.norm(v) - if _jit_isclose(r, 0): + if np.isclose(r, 0): pass - elif _jit_isclose(x**2 + y**2, 0): + elif np.isclose(x**2 + y**2, 0): grad_S_inv[0, 0] = 0.0 grad_S_inv[0, 1] = 0.0 grad_S_inv[0, 2] = 1 diff --git a/src/chemcoord/_cartesian_coordinates/xyz_functions.py b/src/chemcoord/_cartesian_coordinates/xyz_functions.py index 0f3ab44..19c0965 100644 --- a/src/chemcoord/_cartesian_coordinates/xyz_functions.py +++ b/src/chemcoord/_cartesian_coordinates/xyz_functions.py @@ -293,11 +293,6 @@ def concat(cartesians, ignore_index=False, keys=None): return cartesians[0].__class__(new) -@jit(nopython=True, cache=True) -def _jit_isclose(a, b, atol=1e-5, rtol=1e-8): - return np.abs(a - b) <= (atol + rtol * np.abs(b)) - - @jit(nopython=True, cache=True) def _jit_allclose(a, b, atol=1e-5, rtol=1e-8): n, m = a.shape diff --git a/src/chemcoord/_internal_coordinates/_zmat_transformation.py b/src/chemcoord/_internal_coordinates/_zmat_transformation.py index f595c85..37c5181 100644 --- a/src/chemcoord/_internal_coordinates/_zmat_transformation.py +++ b/src/chemcoord/_internal_coordinates/_zmat_transformation.py @@ -18,7 +18,6 @@ get_grad_B, get_ref_pos, ) -from chemcoord._cartesian_coordinates.xyz_functions import _jit_isclose from chemcoord.exceptions import ERR_CODE_OK, ERR_CODE_InvalidReference @@ -26,9 +25,9 @@ def get_S(C, j): S = np.zeros(3) r, alpha, delta = C[:, j] - if _jit_isclose(alpha, np.pi): + if np.isclose(alpha, np.pi): S[2] = r - elif _jit_isclose(alpha, 0): + elif np.isclose(alpha, 0): S[2] = -r else: S[0] = r * sin(alpha) * cos(delta) From 25864887a36704f46fe27c71b973560dd473c20e Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 14:32:53 -0500 Subject: [PATCH 3/9] get rid of custom allclose --- src/chemcoord/_cartesian_coordinates/xyz_functions.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/chemcoord/_cartesian_coordinates/xyz_functions.py b/src/chemcoord/_cartesian_coordinates/xyz_functions.py index 19c0965..3b6a18f 100644 --- a/src/chemcoord/_cartesian_coordinates/xyz_functions.py +++ b/src/chemcoord/_cartesian_coordinates/xyz_functions.py @@ -293,16 +293,6 @@ def concat(cartesians, ignore_index=False, keys=None): return cartesians[0].__class__(new) -@jit(nopython=True, cache=True) -def _jit_allclose(a, b, atol=1e-5, rtol=1e-8): - n, m = a.shape - for i in range(n): - for j in range(m): - if np.abs(a[i, j] - b[i, j]) > (atol + rtol * np.abs(b[i, j])): - return False - return True - - def normalize(vector): """Normalizes a vector""" normed_vector = vector / np.linalg.norm(vector) From d3d0803170444dedf66153ba0d1ffbb872e18b51 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 14:34:06 -0500 Subject: [PATCH 4/9] expose the jitted normalize function --- .../_cartesian_coordinates/_cart_transformation.py | 6 +++--- src/chemcoord/_cartesian_coordinates/xyz_functions.py | 10 ++-------- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py index ebca3c6..d887220 100644 --- a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py +++ b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py @@ -6,7 +6,7 @@ import chemcoord.constants as constants from chemcoord._cartesian_coordinates.xyz_functions import ( - _jit_normalize, + normalize, ) from chemcoord.exceptions import ERR_CODE_OK, ERR_CODE_InvalidReference @@ -57,11 +57,11 @@ def get_B(X, c_table, j): if np.allclose(BA, 0.0): return (ERR_CODE_InvalidReference, B) AD = ref_pos[:, 2] - ref_pos[:, 1] - B[:, 2] = -_jit_normalize(BA) + B[:, 2] = -normalize(BA) N = np.cross(AD, BA) if np.allclose(N, 0.0): return (ERR_CODE_InvalidReference, B) - B[:, 1] = _jit_normalize(N) + B[:, 1] = normalize(N) B[:, 0] = np.cross(B[:, 1], B[:, 2]) return (ERR_CODE_OK, B) diff --git a/src/chemcoord/_cartesian_coordinates/xyz_functions.py b/src/chemcoord/_cartesian_coordinates/xyz_functions.py index 3b6a18f..633cd9d 100644 --- a/src/chemcoord/_cartesian_coordinates/xyz_functions.py +++ b/src/chemcoord/_cartesian_coordinates/xyz_functions.py @@ -293,14 +293,8 @@ def concat(cartesians, ignore_index=False, keys=None): return cartesians[0].__class__(new) -def normalize(vector): - """Normalizes a vector""" - normed_vector = vector / np.linalg.norm(vector) - return normed_vector - - @jit(nopython=True, cache=True) -def _jit_normalize(vector): +def normalize(vector): """Normalizes a vector""" normed_vector = vector / np.linalg.norm(vector) return normed_vector @@ -342,7 +336,7 @@ def _jit_get_rotation_matrix(axis, angle): Returns: Rotation matrix (np.array): """ - axis = _jit_normalize(axis) + axis = normalize(axis) a = m.cos(angle / 2) b, c, d = axis * m.sin(angle / 2) rot_matrix = np.empty((3, 3)) From b16f05da6fe40e8beb19389c73c6e7e0470d3e97 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 14:41:55 -0500 Subject: [PATCH 5/9] use njit instead of jit(nopython=True.. --- .../_cart_transformation.py | 18 +++++++++--------- .../_cartesian_class_core.py | 6 +++--- .../_cartesian_coordinates/xyz_functions.py | 6 +++--- .../_zmat_transformation.py | 18 +++++++++--------- src/chemcoord/constants.py | 4 ++-- 5 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py index d887220..02bd41b 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 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,7 +1130,7 @@ 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: @@ -1141,7 +1141,7 @@ def get_T(X, c_table, j): return err, result -@jit(nopython=True, cache=True) +@njit(cache=True) def get_C(X, c_table): C = np.empty((3, c_table.shape[1])) @@ -1154,7 +1154,7 @@ def get_C(X, c_table): return (ERR_CODE_OK, C) -@jit(nopython=True, cache=True) +@njit(cache=True) def get_grad_C(X, c_table): n_atoms = X.shape[1] grad_C = np.zeros((3, n_atoms, n_atoms, 3)) diff --git a/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py b/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py index 2e97c08..a2ebf84 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 @@ -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. diff --git a/src/chemcoord/_cartesian_coordinates/xyz_functions.py b/src/chemcoord/_cartesian_coordinates/xyz_functions.py index 633cd9d..9ebcfad 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 @@ -293,7 +293,7 @@ def concat(cartesians, ignore_index=False, keys=None): return cartesians[0].__class__(new) -@jit(nopython=True, cache=True) +@njit(cache=True) def normalize(vector): """Normalizes a vector""" normed_vector = vector / np.linalg.norm(vector) @@ -321,7 +321,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 37c5181..46e1736 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 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,14 +135,14 @@ def get_grad_X(C, c_table, chain=True): return grad_X -@jit(nopython=True, cache=True) +@njit(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(cache=True) def remove_translation(grad_X, masses): clean_grad_X = np.empty_like(grad_X) n_atoms = grad_X.shape[1] @@ -152,7 +152,7 @@ def remove_translation(grad_X, masses): return clean_grad_X -@jit(nopython=True, cache=True) +@njit(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 diff --git a/src/chemcoord/constants.py b/src/chemcoord/constants.py index 881695a..58f16b2 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: From 977778dd1fb4b1b61a57e99ecaff917cf4ea19c9 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 14:45:09 -0500 Subject: [PATCH 6/9] Revert "expose the jitted normalize function" This reverts commit d3d0803170444dedf66153ba0d1ffbb872e18b51. --- .../_cartesian_coordinates/_cart_transformation.py | 6 +++--- src/chemcoord/_cartesian_coordinates/xyz_functions.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py index d887220..ebca3c6 100644 --- a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py +++ b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py @@ -6,7 +6,7 @@ import chemcoord.constants as constants from chemcoord._cartesian_coordinates.xyz_functions import ( - normalize, + _jit_normalize, ) from chemcoord.exceptions import ERR_CODE_OK, ERR_CODE_InvalidReference @@ -57,11 +57,11 @@ def get_B(X, c_table, j): if np.allclose(BA, 0.0): return (ERR_CODE_InvalidReference, B) AD = ref_pos[:, 2] - ref_pos[:, 1] - B[:, 2] = -normalize(BA) + B[:, 2] = -_jit_normalize(BA) N = np.cross(AD, BA) if np.allclose(N, 0.0): return (ERR_CODE_InvalidReference, B) - B[:, 1] = normalize(N) + B[:, 1] = _jit_normalize(N) B[:, 0] = np.cross(B[:, 1], B[:, 2]) return (ERR_CODE_OK, B) diff --git a/src/chemcoord/_cartesian_coordinates/xyz_functions.py b/src/chemcoord/_cartesian_coordinates/xyz_functions.py index 633cd9d..3b6a18f 100644 --- a/src/chemcoord/_cartesian_coordinates/xyz_functions.py +++ b/src/chemcoord/_cartesian_coordinates/xyz_functions.py @@ -293,13 +293,19 @@ def concat(cartesians, ignore_index=False, keys=None): return cartesians[0].__class__(new) -@jit(nopython=True, cache=True) def normalize(vector): """Normalizes a vector""" normed_vector = vector / np.linalg.norm(vector) return normed_vector +@jit(nopython=True, cache=True) +def _jit_normalize(vector): + """Normalizes a vector""" + normed_vector = vector / np.linalg.norm(vector) + return normed_vector + + def get_rotation_matrix(axis, angle): """Returns the rotation matrix. @@ -336,7 +342,7 @@ def _jit_get_rotation_matrix(axis, angle): Returns: Rotation matrix (np.array): """ - axis = normalize(axis) + axis = _jit_normalize(axis) a = m.cos(angle / 2) b, c, d = axis * m.sin(angle / 2) rot_matrix = np.empty((3, 3)) From d56642a84b45bd7c06d1927a3e54ac7911f99208 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 15:00:34 -0500 Subject: [PATCH 7/9] use @ instead of np.dot --- dev/proposed_interface_to_chemcoord.py | 3 +-- .../_cart_transformation.py | 16 ++++++++-------- .../_cartesian_class_core.py | 6 +++--- .../asymmetric_unit_cartesian_class.py | 3 +-- 4 files changed, 13 insertions(+), 15 deletions(-) diff --git a/dev/proposed_interface_to_chemcoord.py b/dev/proposed_interface_to_chemcoord.py index 5db1be3..0f23254 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 4ac6252..6ca50fd 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 njit +from numba import njit, prange from numba.extending import overload from numpy import arccos, arctan2, sqrt @@ -1135,17 +1135,17 @@ 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 -@njit(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) @@ -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 a2ebf84..a7854bf 100644 --- a/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py +++ b/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py @@ -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): @@ -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 5494e0f..b1d5111 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) From 616a3b4a2bb60537e6f8988d5a000b26b3e8cb74 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 15:01:56 -0500 Subject: [PATCH 8/9] replaced a few range -> prange --- src/chemcoord/_cartesian_coordinates/_cart_transformation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py index 6ca50fd..985047a 100644 --- a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py +++ b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py @@ -1154,12 +1154,12 @@ def get_C(X, c_table): return (ERR_CODE_OK, C) -@njit(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) From 9cabe1b11ee637613ca2b45e2cf02c8915199a76 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 13 Jan 2025 15:04:20 -0500 Subject: [PATCH 9/9] replaced a couple more range -> prange --- .../_zmat_transformation.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/chemcoord/_internal_coordinates/_zmat_transformation.py b/src/chemcoord/_internal_coordinates/_zmat_transformation.py index 46e1736..b708a06 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 njit +from numba import njit, prange from numpy import cos, cross, sin from numpy.linalg import inv @@ -135,24 +135,24 @@ def get_grad_X(C, c_table, chain=True): return grad_X -@njit(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 -@njit(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 -@njit(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