<repo_name>andrebeu/SEMCSW<file_sep>/cpu_arr_jobsub.cmd
#!/usr/bin/env bash

#SBATCH -t 2:59:00   # runs for 48 hours (max)  
#SBATCH -N 1         # node count 
#SBATCH -c 1         # number of cores 
#SBATCH --mem 4000
#SBATCH -o ./slurms/output.%j.%a.out


# module load pyger/0.9
conda init bash
conda activate sem

# nosplit=${1}
# condition=${2}
# learn_rate=${3}

# alfa=${4}
# lmda=${5}

slurm_arr_idx=${SLURM_ARRAY_TASK_ID}

param_str=`python get_param_jobsub.py ${slurm_arr_idx}`
echo ${param_str}

# submit job
srun python PY-gs-acc.py "${param_str}"

sacct --format="CPUTime,MaxRSS"

<file_sep>/PY-sem_seed_exp.py
# import sys
# import os
# import pandas as pd
# import numpy as np
# import torch as tr
# from CSWSEM import *

# # sweep params
# learn_rate = float(sys.argv[])
# alfa = float(sys.argv[])
# lmda = float(sys.argv[])

# stsize = 25

# ## gridsearch tag
# gs_name = 'gs2'

# ## looping 
# seedL = np.arange(51)
# condL = ['blocked','interleaved']

# dfL = [] # populate from each condition
# for nosplit,condition,seed in itertools.product([0,1],condL,seedL):
#     np.random.seed(seed)
#     tr.manual_seed(seed)
#     ## model tag
#     model_tag = 'nosplit_%i__cond_%s__learnrate_%.3f__alfa_%f__lmbda_%f__seed_%i'%(
#       int(nosplit),condition,learn_rate,alfa,lmda,seed)
#     print('\n\n**',model_tag,'**\n\n')

#     # params
#     exp_kwargs={
#         'condition':condition,
#         'n_train':160,
#         'n_test':40
#     }

#     sem_kwargs={
#         'nosplit':nosplit,
#         'stsize':stsize,
#         'lmda':lmda,
#         'alfa':alfa,
#         'learn_rate':learn_rate,
#         'seed':seed,
#     }

#     # setup
#     task = CSWTask()
#     sem = SEM(**sem_kwargs)

#     # run
#     exp,curr = task.generate_experiment(**exp_kwargs)
#     sem_data = sem.forward_exp(exp)

#     # record curriculum (not ideal, recording with every obs)
#     sem.data.record_exp('curriculum',curr)

#     # save
#     save_fpath = 'gsdata/%s/%s'%(gs_name,model_tag)
#     sem_data_df = pd.DataFrame(sem_data)
#     sem_data_df.to_csv(save_fpath) # just in case
#     dfL.append(sem_data_df)

# ## save again: prefer this
# full_tag = 'learnrate_%.3f__alfa_%f__lmbda_%f'%(learn_rate,alfa,lmda)
# save_fpath = 'gsdata/%s/%s'%(gs_name,full_tag)
# pd.concat(dfL).to_csv(save_fpath)
<file_sep>/README.md
# liketest (03/26/21)
- how to change interleaved without affecting blocked?
  - look into likelihood computation
    - setup 40 trials of interleaved vs single (blocked)
- note RNN size: realized that even when RNN size was set to 1, performance on blocked would not exhibit CI, suggesting learning happening in projection layers. to fix this, inputs are projected down to size 3. 
- refactor argument passing to RNN 
- made batch_sim and BIexp wrapper funs 
- gridsearch over pdim and stsize:
  - small B>I results when pdim=4,stsize=5
- todo:
  - refactor likelihood 
  - check effect of:
    - remove input layer
    - rnn vs gru vs lstm
    - update prior count by 1 vs 5
    - 
- result: good BIseparation on early trials when embed_size == stsize = 10 or 15 (lr=0.05)



# gs2 notes
- change to gridsearch procedure 
  - single job contains multiple seeds
  - parameter passing controlled `get_param_jobsub.py`
  - caution: this gridsearch procedure repeats LSTM sims for all SEM sims


# procedural notes
- NB-run_sem.pynb used for debugging and testing before launching gridsearch
  - code in here should be kept in a state that is amenable to gridsearching




# Model of Curriculum Effect on Schema Learning

## Results

- blocked > interleaved
- early > middle/late

## Modeling procedure

- Fit SEM on blocked/interleaved
- Evaluate on early/middle/late


** MOVING TO DIFFERENT FOLDER: when switching between AB and NF repos, github sometimes bugs out with finder: new files wont load. <file_sep>/code/utils_analysis.py
import os
import numpy as np
import torch as tr

from glob import glob as glob
import pandas as pd

from matplotlib import pyplot as plt

from CSWSEM import *





def make_gsdf_absem(gsname,save=False):
  gs_dir = "gsdata/%s/"%gsname
  fpathL = glob(gs_dir+'*')
  seed_df_L = []
  for fpath in fpathL:
    condition = fpath.split('/')[-1].split('__')[1].split('_')[1]
    seed_df = pd.read_csv(fpath)
    seed_df.loc[:,'model'] = ['LSTM','SEM'][sum(seed_df.loc[:,'nosplit']==1)>0]
    seed_df.loc[:,'condition'] = condition
    seed_df_L.append(seed_df)
  gsdf = pd.concat(seed_df_L)
  gsdf.index = np.arange(len(gsdf))
  gsdf.drop(columns=['Unnamed: 0','like','prior'])
  if save:
      gsdf.to_csv('gsdata/%s.csv'%gsname)
      print('saved %s.csv'%gsname)
  return gsdf


<file_sep>/cpu_jobsub.cmd
#!/usr/bin/env bash

#SBATCH -t 4:00:00   # runs for 48 hours (max)  
#SBATCH -N 1         # node count 
#SBATCH -c 1         # number of cores 
#SBATCH --mem 3000
#SBATCH -o ./slurms/output.%j.%a.out
#SBATCH -e ./slurms/err.%j.%a


# module load pyger/0.9
conda init bash
conda activate sem

model=${1}
cond=${2}
lr=${3}
epoch=${4}
alfa=${5}
lmbda=${6}
seed=${7}

srun python PY-batch_exp.py "${model}" "${cond}" "${lr}" "${epoch}" "${alfa}" "${lmbda}" "${seed}"

sacct --format="CPUTime,MaxRSS"

<file_sep>/code/test_sem_impl.py
from CSWSEM import *

seed = 0
np.random.seed(10)

task = CSWTask()

exp_kwargs={
    'condition':'single',
    'n_train':100,
    'n_test':10
}
exp = task.generate_experiment(**exp_kwargs)


learn_rate = 0.01
sch_kwargs={
    'stsize':10,
    'seed':seed,
    "learn_rate":learn_rate
}
sch = CSWSchema(**sch_kwargs)
event = exp[0]
event_hat = sch.forward(event)



sem_kwargs={
    'lmda':1,
    'alfa':10,
    'nosplit':True
}
sem = SEM(**sem_kwargs)
# sem.forward_trial(event)
sem.forward_exp(exp)













<file_sep>/code/PY-LSTMbatchexp.py
#!/usr/bin/env python
# coding: utf-8

# This script runs one param_set. Outputs results{}.csv and trialxtrial{}.csv

# In[1]:


import sys
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import seaborn as sns
from tqdm import tqdm
from sklearn import metrics
import pandas as pd

from sklearn.metrics import adjusted_rand_score
from sklearn.preprocessing import normalize
from sklearn.linear_model import LogisticRegression
from scipy.special import logsumexp
from scipy.stats import norm
from glob import glob


# In[2]:


from CSWSEM import generate_exp, seed_exp
from vanilla_lstm import VanillaLSTM
from sem.event_models import NonLinearEvent


# ### gridsearch params 

# In[3]:


# parameter search over lr, n_epochs, alpha, lambda
model_type        = str('LSTM')
seed              = int(sys.argv[1])
lr                = float(sys.argv[2])
n_epochs          = int(sys.argv[3])    


# In[4]:


log_alpha         = float(0.0)  # sCRP alpha is set in log scale
log_lambda        = float(0.0)  # sCRP lambda is set in log scale


# In[5]:


condition = 'blocked'
n_train = 160
n_test = 40


# In[6]:


save_dir = 'gsdata/lstm2_aw0/'
model_tag = "%s_cond_%s_lr_%.3f_nepchs_%i_alpha_%.3f-lambda_%.3f_seed_%i"%(
  model_type,condition,lr,n_epochs,log_alpha,log_lambda,seed
)
print(model_tag)


# ### SEM configuration

# In[7]:


optimizer_kwargs = dict(
    lr=lr, beta_1=0.9, beta_2=0.999, epsilon=1e-5, amsgrad=False
) 

f_opts=dict(
  batch_size=25, 
  batch_update=False, 
  dropout=0.0,
  l2_regularization=0.0, 
  n_epochs=n_epochs,
  optimizer_kwargs=optimizer_kwargs
)

# final param dict
sem_kwargs = dict(
  lmda=np.exp(log_lambda), 
  alfa=np.exp(log_alpha), 
  f_opts=f_opts, 
  f_class=VanillaLSTM
)


# # Run model
# 
# main fun call

# In[ ]:


""" 
main fun call
"""
stories_kwargs = {'actor_weight':0.0}
results, trialXtrial, _ = seed_exp( 
              sem_kwargs, stories_kwargs=stories_kwargs,
              model_type=model_type, 
              n_train=n_train, n_test=n_test,
              condition=condition,seed=seed,
)


# In[ ]:


# convert from JSON file format (dict) to pandas df
results = pd.DataFrame(results)
trialXtrial = pd.DataFrame(trialXtrial)


# # save
# 

# In[ ]:


results_fpath = save_dir + "results_" + model_tag + '.csv'
trial_fpath = save_dir + "trial_X_trial_" + model_tag + '.csv'


# In[ ]:


results.to_csv(results_fpath)
trialXtrial.to_csv(trial_fpath)

<file_sep>/get_param_jobsub.py
import sys
import itertools

"""
controls the gridsearch parameter space
given an index, return parameter string
""" 

param_set_idx = int(sys.argv[1])

lr = [0.01, 0.05]
alfa =[0.001, 0.01, 0.1, 1, 10, 100, 1000, 10000, 100000]
lmbda =[0.001, 0.01, 0.1, 1, 10, 100, 1000, 10000, 100000]

itrprod = itertools.product(lr,alfa,lmbda)

for idx,paramL in enumerate(itrprod):
  if idx == param_set_idx:
    print(" ".join([str(i) for i in paramL]))
    break


<file_sep>/loop_jobsub.cmd
#!/usr/bin/env bash

# parameter search over learn_rate, n_epochs, alpha, lambda
# declare -a learn_rate_arr=(0.001 0.005 0.01 0.05)
# declare -a epoch_arr=(1 2 4 8 16 32 64)

declare -a learn_rate_arr=(0.01 0.05 0.1)
declare -a condition_arr=("blocked" "interleaved")
declare -a alfa_arr=(0.01 0.1 1 10 100 1000 10000 100000)
declare -a lmda_arr=(0.01 0.1 1 10 100 1000 10000 100000)

## slurm array idx passed to py script: 
# builds product of params
# takes idx as input
# returns str param set


for learn_rate in "${learn_rate_arr[@]}"; do
  for condition in "${condition_arr[@]}"; do
    # LSTM
    sbatch --array=0-50 cpu_arr_jobsub.cmd "1" "${condition}" "${learn_rate}" "99" "99"
    for alfa in "${alfa_arr[@]}"; do
      for lmda in "${lmda_arr[@]}"; do
        # SEM
        sbatch --array=0-50 cpu_arr_jobsub.cmd "0" "${condition}" "${learn_rate}" "${alfa}" "${lmda}" 
      done
    done
  done
done
<file_sep>/runtest.py
import time
import os
import numpy as np
import torch as tr
from scipy.stats import norm
from CSWSEM import *

sem_kwargs={
  'nosplit':0,
  'alfa':10000000,
  'lmda':0.1,
  'seed':np.random.randint(99),
  'stsize':15,
  'learn_rate':0.05
}


exp_kwargs = {
    'condition':'single',
    'n_train':10,'n_test':2
    }
sem = SEM(**sem_kwargs)
task = CSWTask()
exp,curr = task.generate_experiment(**exp_kwargs)
sem_data = sem.forward_exp(exp,curr)


print('** DONE')<file_sep>/PY-sem_batch_exp.py
import sys
import os
import itertools
import pandas as pd
import numpy as np
import torch as tr
from CSWSEM import *

## input params
param_str = str(sys.argv[1])
learn_rate,alfa,lmda = param_str.split()
learn_rate = float(learn_rate)
alfa = float(alfa)
lmda = float(lmda)


stsize = 25

## gridsearch tag
gs_name = 'gs2'

## looping 
seedL = np.arange(51)
condL = ['blocked','interleaved']

dfL = [] # populate from each condition
for nosplit,condition,seed in itertools.product([0,1],condL,seedL):
    np.random.seed(seed)
    tr.manual_seed(seed)
    ## model tag
    model_tag = 'nosplit_%i__cond_%s__learnrate_%.3f__alfa_%f__lmbda_%f__seed_%i'%(
      int(nosplit),condition,learn_rate,alfa,lmda,seed)
    print('\n\n**',model_tag,'**\n\n')

    # params
    exp_kwargs={
        'condition':condition,
        'n_train':160,
        'n_test':40
    }

    sem_kwargs={
        'nosplit':nosplit,
        'stsize':stsize,
        'lmda':lmda,
        'alfa':alfa,
        'learn_rate':learn_rate,
        'seed':seed,
    }

    # setup
    exp_start_time = time.time()
    task = CSWTask()
    sem = SEM(**sem_kwargs)

    # run
    exp,curr = task.generate_experiment(**exp_kwargs)
    sem_data = sem.forward_exp(exp,curr)

    # record curriculum (not ideal, recording with every obs)
    sem.data.record_exp('condition',condition)
    exp_end_time = time.time()
    sem.data.record_exp('delta_time',exp_end_time-exp_start_time)

    # save
    sem_data_df = pd.DataFrame(sem_data)
    dfL.append(sem_data_df)
    # save_fpath = 'gsdata/%s/%s'%(gs_name,model_tag)
    # sem_data_df.to_csv(save_fpath) # just in case
    

## save again: prefer this
batch_tag = 'batch-learnrate_%.3f__alfa_%f__lmbda_%f'%(learn_rate,alfa,lmda)
save_fpath = 'gsdata/%s/%s'%(gs_name,batch_tag)
pd.concat(dfL).to_csv(save_fpath)
<file_sep>/code/PY-run_simulations.py
#!/usr/bin/env python
# coding: utf-8

# This script runs one param_set. Outputs results{}.csv and trialxtrial{}.csv

import sys

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import seaborn as sns
from tqdm import tqdm
from sklearn import metrics
import pandas as pd

from sklearn.metrics import adjusted_rand_score
from sklearn.preprocessing import normalize
from sklearn.linear_model import LogisticRegression
from scipy.special import logsumexp
from scipy.stats import norm
from glob import glob

from schema_prediction_task_9_8_20 import generate_exp, batch_exp
from vanilla_lstm import VanillaLSTM
from sem.event_models import NonLinearEvent

### simulation information

save_dir = 'data/gridsearch_full1/'


# parameter search over lr, n_epochs, alpha, lambda
model_type        = str(sys.argv[1])
lr                = float(sys.argv[2])
n_epochs          = int(sys.argv[3])    # what does this control?
log_alpha         = float(sys.argv[4])  # sCRP alpha is set in log scale
log_lambda        = float(sys.argv[5])  # sCRP lambda is set in log scale


model_tag = "%s-lr-%.3f-nepchs-%i-alpha-%.3f-lambda-%.3f"%(
  model_type,lr,n_epochs,log_alpha,log_lambda
)
print(model_tag)

# In[4]:

conditions = ['interleaved','blocked','early','middle','late']
n_batch = 50


# ### SEM configuration


# number of trials
n_train = 160
n_test = 40

story_kwargs = dict(seed=None, err=0.2, actor_weight=1.0, instructions_weight=0.0)

optimizer_kwargs = dict(
    lr=lr, beta_1=0.9, beta_2=0.999, epsilon=1e-5, amsgrad=False
) 

f_opts=dict(
  batch_size=25, 
  batch_update=False, 
  dropout=0.0,
  l2_regularization=0.0, 
  n_epochs=n_epochs,
    optimizer_kwargs=optimizer_kwargs
)

f_class = VanillaLSTM # event model class

# final param dict
sem_kwargs = dict(
  lmda=np.exp(log_lambda), 
  alfa=np.exp(log_alpha), 
  f_opts=f_opts, 
  f_class=f_class
)


# toggle between SEM (False) and LSTM (True)
if model_type == 'SEM':
  no_split=False
elif model_type == 'LSTM':
  no_split=True


""" 
batch_exp main fun call
"""

results, trialXtrial, _ = batch_exp(
              sem_kwargs, story_kwargs, no_split=no_split, 
              sem_progress_bar=False, progress_bar=False,
              n_train=n_train, n_test=n_test,
              n_batch=n_batch, conditions=conditions
)

# convert from JSON file format (dict) to pandas df
results = pd.DataFrame(results)
trialXtrial = pd.DataFrame(trialXtrial)


## save
results_fpath = save_dir + "results_" + model_tag + '.csv'
trial_fpath = save_dir + "trial_X_trial_" + model_tag + '.csv'


results.to_csv(results_fpath)
trialXtrial.to_csv(trial_fpath)

<file_sep>/CSWSEM.py
import time
import os
import numpy as np
import torch as tr

from scipy.stats import norm


"""
event (len6): B,L,2,3,4,E
event_hat (len5): Lh,2h,3h,4h,Eh
event_target (len5): Lt,...,Et
"""

## common across all objects
OBSDIM = 10

    

class SEM(object):

    def __init__(self, nosplit, lmda, alfa, stsize, learn_rate, seed):
        """
        """
        self.seed = seed
        tr.manual_seed(seed)
        np.random.seed(seed)
        # SEMBase.__init__()
        self.nosplit = nosplit
        # params
        self.lmda = lmda
        self.alfa = alfa
        self.obsdim = OBSDIM
        self.rnn_kwargs = {
            'stsize':stsize,
            'learn_rate':learn_rate,
            }     
        # collect sem data; locals() returns kwargs dict
        sem_kwargs = locals()
        sem_params = {**sem_kwargs}
        self.data = SEMData(sem_params)
        """
        schlib always has one "inactive" schema  
         similarly, dimensions of prior,likelihood,posterior
         will be len(schlib) == 1+num_active_schemas
        """
        self.active_schema_idx = 0
        self.prev_schema_idx = None
        self._init_schlib()
        self.in_layer = tr.nn.Linear(self.obsdim,self.obsdim)
        return None 

    def _init_schlib(self):
        self.prior = np.array([self.alfa,self.alfa])
        self.schlib = [
            CSWSchema(**self.rnn_kwargs),
            CSWSchema(**self.rnn_kwargs)
            ] 
        self.schema_count = np.array([1,0])
        self.active_schema = self.schlib[0]
        return None

    def get_logpriors(self):
        """ 
        3 cases:
            1) new/inactive schema:
                prior = alfa
            2) old schema, active on previous trial:
                prior = count + lamda
            3) old schema, not active on previous trial:
                prior = count
        - CRP calculation relies on `schema_count`
            which is an array that counts number of times (#obs)
            each schema was used 
        - len(prior) == len(schlib)
        """
        ## init prior
        prior = self.schema_count.copy().astype(float)
        ## active schemas
        if type(self.prev_schema_idx)==int:
            prior[self.prev_schema_idx] += self.lmda
        elif self.prev_schema_idx==None: # tstep0 
            prior[0] = self.alfa
        ## new/inactive schema 
        prior[-1] = self.alfa
        assert len(prior) == len(self.schlib)
        return np.log(prior)

    def get_loglikes(self,event):
        """ wrapper around schema.calc_like
            - active schemas + one inactive schema
        """
        loglikes = -99*np.ones(len(self.schlib))
        for sch_idx,schema in enumerate(self.schlib):
            loglikes[sch_idx] = schema.calc_event_loglike(event)
        return loglikes

    def get_active_schema_idx(self,log_prior,log_like):
        """ 
        splitting SEM uses argmax of posterior log probability
        nonsplitting SEM takes single event
        """
        if self.nosplit:
            return 0
        # edge case first trial
        if self.prev_schema_idx==None:
            return 0
        log_post = log_prior + log_like
        # print(log_prior,log_like,log_post)
        return np.argmax(log_post)

    def select_schema(self,log_prior,log_like):
        """ returns index of active schema
        updates previous_schema_index
        if new schema, append to library
        """
        # select model
        active_schema_idx = self.get_active_schema_idx(log_prior,log_like)
        # update previous schema index
        self.prev_schema_idx = active_schema_idx
        # if new active_schema, update schlib (need to wrap this)
        if active_schema_idx == len(self.schema_count)-1:
            self.schema_count = np.concatenate([self.schema_count,[0]])
            self.schlib.append(CSWSchema(**self.rnn_kwargs))
        self.data.record_trial('active_schema',active_schema_idx)
        # todo: move count update into select_schema
        self.schema_count[active_schema_idx] += 5 
        active_schema = self.schlib[active_schema_idx]
        return active_schema

    # run functions

    def forward_trial(self, event):
        """ 
        wrapper for within trial calculations 
            select schema using argmax(posterior)
            forward and back prop on active schema
        records trial data
        """
        # embed
        event = tr.Tensor(event).unsqueeze(1) # include batch dim
        assert event.shape == (6,1,10)
        # gradient step
        loss = self.active_schema.backprop(event)
        self.data.record_trial('loss',loss)
        # prior & likelihood of each schema
        log_priors = self.get_logpriors() 
        log_likes = self.get_loglikes(event) 
        assert len(log_priors) == len(log_likes) == len(self.schlib)
        # update active schema and update count 
        self.active_schema = self.select_schema(log_priors,log_likes)
        return None

    def forward_exp(self,exp,curr):
        """
        wrapper for run_trial
        """
        # loop over events
        print('active schema of trial t happens at t-1')
        lossL = []
        for trial_num,event in enumerate(exp):
            # print('\n\nTRIAL',trial_num,'nsc',len(self.schlib))
            self.data.new_trial(trial_num)
            self.data.record_trial('curriculum',curr[trial_num])
            self.forward_trial(event)
        # close data
        exp_data = self.data.finalize()
        return exp_data  



class CSWSchema(tr.nn.Module):

    def __init__(self,stsize,learn_rate):
        super().__init__()
        ## network parameters
        self.obsdim = OBSDIM
        self.stsize = stsize
        self.learn_rate = learn_rate
        # setup
        self._build()
        # backprop
        self.lossop = tr.nn.MSELoss()
        self.optiop = tr.optim.Adam(
            self.parameters(), lr=self.learn_rate)
        # init settings
        self.is_active = False
        # loglike params (from NF)
        # NB a function of obs dimension
        self.variance_prior_mode = 1/self.obsdim
        self.var_df0 = 1
        self.var_scale0 = 0.3
        self.sigma = np.ones(self.obsdim)/self.obsdim
        return None

    def _build(self):
        ## architecture setup
        # embedding handled within
        self.elayer = tr.nn.Linear(self.obsdim,self.stsize)
        # self.lstm = tr.nn.LSTM(self.embsize,self.stsize)
        self.gru = tr.nn.GRU(self.stsize,self.stsize)
        self.init_lstm = tr.nn.Parameter(tr.rand(2,1,1,self.stsize),requires_grad=True)
        self.out_layer = tr.nn.Linear(self.stsize,self.obsdim)
        return None

    # supervised 

    def forward(self,event):
        ''' 
        receives full event, returns ehat
            len(event)-1
        main wrapper 
        event is Tensor [tsteps,1,scene_dim]
        returns event_hat
        '''
        eventx = event[:-1]
        # prop
        ehat = eventx
        h_lstm,c_lstm = self.init_lstm 
        ehat = self.elayer(ehat)
        # ehat,(h_lstm,c_lstm) = self.lstm(ehat,(h_lstm,c_lstm))
        ehat,hn = self.gru(ehat,h_lstm)
        ehat = self.out_layer(ehat)
        # numpy mode for like
        assert len(ehat) == len(event)-1
        return ehat

    def backprop(self,event):
        """ update weights 
        """
        ehat = self.forward(event)
        etarget = event[1:]
        assert ehat.shape == etarget.shape == (5,1,10)
        ## 
        self.update_variance(ehat,etarget)
        ## back prop
        loss = 0
        self.optiop.zero_grad()
        for tstep in range(len(ehat)):
            loss += self.lossop(ehat[tstep],etarget[tstep])
            loss.backward(retain_graph=1)
        self.optiop.step()
        self.is_active = True
        # update variance
        return loss.detach().numpy()

    # like funs

    def calc_event_loglike_inactive(self,event):
        """ 
        - evaluate likelihood of each scene under normal_pdf
            summing over components (obsdim)
        - NB currently evaluating Begin,...,4
            i.e. leaveing out end
        returns int, loglike of entire event
        """
        eventx = event[:-1]
        assert self.is_active == False
        loglike_event = 0
        for scene in eventx:
            normal_pdf = norm(0, self.variance_prior_mode ** 0.5)
            scene_loglike = normal_pdf.logpdf(scene).sum()
            loglike_event += scene_loglike
        return loglike_event

    def calc_event_loglike(self,event):
        """ 
        NB self.sigma is a function of prediction error 
            and thus changes over time
        - not fully confident with handling of inactive schema 
        returns int, loglike of entire event
        """
        ## inactive schema:
        # if not self.is_active:
        #     return self.calc_event_loglike_inactive(event)
        # like of schema is function of eventPE
        ehat = self.forward(event).detach().numpy() 
        etarget = event[1:].detach().numpy() # rm BEGIN
        assert ehat.shape == etarget.shape == (5,1,10) # t,batch,obs
        event_pes = np.sum([etarget,-ehat],0).squeeze() # t,obs 
        # loop over scenes of event        
        loglike_event = 0
        for scene_pe in event_pes:
            loglike_event += self.fast_mvnorm_diagonal_logprob(scene_pe, self.sigma)
        return loglike_event

    # NF misc

    def fast_mvnorm_diagonal_logprob(self, x, variances):
        """
        Assumes a zero-mean mulitivariate normal with a diagonal covariance function
        Parameters:
            x: array, shape (D,)
                observations
            variances: array, shape (D,)
                Diagonal values of the covariance function
        output
        ------
            log-probability: float
        """
        log_2pi = np.log(2.0 * np.pi)
        return -0.5 * (log_2pi * np.shape(x)[0] + np.sum(np.log(variances) + (x**2) / variances ))

    def map_variance(self, samples, nu0, var0):
        """
        This estimator assumes an scaled inverse-chi squared prior over the
        variance and a Gaussian likelihood. The parameters d and scale
        of the internal function parameterize the posterior of the variance.
        Taken from Bayesian Data Analysis, ch2 (Gelman)

        samples: N length array or NxD array, where N is the number of 
                 samples and D is the dimensions
        nu0: prior degrees of freedom
        var0: prior scale parameter

        returns: float or D-length array, mode of the posterior

        ## Calculation ##

        the posterior of the variance is thus (Gelman, 2nd edition, page 50):
            
            p(var | y) ~ Inv-X^2(nu0 + n, (nu0 * var0 + n * v) / (nu0 + n) )

        where n is the sample size and v is the empirical variance.  The 
        mode of this posterior simplifies to:

            mode(var|y) = (nu0 * var0 + n * v) / (v0 + n + 2)

        which is just a weighted average of the two modes

        """

        # get n and v from the data
        n = np.shape(samples)[0]
        v = np.var(samples, axis=0)

        mode = (nu0 * var0 + n * v) / (nu0 + n + 2)
        return mode

    def update_variance(self,ehat,etarget):
        ehat = ehat.detach().numpy().squeeze()
        etarget = etarget.detach().numpy().squeeze()
        PE = ehat - etarget
        ##
        self.sigma = self.map_variance(PE, self.var_df0, self.var_scale0)
        return None



class CSWTask():
    """ 
    """

    def __init__(self   ):
        A1,A2,B1,B2 = self._init_paths()
        self.paths = [[A1,A2],[B1,B2]]
        # keep obs dim fixed: NF plate's formula 
        # calculations assumes 10 
        self.obsdim = OBSDIM
        self.tsteps = len(self.paths[0][0])
        self.exp_int = None
        return None

    # helper init
    def _init_paths(self):
        """ 
        begin -> locA -> node11, node 21, node 31, end
        begin -> locA -> node12, node 22, node 32, end
        begin -> locB -> node11, node 22, node 31, end
        begin -> locB -> node12, node 21, node 32, end
        """
        begin,locA,locB = 0,1,2
        node11,node12 = 3,4
        node21,node22 = 5,6
        node31,node32 = 7,8
        end = 9
        A1 = np.array([begin,locA,
            node11,node21,node31,end
            ])
        A2 = np.array([begin,locA,
            node12,node22,node32,end
            ])
        B1 = np.array([begin,locB,
            node11,node22,node31,end
            ])
        B2 = np.array([begin,locB,
            node12,node21,node32,end
            ])
        return A1,A2,B1,B2

    def _init_emat(self):
        self.n_obs = 11
        self.embed_mat = np.random.normal(
            loc=0., scale=1./np.sqrt(self.obsdim), 
            size=(self.n_obs, self.obsdim)
            )
        return None

    # used to generate exp
    def get_curriculum(self,condition,n_train,n_test):
        """ 
        order of events
        NB blocked: ntrain needs to be divisible by 4
        """
        curriculum = []   
        if condition == 'blocked':
            assert n_train%4==0
            curriculum =  \
                [0] * (n_train // 4) + \
                [1] * (n_train // 4) + \
                [0] * (n_train // 4) + \
                [1] * (n_train // 4 )
        elif condition == 'early':
            curriculum =  \
                [0] * (n_train // 4) + \
                [1] * (n_train // 4) + \
                [0, 1] * (n_train // 4)
        elif condition == 'middle':
            curriculum =  \
                [0, 1] * (n_train // 8) + \
                [0] * (n_train // 4) + \
                [1] * (n_train // 4) + \
                [0, 1] * (n_train // 8)
        elif condition == 'late':
            curriculum =  \
                [0, 1] * (n_train // 4) + \
                [0] * (n_train // 4) + \
                [1] * (n_train // 4)
        elif condition == 'interleaved':
            curriculum = [0, 1] * (n_train // 2)
        elif condition == 'single': ## DEBUG
            curriculum =  \
                [0] * (n_train) 
        else:
            print('condition not properly specified')
            assert False
        # 
        curriculum += [int(np.random.rand() < 0.5) for _ in range(n_test)]
        return np.array(curriculum)

    def embed_path(self,path_int):
        """ 
        given exp_int (tsteps)
        returns exp (ntrials,tsteps,obsdim)
        """
        return self.embed_mat[path_int]

    # main wrapper
    def generate_experiment(self,condition,n_train,n_test):
        """ 
        exp arr of events (vec)
        returns [n_train+n_test,tsteps,obsdim]
        """
        self._init_emat()
        # get curriculum
        n_trials = n_train+n_test
        curr = self.get_curriculum(condition,n_train,n_test)
        # generate trials
        exp = -np.ones([n_trials,self.tsteps,self.obsdim])
        self.exp_int = -np.ones([n_trials,self.tsteps],dtype=int)
        for trial_idx in range(n_train+n_test):
            # select A1,A2,B1,B2
            event_type = curr[trial_idx]
            path_type = np.random.randint(2)
            path_int = self.paths[event_type][path_type]
            # embed
            self.exp_int[trial_idx] = path_int
            exp[trial_idx] = self.embed_path(path_int)
        return exp,curr




class SEMData(object):

    def __init__(self, semparams):
        self.semparams = semparams
        self.sem_data = []
        return None

    def new_trial(self,trial_num):
        """ 
        init dict containing trial data 
        NB** trial_data is first appended to sem_data
            and is then modified at runtime.
        """
        self.trial_data = {'trial':trial_num}
        self.sem_data.append(self.trial_data) 

    def record_trial(self,key,value):
        """ populate trial dict 
        NB** trial_data is a dict that has already been 
        appended to sem_data. this method modifies self.trial_data 
        """
        self.trial_data[key] = value

    def record_exp(self,key,value):
        """ record experiment wide keyvalue in sem_data """
        for trial_dict in self.sem_data:
            trial_dict.update({key:value})
        return None

    def record_semparams(self):
        """ update each dict in sem_data
        to include sem semparams
        """
        semparams = {k:v for k,v in self.semparams.items() if k!='self'}
        for trial_dict in self.sem_data:
            trial_dict.update(semparams)
        return None

    def finalize(self):
        self.record_semparams()
        return self.sem_data




<file_sep>/code/PY-gs-acc.py
#!/usr/bin/env python
# coding: utf-8

import sys
import itertools
import os
import numpy as np
import torch as tr
import pandas as pd

from CSWSEM import *


## input params
gs_name = 'gs-acc'

## input params
param_str = str(sys.argv[1])
learn_rate,alfa,lmda = param_str.split()

learn_rate = float(learn_rate)
alfa = float(alfa)
lmda = float(lmda)

stsize = 25



num_seeds = 50
condL = ['blocked','interleaved']

exp_kwargs={
    'n_train':160,
    'n_test':40
}

# NB only SEM - no LSTM
sem_kwargs={
    'nosplit':False,
    'stsize':stsize,
    'lmda':lmda,
    'alfa':alfa,
    'learn_rate':learn_rate,
}


dataL = []
exphat = -np.ones([num_seeds,len(condL),200,5,10]) # (trials,tsteps,obsdim)

## loop over seeds and conditions (Blocked/interleaved)

for seed in np.arange(num_seeds):
  for cidx,cond in enumerate(condL):
    print(seed,cond)
    np.random.seed(seed)
    tr.manual_seed(seed)
    # setup
    sem_kwargs['seed'] = seed
    exp_kwargs['condition'] = cond
    task = CSWTask(seed)
    sem = SEM(**sem_kwargs)
    # run
    exp,curr = task.generate_experiment(**exp_kwargs)
    sem_data = sem.forward_exp(exp,curr)
    ### eval
    exphat[seed,cidx] = np.array([tdata['event_hat'] for tdata in sem_data])
    # record data
    sem.data.record_exp('condition',exp_kwargs['condition'])
    dataL.append(pd.DataFrame(sem_data))


## save

save_fpath = 'gsdata/%s/'%(gs_name)
batch_tag = 'learnrate_%.3f__alfa_%f__lmbda_%f'%(learn_rate,alfa,lmda)
np.save(save_fpath+'exphat__'+batch_tag,exphat)
pd.concat(dataL).to_csv(save_fpath+'semdata__'+batch_tag+'.csv')

