diff --git a/.circleci/config.yml b/.circleci/config.yml index e5114f7..89ac8b3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -36,10 +36,6 @@ run_tests: &run_tests - run: name: Running tests command: python3 -m pytest -Werror --cov=./src/chemcoord tests/ - - run: - name: Upload coverage reports to Codecov - command: | - bash <(curl -s https://codecov.io/bash) - run: name: Prepare documentation command: pip3 install -r docs/requirements.txt diff --git a/.gitignore b/.gitignore index 66f2760..26299c7 100644 --- a/.gitignore +++ b/.gitignore @@ -66,3 +66,6 @@ cealign-0.8-RBS #pytest *.chache/* .cache/ + + +.mypy_cache/ diff --git a/setup.py b/setup.py index 931142e..40cbdb9 100644 --- a/setup.py +++ b/setup.py @@ -17,6 +17,7 @@ "sympy", "six", "pymatgen", + "typing_extensions", ] KEYWORDS = [ "chemcoord", diff --git a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py index 6cfcd7b..3c62ee7 100644 --- a/src/chemcoord/_cartesian_coordinates/_cart_transformation.py +++ b/src/chemcoord/_cartesian_coordinates/_cart_transformation.py @@ -1,6 +1,5 @@ import numba as nb import numpy as np -from numba import njit from numba.extending import overload from numpy import arccos, arctan2, sqrt @@ -8,6 +7,7 @@ from chemcoord._cartesian_coordinates.xyz_functions import ( _jit_normalize, ) +from chemcoord._utilities._decorators import njit from chemcoord.exceptions import ERR_CODE_OK, ERR_CODE_InvalidReference @@ -44,12 +44,12 @@ def f(X, indices): raise AssertionError("Should not be here") -@njit(cache=True) +@njit def get_ref_pos(X, indices): return _stub_get_ref_pos(X, indices) -@njit(cache=True) +@njit 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) -@njit(cache=True) +@njit 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]) @@ -1130,7 +1130,7 @@ def get_grad_S_inv(v): return grad_S_inv -@njit(cache=True) +@njit 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 -@njit(cache=True) +@njit 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) -@njit(cache=True) +@njit 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 642d393..7742451 100644 --- a/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py +++ b/src/chemcoord/_cartesian_coordinates/_cartesian_class_core.py @@ -1,13 +1,16 @@ import collections import copy import itertools +from collections.abc import Sequence from itertools import product +from typing import Any, Union import numba as nb import numpy as np import pandas as pd -from numba import njit +from pandas import DataFrame, Series from sortedcontainers import SortedSet +from typing_extensions import Self import chemcoord._cartesian_coordinates.xyz_functions as xyz_functions import chemcoord.constants as constants @@ -15,8 +18,10 @@ PandasWrapper, ) from chemcoord._generic_classes.generic_core import GenericCore +from chemcoord._utilities._decorators import njit +from chemcoord._utilities.typing import ArithmeticOther, Axes, Matrix, Vector from chemcoord.configuration import settings -from chemcoord.exceptions import IllegalArgumentCombination, PhysicalMeaning +from chemcoord.exceptions import PhysicalMeaning class CartesianCore(PandasWrapper, GenericCore): # noqa: PLW1641 @@ -29,13 +34,10 @@ class CartesianCore(PandasWrapper, GenericCore): # noqa: PLW1641 # overwrites existing method def __init__( self, - frame=None, - atoms=None, - coords=None, - index=None, - metadata=None, - _metadata=None, - ): + frame: DataFrame, + metadata: Union[dict, None] = None, + _metadata: Union[dict, None] = None, + ) -> None: """How to initialize a Cartesian instance. Args: @@ -51,21 +53,8 @@ def __init__( Returns: Cartesian: A new cartesian instance. """ - if bool(atoms is None and coords is None) == bool( - atoms is not None and coords is not None - ): - message = "atoms and coords have to be both None or not None" - raise IllegalArgumentCombination(message) - elif frame is None and atoms is None and coords is None: - message = "Either frame or atoms and coords have to be not None" - raise IllegalArgumentCombination(message) - elif atoms is not None and coords is not None: - dtypes = [("atom", str), ("x", float), ("y", float), ("z", float)] - frame = pd.DataFrame(np.empty(len(atoms), dtype=dtypes), index=index) - frame["atom"] = atoms - frame.loc[:, ["x", "y", "z"]] = coords - elif not isinstance(frame, pd.DataFrame): - raise ValueError("Need a pd.DataFrame as input") + if not isinstance(frame, DataFrame): + raise TypeError("frame has to be a pandas DataFrame") if not self._required_cols <= set(frame.columns): raise PhysicalMeaning( "There are columns missing for a meaningful description of a molecule" @@ -81,15 +70,30 @@ def __init__( else: self._metadata = copy.deepcopy(_metadata) - def _return_appropiate_type(self, selected): - if isinstance(selected, pd.Series): - frame = pd.DataFrame(selected).T + @classmethod + def set_atom_coords( + cls, + atoms: Sequence[str], + coords: Matrix, + index: Union[Axes, None] = None, + ) -> Self: + dtypes = [("atom", str), ("x", float), ("y", float), ("z", float)] + frame = DataFrame(np.empty(len(atoms), dtype=dtypes), index=index) + frame["atom"] = atoms + frame.loc[:, ["x", "y", "z"]] = coords + return cls(frame) + + def _return_appropiate_type( + self, selected: Union[Series, DataFrame] + ) -> Union[Self, Series, DataFrame]: + if isinstance(selected, Series): + frame = DataFrame(selected).T if self._required_cols <= set(frame.columns): selected = frame.apply(pd.to_numeric, errors="ignore") else: return selected - if isinstance(selected, pd.DataFrame) and self._required_cols <= set( + if isinstance(selected, DataFrame) and self._required_cols <= set( selected.columns ): molecule = self.__class__(selected) @@ -99,7 +103,7 @@ def _return_appropiate_type(self, selected): else: return selected - def _test_if_can_be_added(self, other): + def _test_if_can_be_added(self, other: Self) -> None: if not ( set(self.index) == set(other.index) and (self["atom"] == other.loc[self.index, "atom"]).all(axis=None) @@ -110,13 +114,13 @@ def _test_if_can_be_added(self, other): ) raise PhysicalMeaning(message) - def __add__(self, other): + def __add__(self, other: Union[Self, ArithmeticOther]) -> Self: coords = ["x", "y", "z"] new = self.copy() - if isinstance(other, CartesianCore): + if isinstance(other, self.__class__): self._test_if_can_be_added(other) new.loc[:, coords] = self.loc[:, coords] + other.loc[:, coords] - elif isinstance(other, pd.DataFrame): + elif isinstance(other, DataFrame): new.loc[:, coords] = self.loc[:, coords] + other.loc[:, coords] else: try: @@ -126,16 +130,16 @@ def __add__(self, other): new.loc[:, coords] = self.loc[:, coords] + other return new - def __radd__(self, other): + def __radd__(self, other: Union[Self, ArithmeticOther]) -> Self: return self.__add__(other) - def __sub__(self, other): + def __sub__(self, other: Union[Self, ArithmeticOther]) -> Self: coords = ["x", "y", "z"] new = self.copy() - if isinstance(other, CartesianCore): + if isinstance(other, self.__class__): self._test_if_can_be_added(other) new.loc[:, coords] = self.loc[:, coords] - other.loc[:, coords] - elif isinstance(other, pd.DataFrame): + elif isinstance(other, DataFrame): new.loc[:, coords] = self.loc[:, coords] - other.loc[:, coords] else: try: @@ -145,13 +149,13 @@ def __sub__(self, other): new.loc[:, coords] = self.loc[:, coords] - other return new - def __rsub__(self, other): + def __rsub__(self, other: Union[Self, ArithmeticOther]) -> Self: coords = ["x", "y", "z"] new = self.copy() - if isinstance(other, CartesianCore): + if isinstance(other, self.__class__): self._test_if_can_be_added(other) new.loc[:, coords] = other.loc[:, coords] - self.loc[:, coords] - elif isinstance(other, pd.DataFrame): + elif isinstance(other, DataFrame): new.loc[:, coords] = other.loc[:, coords] - self.loc[:, coords] else: try: @@ -161,13 +165,13 @@ def __rsub__(self, other): new.loc[:, coords] = other - self.loc[:, coords] return new - def __mul__(self, other): + def __mul__(self, other: Union[Self, ArithmeticOther]) -> Self: coords = ["x", "y", "z"] new = self.copy() - if isinstance(other, CartesianCore): + if isinstance(other, self.__class__): self._test_if_can_be_added(other) new.loc[:, coords] = self.loc[:, coords] * other.loc[:, coords] - elif isinstance(other, pd.DataFrame): + elif isinstance(other, DataFrame): new.loc[:, coords] = self.loc[:, coords] * other.loc[:, coords] else: try: @@ -177,16 +181,16 @@ def __mul__(self, other): new.loc[:, coords] = self.loc[:, coords] * other return new - def __rmul__(self, other): + def __rmul__(self, other: Union[Self, ArithmeticOther]) -> Self: return self.__mul__(other) - def __truediv__(self, other): + def __truediv__(self, other: Union[Self, ArithmeticOther]) -> Self: coords = ["x", "y", "z"] new = self.copy() - if isinstance(other, CartesianCore): + if isinstance(other, self.__class__): self._test_if_can_be_added(other) new.loc[:, coords] = self.loc[:, coords] / other.loc[:, coords] - elif isinstance(other, pd.DataFrame): + elif isinstance(other, DataFrame): new.loc[:, coords] = self.loc[:, coords] / other.loc[:, coords] else: try: @@ -196,13 +200,13 @@ def __truediv__(self, other): new.loc[:, coords] = self.loc[:, coords] / other return new - def __rtruediv__(self, other): + def __rtruediv__(self, other: Union[Self, ArithmeticOther]) -> Self: coords = ["x", "y", "z"] new = self.copy() - if isinstance(other, CartesianCore): + if isinstance(other, self.__class__): self._test_if_can_be_added(other) new.loc[:, coords] = other.loc[:, coords] / self.loc[:, coords] - elif isinstance(other, pd.DataFrame): + elif isinstance(other, DataFrame): new.loc[:, coords] = other.loc[:, coords] / self.loc[:, coords] else: try: @@ -212,46 +216,49 @@ def __rtruediv__(self, other): new.loc[:, coords] = other / self.loc[:, coords] return new - def __pow__(self, other): + def __pow__(self, other: Union[Self, ArithmeticOther]) -> Self: coords = ["x", "y", "z"] new = self.copy() new.loc[:, coords] = self.loc[:, coords] ** other return new - def __pos__(self): + def __pos__(self) -> Self: return self.copy() - def __neg__(self): + def __neg__(self) -> Self: return -1 * self.copy() - def __abs__(self): + def __abs__(self) -> Self: coords = ["x", "y", "z"] new = self.copy() new.loc[:, coords] = abs(new.loc[:, coords]) return new - def __matmul__(self, other): + def __matmul__(self, other: Matrix) -> Self: return NotImplemented - def __rmatmul__(self, other): + def __rmatmul__(self, other: Matrix) -> Self: coords = ["x", "y", "z"] new = self.copy() new.loc[:, coords] = (np.dot(other, new.loc[:, coords].T)).T return new - def __eq__(self, other): + # Somehow the base class `object` expects the return type to be `bool`` + # but the correct type hint is DataFrame. + # Ignore this override error. + def __eq__(self, other: Self) -> DataFrame: # type: ignore[override] return self._frame == other._frame - def __ne__(self, other): + def __ne__(self, other: Self) -> DataFrame: # type: ignore[override] return self._frame != other._frame - def copy(self): + def copy(self) -> Self: molecule = self.__class__(self._frame) molecule.metadata = self.metadata.copy() molecule._metadata = copy.deepcopy(self._metadata) return molecule - def subs(self, *args): + def subs(self, *args: Any) -> Self: """Substitute a symbolic expression in ``['x', 'y', 'z']`` This is a wrapper around the substitution mechanism of @@ -298,8 +305,12 @@ def subs_function(x): return out @staticmethod - @njit(cache=True) - def _jit_give_bond_array(pos, bond_radii, self_bonding_allowed=False): + @njit + def _jit_give_bond_array( + pos: Matrix[np.floating], + bond_radii: Vector[np.floating], + self_bonding_allowed: bool = False, + ) -> Matrix[np.float64]: """Calculate a boolean array where ``A[i,j] is True`` indicates a bond between the i-th and j-th atom. """ @@ -452,7 +463,7 @@ def complete_calculation(): data = self.add_data([atomic_radius_data, "valency"]) bond_radii = data[atomic_radius_data] if modified_properties is not None: - bond_radii.update(pd.Series(modified_properties)) + bond_radii.update(Series(modified_properties)) bond_radii = bond_radii.values bond_dict = collections.defaultdict(set) for i, j, k in product(*[range(x) for x in fragments.shape]): @@ -742,7 +753,7 @@ def get_bond_lengths(self, indices): :class:`numpy.ndarray`: Vector of angles in degrees. """ coords = ["x", "y", "z"] - if isinstance(indices, pd.DataFrame): + if isinstance(indices, DataFrame): i_pos = self.loc[indices.index, coords].values b_pos = self.loc[indices.loc[:, "b"], coords].values else: @@ -772,7 +783,7 @@ def get_angle_degrees(self, indices): :class:`numpy.ndarray`: Vector of angles in degrees. """ coords = ["x", "y", "z"] - if isinstance(indices, pd.DataFrame): + if isinstance(indices, DataFrame): i_pos = self.loc[indices.index, coords].values b_pos = self.loc[indices.loc[:, "b"], coords].values a_pos = self.loc[indices.loc[:, "a"], coords].values @@ -812,7 +823,7 @@ def get_dihedral_degrees(self, indices, start_row=0): :class:`numpy.ndarray`: Vector of angles in degrees. """ coords = ["x", "y", "z"] - if isinstance(indices, pd.DataFrame): + if isinstance(indices, DataFrame): i_pos = self.loc[indices.index, coords].values b_pos = self.loc[indices.loc[:, "b"], coords].values a_pos = self.loc[indices.loc[:, "a"], coords].values @@ -982,7 +993,7 @@ def get_without(self, fragments, use_lookup=None): return sorted(missing_part, key=len, reverse=True) @staticmethod - @njit(cache=True) + @njit 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/_cartesian_class_io.py b/src/chemcoord/_cartesian_coordinates/_cartesian_class_io.py index 285bb32..0f54fba 100644 --- a/src/chemcoord/_cartesian_coordinates/_cartesian_class_io.py +++ b/src/chemcoord/_cartesian_coordinates/_cartesian_class_io.py @@ -277,7 +277,7 @@ def from_pyscf_molecule(cls, mol): Returns: Cartesian: """ - return cls( + return cls.set_atom_coords( atoms=mol.elements, coords=mol.atom_coords(unit="Angstrom"), ) @@ -458,7 +458,7 @@ def from_pymatgen_molecule(cls, molecule): Returns: Cartesian: """ - return cls( + return cls.set_atom_coords( atoms=[el.value for el in molecule.species], coords=molecule.cart_coords ) diff --git a/src/chemcoord/_cartesian_coordinates/_cartesian_class_pandas_wrapper.py b/src/chemcoord/_cartesian_coordinates/_cartesian_class_pandas_wrapper.py index ccfb40c..d622966 100644 --- a/src/chemcoord/_cartesian_coordinates/_cartesian_class_pandas_wrapper.py +++ b/src/chemcoord/_cartesian_coordinates/_cartesian_class_pandas_wrapper.py @@ -4,7 +4,7 @@ from chemcoord.exceptions import PhysicalMeaning -class PandasWrapper(object): +class PandasWrapper: """This class provides wrappers for :class:`pandas.DataFrame` methods. It has the same behaviour as the :class:`~pandas.DataFrame` diff --git a/src/chemcoord/_cartesian_coordinates/_indexers.py b/src/chemcoord/_cartesian_coordinates/_indexers.py index 46ea499..2044807 100644 --- a/src/chemcoord/_cartesian_coordinates/_indexers.py +++ b/src/chemcoord/_cartesian_coordinates/_indexers.py @@ -3,7 +3,7 @@ from chemcoord._utilities._temporary_deprecation_workarounds import is_iterable -class _generic_Indexer(object): +class _generic_Indexer: def __init__(self, molecule): self.molecule = molecule diff --git a/src/chemcoord/_cartesian_coordinates/py.typed b/src/chemcoord/_cartesian_coordinates/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/src/chemcoord/_cartesian_coordinates/xyz_functions.py b/src/chemcoord/_cartesian_coordinates/xyz_functions.py index 353a1dc..1e6f7f7 100644 --- a/src/chemcoord/_cartesian_coordinates/xyz_functions.py +++ b/src/chemcoord/_cartesian_coordinates/xyz_functions.py @@ -8,8 +8,8 @@ import numpy as np import pandas as pd import sympy -from numba import njit +from chemcoord._utilities._decorators import njit from chemcoord.configuration import settings @@ -299,7 +299,7 @@ def normalize(vector): return normed_vector -@njit(cache=True) +@njit 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) -@njit(cache=True) +@njit def _jit_get_rotation_matrix(axis, angle): """Returns the rotation matrix. diff --git a/src/chemcoord/_generic_classes/generic_IO.py b/src/chemcoord/_generic_classes/generic_IO.py index 633a840..b847751 100644 --- a/src/chemcoord/_generic_classes/generic_IO.py +++ b/src/chemcoord/_generic_classes/generic_IO.py @@ -2,7 +2,7 @@ import sympy -class GenericIO(object): +class GenericIO: def _sympy_formatter(self): def formatter(x): if isinstance(x, sympy.Basic): diff --git a/src/chemcoord/_generic_classes/generic_core.py b/src/chemcoord/_generic_classes/generic_core.py index 1c8914b..a2a4fb1 100644 --- a/src/chemcoord/_generic_classes/generic_core.py +++ b/src/chemcoord/_generic_classes/generic_core.py @@ -3,7 +3,7 @@ import chemcoord.constants as constants -class GenericCore(object): +class GenericCore: def add_data(self, new_cols=None): """Adds a column with the requested data. diff --git a/src/chemcoord/_internal_coordinates/_indexers.py b/src/chemcoord/_internal_coordinates/_indexers.py index 9184d95..26cb2ee 100644 --- a/src/chemcoord/_internal_coordinates/_indexers.py +++ b/src/chemcoord/_internal_coordinates/_indexers.py @@ -4,7 +4,7 @@ from chemcoord.exceptions import InvalidReference -class _generic_Indexer(object): +class _generic_Indexer: def __init__(self, molecule): self.molecule = molecule diff --git a/src/chemcoord/_internal_coordinates/_zmat_class_pandas_wrapper.py b/src/chemcoord/_internal_coordinates/_zmat_class_pandas_wrapper.py index 7f95c7b..3b38011 100644 --- a/src/chemcoord/_internal_coordinates/_zmat_class_pandas_wrapper.py +++ b/src/chemcoord/_internal_coordinates/_zmat_class_pandas_wrapper.py @@ -1,4 +1,4 @@ -class PandasWrapper(object): +class PandasWrapper: """This class provides wrappers for :class:`pandas.DataFrame` methods. It has the same behaviour as the :class:`~pandas.DataFrame` diff --git a/src/chemcoord/_internal_coordinates/_zmat_transformation.py b/src/chemcoord/_internal_coordinates/_zmat_transformation.py index 8ee0f4b..29d6a18 100644 --- a/src/chemcoord/_internal_coordinates/_zmat_transformation.py +++ b/src/chemcoord/_internal_coordinates/_zmat_transformation.py @@ -1,6 +1,6 @@ import numba as nb import numpy as np -from numba import njit, prange +from numba import prange from numpy import cos, cross, sin from numpy.linalg import inv @@ -10,10 +10,11 @@ get_grad_B, get_ref_pos, ) +from chemcoord._utilities._decorators import njit from chemcoord.exceptions import ERR_CODE_OK, ERR_CODE_InvalidReference -@njit(cache=True) +@njit def get_S(C, j): S = np.zeros(3) r, alpha, delta = C[:, j] @@ -28,7 +29,7 @@ def get_S(C, j): return S -@njit(cache=True) +@njit def get_grad_S(C, j): grad_S = np.empty((3, 3), dtype=nb.f8) r, alpha, delta = C[:, j] @@ -50,7 +51,7 @@ def get_grad_S(C, j): return grad_S -@njit(cache=True) +@njit def get_X(C, c_table): X = np.empty_like(C) n_atoms = X.shape[1] @@ -62,7 +63,7 @@ def get_X(C, c_table): return (ERR_CODE_OK, j, X) # pylint:disable=undefined-loop-variable -@njit(cache=True) +@njit def chain_grad(X, grad_X, C, c_table, j, l): """Chain the gradients. @@ -114,7 +115,7 @@ def chain_grad(X, grad_X, C, c_table, j, l): return new_grad_X -@njit(cache=True) +@njit def get_grad_X(C, c_table, chain=True): n_atoms = C.shape[1] grad_X = np.zeros((3, n_atoms, n_atoms, 3)) @@ -127,14 +128,14 @@ def get_grad_X(C, c_table, chain=True): return grad_X -@njit(cache=True) +@njit def to_barycenter(X, masses): M = masses.sum() v = (X * masses).sum(axis=1).reshape((3, 1)) / M return X - v -@njit(parallel=True, cache=True) +@njit(parallel=True) def remove_translation(grad_X, masses): clean_grad_X = np.empty_like(grad_X) n_atoms = grad_X.shape[1] @@ -144,7 +145,7 @@ def remove_translation(grad_X, masses): return clean_grad_X -@njit(parallel=True, cache=True) +@njit(parallel=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/_internal_coordinates/zmat_functions.py b/src/chemcoord/_internal_coordinates/zmat_functions.py index b1e9df7..5518902 100644 --- a/src/chemcoord/_internal_coordinates/zmat_functions.py +++ b/src/chemcoord/_internal_coordinates/zmat_functions.py @@ -4,7 +4,7 @@ from chemcoord._internal_coordinates.zmat_class_main import Zmat -class DummyManipulation(object): +class DummyManipulation: """Contextmanager that controls the behaviour of :meth:`~chemcoord.Zmat.safe_loc` and :meth:`~chemcoord.Zmat.safe_iloc`. @@ -41,7 +41,7 @@ def __exit__(self, exc_type, exc_value, traceback): self.cls.dummy_manipulation_allowed = self.old_value -class TestOperators(object): +class TestOperators: """Switch the validity testing of zmatrices resulting from operators. The following examples is done with ``+`` @@ -67,7 +67,7 @@ def __exit__(self, exc_type, exc_value, traceback): self.cls.test_operators = self.old_value -class PureInternalMovement(object): +class PureInternalMovement: """Remove the translational and rotational degrees of freedom. When doing assignments to the z-matrix:: @@ -122,4 +122,6 @@ def apply_grad_cartesian_tensor(grad_X, zmat_dist): Cartesian, ) - return Cartesian(atoms=zmat_dist["atom"], coords=cart_dist, index=zmat_dist.index) + return Cartesian.set_atom_coords( + atoms=zmat_dist["atom"], coords=cart_dist, index=zmat_dist.index + ) diff --git a/src/chemcoord/_utilities/_decorators.py b/src/chemcoord/_utilities/_decorators.py index 42dd7af..d9724dc 100644 --- a/src/chemcoord/_utilities/_decorators.py +++ b/src/chemcoord/_utilities/_decorators.py @@ -1,13 +1,17 @@ # The following code was taken from the MIT licensed pandas project # and modified. http://pandas.pydata.org/ +from collections.abc import Callable from textwrap import dedent +from typing import TypeVar, Union, overload + +import numba as nb # # Substitution and Appender are derived from matplotlib.docstring (1.1.0) # # module http://matplotlib.org/users/license.html # -class Substitution(object): +class Substitution: """ A decorator to take a function's docstring and perform string substitution on it. @@ -63,7 +67,7 @@ def from_params(cls, params): return result -class Appender(object): +class Appender: """ A function decorator that will append an addendum to the docstring of the target function. @@ -103,3 +107,40 @@ def indent(text, indents=1): return "" jointext = "".join(["\n"] + [" "] * indents) return jointext.join(text.split("\n")) + + +Function = TypeVar("Function", bound=Callable) + + +@overload +def njit(f: Function, **kwargs) -> Function: ... +@overload +def njit(**kwargs) -> Callable[[Function], Function]: ... + + +def njit( + f: Union[Function, None] = None, **kwargs +) -> Union[Function, Callable[[Function], Function]]: + """Type-safe jit wrapper that caches the compiled function + + With this jit wrapper, you can actually use static typing together with numba. + The crucial declaration is that the decorated function's interface is preserved, + i.e. mapping :class:`Function` to :class:`Function`. + Otherwise the following example would not raise a type error: + + .. code-block:: python + + @numba.njit + def f(x: int) -> int: + return x + + f(2.0) # No type error + + While the same example, using this custom :func:`njit` would raise a type error. + + In addition to type safety, this wrapper also sets :code:`cache=True` by default. + """ + if f is None: + return nb.njit(cache=True, **kwargs) + else: + return nb.njit(f, cache=True, **kwargs) diff --git a/src/chemcoord/_utilities/typing.py b/src/chemcoord/_utilities/typing.py new file mode 100644 index 0000000..a15c159 --- /dev/null +++ b/src/chemcoord/_utilities/typing.py @@ -0,0 +1,59 @@ +"""Define some types that do not fit into one particular module + +In particular it enables barebone typechecking for the shape of numpy arrays + +Inspired by +https://stackoverflow.com/questions/75495212/type-hinting-numpy-arrays-and-batches + +Note that most numpy functions return :python:`ndarray[Any, Any]` +i.e. the type is mostly useful to document intent to the developer. +""" + +import os +from collections.abc import Sequence +from typing import Any, Dict, Tuple, TypeVar, Union + +import numpy as np + +# Reexpose some pandas types +from pandas._typing import Axes # noqa: F401 +from typing_extensions import TypeAlias + +# We want the dtype to behave covariant, i.e. if a +# Vector[float] is allowed, then the more specific +# Vector[float64] should also be allowed. +# Also see here: +# https://stackoverflow.com/questions/61568462/what-does-typevara-b-covariant-true-mean +#: Type annotation of a generic covariant type. +T_dtype_co = TypeVar("T_dtype_co", bound=np.generic, covariant=True) + +# Currently we can define :code:`Matrix` and higher order tensors +# only with shape :code`Tuple[int, ...]` because of +# https://github.com/numpy/numpy/issues/27957 +# make the typechecks more strict over time, when shape checking finally comes to numpy. + +#: Type annotation of a vector. +Vector = np.ndarray[Tuple[int], np.dtype[T_dtype_co]] +#: Type annotation of a matrix. +Matrix = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] +#: Type annotation of a tensor. +Tensor3D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] +#: Type annotation of a tensor. +Tensor4D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] +#: Type annotation of a tensor. +Tensor5D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] +#: Type annotation of a tensor. +Tensor6D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] +#: Type annotation of a tensor. +Tensor7D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] +#: Type annotation of a tensor. +Tensor = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] + +#: Type annotation for pathlike objects. +PathLike: TypeAlias = Union[str, os.PathLike] +#: Type annotation for dictionaries holding keyword arguments. +KwargDict: TypeAlias = Dict[str, Any] + +ArithmeticOther = Union[ + float, np.floating, Sequence, Sequence[Sequence], Vector, Matrix +] diff --git a/src/chemcoord/constants.py b/src/chemcoord/constants.py index 58f16b2..59948e3 100644 --- a/src/chemcoord/constants.py +++ b/src/chemcoord/constants.py @@ -15,7 +15,8 @@ import numpy as np import pandas as pd -from numba import njit + +from chemcoord._utilities._decorators import njit keys_below_are_abs_refs = -sys.maxsize + 100 int_label = { @@ -39,7 +40,7 @@ } -@njit(cache=True) +@njit def _jit_absolute_refs(j): # Because dicts are not supported in numba :( if j == -sys.maxsize - 1: diff --git a/src/chemcoord/py.typed b/src/chemcoord/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/tests/cartesian_coordinates/Cartesian/test_core.py b/tests/cartesian_coordinates/Cartesian/test_core.py index 6584bb0..32bf9d7 100644 --- a/tests/cartesian_coordinates/Cartesian/test_core.py +++ b/tests/cartesian_coordinates/Cartesian/test_core.py @@ -96,7 +96,7 @@ def get_complete_path(structure): def test_init(): - with pytest.raises(ValueError): + with pytest.raises(TypeError): cc.Cartesian(5) with pytest.raises(PhysicalMeaning): cc.Cartesian(molecule.loc[:, ["atom", "x"]])