Merge pull request #78 from JHay0112/autodiff-grad

Autodiff grad

Author: JHay0112

jmath/autodiff.py

@@ -1,5 +1,20 @@
 '''
     Automatic Differentiation
+
+    Examples
+    --------
+
+    >>> from jmath.autodiff import x, y
+    >>> f = 6*x*y*2 + y + 9
+    >>> f_prime = f.d(x)
+    >>> f_prime(x = 2, y = 1)
+    6
+
+    >>> from jmath.autodiff import x, y
+    >>> f = x*y
+    >>> grad_f = f.d(x, y)
+    >>> grad_f(x = 3, y = 2)
+    Vector(2, 3)
 '''
 
 # - Imports
@@ -10,11 +25,12 @@
 from functools import wraps
 from types import FunctionType
 from .uncertainties import Uncertainty
+from .linearalgebra import Vector
 from typing import Any, Union, Callable, Tuple
 
 # - Typing
 
-Supported = Union[int, float, Uncertainty, 'Function', 'Variable']
+Supported = Union[int, float, Uncertainty, 'Function', 'Variable', Vector]
 Numeric = Union[int, float, Uncertainty]
 
 # - Classes
@@ -63,7 +79,10 @@ def __str__(self):
         # Standard function
         return f"{self.func.__name__}{str(params)[:-2]})"
 
-    def __call__(self, **kwargs):
+    def __call__(self, *args, **kwargs):
+
+        if len(args) != 0:
+            raise AttributeError("Inputs must be assigned to a variable e.g. f(x = 3) rather than f(3)!")
 
         if not isinstance(self.func, Callable):
             return self.func
@@ -225,7 +244,7 @@ def register(self, *inputs: 'Function'):
         '''
         self.inputs = inputs
 
-    def differentiate(self, wrt: Union['Variable', str]) -> 'Function':
+    def differentiate(self, *wrt: Union['Variable', str]) -> Union['Function', Vector]:
         '''
             Differentiates the function with respect to a variable.
 
@@ -235,6 +254,15 @@ def differentiate(self, wrt: Union['Variable', str]) -> 'Function':
             wrt
                 The variable to differentiate with respect to.
         '''
+        if len(wrt) == 1:
+            # Single differential case
+            wrt = wrt[0]
+        else:
+            # Multiple case
+            results = []
+            for var in wrt:
+                results.append(self.differentiate(var))
+            return Vector(results)
         # The differentiated function
         func = 0
         # Move across inputs
@@ -248,8 +276,8 @@ def differentiate(self, wrt: Union['Variable', str]) -> 'Function':
         return func
 
     @wraps(differentiate)
-    def d(self, wrt: Union['Variable', str]) -> 'Function':
-        return self.differentiate(wrt)
+    def d(self, *wrt: Union['Variable', str]) -> 'Function':
+        return self.differentiate(*wrt)
 
 class Variable(Function):
     '''
@@ -272,9 +300,12 @@ def __str__(self):
 
         return self.id
 
-    def __call__(self, input: Any) -> Any:
+    def __call__(self, input: Any = None, **kwargs) -> Any:
 
-        return input
+        if input is not None:
+            return input
+        else:
+            return kwargs[self.id]
 
     def differentiate(self, wrt: 'Variable') -> int:
 

jmath/linearalgebra/vectors.py

@@ -93,6 +93,17 @@ def inner(*args, **kwargs):
 
         return inner
 
+    def __call__(self, *args, **kwargs) -> 'Vector':
+        """Call vector of functions."""
+        results = []
+        for component in self.components:
+            # Only call callable functions
+            if isinstance(component, Callable):
+                results.append(component(*args, **kwargs))
+            else:
+                results.append(component)
+        return Vector(results)
+
     def __eq__(self, vector: "Vector") -> bool:
         """Tests equality of vectors"""
         if isinstance(vector, Vector) or isinstance(vector, Point):
@@ -213,6 +224,34 @@ def angle_between(self, vector: "Vector") -> float:
             vector = vector.vector
 
         return round(math.acos((self @ vector)/(self.magnitude() * vector.magnitude())), 5)
+    
+    def differentiate(self, *wrt) -> 'Vector':
+        '''
+            Differentiate functions in vector with respect to a variable.
+
+            Parameters
+            ----------
+
+            wrt
+                The variables to differentiate with respect to.
+        '''
+        results = []
+        for component in self.components:
+            # Check if is callable
+            if isinstance(component, Callable):
+                # Then differentiate it
+                results.append(component.d(*wrt))
+            else:
+                # Non-differentiable
+                # So just append a zero
+                results.append(0)
+        return Vector(results)
+
+    def d(self, *wrt) -> 'Vector':
+        '''
+            Differentiate short hand.
+        '''
+        return self.differentiate(*wrt)
 
 class Point(Vector):
     """

tests/test_unit_spaces.py

@@ -31,7 +31,7 @@ def test_conversion_leak():
     space = UnitSpace()
     space.a = Unit("a")
     space.b = Unit("b")
-    space.define_conversion(space.a, space.b, random_integer())
+    space.define_conversion(space.a, space.b, random_integer(non_zero = True))
     assert space.a == space.a.convert_to(space.b).convert_to(space.a)
     assert universal.units == {}
 

tests/test_units.py

@@ -50,7 +50,7 @@ def test_conversion():
     """Tests that unit conversion happens as expected."""
     a = Unit("a")
     b = Unit("b")
-    coeffecient = random_integer()
+    coeffecient = random_integer(non_zero = True)
     define_conversion(a, b, coeffecient)
 
     assert a.convert_to(b).value == (a*coeffecient).value