<jupyter_start><jupyter_text>Generate DataMake sure you have also copied '/home/mcarilli/CME/2D_CME/learning_environment/train_2D_rt_v1.py' and'/home/mcarilli/CME/2D_CME/learning_environment/tools_2D_rt_v1.py' to your file location. If you want to load in pre-generated data, there are 20 batches of 5,120 are already stored in:1. '/home/mcarilli/CME/2D_CME/learning_environment/training_data_quadvec/'2. '/home/mcarilli/CME/2D_CME/learning_environment/training_data_fixedquad/'There are also 10 batches of 512 parameters in 1 (_fixedquad/_) and 10 batches of 256 parameterc in 2 (_quadvec/_). generated using quad_vec and fixed_quad, respectively. If you want to make and store other data files, here is some code:<jupyter_code># define your data path

path =  './training_data/'
set_size = 256
num_files = 10
N = num_files*set_size


params = tr.generate_param_vectors(N)
# define your xmax function!! 

def xmax_fun(xmax):

    return(4.5*xmax)
#generate_sets(set_size,num_files,param_vectors,method,xmax_fun,NCOR,path_to_directory)
tr.generate_sets_pcme(set_size = 256,num_files=num_files,param_vectors = params,
                 method = 'quad_vec',
                 xmax_fun=xmax_fun,
                 NCOR = 60,
                 path_to_directory=path)<jupyter_output><empty_output><jupyter_text>Train ! Here are the generating grid and get_ypred_at_RT functions! You can choose to generate the grid using NORM_nas, NORM_mat quantiles (linear right now, rectangular grid), or performin an affine transform with the covariance matrix and means. To change, merely change the generate_grid function in get_ypred_at_RT.<jupyter_code>def generate_grid(npdf,VAR,MU,p,quantiles=None):
    
    if quantiles=='PRESET':
        logstd = torch.sqrt(torch.log((VAR/MU**2)+1))
        logmean = torch.log(MU**2/torch.sqrt(VAR+MU**2))
        translin_0 = torch.exp(logmean[0]+logstd[0]*NORM_nas)
        translin_1 = torch.exp(logmean[1]+logstd[1]*NORM_mat)
        
        
        GRID_nas,GRID_mat = torch.meshgrid(translin_0,translin_1,indexing='ij')
        
        return GRID_nas,GRID_mat

    
def generate_grid_affine(npdf,VAR,MU,p,quantiles=None):
    if quantiles=='PRESET':
        logvar = torch.log((VAR/MU**2)+1)
        logmean = torch.log(MU**2/torch.sqrt(VAR+MU**2))

        var_ij = p[0]**2/ (p[1]+p[2]) 
        logst_ij = torch.log( var_ij * np.exp(-(logmean.sum()+logvar.sum()/2)) +1 ) #set this to zero if you want to 
        #just use a grid with zero correlation
        
        xv, yv = torch.meshgrid(NORM_nas,NORM_mat,indexing='ij')

        a = torch.sqrt(logvar[0])
        b = logst_ij/a
        c = torch.sqrt(logvar[1]-b**2)

        GRID_nas = torch.exp(a*xv + logmean[0]) #this is correct
        GRID_mat = torch.exp(b*xv+c*yv + logmean[1])
        
        return GRID_nas,GRID_mat

def get_ypred_at_RT(p,npdf,w,hyp,quantiles='PRESET',
                   first_special=False,special_std='tail_prob',poisson_kernel = False,
                   kernel_type='independent',use_global_corr = True):
    p = 10**p
    MU, VAR, STD, xmax = tr.get_moments(p)
    
    GRID_nas,GRID_mat = generate_grid_affine(npdf,VAR,MU,p,quantiles=quantiles)

    
   
    xgrid_nas = torch.arange(xmax[0]+1)
    xgrid_mat = torch.arange(xmax[1]+1)
    
    gammaln_xgrid_nas = lnfactorial[1:(xmax[0]+2)]
    gammaln_xgrid_mat = lnfactorial[1:(xmax[1]+2)] 

    Y = torch.zeros((xmax[0]+1,xmax[1]+1))
    if kernel_type=='correlated':
        eps=1e-8
        f = p[1]/(p[1]+p[2])
        if use_global_corr:
            rho_ = p[0]*np.sqrt(f*(1-f)/((1+p[0])*(1+p[0]*f)))
        else:
            rho_ = (hyp-1)/5 #hmmm
        rho2 = rho_**2

        a = (torch.sqrt(rho2* (GRID_nas - GRID_mat)**2 
                + 4 * GRID_nas *  GRID_mat)/rho_
            - GRID_nas - GRID_mat)/2
        log_theta_mu_eps = torch.log(a + GRID_nas  + GRID_mat + eps) 
    else:

        s_nas = torch.zeros((GRID_nas.shape))
        s_mat =  torch.zeros((GRID_mat.shape))

        spec = 0 if first_special else -1
        if first_special:
            s_nas[1:,:] = torch.diff(GRID_nas,axis=0)
            s_mat[:,1:] = torch.diff(GRID_mat)
        else: #last special... for now
            s_nas[:-1,:] = torch.diff(GRID_nas,axis=0)
            s_mat[:,:-1] = torch.diff(GRID_mat)

        if special_std == 'mean':
            s_nas[spec,:] = GRID_nas[spec,:] #not so sure about these -- need to be fixed......
            s_mat[:,spec] = GRID_mat[:,spec]
        elif special_std == 'neighbor': #assign_neighbor_to_special
            s_nas[spec,:] = s_nas[1,:] if first_special else s_nas[-2,:]
            s_mat[:,spec] = s_mat[:,1] if first_special else s_mat[:,-2]

        elif special_std == 'tail_prob':
            if first_special:
                print('If you are using this setting, you are doing something wrong.')
            t_max = torch.log(p[1]/p[2])/(p[1] - p[2])
            f = (torch.exp(-p[2]*t_max) - torch.exp(-p[1]*t_max)) * p[1]/(p[1] - p[2]) * p[0]
            tailratio = 1/(1+1/f) #the mature tail ratio
            s_mat[:,spec] = torch.sqrt(GRID_mat[:,spec] / (1-tailratio))
            tailratio = p[0]/(1+p[0]) #the nascent tail ratio
            s_nas[spec,:] = torch.sqrt(GRID_nas[spec,:] / (1-tailratio))
        elif special_std == 'tail_prob_all': #not sure this is gonna work
            t_max = torch.log(p[1]/p[2])/(p[1] - p[2])
            f = (torch.exp(-p[2]*t_max) - torch.exp(-p[1]*t_max)) * p[1]/(p[1] - p[2]) * p[0]
            tailratio = 1/(1+1/f) #the mature tail ratio
            s_mat = torch.sqrt(GRID_mat / (1-tailratio))
            tailratio = p[0]/(1+p[0]) #the nascent tail ratio
            s_nas = torch.sqrt(GRID_nas / (1-tailratio))
        else:
            print('did not specify a standard deviation convention!')

        if special_std != 'tail_prob_all':
            s_nas *= hyp
            s_mat *= hyp
        v_nas = s_nas**2
        v_mat = s_mat**2

        r_nas = GRID_nas**2/(v_nas-GRID_nas)
        p_nas = 1-GRID_nas/v_nas 
        r_mat = GRID_mat**2/(v_mat-GRID_mat)
        p_mat = 1-GRID_mat/v_mat 

        if kernel_type=='Famoye': #should use tail_prob_all... at least for now
            Z_nas = r_nas / (r_nas + GRID_nas * (1-np.exp(-1)))
            Z_mat = r_mat / (r_mat + GRID_mat * (1-np.exp(-1)))
            c_nas = (Z_nas)**r_nas
            c_mat = (Z_mat)**r_mat
            A_nas = GRID_nas * (Z_nas*np.exp(-1)-1)
            A_mat = GRID_mat * (Z_mat*np.exp(-1)-1)
            Lam = hyp/(c_nas*c_mat*A_nas*A_mat/(s_nas*s_mat))
#             print(c_nas*c_mat*A_nas*A_mat/(s_nas*s_mat)*Lam)
#             raise
    for i in range(npdf[0]):
        for j in range(npdf[1]):
 
            if kernel_type != 'correlated':
                lnas = -GRID_nas[i,j] + xgrid_nas * torch.log(GRID_nas[i,j]) - gammaln_xgrid_nas

                if p_nas[i,j] > 1e-10 and not poisson_kernel: #hmm doesn't seem to work
                    lnas += torch.special.gammaln(xgrid_nas+r_nas[i,j]) - torch.special.gammaln(r_nas[i,j]) \
                    - xgrid_nas*torch.log(r_nas[i,j] + GRID_nas[i,j]) + GRID_nas[i,j] \
                    + r_nas[i,j]*torch.log(1-p_nas[i,j])

                lmat =  - GRID_mat[i,j] + xgrid_mat * torch.log(GRID_mat[i,j]) - gammaln_xgrid_mat
                if p_mat[i,j] > 1e-10 and not poisson_kernel:
                    lmat += torch.special.gammaln(xgrid_mat+r_mat[i,j]) - torch.special.gammaln(r_mat[i,j]) \
                    - xgrid_mat*torch.log(r_mat[i,j] + GRID_mat[i,j]) + GRID_mat[i,j] \
                    + r_mat[i,j]*torch.log(1-p_mat[i,j]) #wasteful: we're recomputing a lot of stuff.
                if kernel_type == 'Famoye':
                    C = 1+Lam[i,j] * (torch.exp(-xgrid_nas[:,None]) - c_nas[i,j]) \
                        * (torch.exp(-xgrid_mat[None,:]) - c_mat[i,j])
                    Y += w[i*npdf[1] + j] * torch.exp(lnas[:,None] + lmat[None,:]) * C
                else:
                    Y += w[i*npdf[1] + j] * torch.exp(lnas[:,None] + lmat[None,:])
            elif kernel_type == 'correlated':
                Z = a[i,j] * (torch.log(a[i,j] + eps) - log_theta_mu_eps[i,j]) \
                    + xgrid_nas[:,None] * (torch.log(GRID_nas[i,j]  + eps) - log_theta_mu_eps[i,j]) \
                    + xgrid_mat[None,:] * (torch.log(GRID_mat[i,j]  + eps) - log_theta_mu_eps[i,j]) \
                    + torch.lgamma(a[i,j] + xgrid_nas[:,None] + xgrid_mat[None,:]) \
                    - torch.lgamma(a[i,j]) \
                    - gammaln_xgrid_nas[:,None] \
                    - gammaln_xgrid_mat[None,:]
#                 Y += w[i*npdf[1] + j] * Z
                Y += w[i*npdf[1] + j] * torch.exp(Z)

#     X = torch.exp(Y)
#     X[X<1e-16]=1e-16
    Y[Y<1e-16]=1e-16
#     raise
            #print('Y shape: ',Y.shape)
            #note convention change. Y = the predicted PMF is now returned in the same shape as the original histogram.
            #this is fine bc Y is flattened anyway later on down the line.
    return Y



# define NORM and YPRED_FUN
npdf = np.array([5,5])
cheb=True

if cheb:
    n = np.arange(npdf[0])
    q = np.flip((np.cos((2*(n+1)-1)/(2*npdf[0])*np.pi)+1)/2)
else:
    q = np.linspace(0,1,npdf[0]+2)[1:-1]
    
NORM_nas = torch.tensor(stats.norm.ppf(q))

if cheb:
    n = np.arange(npdf[1])
    q = np.flip((np.cos((2*(n+1)-1)/(2*npdf[1])*np.pi)+1)/2)
else:
    q = np.linspace(0,1,npdf[1]+2)[1:-1]

NORM_mat = torch.tensor(stats.norm.ppf(q))
  
# q = np.linspace(0,1,npdf[1]+2)[1:-1]
# n = np.arange(npdf[1])
# q = np.flip((np.cos((2*(n+1)-1)/(2*npdf[1])*np.pi)+1)/2)
#     print(q)
    
lnfactorial = torch.special.gammaln(torch.arange(1003))<jupyter_output><empty_output><jupyter_text>Load in how many training and testing files you want. Get data now also needs set size<jupyter_code>train_list, test_list = tr.get_data(set_size = 256,
                                    number_of_training_files=2,
                                 number_of_testing_files=1,total_files=3,file_path=path)<jupyter_output><empty_output><jupyter_text>More useful to define this stuff in a single place.`neighbor` special_std seem to break a lot. `tail_prob` works consistently. But it's suboptimal.<jupyter_code># get_ypred = lambda p,npdf,w,hyp : get_ypred_at_RT(p,npdf,w,hyp,quantiles='PRESET',\
#                                                   first_special=False,special_std='tail_prob',\
#                                                   poisson_kernel=False,kernel_type='correlated')
get_ypred = lambda p,npdf,w,hyp : get_ypred_at_RT(p,npdf,w,hyp,quantiles='PRESET',\
                                                  first_special=False,special_std='tail_prob_all',\
                                                  poisson_kernel=False,kernel_type='Famoye',\
                                                  use_global_corr=False)
t1 = time.time()
e,b,t,model = tr.train_MLP(train_list[:100],test_list[:40],num_epochs=2,
                           npdf=npdf,batchsize=20,get_ypred_at_RT=get_ypred,
                           metric='kld',learning_rate=1e-3,MLP=1,softmax_type='vanilla')
t2 = time.time()

print(t2-t1)



# IF YOU WANT TO SAVE
#path1 = 'DEFINE/PATH/TO/STORE'
#np.save(path1+'e_t',[e,t])
#torch.save(model.state_dict(),'./test_MODEL')<jupyter_output>Epoch Number: 0
Epoch Number: 1
33.7603235244751<jupyter_text>we gotta start saving some of the metadatai recommend using YYMMDD format, because that's sortable whereas MMDDYY is not<jupyter_code># torch.save(model.state_dict(),'./models/220323_10x11pdf_512params_5e_rad1softmax_cheb_logw_MODEL')<jupyter_output><empty_output><jupyter_text>Plotting!If you have already stored your model, load it in!<jupyter_code># pre-loaded model

#path = './dir_03032022/'
#npdf = [10,11]
#model = tr.my_MLP1(3,npdf[0]*npdf[1])
#model.load_state_dict(torch.load(path+'MODEL_kld'))
#model.eval()
#array = np.load(path+ 'mselog.npy',allow_pickle=True)
#e = array[0]
#t = array[1]
t
metric = 'kld'

tt.plot_training(e,t,metric)<jupyter_output><empty_output><jupyter_text>Metric Histogram<jupyter_code>n__ = 300

# i know the function still says test_klds, but now it calculates whatever metric you give it

metrics, metric_mean = tr.calculate_test_metrics(test_list=train_list[:n__],npdf=npdf,
                                        model=model,get_ypred_at_RT=get_ypred,
                                             metric = metric)
tt.plot_histogram(metrics,bins=10)
tt.plot_CDF(metrics)<jupyter_output><empty_output><jupyter_text>Plot Param Quantiles<jupyter_code>tt.plot_param_quantiles(metrics,train_list[:n__])<jupyter_output><empty_output><jupyter_text>Plot PMFs<jupyter_code>tt.plot_PMF_grid(file_list=train_list[0:4],npdf=npdf,nrows=2,
                 ncols=2,model=model,get_ypred_at_RT=get_ypred)
def plot_PMF_grid(file_list,npdf,nrows,ncols,model,get_ypred_at_RT,kld=True,plot_points=True,log=True):
    
    p_list,y_list = tr.load_training_data(file_list)
#     rand = np.zeros(nrows*ncols) 
    
# #     print(p_list)
#     for i in range(nrows*ncols):
#         rand[i] = random.randint(0,len(y_list))
    rand_inds = np.random.choice(len(y_list),nrows*ncols,replace=False)
    y = []
    Y = []
    
    for r in rand_inds:
        r = int(r)
        y_pred = tr.get_predicted_PMF(p_list=p_list,
                                   npdf=npdf,position=r,model=model,get_ypred_at_RT = get_ypred_at_RT)
        
        y.append(y_list[r])
        Y.append(y_pred)
    
    Y = [Y_.detach().numpy() for Y_ in Y]
    y = [y_.detach().numpy() for y_ in y]
    
    Y = [Y_.reshape(y[i].shape) for i,Y_ in enumerate(Y)]

    fig, ax1 = plt.subplots(nrows=nrows, ncols=2*ncols, figsize=(15, 15))
    k = 0
#     print(rand_inds)
    j_num = np.arange(0,ncols*2,2)
    
    for i in range(nrows):
        for j in j_num:
            y_ = y[k]
            Y_ = Y[k]
            cm='viridis'
            
            vmin = np.minimum(y_.min(),Y_.min())
            vmax = np.maximum(y_.max(),Y_.max())
            print(vmin)
            if log:
                vmin = np.log10(vmin)
                vmax = np.log10(vmax)
                ax1[i,j].imshow(np.log10(y_).T,cmap=cm,aspect='auto',vmin=vmin,vmax=vmax)
                ax1[i,j+1].imshow(np.log10(Y_).T,cmap=cm,aspect='auto',vmin=vmin,vmax=vmax)
            else:
                ax1[i,j].imshow((y_).T,cmap=cm,aspect='auto',vmin=vmin,vmax=vmax)
                ax1[i,j+1].imshow((Y_).T,cmap=cm,aspect='auto',vmin=vmin,vmax=vmax)
            
            if plot_points:
                p__ = 10**p_list[rand_inds[k]]
                MU, VAR, STD, xmax = tr.get_moments(p__)
                GRID_nas,GRID_mat = generate_grid_affine(npdf,VAR,MU,p__,quantiles='PRESET')
                ax1[i,j+1].scatter(GRID_nas[:],GRID_mat[:],s=5,c='k')
            ax1[i,j].invert_yaxis()
            ax1[i,j].set_title('True log-PMF & basis locations')
            ax1[i,j+1].invert_yaxis()
            ax1[i,j+1].set_title('Reconstructed log-PMF')
            
            if kld == True:
                kld_ = -np.sum(y_.flatten()*np.log(Y_.flatten()/y_.flatten()))
                ax1[i,j].title.set_text(f'KLD: {kld_}')
            k = k + 1

        
    fig.tight_layout()<jupyter_output><empty_output><jupyter_text>it would be nice to have these plotted in a deterministic order<jupyter_code>plot_PMF_grid(file_list=train_list[0:10],npdf=npdf,nrows=5,
                 ncols=2,model=model,get_ypred_at_RT=get_ypred,log=True,plot_points=False)
plot_PMF_grid(file_list=train_list[0:10],npdf=npdf,nrows=5,
                 ncols=2,model=model,get_ypred_at_RT=get_ypred,log=False,plot_points=False)
# train_list[0:10][0]<jupyter_output><empty_output><jupyter_text>fwiw it'd be even nicer to have it generate a full report every time it trains a model...<jupyter_code>params, ys = tr.load_training_data(train_list)
# plot individual PMF

p_num = 21

params[p_num]

tt.plot_PMF(params[p_num:p_num+1],ys[p_num],npdf,model,get_ypred_at_RT_affine,metric)<jupyter_output><empty_output><jupyter_text>To do : move in nnlsq code from `projects/gg220304_ml_cme_2d/learning_env.ipynb`<jupyter_code>params,ys = tr.load_training_data(train_list)


import scipy
from numpy import linalg

import oct2py
from oct2py import octave
oc = oct2py.Oct2Py()
nr=2
# npdf=np.asarray([10,11])
fig1,ax1=plt.subplots(nrows = nr,ncols=2,figsize=(8,nr*3))
for j in range(nr):
    ax1[j,0].imshow(np.log10(ys[j]).T,aspect='auto')
    ax1[j,0].invert_yaxis()
    Y_ = np.zeros((np.prod(ys[j].shape),np.prod(npdf)))
    for i in range(np.prod(npdf)):
        w_ = torch.zeros(np.prod(npdf))
        w_[i] = 1
        Y_[:,i] = get_ypred(params[j],npdf,w_,0.5).flatten()
        
#     c=np.matmul(Y_.T,Y_)
#     q=-np.dot(ys[j].flatten(),Y_)
    
#     lb = np.zeros(npdf.prod())
#     ub = np.ones(npdf.prod())
    
#     A = np.ones(npdf.prod())[None,:]
#     b = np.asarray([1])

#     x0 = np.random.rand(npdf.prod())
#     x0 /= np.sum(x0)
    
#     w_nnl = oc.qp(x0, c, q, A, b, lb, ub)
#     w_nnl[w_nnl<0]=0
    w_nnl = scipy.optimize.nnls(Y_,ys[j].flatten())[0]
    w_nnl /= w_nnl.sum()
#     w_nll = 
#     w_nnl = linalg.lstsq(Y_,ys[j].flatten(),rcond=None)[0] #not actually nonnegative least sq

    ypred_nnl = np.dot(Y_,w_nnl).reshape(ys[j].shape)
#     ypred_nnl[ypred_nnl<1e-16]=1e-16

    ax1[j,1].imshow(np.log10(ypred_nnl.T),aspect='auto')
    ax1[j,1].invert_yaxis()

#     p__ = 10**p_list[rand_inds[k]]
#     MU, VAR, STD, xmax = tr.get_moments(10**params[j])
#     GRID_nas,GRID_mat = generate_grid_affine(npdf,VAR,MU,10**params[j],quantiles='PRESET')
#     ax1[j,1].scatter(GRID_nas[:],GRID_mat[:],s=5,c='k')

    print(j)

ys[0].shape
fig1,ax1=plt.subplots(nrows = npdf[0],ncols=npdf[1],figsize=(15,15))
for i in range(npdf[0]):
    for j in range(npdf[1]):
        ker = Y_[:,i*npdf[1] + j].reshape(ys[nr-1].shape)
        ker = np.log10(ker)
        ax1[i,j].imshow(ker.T,aspect='auto')
        ax1[i,j].invert_yaxis()
#         ax1[i,j].plot(ker.sum(0))
# 
ker
xgrid_nas = np.arange(20)
xgrid_mat = np.arange(30)

# mu_nas = 10
# mu_mat = 20
p_nas = 0.4
p_mat = 0.2

gammaln_xgrid_nas = lnfactorial[1:(xgrid_nas[-1]+2)].detach().numpy()
gammaln_xgrid_mat = lnfactorial[1:(xgrid_mat[-1]+2)].detach().numpy()

lmat
VAR
MU
p = 10**params[0]
MU, VAR, STD, xmax = tr.get_moments(p)
p = p.detach().numpy()

GRID_nas,GRID_mat = generate_grid_affine(npdf,VAR,MU,p,quantiles='PRESET')

hyp = 0.4

fig1,ax1=plt.subplots(nrows = npdf[0],ncols=npdf[1],figsize=(15,15))
for i in range(npdf[0]):
    for j in range(npdf[1]):
        mu_nas = GRID_nas[i,j].detach().numpy()
        mu_mat = GRID_mat[i,j].detach().numpy()
        
        t_max = np.log(p[1]/p[2])/(p[1] - p[2])
        f = (np.exp(-p[2]*t_max) - np.exp(-p[1]*t_max)) * p[1]/(p[1] - p[2]) * p[0]
        tailratio = 1/(1+1/f) #the mature tail ratio
        v_mat = mu_mat/(1-tailratio)
        tailratio = p[0]/(1+p[0]) #the nascent tail ratio
        v_nas = mu_nas / (1-tailratio)

        s_nas = np.sqrt(v_nas)
        s_mat = np.sqrt(v_mat)
        
#         v_nas = s_nas**2
#         v_mat = s_mat**2

        r_nas = mu_nas**2/(v_nas-mu_nas)
        p_nas = 1-mu_nas/v_nas 
        r_mat = mu_mat**2/(v_mat-mu_mat)
        p_mat = 1-mu_mat/v_mat 
#         print('hello')
#         print(r_nas)
#         print(r_mat)
#         print(mu_nas)
#         print(r_nas)
#         print('hello')
#         print(mu_mat)
#         print(v_mat)
#         Z_nas = r_nas / (r_nas + mu_nas * (1-np.exp(-1)))
#         Z_mat = r_mat / (r_mat + mu_mat * (1-np.exp(-1)))
#         c_nas = (Z_nas)**r_nas
        c_nas = ((1-p_nas)/(1-p_nas * np.exp(-1)))**r_nas
#         c_mat = (Z_mat)**r_mat
        c_mat = ((1-p_mat)/(1-p_mat * np.exp(-1)))**r_mat
#         A_nas = mu_nas * (Z_nas*np.exp(-1)-1)
#         print(A_nas)
        A_nas = r_nas * p_nas * np.exp(-1) / (1-p_nas*np.exp(-1)) - r_nas * p_nas /(1-p_nas)
#         print(A_nas)
#         A_mat = mu_mat * (Z_mat*np.exp(-1)-1)
#         print(A_mat)
        A_mat = r_mat * p_mat * np.exp(-1) / (1-p_mat*np.exp(-1)) - r_mat * p_mat /(1-p_mat)
#         print(A_mat)
        Lam = hyp/(c_nas*c_mat*A_nas*A_mat/(s_nas*s_mat))

#         ker = Y_[:,i*npdf[1] + j].reshape(ys[nr-1].shape)

        lnas = -mu_nas + xgrid_nas * np.log(mu_nas) - gammaln_xgrid_nas

        if p_nas > 1e-10: #hmm doesn't seem to work
            lnas += scipy.special.gammaln(xgrid_nas+r_nas) - scipy.special.gammaln(r_nas) \
            - xgrid_nas*np.log(r_nas + mu_nas) +mu_nas \
            + r_nas*np.log(1-p_nas)

        lmat =  - mu_mat + xgrid_mat * np.log(mu_mat) - gammaln_xgrid_mat
#         print(lmat)
        if p_mat > 1e-10:
            lmat += scipy.special.gammaln(xgrid_mat+r_mat) - scipy.special.gammaln(r_mat) \
            - xgrid_mat*np.log(r_mat + mu_mat) + mu_mat \
            + r_mat*np.log(1-p_mat) #wasteful: we're recomputing a lot of stuff.
        ker1 = np.exp(lnas[:,None] + lmat[None,:])
        
#         print(type(Lam))
#         c_nas = ((1-p_nas)/(1-p_nas*np.exp(-1)))**r_nas
#         c_mat = ((1-p_mat)/(1-p_mat*np.exp(-1)))**r_mat
#         print(c_mat)
#         Lam = 0.6
#         print(c_nas*c_mat*A_nas*A_mat/(s_nas*s_mat)*Lam)
        C = 1 + Lam * (np.exp(-xgrid_nas[:,None]) - c_nas) * (np.exp(-xgrid_mat[None,:]) - c_mat)
    
#         print(c_mat)
#         raise
        ker = ker1*C

        
        print(ker.min()) #THIS SHOULD NEVER GO BELOW ZERO
#         ker = np.log10(ker)
#         ax1[i,j].imshow(ker.T,aspect='auto')
#         ax1[i,j].invert_yaxis()
#         raise
        ax1[i,j].plot(ker.sum(0))
        ax1[i,j].plot(ker.sum(1))
fig1.tight_layout()
# ker<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1596933591946061, 'fraction_numerical': 0.026045996120809087, 'mean_word_length': 2.5663702239789195, 'pattern_counts': {'":': 0, '<': 52, '<?xml version=': 0, '>': 52, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '136243', 'license': 'MIT', 'n_tokens_mistral': 9141, 'n_tokens_neox': 8793, 'n_words': 1316, 'path': '2D/gg220323_learning_env_affine.ipynb', 'repo_name': 'mtcarilli/CME_approximations', 'size': '21654'}