# 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...
# -------------------
# MODIFIED on 2002-8-29 to test ddasrt instead [plus added G]...

import Numeric
import ddasrt

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

## added for ddasrt:
def G(neq,t,Y,ng,GOUT,RPAR,IPAR):
    """G(neq,t,Y,ng,GOUT,RPAR,IPAR) is the function to be used to return the monitoring function
    results.  It is optional [set ng to 0 when calling ddasrt]
    neq -> number of equations
    t -> scalar (time)
    Y is array of size of number of equations [state of system]
    ng -> number of monitoring functions
    GOUT -> array of size ng, output of our monitoring functions
    RPAR and IPAR are real and integer parameters array (see RES.__doc__ for more info)

    In our example, we use G = Y[0] - 10 so that we detect when Y[0] = 10
    """
    GOUT[0] = Y[0] - 10
    
    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 [same as ddassl]
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)
NG = 1  # we use one monitoring function

# 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
    # ddasrt returns array JROOT, where JROOT[i] = 1 if G[i] has a root at T; 0 otherwise
    JROOT = ddasrt.ddasrt(RES,NEQ,T,Y,YDOT,TOUT,INFO,RTOL,ATOL,IDID,RWORK,LRW,IWORK,LIW,RPAR,IPAR,JAC,G,NG)

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