""" lorenz.py

This file may be imported into other scripts to provide a python
interface to gsl for integrating the lorenz system.  It may also be
called as "main" to make data files.  Here is the Lorenz system:

   \dot x = s(y-x)
   \dot y = rx -xz -y
   \dot z = xy -bz

Options:

--data_dir=  Specify directory that will hold created files
--TSintro    Make: TSintro_1_vfine TSintro_2_pu TSintro_3_q
--L20K       Make: lorenz.xyz lorenz.4
"""
import numpy, scipy, os

extra_code = \
"""
extern "C" int Lsteps(double *y, struct ParPac *pac, double ts, int Nt,
    float *result);
extern "C" int Ltansteps( double *y, struct ParPac *pac, double ts,
    int Nt, double *result);                                     
struct ParPac {double s,b,r;};
"""
def Lsteps(IC,      # IC[0:3] is the initial condition
           s, b, r, # These are the Lorenz parameters
           T_step,  # The time between returned samples
           N_steps  # N_steps The number of returned samples
           ):

    import scipy.weave as weave
    Lsteps_w = \
"""
struct ParPac pac;             /* Paramters structure */
pac.s = s; pac.b = b;pac.r = r;
Lsteps( IC_A, &pac, T_step, N_steps, result_R);
"""
    result = numpy.ones((N_steps,3),numpy.float32)
    result_R = result.reshape(-1) # I don't know weave can use a 2d array
    IC_A = numpy.array(IC,numpy.float64)
    # The library "lorenz" has the compiled c code to integrate the equations
    rv = weave.inline(Lsteps_w,
                      ['IC_A', 's', 'b', 'r', 'T_step', 'N_steps', 'result_R'],
                      libraries=['lorenz','gsl', 'gslcblas', 'm'],
                      support_code=extra_code,
                      library_dirs=[os.path.abspath('code/python')]
                 )
    return result
def Ltan_steps(IC,       # IC[0:12] is the initial condition
               s, b, r,  # These are the Lorenz parameters
               T_step,   # The time between returned samples
               N_steps
               ):
    """ Return RV[0:N_steps+1,0:4,0:3) where RV[t,0,0:3] is the initial x
    vector and RV[t,1:4,0:3] is the derivative after integrating t time
    steps.
    """

    import scipy.weave as weave
    Ltan_steps_w = \
"""
struct ParPac pac;             /* Paramters structure */
pac.s = s; pac.b = b;pac.r = r;
Ltansteps( IC_A, &pac, T_step, N_steps, result_R);
"""
    RV = numpy.ones((N_steps,4,3),numpy.float64)
    result_R = RV.reshape(-1) # I don't know weave can use a 3d array
    IC_A = numpy.array(IC,numpy.float64)
    rv = weave.inline(Ltan_steps_w,
                      ['IC_A', 's', 'b', 'r', 'T_step', 'N_steps', 'result_R'],
                      libraries=['lorenz','gsl', 'gslcblas', 'm'],
                      support_code=extra_code,
                      library_dirs=[os.path.abspath('code/python')]
                 )
    return RV
def Ltan_one(IC,       # IC[0:3] is the initial condition
             s, b, r,  # These are the Lorenz parameters
             T_step    # The time step
             ):
    """ Integrate initial condition IC and the tangent (initialize here as
    identity) one step.  Return a pair consisting of the final
    location x(t) and the derivative D.
    """
    IC_4_3 = numpy.zeros((4,3),numpy.float64)
    IC_4_3[0,0:3] = IC
    IC_4_3[1,0] = 1.0
    IC_4_3[2,1] = 1.0
    IC_4_3[3,2] = 1.0
    RV = Ltan_steps(IC_4_3.reshape(-1),s,b,r,T_step,2)
    X = RV[1,0,0:3]
    D = RV[1,1:4,0:3]
    return X,D
################ Begin code for stand alone script #############
if __name__ == '__main__':

    s = 10.0
    r = 28.0
    b = 8.0/3
    import optparse, os
    cwd = os.getcwd()
    OP = optparse.OptionParser(
    usage=
    """Usage: %prog [options]""",
    description=
    """Integrate the Lorenz system and write results to files.""")
    OP.add_option('--TSintro', action="store_true", default=False, help='')
    OP.add_option('--L20K', action="store_true", default=False, help='')
    OP.add_option('--data_dir', default=cwd, help='')
    opt,pargs = OP.parse_args()
    def relax():
        # Relax to the attractor:
        global s,r,b
        IC = [1.0,1.0,1.0]
        ts = 7.5
        Nt = 3
        return Lsteps(IC,s,b,r,ts,Nt)[-1,:]
    if opt.TSintro:
        # Make vfine with time step .003 for TSintro and lorenz
        ts    = 0.003                           # Time step
        Nt    = 2000                            # Number of samples
        A = Lsteps(relax(),s,b,r,ts,Nt)
        f = open(opt.data_dir + "/TSintro_1_vfine", 'w')
        for i in xrange(0,Nt):
            t = ts * i
            print >> f, t, A[i][0]
        # Decimate vfine to make pu with time step 0.15
        pu  = A[0:2000:50, 0]               # Decimate by 50
        ts *= 50
        f = open(opt.data_dir + "/TSintro_2_pu", 'w')
        for i in xrange(0, Nt/50):
            t = ts * i
            print >> f, t, pu[i]
        # Quantize pu to make q
        q = numpy.ceil(pu / 10 + 2)  # Quantize by mapping
                                     # {(-20,-10],(-10,0],(0,10],(10,20]}
                                     # to {1,2,3,4} respectively
        f = open(opt.data_dir + "/TSintro_3_q", 'w')
        for i in xrange(0, Nt/50):
            t = ts * i
            print >> f, i, q[i]
    if opt.L20K:
        # Make 20,000 point trajectory, and write it to "lorenz.xyz"
        ts     = 0.15                           # Time step
        Nt     = 20000                          # Number of samples
        A = Lsteps(relax(),s,b,r,ts,Nt)
        f = open(opt.data_dir + "/lorenz.xyz", 'w')
        for xyz in A:
            print >> f, 3*'%7.3f '%tuple(xyz)
        # Quantize the trajectory and write it to "lorenz.4"
        f = open(opt.data_dir + "/lorenz.4", 'w')
        for i in  numpy.ceil(A[:, 0] / 10 + 2):
            print >> f, int(i)

# Local Variables:
# mode: python
# End: