# Automatic first-order derivatives
#
# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
# last revision: 2010-4-14
#

"""
Automatic differentiation for functions with any number of variables

Instances of the class DerivVar represent the values of a function and
its partial X{derivatives} with respect to a list of variables. All
common mathematical operations and functions are available for these
numbers. There is no restriction on the type of the numbers fed into
the code; it works for real and complex numbers as well as for any
Python type that implements the necessary operations.

This module is as far as possible compatible with the n-th order
derivatives module Derivatives. If only first-order derivatives
are required, this module is faster than the general one.

Example::

  print sin(DerivVar(2))

produces the output::

  (0.909297426826, [-0.416146836547])

The first number is the value of sin(2); the number in the following
list is the value of the derivative of sin(x) at x=2, i.e. cos(2).

When there is more than one variable, DerivVar must be called with
an integer second argument that specifies the number of the variable.

Example::

    >>>x = DerivVar(7., 0)
    >>>y = DerivVar(42., 1)
    >>>z = DerivVar(pi, 2)
    >>>print (sqrt(pow(x,2)+pow(y,2)+pow(z,2)))

    produces the output

    >>>(42.6950770511, [0.163953328662, 0.98371997197, 0.0735820818365])

The numbers in the list are the partial derivatives with respect
to x, y, and z, respectively.

Note: It doesn't make sense to use DerivVar with different values
for the same variable index in one calculation, but there is
no check for this. I.e.::

    >>>print DerivVar(3, 0)+DerivVar(5, 0)

    produces

    >>>(8, [2])

but this result is meaningless.
"""


from Scientific import N; Numeric = N


# The following class represents variables with derivatives:

class DerivVar:

    """
    Numerical variable with automatic derivatives of first order
    """

    def __init__(self, value, index=0, order=1):
        """
        @param value: the numerical value of the variable
        @type value: number
        @param index: the variable index, which serves to
            distinguish between variables and as an index for
            the derivative lists. Each explicitly created
            instance of DerivVar must have a unique index.
        @type index: C{int}
        @param order: the derivative order, must be zero or one
        @type order: C{int}
        @raise ValueError: if order < 0 or order > 1
        """
        if order < 0 or order > 1:
            raise ValueError('Only first-order derivatives')
        self.value = value
        if order == 0:
            self.deriv = []
        elif type(index) == type([]):
            self.deriv = index
        else:
            self.deriv = index*[0] + [1]

    def __getitem__(self, order):
        """
        @param order: derivative order
        @type order: C{int}
        @return: a list of all derivatives of the given order
        @rtype: C{list}
        @raise ValueError: if order < 0 or order > 1
        """
        if order < 0 or order > 1:
            raise ValueError('Index out of range')
        if order == 0:
            return self.value
        else:
            return self.deriv

    def __repr__(self):
        return `(self.value, self.deriv)`

    def __str__(self):
        return str((self.value, self.deriv))

    def __coerce__(self, other):
        if isDerivVar(other):
            return self, other
        else:
            return self, DerivVar(other, [])

    def __cmp__(self, other):
        return cmp(self.value, other.value)

    def __neg__(self):
        return DerivVar(-self.value, map(lambda a: -a, self.deriv))

    def __pos__(self):
        return self

    def __abs__(self): # cf maple signum # derivate of abs
        absvalue = abs(self.value)
        return DerivVar(absvalue, map(lambda a, d=self.value/absvalue:
                                      d*a, self.deriv))
    def __nonzero__(self):
        return self.value != 0

    def __add__(self, other):
        return DerivVar(self.value + other.value,
                        _mapderiv(lambda a,b: a+b, self.deriv, other.deriv))
    __radd__ = __add__

    def __sub__(self, other):
        return DerivVar(self.value - other.value,
                        _mapderiv(lambda a,b: a-b, self.deriv, other.deriv))

    def __rsub__(self, other):
        return DerivVar(other.value - self.value,
                        _mapderiv(lambda a,b: a-b, other.deriv, self.deriv))

    def __mul__(self, other):
        return DerivVar(self.value*other.value,
                        _mapderiv(lambda a,b: a+b,
                                  map(lambda x,f=other.value:f*x, self.deriv),
                                  map(lambda x,f=self.value:f*x, other.deriv)))
    __rmul__ = __mul__

    def __div__(self, other):
        if not other.value:
            raise ZeroDivisionError('DerivVar division')
        inv = 1./other.value
        return DerivVar(self.value*inv,
                        _mapderiv(lambda a,b: a-b,
                                  map(lambda x,f=inv: f*x, self.deriv),
                                  map(lambda x,f=self.value*inv*inv: f*x,
                                      other.deriv)))
    def __rdiv__(self, other):
        return other/self

    __truediv__ = __div__

    def __pow__(self, other, z=None):
        if z is not None:
            raise TypeError('DerivVar does not support ternary pow()')
        val1 = pow(self.value, other.value-1)
        val = val1*self.value
        deriv1 = map(lambda x, f=val1*other.value: f*x, self.deriv)
        if isDerivVar(other) and len(other.deriv) > 0:
            deriv2 = map(lambda x, f=val*Numeric.log(self.value): f*x,
                         other.deriv)
            return DerivVar(val, _mapderiv(lambda a,b: a+b, deriv1, deriv2))
        else:
            return DerivVar(val, deriv1)

    def __rpow__(self, other):
        return pow(other, self)

    def exp(self):
        v = Numeric.exp(self.value)
        return DerivVar(v, map(lambda x, f=v: f*x, self.deriv))

    def log(self):
        v = Numeric.log(self.value)
        d = 1./self.value
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def log10(self):
        v = Numeric.log10(self.value)
        d = 1./(self.value * Numeric.log(10))
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def sqrt(self):
        v = Numeric.sqrt(self.value)
        d = 0.5/v
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def sign(self):
        if self.value == 0:
            raise ValueError("can't differentiate sign() at zero")
        return DerivVar(Numeric.sign(self.value), 0)

    def sin(self):
        v = Numeric.sin(self.value)
        d = Numeric.cos(self.value)
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def cos(self):
        v = Numeric.cos(self.value)
        d = -Numeric.sin(self.value)
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def tan(self):
        v = Numeric.tan(self.value)
        d = 1.+pow(v,2)
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def sinh(self):
        v = Numeric.sinh(self.value)
        d = Numeric.cosh(self.value)
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def cosh(self):
        v = Numeric.cosh(self.value)
        d = Numeric.sinh(self.value)
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def tanh(self):
        v = Numeric.tanh(self.value)
        d = 1./pow(Numeric.cosh(self.value),2)
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def arcsin(self):
        v = Numeric.arcsin(self.value)
        d = 1./Numeric.sqrt(1.-pow(self.value,2))
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def arccos(self):
        v = Numeric.arccos(self.value)
        d = -1./Numeric.sqrt(1.-pow(self.value,2))
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def arctan(self):
        v = Numeric.arctan(self.value)
        d = 1./(1.+pow(self.value,2))
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

    def arctan2(self, other):
        den = self.value*self.value+other.value*other.value
        s = self.value/den
        o = other.value/den
        return DerivVar(Numeric.arctan2(self.value, other.value),
                        _mapderiv(lambda a, b: a-b,
                                  map(lambda x, f=o: f*x, self.deriv),
                                  map(lambda x, f=s: f*x, other.deriv)))

    def gamma(self):
        from transcendental import gamma, psi
        v = gamma(self.value)
        d = v*psi(self.value)
        return DerivVar(v, map(lambda x, f=d: f*x, self.deriv))

# Type check

def isDerivVar(x):
    """
    @param x: an arbitrary object
    @return: True if x is a DerivVar object, False otherwise
    @rtype: bool
    """
    return hasattr(x,'value') and hasattr(x,'deriv')

# Map a binary function on two first derivative lists

def _mapderiv(func, a, b):
    nvars = max(len(a), len(b))
    a = a + (nvars-len(a))*[0]
    b = b + (nvars-len(b))*[0]
    return map(func, a, b)


# Define vector of DerivVars

def DerivVector(x, y, z, index=0):

    """
    @param x: x component of the vector
    @type x: number
    @param y: y component of the vector
    @type y: number
    @param z: z component of the vector
    @type z: number
    @param index: the DerivVar index for the x component. The y and z
                  components receive consecutive indices.
    @type index: C{int}
    @return: a vector whose components are DerivVar objects
    @rtype: L{Scientific.Geometry.Vector}
    """

    from Scientific.Geometry.VectorModule import Vector
    if isDerivVar(x) and isDerivVar(y) and isDerivVar(z):
        return Vector(x, y, z)
    else:
        return Vector(DerivVar(x, index),
                      DerivVar(y, index+1),
                      DerivVar(z, index+2))

# Functions with derivatives

class DerivFn(object):

    """
    Function with derivatives, applicable to DerivVar objects
    """

    def __init__(self, fn, *deriv_fns):
        """
        @param fn: a Python function from numbers to numbers
        @param *deriv_fns: Python functions from numbers to numbers defining
                           the partial derivatives of the defined function
                           with respect to its arguments. There must be as many
                           derivative functions as the main function has
                           arguments.
        """
        self.fn = fn
        self.deriv_fns = deriv_fns

    def __call__(self, *args):
        """
        Apply the function to the supplied arguments. The number of arguments
        must be equal to the number of derivative functions given in the
        constructor. The arguments can be plain numbers or DerivVar objects.
        The return value is a DerivVar object.
        """
        assert len(args) == len(self.deriv_fns)
        values = []
        derivs = []
        nderivs = 0
        for x in args:
            if isinstance(x, DerivVar):
                values.append(x.value)
                derivs.append(x.deriv)
                nderivs = max(nderivs, len(x.deriv))
            else:
                values.append(x)
                derivs.append([])
        derivs = [d+(nderivs-len(d))*[0] for d in derivs]
        v = self.fn(*values)
        d = [f(*values) for f in self.deriv_fns]
        rderivs = nderivs*[0]
        for i in range(len(d)):
            rderivs = [r+d[i]*x for (r, x) in zip(rderivs, derivs[i])]
        if rderivs:
            return DerivVar(v, rderivs)
        else:
            return v

class NumDerivFn(DerivFn):

    """
    Function with derivatives evaluated by numerical approximation,
    applicable to DerivVar objects
    """
    
    def __init__(self, fn, *hs):
        """
        @param fn: a Python function from numbers to numbers
        @param hs: one step value h for each argument of the function.
                   The ith partial derivative is calculated
                   symmetrically as
                   (f(xs[i]+hs[i])-f(xs[i]-hs[i]))/(2.*hs[i]).
        """
        deriv_fns = []
        for i in range(len(hs)):
            h = hs[i]
            deriv_fns.append(self._makeDerivFn(fn, i, h))
        DerivFn.__init__(self, fn, *deriv_fns)

    def _makeDerivFn(self, fn, i, h):
        def deriv_fn(*args):
            args_p = args[:i] + (args[i]+h,) + args[i+1:]
            args_m = args[:i] + (args[i]-h,) + args[i+1:]
            return (fn(*args_p)-fn(*args_m))/(2.*h)
        return deriv_fn
