""" Sparse_hmm_lor.py

For making the "likeLor" figure in that plots cross entropy vs number
of states, going up to millions of states.

Builds discrete output HMMs of quantized observations of the Lorenz
system by using information about the continuous states of the system.
Goal is to build an HMM that gets close to the entropy bound
estimated via Lyapunov exponents.
"""

import  sys, numpy, scipy.sparse, math, time, lorenz

def F_step(Pyt,PN,a,alpha,ScS,YcS):
    import scipy.weave as weave
    ScS_indptr = ScS.indptr     # I would like to avoid these translations
    ScS_indices = ScS.indices
    ScS_data = ScS.data
    YcS_indptr = YcS.indptr
    YcS_indices = YcS.indices
    YcS_data = YcS.data
    return weave.inline("""
int i,j,k;
int yt = Pyt;            /* Weave chokes without these cast like things */
int N = PN;
double alpha_i,gamma;
for(i=0;i<N;i++){
    a[i] = 0;
}
for(i=0;i<N;i++){
    alpha_i = alpha[i];
    if (alpha_i < 1e-20) continue;
    for(j=ScS_indptr[i];j<ScS_indptr[i+1];j++){
        k = ScS_indices[j];
        a[k] += alpha_i*ScS_data[j];
    }
}
for(i=0;i<N;i++){
    alpha[i] = 0;
}
for(j=YcS_indptr[yt];j<YcS_indptr[yt+1];j++){
    k = YcS_indices[j];
    alpha[k] = a[k]*YcS_data[j];
}
gamma = 0.0;
for(i=0;i<N;i++)
    gamma += alpha[i];
for(i=0;i<N;i++)
    alpha[i] = alpha[i]/gamma;
return_val = gamma;
""",
                      ['Pyt', 'PN', 'a', 'alpha',
                       'ScS_indptr', 'ScS_indices', 'ScS_data',
                       'YcS_indptr', 'YcS_indices', 'YcS_data',
                       ],compiler='gcc')
class Sparse:
    """ To use scipy.sparse efficiently, I've resorted to doing the key
    manipulations in C.
    """
    def __init__(self, P_S0,P_ScS,P_YcS):
        """Build a new Hidden Markov Model"""
        self.P_S0 = P_S0
        self.N = len(self.P_S0)
        self.P_ScS = P_ScS.tocsr()
        self.P_YcS = P_YcS.tocsr()
        return # End of __init__()
    def like(self,
             Y
             ):
        """
        alpha   Updated state distribution as a dense vector
        a       Forecast, viz, ScS*alpha as a dense vector
        """
        N = self.N
        ScS = self.P_ScS.T
        YcS = self.P_YcS
        alpha = self.P_S0
        a = numpy.empty((N,))
        rv = 0.0
        for t in xrange(len(Y)):
            gamma = F_step(Y[t],N,a,alpha,ScS,YcS)
            assert gamma > 0,'gamma=%f'%gamma
            rv += math.log(gamma)
        return rv
# Parameters of the Lorenz system
s = 10.0
r = 28.0
b = 8.0/3
# Other parameters
Y_res = 10.0         # Quantization resolution for observations
Trelax = 7.5         # Time step for relaxation.
ts = 0.15            # Time step
Nrelax = 3           # Number of samples for relaxation
Nbuild = 8000000     # Number of samples for building initial model
Ntest =    10000     # Number of samples for testing

# Relax to the attractor:
result = lorenz.Lsteps([1,1,1],s,b,r,Trelax,Nrelax)
ic = result[-1,:]                              # reset ICs

# Make data for building models: build_x and build_y
result = lorenz.Lsteps(ic,s,b,r,ts,Nbuild)
build_x = result
ic = result[-1,:]  # Save ic for making test data
y = numpy.ceil(result[:,0]/Y_res).astype(numpy.int32)
# Next make Y_dict a dict with keys being all the different values in
# y.  Y_dict[y] is a positive integer.  In other words, Y_dict maps y
# values to a small interval of positive integers
Y_dict = {}
for t in xrange(Nbuild):
    Y_key = y[t]
    if not Y_dict.has_key(Y_key):
        Y_dict[Y_key] = len(Y_dict)
    y[t] = Y_dict[Y_key]
CardY = len(Y_dict)
build_y = y

# Make testing data: test_y
result = lorenz.Lsteps(ic,s,b,r,ts,Ntest)
y = numpy.ceil(result[:,0]/Y_res).astype(numpy.int32)
for t in xrange(Ntest):  # Better not have any y values in test that
    y[t] = Y_dict[y[t]]  # didn't appear in build
test_y = y

outfile = open(sys.argv[1],'w')
print >>outfile,  'ts=',ts, 'Trelax=', Trelax, 'Nrelax=', Nrelax
print >>outfile,  'Y_res=' ,Y_res, 'CardY=', CardY
print >>outfile,  'Nbuild=',Nbuild,'Ntest=',Ntest
print >>outfile,  '%6s %7s %7s %6s %6s %9s %-13s'%('Tprep','S_res','Nstates','Tlil',
      'h_raw','Tlike', 'Total time')
def sparse_norm(A,dim):
    A = A.tocoo()
    norm = A.sum(dim).A.reshape(-1) # sum over selected index
    for k in xrange(len(A.data)):
        if dim == 1:
            A.data[k] /= norm[A.row[k]]
        else:
            A.data[k] /= norm[A.col[k]]
    return A,norm
# Loop over quant resolution for states
#for exponent in numpy.arange(5,-4.51,-.5):
for exponent in numpy.arange(5,-4.51,-.5):
#for exponent in numpy.arange(5,-2.51,-.5):
    Tstart = time.time()
    S_res = 2.0** exponent
    # Quantize the trajectory to get outputs and states
    q = numpy.ceil(build_x/S_res).astype(numpy.int32)
    # State_dict is a dictionary.  The keys are quantized state vectors
    # and the values are indices assigned in the order in which keys
    # appear for the first time
    st = numpy.empty((Nbuild),numpy.int32)
    State_dict = {}
    # st is state index as a function of time
    for t in xrange(Nbuild):
        S_key = tuple(q[t,:])
        if not State_dict.has_key(S_key):
            State_dict[S_key] = int(len(State_dict))
        st[t] = State_dict[S_key]
    Nstates = len(State_dict)
    q = 0          #  Memory management?
    State_dict = 0 #  Memory management?
    result = 0
    Tlil = time.time()
    P_ScS=scipy.sparse.lil_matrix((Nstates,Nstates))
    P_YcS=scipy.sparse.lil_matrix((CardY,Nstates))
    Ycount = numpy.zeros((CardY,Nstates),scipy.int32)
    Sdicts = Nstates*[None]
    for s in xrange(Nstates):
        Sdicts[s] = {}
    for t in xrange(Nbuild-1):
        Ycount[build_y[t],st[t]] += 1
        if Sdicts[st[t]].has_key(st[t+1]):
            Sdicts[st[t]][st[t+1]] += 1
        else:
            Sdicts[st[t]][st[t+1]] = 1
    for s in xrange(Nstates):
        for n in Sdicts[s].keys():
            P_ScS[s,n] = Sdicts[s][n]
        for y in xrange(CardY):
            if Ycount[y,s] > 0:
                P_YcS[y,s] = Ycount[y,s]
    Tlike = time.time()
    P_YcS,norm = sparse_norm(P_YcS,0)
    P_ScS,norm = sparse_norm(P_ScS,1)
    P_S0 = norm/norm.sum()
    model = Sparse(P_S0,P_ScS,P_YcS)

    h_raw = model.like(test_y)/Ntest
    Tend     = time.time()
    print >>outfile, '%6.1f'%(Tlil-Tstart), '%7.4f'%S_res, '%7d'%Nstates,\
          '%6.1f'%(Tlike-Tlil), '%6.3f'%-h_raw, '%9.2f'%(Tend-Tlike),\
          '%9.1f'%(Tend - Tstart)

#Local Variables:
#mode:python
#End: