# Automatic nth-order derivatives
#
# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
# last revision: 2000-6-5
#

"""This module provides automatic differentiation for functions with
any number of variables up to any order. An instance of the class
DerivVar represents the value of a function and the values of 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 significantly 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.

Example:

  >>>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 has the lower one of the two orders. An exception are
zeroth-order variables, which are treated as constants.

Caution: 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 multiple DerivVar objects 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 Numeric


# The following class represents variables with derivatives:

class DerivVar:

    """Variable with derivatives

    Constructor: DerivVar(|value|, |index| = 0, |order| = 1)

    Arguments:

    |value| -- the numerical value of the variable

    |index| -- the variable index (an integer), 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.

    |order| -- the derivative order

    Indexing with an integer yields the derivatives of the corresponding
    order.
    """

    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):
        "Returns a DerivVar object with a lower derivative 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 ValueError, '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 self.__neg__()
	else:
	    raise ValueError, "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 Numeric.exp(Numeric.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(Numeric.exp, Numeric.exp)

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

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

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

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

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

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

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

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

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

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

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

# Type check

def isDerivVar(x):
    "Returns 1 if |x| is a DerivVar object."
    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 ValueError, 'Internal error'
    return d

# Define vector of DerivVars

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

    """Returns a vector whose components are DerivVar objects.

    Arguments:

    |x|, |y|, |z| -- vector components (numbers)

    |index| -- the DerivVar index for the x component. The y and z
               components receive consecutive indices.

    |order| -- the derivative order
    """

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