diff --git a/autode/opt/coordinates/autodiff.py b/autode/opt/coordinates/autodiff.py index ef18c27f2..dc7d51ca3 100644 --- a/autode/opt/coordinates/autodiff.py +++ b/autode/opt/coordinates/autodiff.py @@ -2,6 +2,7 @@ Automatic differentiation routines in pure Python """ from typing import Union, Callable, Sequence, Optional +from enum import Enum from copy import deepcopy import numpy as np import math @@ -10,10 +11,18 @@ numeric_type = Union[float, int] +class DerivativeOrder(Enum): + """Order of derivative""" + + zeroth = 0 + first = 1 + second = 2 + + def get_differentiable_vars( values: Sequence[numeric_type], symbols: Sequence[str], - deriv_order: int = 2, + deriv_order: DerivativeOrder = DerivativeOrder.second, ): """ Obtain differentiable variables from a series of numbers @@ -74,7 +83,7 @@ def __init__( # load the derivatives with sanity checks self._first_der: Optional[np.ndarray] = None self._second_der: Optional[np.ndarray] = None - self._order = 0 + self._order = DerivativeOrder.zeroth if first_der is None: return assert isinstance(first_der, np.ndarray) @@ -85,7 +94,7 @@ def __init__( f" shape of derivative array {first_der.shape}" ) self._first_der = first_der.astype(float) - self._order = 1 + self._order = DerivativeOrder.first if second_der is None: return @@ -96,7 +105,7 @@ def __init__( f" shape of second derivative matrix {first_der.shape}" ) self._second_der = second_der.astype(float) - self._order = 2 + self._order = DerivativeOrder.second @property def n_vars(self) -> int: @@ -132,7 +141,11 @@ def value(self, value): @classmethod def from_variable( - cls, value: float, symbol: str, all_symbols: Sequence[str], order: int + cls, + value: float, + symbol: str, + all_symbols: Sequence[str], + order: DerivativeOrder, ): """ Create a hyper-dual number from one variable, requires @@ -159,10 +172,10 @@ def from_variable( n = len(all_symbols) idx = list(all_symbols).index(symbol) - if order >= 1: + if order == DerivativeOrder.first or order == DerivativeOrder.second: first_der = np.zeros(shape=n, dtype=float) first_der[idx] = 1.0 - if order == 2: + if order == DerivativeOrder.second: second_der = np.zeros(shape=(n, n), dtype=float) return VectorHyperDual(val, all_symbols, first_der, second_der) @@ -187,7 +200,7 @@ def differentiate_wrt( if symbol1 not in self._symbols: return None - if self._order == 0: + if self._order == DerivativeOrder.zeroth: return None idx_1 = self._symbols.index(symbol1) @@ -200,7 +213,7 @@ def differentiate_wrt( return None idx_2 = self._symbols.index(symbol2) # check if second derivs are available - if self._order == 1: + if self._order == DerivativeOrder.first: return None assert self._second_der is not None return self._second_der[idx_1, idx_2] @@ -221,13 +234,13 @@ def __add__( self._check_compatible(other) val = self._val + other._val - if self._order == 0: + if self._order == DerivativeOrder.zeroth: return VectorHyperDual(val, self._symbols) assert self._first_der is not None assert other._first_der is not None first_der = self._first_der + other._first_der - if self._order == 1: + if self._order == DerivativeOrder.first: return VectorHyperDual(val, self._symbols, first_der) assert self._second_der is not None @@ -246,9 +259,10 @@ def __neg__(self) -> "VectorHyperDual": """Unary negative operation""" new = self.copy() new._val = -new._val - if self._order >= 1: + if self._order == DerivativeOrder.first: + new._first_der = -new._first_der + elif self._order == DerivativeOrder.second: new._first_der = -new._first_der - if self._order == 2: new._second_der = -new._second_der return new @@ -272,7 +286,7 @@ def __mul__(self, other) -> "VectorHyperDual": self._check_compatible(other) val = self._val * other._val - if self._order == 0: + if self._order == DerivativeOrder.zeroth: return VectorHyperDual(val, self._symbols) assert self._first_der is not None @@ -280,7 +294,7 @@ def __mul__(self, other) -> "VectorHyperDual": first_der = ( self._val * other._first_der + other._val * self._first_der ) - if self._order == 1: + if self._order == DerivativeOrder.first: return VectorHyperDual(val, self._symbols, first_der) assert self._second_der is not None @@ -340,14 +354,14 @@ def apply_operation( val = operator(num._val) - if num._order == 0: + if num._order == DerivativeOrder.zeroth: return VectorHyperDual(val, num._symbols) assert num._first_der is not None f_dash_x0 = operator_first_deriv(num._val) first_der = num._first_der * f_dash_x0 - if num._order == 1: + if num._order == DerivativeOrder.first: return VectorHyperDual(val, num._symbols, first_der) assert num._second_der is not None diff --git a/autode/opt/coordinates/primitives.py b/autode/opt/coordinates/primitives.py index 3c3c69867..6aea394c8 100644 --- a/autode/opt/coordinates/primitives.py +++ b/autode/opt/coordinates/primitives.py @@ -1,10 +1,11 @@ import numpy as np from abc import ABC, abstractmethod -from typing import Tuple, TYPE_CHECKING, List, Optional +from typing import Tuple, TYPE_CHECKING, List from autode.opt.coordinates.autodiff import ( get_differentiable_vars, DifferentiableMath, + DerivativeOrder, ) if TYPE_CHECKING: @@ -50,7 +51,7 @@ def _cross_vec3( def _get_vars_from_atom_idxs( *args: int, x: "CartesianCoordinates", - deriv_order: int, + deriv_order: DerivativeOrder, ) -> List["VectorHyperDual"]: """ Obtain differentiable variables from the Cartesian components @@ -88,7 +89,10 @@ def __init__(self, *atom_indexes: int): """A primitive internal coordinate that involves a number of atoms""" self._atom_indexes = atom_indexes - def _evaluate(self, x: "CartesianCoordinates", deriv_order: int): + @abstractmethod + def _evaluate( + self, x: "CartesianCoordinates", deriv_order: DerivativeOrder + ): """ The function that performs the main evaluation of the PIC, and optionally returns derivative or second derivatives. @@ -106,7 +110,7 @@ def _evaluate(self, x: "CartesianCoordinates", deriv_order: int): def __call__(self, x: "CartesianCoordinates") -> float: """Return the value of this PIC given a set of cartesian coordinates""" _x = x.ravel() - res = self._evaluate(_x, deriv_order=0) + res = self._evaluate(_x, deriv_order=DerivativeOrder.zeroth) return res.value def derivative( @@ -133,7 +137,7 @@ def derivative( (np.ndarray): Derivative array of shape (N, ) """ _x = x.ravel() - res = self._evaluate(_x, deriv_order=1) + res = self._evaluate(_x, deriv_order=DerivativeOrder.first) derivs = np.zeros_like(_x, dtype=float) for i in range(_x.shape[0]): dqdx_i = res.differentiate_wrt(str(i)) @@ -167,7 +171,7 @@ def second_derivative( """ _x = x.ravel() x_n = _x.shape[0] - res = self._evaluate(_x, deriv_order=2) + res = self._evaluate(_x, deriv_order=DerivativeOrder.second) derivs = np.zeros(shape=(x_n, x_n), dtype=float) for i in range(x_n): for j in range(x_n): @@ -256,7 +260,7 @@ class PrimitiveInverseDistance(_DistanceFunction): """ def _evaluate( - self, x: "CartesianCoordinates", deriv_order: int + self, x: "CartesianCoordinates", deriv_order: DerivativeOrder ) -> "VectorHyperDual": """1 / |x_i - x_j|""" i_0, i_1, i_2, j_0, j_1, j_2 = _get_vars_from_atom_idxs( @@ -278,7 +282,7 @@ class PrimitiveDistance(_DistanceFunction): """ def _evaluate( - self, x: "CartesianCoordinates", deriv_order: int + self, x: "CartesianCoordinates", deriv_order: DerivativeOrder ) -> "VectorHyperDual": """|x_i - x_j|""" i_0, i_1, i_2, j_0, j_1, j_2 = _get_vars_from_atom_idxs( @@ -339,7 +343,9 @@ def __eq__(self, other) -> bool: and other._ordered_idxs == self._ordered_idxs ) - def _evaluate(self, x: "CartesianCoordinates", deriv_order: int): + def _evaluate( + self, x: "CartesianCoordinates", deriv_order: DerivativeOrder + ): """m - o - n angle""" variables = _get_vars_from_atom_idxs( self.m, self.o, self.n, x=x, deriv_order=deriv_order @@ -411,7 +417,7 @@ def __eq__(self, other) -> bool: ) def _evaluate( - self, x: "CartesianCoordinates", deriv_order: int + self, x: "CartesianCoordinates", deriv_order: DerivativeOrder ) -> "VectorHyperDual": """Dihedral m-o-p-n""" # https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics