Skip to content

Commit

Permalink
Use numpy functions (#116)
Browse files Browse the repository at this point in the history
There are a few numpy functions such as `numpy.cross` which were not available for numba when I wrote the package.
With this PR the corresponding numpy functions are used instead of the custom ones.

Fixes #113
  • Loading branch information
mcocdawc authored Jan 13, 2025
1 parent 9165c76 commit eae7421
Show file tree
Hide file tree
Showing 3 changed files with 9 additions and 37 deletions.
16 changes: 7 additions & 9 deletions src/chemcoord/_cartesian_coordinates/_cart_transformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,6 @@

import chemcoord.constants as constants
from chemcoord._cartesian_coordinates.xyz_functions import (
_jit_cross,
_jit_isclose,
_jit_normalize,
)
from chemcoord.exceptions import ERR_CODE_OK, ERR_CODE_InvalidReference
Expand Down Expand Up @@ -56,15 +54,15 @@ 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 = _jit_cross(AD, BA)
if _jit_isclose(N, 0.0).all():
N = np.cross(AD, BA)
if np.allclose(N, 0.0):
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)


Expand All @@ -77,7 +75,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)
Expand Down Expand Up @@ -1107,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
Expand Down
25 changes: 0 additions & 25 deletions src/chemcoord/_cartesian_coordinates/xyz_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import warnings
from threading import Thread

import numba as nb
import numpy as np
import pandas as pd
import sympy
Expand Down Expand Up @@ -294,30 +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
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


@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)
Expand Down
5 changes: 2 additions & 3 deletions src/chemcoord/_internal_coordinates/_zmat_transformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,16 @@
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


@jit(nopython=True, cache=True)
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)
Expand Down

0 comments on commit eae7421

Please sign in to comment.