# Automatic nth-order derivatives
#
# Written by Konrad Hinsen <hinsenk@ere.umontreal.ca>
# last revision: 1996-3-5
#

"""This module provides automatic differentiation for functions with
any number of variables up to any order. Instances of the class
DerivVar represent the values of a function and its partial
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.

If only first-order derivatives are required, the module
'FirstDerivatives' should be used. It is compatible to this
one, but faster.

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.

Higher-order derivatives are requested with an optional third
argument to DerivVar:

  x = DerivVar(3., 0, 3)
  y = DerivVar(5., 1, 3)
  print sqrt(x*y)

produces the output

  (3.87298334621,
     [0.645497224368, 0.387298334621],
       [[-0.107582870728, 0.0645497224368],
        [0.0645497224368, -0.0387298334621]],
          [[[0.053791435364, -0.0107582870728],
            [-0.0107582870728, -0.00645497224368]],
           [[-0.0107582870728, -0.00645497224368],
            [-0.00645497224368, 0.0116189500386]]])

The individual orders can be extracted by indexing:

  print sqrt(x*y)[0]
    3.87298334621
  print sqrt(x*y)[1]
    [0.645497224368, 0.387298334621]

An n-th order derivative is represented by a nested list of
depth n.

When variables with different differentiation orders are mixed,
the result gets the lower of the two orders. An exception are
zeroth-order variables, which are treated as constants.

Warning: Higher-order derivatives are implemented by recursively
using DerivVars to represent derivatives. This makes the code
very slow for high orders.

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.
"""


import umath, Vector


# Error type

DerivError = 'DerivError'


# The following class represents variables with derivatives:

class DerivVar:

    def __init__(self, value, index=0, order = 1, recursive = None):
	self.value = value
	if recursive:
	    d = 0
	else:
	    d = 1
	if type(index) == type([]):
	    self.deriv = index
	elif order == 0:
	    self.deriv = []
	elif order == 1:
	    self.deriv = index*[0] + [d]
	else:
	    self.deriv = []
	    for i in range(index):
		self.deriv.append(DerivVar(0, index, order-1, 1))
	    self.deriv.append(DerivVar(d, index, order-1, 1))
	self.order = order

    def toOrder(self, order):
	if self.order <= order:
	    return self
	if order == 0:
	    return self.value
	return DerivVar(self.value, map(lambda x, o=order-1: x.toOrder(o),
					self.deriv), order)

    def __getitem__(self, item):
	if item < 0 or item > self.order:
	    raise DerivError, 'Index out of range'
	if item == 0:
	    return self.value
	else:
	    return map(lambda d, i=item-1: _indexDeriv(d,i), self.deriv)

    def __repr__(self):
	return repr(tuple(map(lambda n, x=self: x[n], range(self.order+1))))

    def __str__(self):
	return str(tuple(map(lambda n, x=self: x[n], range(self.order+1))))

    def __coerce__(self, other):
	if isDerivVar(other):
	    if self.order == other.order or self.order == 0 or other.order == 0:
		return self, other
	    order = min(self.order, other.order)
	    return self.toOrder(order), other.toOrder(order)
	else:
	    return self, DerivVar(other, [], 0)

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

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

    def __pos__(self):
	return self

    def __abs__(self):
	if self.value > 0:
	    return self
	elif self.value < 0:
	    return __neg__(self)
	else:
	    raise DerivError, "can't differentiate abs() at zero"

    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),
			max(self.order, other.order))
    __radd__ = __add__

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

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

    def __mul__(self, other):
	if self.order < 2:
	    s1 = self.value
	else:
	    s1 = self.toOrder(self.order-1)
	if other.order < 2:
	    o1 = other.value
	else:
	    o1 = other.toOrder(other.order-1)
	return DerivVar(self.value*other.value,
			_mapderiv(lambda a,b: a+b,
				  map(lambda x,f=o1: f*x, self.deriv),
				  map(lambda x,f=s1: f*x, other.deriv)),
			max(self.order, other.order))
    __rmul__ = __mul__

    def __div__(self, other):
	if not other.value:
	    raise ZeroDivisionError, 'DerivVar division'
	if self.order < 2:
	    s1 = self.value
	else:
	    s1 = self.toOrder(self.order-1)
	if other.order < 2:
	    o1i = 1./other.value
	else:
	    o1i = 1./other.toOrder(other.order-1)
	return DerivVar(_toFloat(self.value)/other.value,
			_mapderiv(lambda a,b: a-b,
				  map(lambda x,f=o1i: x*f, self.deriv),
				  map(lambda x,f=s1*pow(o1i, 2): f*x,
				      other.deriv)),
				  max(self.order, other.order))
    def __rdiv__(self, other):
	return other/self

    def __pow__(self, other, z=None):
	if z is not None:
	    raise TypeError, 'DerivVar does not support ternary pow()'
	if len(other.deriv) > 0:
	    return umath.exp(umath.log(self)*other)
	else:
	    if self.order < 2:
		ps1 = other.value*pow(self.value, other.value-1)
	    else:
		ps1 = other.value*pow(self.toOrder(self.order-1), other.value-1)
	    return DerivVar(pow(self.value, other.value),
			    map(lambda x,f=ps1: f*x, self.deriv),
			    max(self.order, other.order))

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

    def _mathfunc(self, f, d):
	if self.order < 2:
	    fd = d(self.value)
	else:
	    fd = d(self.toOrder(self.order-1))
	return DerivVar(f(self.value), map(lambda x, f=fd: f*x, self.deriv),
			self.order)

    def exp(self):
	return self._mathfunc(umath.exp, umath.exp)

    def log(self):
	return self._mathfunc(umath.log, lambda x: 1./x)

    def sqrt(self):
	return self._mathfunc(umath.sqrt, lambda x: 0.5/umath.sqrt(x))

    def sin(self):
	return self._mathfunc(umath.sin, umath.cos)

    def cos(self):
	return self._mathfunc(umath.cos, lambda x: -umath.sin(x))

    def tan(self):
	return self._mathfunc(umath.tan, lambda x: 1.+pow(umath.tan(x),2))

    def sinh(self):
	return self._mathfunc(umath.sinh, umath.cosh)

    def cosh(self):
	return self._mathfunc(umath.cosh, umath.sinh)

    def tanh(self):
	return self._mathfunc(umath.tanh, lambda x: 1./pow(umath.cosh(x),2))

    def arcsin(self):
	return self._mathfunc(umath.arcsin, lambda x:
			      1./umath.sqrt(1.-pow(x,2)))

    def arccos(self):
	return self._mathfunc(umath.arccos, lambda x:
			      -1./umath.sqrt(1.-pow(x,2)))

    def arctan(self):
	return self._mathfunc(umath.arctan, lambda x:
			      1./(1+pow(x,2)))

# Type check

def isDerivVar(x):
    return hasattr(x,'value') and hasattr(x,'deriv') and hasattr(x,'order')


# 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)

# Convert argument to float if it is integer

def _toFloat(x):
    if type(x) == type(0):
	return float(x)
    return x

# Subscript for DerivVar or ordinary number

def _indexDeriv(d, i):
    if isDerivVar(d):
	return d[i]
    if i != 0:
	raise DerivError, 'Internal error'
    return d

# Define vector of DerivVars

def DerivVector(x, y, z, index=0, order = 1):
    if isDerivVar(x) and isDerivVar(y) and isDerivVar(z):
	return Vector.Vector(x, y, z)
    else:
	return Vector.Vector(DerivVar(x, index, order),
			     DerivVar(y, index+1, order),
			     DerivVar(z, index+2, order))