Use enum for derivative order

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,15 +286,15 @@ 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
             assert other._first_der is not None
             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

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