# written by SLJ on 2002/08/06
# corrected on 2002-08-09 (for info array)
# this is to test DDASSL...
# a simple DAE representing a free falling body...
# Changed the IPAR array on 2002-8-15
# Added the JAC routine on 2002-8-15 22h50
# 2002-8-29: better interface; put the run method...

import Numeric
import ddassl

g = 9.8

def RES(t,Y,YDOT,DELTA,ires,RPAR,IPAR):
    """ RES(t,Y,YDOT,DELTA,ires,RPAR,IPAR) is the method representing the DAE
    t -> scalar (time)
    Y, YDOT, DELTA are arrays of size of number of equation;
        DELTA = G(t,Y,YDOT) (in place)
    ires -> integer array of size 1
    RPAR, IPAR -> real and integer parameters array.
    *** Note that IPAR[0] = dim(Y), IPAR[1] = dim(IPAR) and IPAR[2] = dim(RPAR)

    In our case, the function is simply representing the DAE system:
    dx/dt - v = 0 and dv/dt - g = 0 (free fall body)
    where Y = [x,v].  I.e. G(t,Y,YDOT) = [x_dot - v, v_dot - g]"""
#    print "Entering RES..."
    DELTA[0] = YDOT[0] - Y[1]
    DELTA[1] = YDOT[1] - g
#    print "Exiting RES..."
    return

def JAC(t,Y,YDOT,PD,cj,RPAR,IPAR):
    """JAC(t,Y,YDOT,PD,cj,RPAR,IPAR) is the method to be used to return the Jacobian
    of the DAE.  It is optional...
    t -> scalar (time)
    Y and YDOT are arrays of size of numer of equations;
        PD is the jacobian array of the equation (size NEQ*NEQ) with the following shape:
        PD(i,j) = dG(i)/dY(j) + cj * dG(j)/dYDOT(j)
        ***NOTE: PD is always set to zero before the call, so JAC needs only to set
        the nonzero values...
    cj -> scalar provided by user
    RPAR and IPAR are real and integer parameters array (see RES.__doc__ for more info)

    In our DAE system, we have PD = [[0,-1],[0,0]] + cj*identity = [[cj,-cj],[0,cj]]
    Note that the part with cj will always be the identity for ODE's...
    """
#    print "Entering JAC..."
    # note: setting only nonzero values
    PD[0,0] = cj
    PD[0,1] = -cj
    PD[1,1] = cj
#    print "Exiting JAC..."
    
    return

NEQ = 2 # number of equations
INFO = Numeric.zeros(15) # to tell DDASSL how to compute the solution
INFO[4] = 1  # so that the code uses our jacobian routine JAC.


RTOL = Numeric.array([0.0]) # won't use relative error
ATOL = Numeric.array([0.001]) # absolute tolerance

IDID = Numeric.array([0])  # to report what the code did (will be changed in place)

LRW = 50 + 10*NEQ+NEQ**2   # size of real work array
LIW = 20 + 2*NEQ  # size of integer work array
RWORK = Numeric.array([0.0]*LRW)
IWORK = Numeric.zeros(LIW)

# parameters array (not used in our code;
# but still need to be array for type consistency...)
RPAR = Numeric.array([0.0, 0])
IPAR = Numeric.array([2, 3, 1])  # [size of y, size of IPAR, size of RPAR]
				#(my convention)

# initial conditions: (note that they are consistent!!!)
T = Numeric.array([0.0])
Y = Numeric.array([0.0,1.0])
YDOT = Numeric.array([1.0,g])
TOUT = Numeric.array([10.0])

# some tests!...
def run():
    print "Testing RES..."
    RES(T,Y,YDOT,RWORK,0,RPAR,IPAR)
    print 'RES(T,Y,YDOT,...) = ', RWORK[0], RWORK[1]

    print 'Y before ddassl call: ', Y

    ddassl.ddassl(RES,NEQ,T,Y,YDOT,TOUT,INFO,RTOL,ATOL,IDID,RWORK,LRW,IWORK,LIW,RPAR,IPAR,JAC)

    print 'Y after: ', Y, 'at T = :', T
    print 'with IDID = ', IDID