gsc2.py

# -*- coding: utf-8 -*-
"""
GSC Software: 
Author: Pyeong Whan Cho
Department of Cognitive Science, Johns Hopkins University
Summer 2015
(based on the MATLAB program, LDNet2.0)
"""
from math import *
import numpy as np
from matplotlib import pyplot as plt
from scipy import optimize as optim
import pandas as pd
import pylab as P
from scipy import linalg as la
from numbers import Number
import sys
import copy
import itertools
import pprint


# Nested dictionary
class Vividict(dict):
    #http://stackoverflow.com/questions/16333296/how-do-you-create-nested-dict-in-python
    def __missing__(self, key):
        value = self[key] = type(self)()
        return value

class RefSet():
    def __init__(self):
        self.dict = Vividict()
        
    def add(self, key, val, isSym=True, isTree=False, isSpanRole=False):
        if isTree:
            if isSpanRole:
                sys.exit('Not yet implemented. Check tree.py.')
            else:
                self.dict[key]['sym'] = ['/'.join(binding) for binding in val.recursiveFRtoString]
        else:
            if isSym:
                # val: a list of f/r bindings
                self.dict[key]['sym'] = val
            else:
                # val: a state vector (1d-array) in the pattern coordinate
                self.dict[key]['num'] = val
    
    def disp(self):
        pprint.pprint(self.dict)



class GscNet():
    def __init__(self, filler_names, role_names, grammar=None, 
                 WGC=None, bGC=None, extC=None, prev_extC=None, z=0.5, beta=4, 
                 q_init=0, q_rate=4, q_max=50, q_fun='plinear', c=0.5,  # check q_max
                 T_init=0, dt=0.001, reptype_r='local', reptype_f='local',
                 dp_r=0, dp_f=0, ndim_r=None, ndim_f=None, 
                 T_decay_rate=0, T_min=0, quant_list=None, grid_points=None,
                 F=None, R=None, getGPset=False):
        '''Construct an instance of GscNet.'''
        # filler_names: a list of filler names
        # role_names: a list of role names
        # hgRules: a list of harmonic grammar rules
        
        # Local representation        
        self.binding_names = bind(filler_names, role_names)
        self.filler_names = filler_names
        self.role_names = role_names
        self.nbindings = len(self.binding_names)    # number of bindings
        self.nfillers = len(filler_names)           # number of fillers
        self.nroles = len(role_names)               # number of roles
        
        # Network parameters (C-space) --- local representation
        if WGC is None:
            WGC = np.zeros(self.nbindings**2).reshape(self.nbindings, self.nbindings, order='F')
        if bGC is None:
            bGC = np.zeros(self.nbindings)#.reshape(self.nbindings)
        if extC is None:
            extC = np.zeros(self.nbindings)#.reshape(self.nbindings)
        if prev_extC is None:
            prev_extC = np.zeros(self.nbindings)#.reshape(self.nbindings)
        
        self.WGC = WGC            # weight matrix
        self.bGC = bGC            # bias vector
        self.extC = extC        # external input
        self.prev_extC = prev_extC
        
        # Network parameters (S-space) --- distributed representation
        # If 'local' is chosen, these parameters would be same as the above parameters.
        if ndim_r is None:
            ndim_r = self.nroles
        if ndim_f is None:
            ndim_f = self.nfillers
        
        self.reptype_r = reptype_r
        self.reptype_f = reptype_f
        self.ndim_r = ndim_r
        self.ndim_f = ndim_f
        
        # R: a matrix whose column vectors correspond to roles.
        # F: a matrix whose column vectors correspond to fillers.
        # TP: a matrix whose column vectors correspond to bindings.
        #     it is used for a state conversion from C- to S-space.
        # TPinv: the inverse of TP. Used for a state conversion from S- to C-Space.
        #        for now, only works for symmetric TP.
        self.R = encode_symbols(self.nroles, reptype=reptype_r, dp=dp_r, ndim=ndim_r)
        self.F = encode_symbols(self.nfillers, reptype=reptype_f, dp=dp_f, ndim=ndim_f)
        
        # If either R or F is provided, overwrite self.R or self.F.
        if R is not None:
            self.R = R
        if F is not None:
            self.F = F
            
        self.TP, self.TPinv = compute_TPmat(self.R, self.F)  # Pay attention to the argument order.
        
        self.nunits = ndim_r * ndim_f       # CHECK: should be same as the number of rows of TP
        self.act = np.zeros(self.nunits)    # activation state in S-space.
        self.actC = self.S2C()
        
#        self.W = compute_sspace_weights(self.WC, self)
#        self.b = compute_sspace_biases(self.bC, self)
#        self.ext = compute_sspace_biases(self.extC, self) # CHECK:
        
        # Add the bowl with bowl parameters (C-space is assumed)
        self.z = z                          # bowl center: vector or scalar
        self.beta = beta                    # bowl strength
        self.WBC = -self.beta * np.eye(self.nbindings)
        self.bBC = self.beta * self.z * np.ones(self.nbindings)
        #self.check_beta()
        
        self.WC = self.WGC + self.WBC
        self.bC = self.bGC + self.bBC
        
        self.WG = compute_sspace_weights(self.WGC, self)
        self.WB = compute_sspace_weights(self.WBC, self)
        self.bG = compute_sspace_biases(self.bGC, self)
        self.bB = compute_sspace_biases(self.bBC, self)
        
        self.W = compute_sspace_weights(self.WC, self)
        self.b = compute_sspace_biases(self.bC, self)
        self.zeta = self.TP.dot(self.z * np.ones(self.nbindings))
        self.ext = compute_sspace_biases(self.extC, self)
        self.prev_ext = compute_sspace_biases(self.prev_extC, self)
        
        self.check_beta()
        
        # Quantization dynamics parameter        
        self.q_init = q_init                # quantization strength
        self.q = copy.deepcopy(self.q_init) 
        self.q_rate = q_rate
        self.q_fun = q_fun
        self.q_max = q_max                  # used when qfun = 'linear'
        self.q_time = 0
        self.c = c                          # quantization parameter

        # Update parameters        
        self.T_init = T_init
        self.T = copy.deepcopy(self.T_init) # temperature
        self.T_decay_rate = T_decay_rate    # temperature decay rate
        self.T_min = T_min        
        self.dt = dt                        # time step constant
        
        # Clamping state
        self.clamped = False
        self.act_trace = None
        self.quant_list = quant_list

        # Test if W is symmetric
        if np.allclose(self.W, self.W.T) == False:
            sys.exit("The weight matrix (2D array) is not symmetric. Please check it.")

        # unit labels
        ndigits = len(str(self.nunits))
        self.unit_names = ['U' + str(ii+1).zfill(ndigits) for ii in list(range(self.nunits))]

        self.speed = None
        self.ema_speed = None
        self.grid_points = grid_points
        
        self.getGPset = getGPset
        if getGPset:
            self.all_grid_points()
        
    def randomize_state(self, minact=0, maxact=1):
        '''Set the activation state to a random vector inside a unit hypercube'''
        actC = np.random.uniform(minact, maxact, self.nbindings)
        self.act = self.C2S(actC)

    def set_init_state(self, mu=0.2, sd=0.2):
        actC = np.random.normal(loc=mu, scale=sd, size=self.nbindings)
        self.act = self.C2S(actC)

    def randomize_weights(self, mu=0, sigma=1):
        '''Randomize WC'''
        WC = sigma*np.random.randn(self.nbindings**2).reshape(self.nbindings, self.nbindings, order='F') + mu
        WC = (WC + WC.T)/2
        self.WC = WC
        self.W = compute_sspace_weights(self.WC, self)

    def all_grid_points(self):
        if self.quant_list is None:
            quant_list = [None] * self.nroles
            for rind, role in enumerate(self.role_names):
                quant_list[rind] = [self.binding_names[ii] for ii in self.find_roles(role)]        
        else:
            quant_list = self.quant_list
    
        gpset = list(itertools.product(*quant_list))
        gpset = [list(gp) for ii, gp in enumerate(gpset)]
        
        #all_grid_points = [None] * len(temp)
        gpset_gh = np.zeros(len(gpset))
        for ii, gp in enumerate(gpset):
            actC = np.zeros(self.nbindings)
            #actC[self.find_bindings(list(gp))] = 1.0
            actC[self.find_bindings(gp)] = 1.0
            gpset_gh[ii] = self.Hg(act=self.C2S(actC))
    
        idx = np.argsort(gpset_gh)[::-1]
        self.gpset = [gpset[ii] for ii in idx]
        self.gpset_gh = gpset_gh[idx]
        
    ##########################################################################
    ###
    ### set weights and biases 
    ###        
    ##########################################################################
        
    def set_weight(self, binding_name1, binding_name2, weight, symmetric=True):
        '''Set the weight of a connection between binding1 and binding2.
        When symmetric is set to True (default), the connection weight from 
        binding2 to binding1 is set to the same value.'''
        idx1 = self.find_bindings(binding_name1)
        idx2 = self.find_bindings(binding_name2)
        if symmetric:
            self.WGC[idx1, idx2] = self.WGC[idx2, idx1] = weight
        else:
            self.WGC[idx2, idx1] = weight
        self.WC = self.WGC + self.WBC
        self.WG = compute_sspace_weights(self.WGC, self)
        self.W = compute_sspace_weights(self.WC, self)

    def set_bias(self, binding_name, bias):
        '''Set the bias of a binding to [bias].'''
        idx = self.find_bindings(binding_name)
        self.bGC[idx] = bias
        self.bC = self.bGC + self.bBC
        self.bG = compute_sspace_biases(self.bGC, self)
        self.b = compute_sspace_biases(self.bC, self)
        
    def set_role_bias(self, role_name, bias):
        '''Set the bias of bindings of all fillers with particular roles to [bias].'''
        role_list = [bb.split('/')[1] for bb in self.binding_names]
        if not isinstance(role_name, list):
            role_name = [role_name]
        for jj, role in enumerate(role_name):
            idx = [ii for ii, rr in enumerate(role_list) if role == rr]
            self.bGC[idx] = bias
            
        self.bC = self.bGC + self.bBC
        self.bG = compute_sspace_biases(self.bGC, self)
        self.b = compute_sspace_biases(self.bC, self)
        
    def set_filler_bias(self, filler_name, bias):
        '''Set the bias of bindings of all roles with particular fillers to [bias].'''
        filler_list = [bb.split('/')[0] for bb in self.binding_names]
        if not isinstance(filler_name, list):
            filler_name = [filler_name]
        for jj, filler in enumerate(filler_name):
            idx = [ii for ii, ff in enumerate(filler_list) if filler == ff]
            self.bGC[idx] = bias
        self.bC = self.bGC + self.bBC
        self.bG = compute_sspace_biases(self.bGC, self)
        self.b = compute_sspace_biases(self.bC, self)

    # set_state()
    def set_state(self, binding_names, vals=1.0):
        idx = self.find_bindings(binding_names)
        self.actC = np.zeros(self.nbindings)
        self.actC[idx] = vals
        self.act = self.C2S()
    
    ##########################################################################
    ###
    ### Etc
    ###        
    ##########################################################################

    def vec2mat(self, act=None):
        if act is None:
            act = self.S2C()
        return act.reshape(self.nfillers, self.nroles, order='F')
        
    def C2S(self, C=None):
        '''Change basis: from C- to S-space.'''        
        if C is None:
            C = self.actC
        return self.TP.dot(C)
    
    def S2C(self, S=None):
        '''Change basis: from S- to C-space.'''        
        if S is None:
            S = self.act
        return self.TPinv.dot(S)

    def find_bindings(self, binding_names):
        '''Find the indices of the bindings from the list of binding names.'''
        
        if not isinstance(binding_names, list):
            binding_names = [binding_names]
        return [self.binding_names.index(bb) for bb in binding_names]
        
    def find_roles(self, role_name):
        
        if not isinstance(role_name, list):
            role_name = [role_name]
            
        role_list = [bb.split('/')[1] for bb in self.binding_names]
        role_idx = []
        for jj, role in enumerate(role_name):
            #idx = [ii for ii, rr in enumerate(role_list) if role_name == rr]
            idx = [ii for ii, rr in enumerate(role_list) if role == rr]
            role_idx += idx
        #idx_all = [[ii for ii, rr in enumerate(role_list) if role == rr] for jj, role in enumerate(role_name)]
        return role_idx
    
    def find_fillers(self, filler_name):
        
        if not isinstance(filler_name, list):
            filler_name = [filler_name]
            
        filler_list = [bb.split('/')[0] for bb in self.binding_names]
        filler_idx = []
        for jj, filler in enumerate(filler_name):
            idx = [ii for ii, ff in enumerate(filler_list) if filler == ff]
            filler_idx += idx
        return filler_idx
        
    def read_state(self, act=None):
        '''Print the current state (C-SPACE) in a readable format. Pandas should be installed.'''
        
        if act is None:
            act = self.act            
        #actC = self.S2C().reshape(self.nfillers, self.nroles, order='F')
        actC = self.vec2mat(self.S2C(act))
        print(pd.DataFrame(actC, index=self.filler_names, columns=self.role_names))
    
    def read_grid_point(self, act=None, disp=True):
        '''Print a grid point close to the current state. The grid point will be
        chosen by the snapping method (winner-takes-it-all).'''
        
        if act is None:
            act = self.act
            
        if self.quant_list is None:
            # By default, it is assumed that a role can be bound with only one filler.
            actC = self.vec2mat(self.S2C(act))
            winner_idx = np.argmax(actC, axis=0)
            winners = [self.filler_names[ii] for ii in winner_idx]
            winners = ["%s/%s" % bb for bb in zip(winners, self.role_names)]
        else:
            # quant_list: a list of lists of binding names. 
            actC = self.S2C(act)  # pattern coordinate
            winners = []
            for kk, group in enumerate(self.quant_list):
                idx = self.find_bindings(group)
                winner_idx = [idx[ii] for ii, jj in enumerate(actC[idx]) if jj == max(actC[idx])]
                winner = [self.binding_names[jj] for ii, jj in enumerate(winner_idx)]
                winners.extend(winner)
        if disp:
            print(winners)
        return winners
            
    def read_weight(self, which='WGC'):
        '''Print the weight matrix in a readable format (in the pattern coordinate).'''
        if which[-1] == 'C':
            print(pd.DataFrame(getattr(self, which), index=self.binding_names, columns=self.binding_names))
        else:
            print(pd.DataFrame(getattr(self, which), index=self.unit_names, columns=self.unit_names))
    
    def read_bias(self, which='bGC', print_vertical=True):
        '''Print the bias vector (in the pattern coordinate).'''
        if which[-1] == 'C':
            if print_vertical:
                print(pd.DataFrame(getattr(self, which).reshape(self.nbindings, 1), index=self.binding_names, columns=["bias"]))
            else:
                print(pd.DataFrame(getattr(self, which).reshape(1, self.nbindings), index=["bias"], columns=self.binding_names))
        else:
            if print_vertical:
                print(pd.DataFrame(getattr(self, which).reshape(self.nbindings, 1), index=self.unit_names, columns=["bias"]))
            else:
                print(pd.DataFrame(getattr(self, which).reshape(1, self.nbindings), index=["bias"], columns=self.unit_names))

    def hinton(self, which='WGC', label=True):
        '''Draw a hinton diagram'''
        # CHECK: which of W or W-beta*I should be visualized?
        if label:
            if which[-1] == 'C':
                labels=self.binding_names
            else:
                labels=self.unit_names
        else:
            labels=None
        hinton(getattr(self, which), xlabels=labels, ylabels=labels)
        
    def act_clamped(self, act=None):
        '''S-space'''
        if act is None:
            act = self.act
        return self.projmat.dot(act) + self.clampvec
    
    def check_beta(self, disp=False):
        'Compute and print the recommended beta value given the weights and biases in the C-space.'''
        eigvals, eigvecs = la.eigh(self.WGC) # WGC should be a symmetric matrix. So eigh() was used instead of eig()
        eig_max = max(eigvals)               # Condition 1: beta > eig_max to be stable
        if self.nbindings == 1:
            beta1 = -(self.bGC+self.extC)/self.z
            beta2 = (self.bGC+self.extC+eig_max)/(1-self.z)
        else:
            beta1 = -min(self.bGC+self.extC)/self.z # Condition 2: beta > beta1
            beta2 = (max(self.bGC+self.extC)+eig_max)/(1-self.z) # Condition 3: beta > beta2
        beta_min = max(eig_max, beta1, beta2)
        if self.beta <= beta_min:
            sys.exit("Beta (bowl strength) should be greater than %.4f." % beta_min)
        if disp:
            print('Recommended beta min = %.3f' % beta_min)

    def info(self):
        '''Print the network information.'''
        print('Fillers = ', self.filler_names)
        print('Roles = ', self.role_names)
        print('Num_of_units = ', self.nunits)
        print('Current T = ', self.T)
        # Add more information ...

    def compute_dist(self, ref_point, norm_ord=2, space='S'):
        """
        Compute the distance of the current state from a grid point. 
        [grid point] is a set of bindings.
        """
        
        idx = self.find_bindings(ref_point)
        destC = np.zeros(self.nbindings)
        destC[idx] = 1.0
        
        if space == 'S':
            state1 = self.act
            state2 = self.C2S(destC)
        elif space == 'C':
            state1 = self.S2C(self.act)
            state2 = destC
            
        return np.linalg.norm(state1-state2, ord=norm_ord)


    ##########################################################################
    ###
    ### Set or clear external input / Clamp or unclamp 
    ###        
    ##########################################################################
        
    def set_input(self, binding_names, ext_vals, inhib_comp=False):
        '''Set external input.'''

        # priming effect [CHECK]
        # self.prev_extC = copy.deepcopy(self.extC) # move to run()
        
        if not isinstance(ext_vals, list):
            ext_vals = [ext_vals]
        if not isinstance(binding_names, list):
            binding_names = [binding_names]
        if len(ext_vals) > 1:
            if len(binding_names) != len(ext_vals):
                sys.exit("binding_names and ext_vals have different lengths.")
                
        self.clear_input()  # Consider removing this line.
        if inhib_comp:
            # Check qunat_list
            if self.quant_list is None:
                # Assume filler competition in each role.
                # Find the current role.
                role_names = [b.split('/')[1] for b in binding_names]
                idx = self.find_roles(role_names)
                self.extC[idx] = -np.asarray(ext_vals) # -ext_vals (list object)
            else:
                for ii, bb in enumerate(binding_names):
                    group_id = [jj for jj, group in enumerate(self.quant_list) if bb in group]
                    for kk, gid in enumerate(group_id):
                        curr_bindings = self.quant_list[gid]
                        idx = self.find_bindings(curr_bindings)
                        self.extC[idx] = -np.asarray(ext_vals)   # -ext_vals (list object)                     
            
        idx = self.find_bindings(binding_names)
        self.extC[idx] = ext_vals        
        self.ext = compute_sspace_biases(self.extC, self)
        
    def clear_input(self):
        '''Remove external input.'''
        self.extC = np.zeros(self.nbindings)
        self.ext = compute_sspace_biases(self.extC, self)
        
    def clamp(self, binding_names, clamp_vals=1.0, clamp_comp=False):
        '''Clamp f/r bindings'''
        # A matrix of basic vectors each of which corresponds to a filler/role
        # binding whose activation state can change.
        if not isinstance(clamp_vals, list):
            clamp_vals = [clamp_vals]
        if not isinstance(binding_names, list):
            binding_names = [binding_names]
        if len(clamp_vals) > 1:
            if len(clamp_vals) != len(binding_names):
                sys.exit('The number of bindings clamped is not equal to the number of values provided.')
             
        self.clamped = True
        self.binding_names_clamped = binding_names
        clampvecC = np.zeros(self.nbindings)
        
        # ADD
        if clamp_comp:
            # Check qunat_list
            if self.quant_list is None:
                # Assume filler competition in each role.
                # Find the current role.
                role_names = [b.split('/')[1] for b in binding_names]
                idx1 = self.find_roles(role_names)
                clampvecC[idx1] = 0.0                
            else:
                for ii, bb in enumerate(binding_names):
                    group_id = [jj for jj, group in enumerate(self.quant_list) if bb in group]
                    for kk, gid in enumerate(group_id):
                        curr_bindings = self.quant_list[gid]
                        idx1 = self.find_bindings(curr_bindings)
                        clampvecC[idx1] = 0.0                       
         
        idx = self.find_bindings(binding_names)
        clampvecC[idx] = clamp_vals
        self.clampvecC = clampvecC
        
        if clamp_comp:
            idx += idx1 # CHECK        
            idx.sort()
        
        # Choose unclamped bindings. --- free to vary.
        idx0 = [bb for bb in np.arange(self.nbindings) if bb not in idx]
        A = self.TP[:, idx0] # complement set of idx (basis vectors of the subspace)
        if len(idx0) > 0:
            self.projmat = compute_projmat(A)
        else:
            self.projmat = np.zeros((self.nunits, self.nunits))
        self.clampvec = self.C2S(clampvecC)
        self.act = self.act_clamped(self.act)
        self.actC = self.S2C()
        
    def unclamp(self):
        if self.clamped is True:
            del self.clampvec
            del self.clampvecC
            del self.projmat
            del self.binding_names_clamped
            self.clamped = False
            
    ##########################################################################
    ###
    ### Harmony and harmony gradient
    ###        
    ##########################################################################        
        
#    def harmony(self, act=None, quant_list=None):
#        '''Compute the total harmony of the current state'''
#        if act is None:
#            act = self.act
#            
#        if quant_list is None:
#            quant_list = self.quant_list
#            
#        return harmony(net=self, act=act, quant_list=quant_list)

#    def H(self, act=None, quant_list=None):
#        '''Compute the total harmony of the current state'''
#        if act is None:
#            act = self.act            
#        if quant_list is None:
#            quant_list = self.quant_list
#        return H(net=self, act=act, quant_list=quant_list)
#
#    def HGrad(self, act=None, quant_list=None):
#        '''Compute the harmony gradient evaluated at the current state'''
#        if act is None:
#            act = self.act
#        if quant_list is None:
#            quant_list = self.quant_list
#        return HGrad(net=self, act=act, quant_list=quant_list)

    def H(self, act=None, quant_list=None):
        '''Compute the total harmony of the current state'''
        if act is None:
            act = self.act            
        if quant_list is None:
            quant_list = self.quant_list
        return Hg(net=self, act=act) + self.q * Q(net=self, n=act)

    def Hg(self, act=None):
        '''Compute the grammar harmony of the current state'''
        if act is None:
            act = self.act
        return Hg(net=self, act=act)

    def Hq(self, act=None, quant_list=None):
        if act is None:
            act = self.act            
        if quant_list is None:
            quant_list = self.quant_list
        return Q(net=self, n=act)

    def H0(self, act=None):
        '''Compute the grammar harmony of the current state'''
        if act is None:
            act = self.act
        return H0(net=self, act=act)

    def H1(self, act=None):
        '''Compute the grammar harmony of the current state'''
        if act is None:
            act = self.act
        return H1(net=self, act=act)

    def HGrad(self, act=None, quant_list=None):
        '''Compute the harmony gradient evaluated at the current state'''
        if act is None:
            act = self.act
        if quant_list is None:
            quant_list = self.quant_list
        return HgGrad(net=self, act=act) + self.q * QGrad(net=self, n=act)
        
    def HgGrad(self, act=None):
        '''Compute the grammar harmony of the current state'''
        if act is None:
            act = self.act
        return HgGrad(net=self, act=act)

    def H0Grad(self, act=None, quant_list=None):
        '''Compute the harmony gradient evaluated at the current state'''
        if act is None:
            act = self.act
        if quant_list is None:
            quant_list = self.quant_list
        return H0Grad(net=self, act=act, quant_list=quant_list)
        
    def H1Grad(self, act=None):
        '''Compute the grammar harmony of the current state'''
        if act is None:
            act = self.act
        return H1Grad(net=self, act=act)
        
    def HqGrad(self, act=None, quant_list=None):
        if act is None:
            act = self.act            
        if quant_list is None:
            quant_list = self.quant_list
        return QGrad(net=self, n=act)

#    def oharmony(self, act=None):
#        '''Compute the grammar harmony of the current state'''
#        if act is None:
#            act = self.act
#        return oharmony(net=self, act=act)

#    def qharmony(self, act=None, quant_list=None):
#        '''Compute the grammar harmony of the current state'''
#        if act is None:
#            act = self.act            
#        if quant_list is None:
#            quant_list = self.quant_list            
#        return q2harmony(net=self, act=act, quant_list=quant_list)

#    def Hq(self, act=None, quant_list=None):
#        '''Compute the grammar harmony of the current state'''
#        if act is None:
#            act = self.act            
#        if quant_list is None:
#            quant_list = self.quant_list            
#        return Hq(net=self, act=act, quant_list=quant_list)

#    def gharmony(self, act=None):
#        '''Compute the grammar harmony of the current state'''
#        if act is None:
#            act = self.act
#        return gharmony(net=self, act=act)

#    def bharmony(self, act=None):
#        '''Compute the grammar harmony of the current state'''
#        if act is None:
#            act = self.act
#        return bharmony(net=self, act=act)

#    def hgrad(self, act=None, quant_list=None):
#        '''Compute the harmony gradient evaluated at the current state'''
#        if act is None:
#            act = self.act
#        if quant_list is None:
#            quant_list = self.quant_list  
#        return hgrad(net=self, act=act, quant_list=quant_list)
    
    ##########################################################################
    ###
    ### Update state and parameters
    ###        
    ##########################################################################        
        
    #def update(self, update_T=True, update_q=True, q_fun='log', quant_list=None):
    def update(self, update_T=True, update_q=True, q_fun=None, space='S', norm_ord=2):
        
        if q_fun is None:
            q_fun = self.q_fun
        
        self.prev_act = copy.deepcopy(self.act)
        self.update_state()
        self.update_speed(space=space, norm_ord=norm_ord)
        if self.grid_points is not None:
            self.update_dist(self.grid_points, space=space, norm_ord=norm_ord)
            
        if update_T:
            self.update_T()
        if update_q:
            self.update_q(fun=q_fun)
        
    def update_state(self):
        '''Update the current state (with noise)'''
        
        #self.act += self.dt * self.hgrad(quant_list=self.quant_list) # update state (deterministic change)
        self.act += self.dt * self.HGrad(quant_list=self.quant_list) # update state (deterministic change)
        self.add_noise()                   # add noise (stochastic part)
        if self.clamped:
            self.act = self.act_clamped()
        self.actC = self.S2C()

    def add_noise(self):
        '''Add a noise to the current activation state.'''        
        noise = sqrt(2 * self.T * self.dt) * np.random.randn(self.nunits)
        self.act += noise
    
    def update_T(self, method='exponential'):
        '''Update the current temperature'''        
        # In the LDNet, the decay rate per timestep is adjusted based on the time step constant.
        # For simplicity, the decay rate is assumed to be per-step.
        if method == 'exponential':
            self.T = (1-self.T_decay_rate) * (self.T - self.T_min) + self.T_min
            
    def update_q(self, fun='plinear'):
        
        if fun == 'log':
            self.q_time += 1
            self.q = self.q_rate * log(self.q_time)
            # self.q = self.q_rate * log(exp(self.q/self.q_rate)+1)  # without q_time
        if fun == 'linear':
            self.q += self.q_rate
        if fun == 'plinear': 
            # piecewise linear
            if self.q <= self.q_max:
                self.q += self.q_rate
            
    def reset(self):
        
        self.q = self.q_init
        self.q_time = 0
        self.T = self.T_init
        self.randomize_state()
        
        # [CHECK]
        self.prev_extC = np.zeros(self.nbindings)
        self.prev_ext = compute_sspace_biases(self.prev_extC, self)
        self.unclamp()
        self.clear_input()

    def update_speed(self, ema_factor=0.1, norm_ord=2, space='S'):
        if space == 'S':
            diff = self.act - self.prev_act
        elif space == 'C':
            diff = self.S2C(self.act) - self.S2C(self.prev_act)
            
        self.speed = np.linalg.norm(diff, ord=norm_ord) #/ net.dt
        if self.ema_speed is None:
            self.ema_speed = self.speed
        else:
            self.ema_speed = ema_factor * self.ema_speed + (1-ema_factor) * self.speed
        #return self.ema_speed

    def update_dist(self, grid_points, norm_ord=2, space='S'):
        
        if not any(isinstance(tt, list) for tt in grid_points): 
            grid_points = [grid_points]

        dist = np.zeros(len(grid_points)) 
        for ii, grid_point in enumerate(grid_points):
            dist[ii] = self.compute_dist(ref_point=grid_point, norm_ord=norm_ord, space=space)
            
        self.dist = dist
        
    def update_traces(self):
        
        self.act_trace[self.step, :] = self.act
#        self.h_trace[self.step] = self.harmony()
#        self.oh_trace[self.step] = self.oharmony()
#        self.qh_trace[self.step] = self.qharmony()
#        self.gh_trace[self.step] = self.gharmony()
#        self.bh_trace[self.step] = self.bharmony()
        self.H_trace[self.step] = self.H()
        self.Hg_trace[self.step] = self.Hg()
        self.Hq_trace[self.step] = self.Hq()
        self.H0_trace[self.step] = self.H0()
        self.H1_trace[self.step] = self.H1()
        self.speed_trace[self.step] = self.speed
        self.ema_speed_trace[self.step] = self.ema_speed
        self.q_trace[self.step] = self.q
        self.T_trace[self.step] = self.T
        self.gp_trace[self.step] = self.read_grid_point(disp=False)
        # Consider storing winner indices instead of the F/R binding strings.
        # It will reduce the data size a lot
        
        if self.grid_points is not None:
            self.dist_trace[self.step, :] = self.dist
                
    # def update_input(self):
                
    ##########################################################################
    ###
    ### Run 
    ###        
    ##########################################################################        

    def run(self, maxstep, plot=False, grayscale=False, colorbar=True, 
            tol=None, ema_factor=0.1, update_T=True, update_q=True, q_fun=None, 
            grid_points=None, testvar='ema_speed', space='S', norm_ord=2, 
            ext_overlap=False, ext_overlap_steps=0):
        '''
        Run simulations for [maxstep] steps.
        '''
        if q_fun is None:
            q_fun = self.q_fun
            
        if grid_points is None:
            grid_points = self.grid_points

        # initialize 
        self.converged = False
        self.act_trace = np.zeros((maxstep+1, self.nunits))
#        self.h_trace = np.zeros(maxstep+1)
#        self.oh_trace = np.zeros(maxstep+1)
#        self.qh_trace = np.zeros(maxstep+1)
#        self.gh_trace = np.zeros(maxstep+1)
#        self.bh_trace = np.zeros(maxstep+1)
        self.H_trace = np.zeros(maxstep+1)
        self.Hg_trace = np.zeros(maxstep+1)
        self.Hq_trace = np.zeros(maxstep+1)
        self.H0_trace = np.zeros(maxstep+1)
        self.H1_trace = np.zeros(maxstep+1)
        self.speed_trace = np.zeros(maxstep+1)
        self.ema_speed_trace = np.zeros(maxstep+1)
        self.q_trace = np.zeros(maxstep+1)
        self.T_trace = np.zeros(maxstep+1)
        self.gp_trace = [[] for _ in range(maxstep+1)]
        
        if grid_points is not None:
            if not any(isinstance(pp, list) for pp in grid_points): 
                grid_points = [grid_points]
            n_grid_points = len(grid_points)
            self.dist_trace = np.zeros((maxstep+1, len(grid_points)))
        
        # Log initial state
        self.step = 0
        if grid_points is not None:
            self.dist = np.empty(n_grid_points)
            self.dist[:] = np.NAN
        self.update_traces()
        
        prev_extC = copy.deepcopy(self.prev_extC)
        curr_extC = copy.deepcopy(self.extC)
        for tt in np.arange(maxstep)+1:
            self.step = tt
            if ext_overlap and (tt <= ext_overlap_steps):
                # Linear overlap
                # update extC
                self.extC = (1 - tt/ext_overlap_steps) * prev_extC + (tt/ext_overlap_steps) * curr_extC
                self.ext = compute_sspace_biases(self.extC, self)
                
            self.update(update_T=update_T, update_q=update_q, q_fun=q_fun, 
                        space=space, norm_ord=norm_ord)
            self.update_traces()

            if tol is not None:
                self.check_convergence(tol=tol, testvar=testvar)
                if self.converged:
                    break
        
        self.rt = self.step
        #self.prev_extC = copy.deepcopy(self.extC) # this is done in the sim method.
        
        if plot:
            actC_trace = self.S2C(self.act_trace[:(self.rt+1),:].T).T
            heatmap(actC_trace.T, xlabel="Timestep", ylabel="Bindings", 
                    yticklabels=self.binding_names, grayscale=grayscale, colorbar=colorbar)

    def check_convergence(self, tol, testvar='ema_speed'):
        '''Check if the convergence criterion (distance vs. ema_speed) has been satisfied.'''
        
        if testvar == 'dist':               
            if (self.dist < tol).any():
                self.converged = True
                
        if testvar == 'ema_speed':
            if self.ema_speed < tol:
                self.converged = True

    def simulate(self):
        
        return 0
        
    def plot_act_trace(self, timesteps=None):
        
        if hasattr(self, 'act_trace'):
            if timesteps is None:
                timesteps = list(range(self.rt+1))
            
            actC_trace = self.S2C(self.act_trace[timesteps,:].T).T
            heatmap(actC_trace.T, xlabel="Timestep", ylabel="Bindings", 
                    xticklabels=None, yticklabels=self.binding_names, grayscale=False)
        else:
            sys.exit("There is no [act_trace] attribute in the current object.")

    def plot_trace(self, varname, timesteps=None):
        if timesteps is None:
            timesteps = list(range(self.rt+1))
        curr_trace = getattr(self, varname+'_trace')
        plt.plot(curr_trace[timesteps])
        plt.xlabel('Timestep', fontsize=16)
        plt.ylabel(varname, fontsize=16)
        plt.show()
        
#    def plot_dist_trace(self, timesteps=None):
#        if timesteps is None:
#            timesteps = list(range(self.rt+1))
#        plt.plot(self.dist_trace[timesteps, :])
#        plt.xlabel('Timestep', fontsize=16)
#        plt.ylabel('Distance', fontsize=16)
#        plt.show()
        
    ##########################################################################
    ###
    ### Update state and parameters
    ###        
    ##########################################################################        

    def optim(self, initvals=[], method='nelder-mead', options={'xtol': 1e-8, 'disp': True}):
        '''Find a local optimum given an initial guess (by default, the current state is used)'''
        # CHECK: does this work correctly when some bindings are clamped?
        if initvals==[]:
            initvals = self.act
        res = optim.minimize(lambda x: -harmony(self, x), initvals, method='nelder-mead', options=options)
        res.fun = -res.fun # Convert the (locally) minimal negative harmony (energy) to the (locally) maximal harmony
        return res

    def harmony_landscape(self, binding_name, minact, maxact):
        '''1D harmony landscape'''
        # To-do: Currently it works only with local representation.
        if isinstance(binding_name, list):
            if len(binding_name) > 1:
                sys.exit("Choose only one binding.!")
                
        idx = self.binding_names.index(binding_name)
        agrid = np.linspace(minact, maxact, 1000)
        hgrid = np.zeros(len(agrid))
        actC = self.S2C()
        for ii, aa in enumerate(agrid):
            act0 = actC
            act0[idx] = aa
            act = self.C2S(act0)
            hgrid[ii] = harmony(net=self, act=act)
        return agrid, hgrid

 
 
class WeightTP():
    
    def __init__(self, binding_names):
        self.binding_names = binding_names
        self.nbindings = len(binding_names)
        self.WGC = np.zeros((self.nbindings, self.nbindings))
        
    def set_weight(self, binding1, binding2, weight):
        idx1 = self.binding_names.index(binding1)
        idx2 = self.binding_names.index(binding2)
        self.WGC[idx1, idx2] = self.WGC[idx2, idx1] = weight
    
    def show(self):
        print(pd.DataFrame(self.WGC, index=self.binding_names, columns=self.binding_names))

class BiasTP():
    def __init__(self, binding_names):
        self.binding_names = binding_names
        self.nbindings = len(binding_names)
        self.bGC = np.zeros(self.nbindings)
    
    def set_bias(self, binding, bias):
        idx = self.binding_names.index(binding)
        self.bGC[idx] = bias
        
    def set_role_bias(self, role, bias):
        role_list = [b.split('/')[1] for b in self.binding_names]
        idx = [ii for ii, rr in enumerate(role_list) if role == rr]
        self.bGC[idx] = bias

    def set_filler_bias(self, filler, bias):
        filler_list = [b.split('/')[0] for b in self.binding_names]
        idx = [ii for ii, ff in enumerate(filler_list) if filler == ff]
        self.bGC[idx] = bias
        
    def show(self):
        print(pd.DataFrame(self.bGC, index=self.binding_names, columns=["bias"]))


# wh-question
# To-do: quant_list should be considered in read_grid_point()
class GscNet0():
    # local representation first.
    def __init__(self, net_list, group_names):
        self.ngroups = len(net_list)
        self.group_names = group_names
        self.groups = net_list
        
        self.binding_names = []
        for idx, net in enumerate(self.groups):
            self.binding_names = self.binding_names + [b + ':' + group_names[idx] for b in net.binding_names]
                    
        self.nbindings = len(self.binding_names)
        
        self.WGC = np.zeros((self.nbindings, self.nbindings))
        self.WBC = np.zeros((self.nbindings, self.nbindings))
        for ii, net in enumerate(self.groups):
            curr_bindings = [b + ':' + group_names[ii] for b in net.binding_names]
            idx = [self.binding_names.index(bb) for bb in curr_bindings]
            self.WGC[np.ix_(idx,idx)] = net.WGC
            self.WBC[np.ix_(idx,idx)] = net.WBC
        self.WC = self.WGC + self.WBC

        self.bGC = np.zeros(self.nbindings)
        self.bBC = np.zeros(self.nbindings)
        self.extC = np.zeros(self.nbindings)
        nunits = 0
        self.group_unit_idx = list()
        self.group_binding_idx = list()
        for ii, net in enumerate(self.groups):
            curr_bindings = [b + ':' + group_names[ii] for b in net.binding_names]
            idx = [self.binding_names.index(bb) for bb in curr_bindings]
            self.bGC[np.ix_(idx)] = net.bGC
            self.bBC[np.ix_(idx)] = net.bBC
            self.extC[np.ix_(idx)] = net.extC
            unit_idx = np.arange(net.nunits) + nunits
            self.group_unit_idx.append(unit_idx) # WRONG ....
            self.group_binding_idx.append(idx)
            nunits = nunits + net.nunits
        self.bC = self.bGC + self.bBC
        
        self.nunits = nunits
        
        self.W = self.compute_sspace_weights()
        self.b, self.ext = self.compute_sspace_b_and_ext()
        self.act = np.zeros(self.nunits)

    def compute_sspace_b_and_ext(self):
        b_distributed = np.zeros(self.nunits)
        ext_distributed = np.zeros(self.nunits)
        for ii, net in enumerate(self.groups):
            b_distributed[self.group_unit_idx[ii]] = net.b
            ext_distributed[self.group_unit_idx[ii]] = net.ext
        return b_distributed, ext_distributed
        
    # The program does not use this:
    def compute_sspace_weights(self):
        W_distributed = np.zeros((self.nunits, self.nunits))
        for ii, net1 in enumerate(self.groups):
            idx1 = self.group_unit_idx[ii]
            bidx1 = self.group_binding_idx[ii]
            for jj, net2 in enumerate(self.groups):
                idx2 = self.group_unit_idx[jj]
                bidx2 = self.group_binding_idx[jj]
                W_local = self.WC[np.ix_(bidx1, bidx2)]
                if ii == jj:
                    W_distributed[np.ix_(idx1, idx1)] = net1.W
                else:
                    W_temp = np.zeros((net1.nunits, net2.nunits)) # from net2 to net1
                    for kk in np.arange(0, net1.nbindings):
                        si = net1.TP[:, kk]
                        for ll in np.arange(0, net2.nbindings):
                            sj = net2.TP[:, ll]
                            #W_temp = W_temp + W_local[kk,ll] * (np.outer(si, sj) + np.outer(sj, si)) / (np.dot(si, si) * np.dot(sj, sj))
                            W_temp = W_temp + W_local[kk,ll] * np.outer(si, sj) / (np.dot(si, si) * np.dot(sj, sj))
                    W_distributed[np.ix_(idx1, idx2)] = W_temp
        return W_distributed

    def randomize_state(self, minact=0, maxact=1):
        '''Set the activation state to a random vector inside a unit hypercube'''
        for ii, net in enumerate(self.groups):
            net.randomize_state()
        
#    def randomize_weights(self, mu=0, sigma=1):
#        '''Randomize WC'''
#        WC = sigma*np.random.randn(self.nbindings**2).reshape(self.nbindings, self.nbindings, order='F') + mu
#        WC = (WC + WC.T)/2
#        self.WC = WC
#        self.W = compute_sspace_weights(self.WC, self)
        
    ##########################################################################
    ###
    ### set weights and biases (Use this only for setting the weights between networks)
    ###        
    ##########################################################################
        
    def set_weight(self, binding_name1, binding_name2, weight, symmetric=True):
        idx1 = self.binding_names.index(binding_name1)
        idx2 = self.binding_names.index(binding_name2)
        if symmetric:
            self.WGC[idx1, idx2] = self.WGC[idx2, idx1] = weight
        else:
            self.WGC[idx2, idx1] = weight
        self.WC = self.WGC + self.WBC
        self.W = self.compute_sspace_weights()

    def set_bias(self, binding_name, bias):
        idx = self.binding_names.index(binding_name)
        self.bGC[idx] = bias
        self.bC = self.bGC + self.bBC
        self.b, self.ext = self.compute_sspace_b_and_ext()
        
    def set_role_bias(self, role_name, bias):
        role_list = [bb.split('/')[1] for bb in self.binding_names]
        idx = [ii for ii, rr in enumerate(role_list) if role_name in rr]
        self.bGC[idx] = bias
        self.bC = self.bGC + self.bBC
        self.b, self.ext = self.compute_sspace_b_and_ext()
        
    def set_filler_bias(self, filler_name, bias):
        filler_list = [bb.split('/')[0] for bb in self.binding_names]
        idx = [ii for ii, ff in enumerate(filler_list) if filler_name in ff]
        self.bGC[idx] = bias
        self.bC = self.bGC + self.bBC
        self.b, self.ext = self.compute_sspace_b_and_ext()
        
    ##########################################################################
    ###
    ### Etc --- CHECK if it works
    ###        
    ##########################################################################

    def set_seed(self, num):        
        np.random.seed(num)
        
    def find_bindings(self, binding_names):
        '''Find the indices of the bindings from the list of binding names.'''
        
        if not isinstance(binding_names, list):
            binding_names = [binding_names]
        return [self.binding_names.index(bb) for bb in binding_names]
        
    def read_state(self):
        '''Print the current state (C-SPACE) in a readable format. Pandas should be installed.'''
        for ii, net in enumerate(self.groups):
            print("Group: %s" % self.group_names[ii])
            net.read_state()
            print("\n")
    
    def read_grid_point(self):
        # CHECK: current method assumes filler competition in each role.
        # quant_list should be considered
        for ii, net in enumerate(self.groups):
            print("Group: %s" % self.group_names[ii])
            print(net.read_grid_point(disp=False))
            print("\n")
        
    def read_weight(self):
        print(pd.DataFrame(self.WC, index=self.binding_names, columns=self.binding_names))
        print(pd.DataFrame(self.bC.reshape(1, self.nbindings), index=["bias"], columns=self.binding_names))
    
    def hinton(self, label=True):
        '''Draw a hinton diagram'''
        # CHECK: which of W or W-beta*I should be visualized?
        if label:
            labels=self.binding_names
        else:
            labels=None            
        hinton(self.WC, xlabels=labels, ylabels=labels)
        
#    def act_clamped(self, act=None):
#        # Not used. Consider removing this part.
#        if act is None:
#            act = self.act
#        return self.projmat.dot(act) + self.clampvec
#    
#    def check_beta(self):
#        # Unused.
#        # To-do: eig() vs. eigh()
#        eigvals, eigvecs = la.eigh(self.W) # W should be a symmetric matrix.
#        eig_max = max(eigvals)               # Condition 1: beta > eig_max to be stable
#        if self.nbindings == 1:
#            beta1 = -(self.b+self.ext)/self.z
#            beta2 = (self.b+self.ext+eig_max)/(1-self.z)
#        else:
#            beta1 = -min(self.b+self.ext)/self.z # Condition 2: beta > beta1
#            beta2 = (max(self.b+self.ext)+eig_max)/(1-self.z) # Condition 3: beta > beta2
#        beta_min = max(eig_max, beta1, beta2)
#        if self.beta <= beta_min:
#            sys.exit("Beta (bowl strength) should be greater than %.4f." % beta_min)
#
#    def info(self):
#        '''Print the network information.'''
#        print('Fillers = ', self.fillers)
#        print('Roles = ', self.roles)
#        print('Num_of_units = ', self.nunits)
#        print('Current T = ', self.T)
#        # Add more information ...
                
    ##########################################################################
    ###
    ### Set or clear external input / Clamp or unclamp 
    ###        
    ##########################################################################
        
    def set_input(self, binding_names, ext_vals):
        if len(binding_names) != len(ext_vals):
            sys.exit("binding_names and ext_vals have different lengths.")
        self.extC = np.zeros(self.nbindings)
        ind = self.find_bindings(binding_names)
        self.extC[ind] = ext_vals
        self.ext = compute_sspace_biases(self.extC, self)
        
    def clear_input(self):
        self.extC = np.zeros(self.nbindings)
        self.ext = compute_sspace_biases(self.extC, self)
        
    def clamp(self, binding_names, clamp_vals):
        '''Clamp f/r bindings'''
        # A matrix of basic vectors each of which corresponds to a filler/role
        # binding whose activation state can change.
        if len(clamp_vals) != len(binding_names):
            sys.exit('The number of bindings clamped is not equal to the number of values provided.')
            
        self.clamped = True
#        self.binding_names_clamped = binding_names
#        
#        idx = self.find_bindings(binding_names)
#        clampvecC = np.zeros(self.nbindings)
#        clampvecC[idx] = clamp_vals
#        self.clampvecC = clampvecC
#        
#        # Choose unclamped bindings. --- free to vary.
#        idx0 = [bb for bb in np.arange(self.nbindings) if bb not in idx]
#        A = self.TP[:, idx0] # complement set of idx (basis vectors of the subspace)
#        self.projmat = compute_projmat(A)
#        self.clampvec = self.TP.dot(clampvecC)
#        self.act = self.act_clamped(self.act)
        
        # CHECK: clamping each group
        temp = [b.split(':') for b in binding_names]
        group_num = [self.group_names.index(b[1]) for b in temp]
        bb_names = [b[0] for b in temp]
        
        for ii, gg in enumerate(set(group_num)):
            idx1 = [jj for jj, val in enumerate(group_num) if val==gg]
            curr_bindings = [bb_names[jj] for jj in idx1]
            curr_vals = [clamp_vals[jj] for jj in idx1]
            self.groups[gg].clamp(binding_names=curr_bindings, clamp_vals=curr_vals)
        
    def unclamp(self):
        self.clamped = False
        for ii, net in enumerate(self.groups):
            if net.clamped:
                net.unclamp()
        
#    def clamp_group(self, group_name, clamp_vals):
#        self.group_clamped = group_name
#        group_idx = self.group_names.index(group_name)
#        binding_names = [b + ':' + group_name for b in self.groups[group_idx].binding_names]
#        self.clamp(binding_names, clamp_vals)

    ##########################################################################
    ###
    ### Update state and parameters
    ###        
    ##########################################################################        
        
    def update(self, update_T=True, update_q=True):
        '''Update the current state (with noise) for each group'''
        for ii, net1 in enumerate(self.groups):     # to
            extC = np.zeros(net1.nbindings)
            binding_names1 = [b + ':' + self.group_names[ii] for b in net1.binding_names]
            idx1 = [self.binding_names.index(b) for b in binding_names1]
            for jj, net2 in enumerate(self.groups): # from
                # between group weights 
                # curr_W = self.W[idx1, idx2]
                if not (ii==jj):
                    binding_names2 = [b + ':' + self.group_names[jj] for b in net2.binding_names]            
                    idx2 = [self.binding_names.index(b) for b in binding_names2]
                    curr_WGC = self.WGC[np.ix_(idx1, idx2)]
                    extC = extC + curr_WGC.dot(net2.S2C(net2.act))
            net1.ext = net1.C2S(extC)
            
        for ii, net in enumerate(self.groups):
            net.update(update_T = update_T, update_q = update_q)
            self.act[self.group_unit_idx[ii]] = net.act
                    
#    def update_state(self):
#        '''Update the current state (with noise)'''
#        self.act += self.dt * self.hgrad() # update state (deterministic change)
#        self.add_noise()                   # add noise (stochastic part)
#        if self.clamped:
#            self.act = self.act_clamped()
#
#    def add_noise(self):
#        '''Add a noise to the current activation state.'''
#        noise = sqrt(2 * self.T * self.dt) * np.random.randn(self.nunits)
#        self.act += noise
    
#    def update_T(self, method='exponential'):
#        '''Update the current temperature'''
#        # In the LDNet, the decay rate per timestep is adjusted based on the time step constant.
#        # For simplicity, the decay rate is assumed to be per-step.
#        if method == 'exponential':
#            self.T = (1-self.T_decay_rate) * (self.T - self.T_min) + self.T_min
#                
#    def update_q(self, method='log'):
#        if method == 'log':
#            self.q_time += 1
#            self.q = self.q_factor * log(self.q_time)
            
    def reset(self):
        for ii, net in enumerate(self.groups):
            net.reset()

    ##########################################################################
    ###
    ### Harmony and harmony gradient
    ###        
    ##########################################################################        
        
    def harmony(self, act=None, quant_list=None):
        '''Compute the total harmony of the current state'''
        if act is None:
            act = self.act
        return harmony(net=self, act=act, quant_list=quant_list)
        
    def gharmony(self, act=None):
        '''Compute the grammar harmony of the current state'''
        if act is None:
            act = self.act
        return gharmony(net=self, act=act)

    def hgrad(self, act=None, quant_list=None):
        '''Compute the harmony gradient evaluated at the current state'''
        if act is None:
            act = self.act
        return hgrad(net=self, act=act, quant_list=quant_list)


    ##########################################################################
    ###
    ### Run 
    ###        
    ##########################################################################        

    def run(self, maxstep, tol=None):
        self.rt, self.final_state, self.act_trace, self.ema_speed_trace = run_GscNet0(self, maxstep=maxstep, tol=tol)

    def plot_act_trace(self):
        for ii, net in enumerate(self.groups):
            curr_act_trace = self.act_trace[:, self.group_unit_idx[ii]]
            actC_trace = net.TPinv.dot(curr_act_trace.T).T
            print("Group: %s" % self.group_names[ii])
            heatmap(actC_trace.T, xlabel="Timestep", ylabel="Bindings", yticklabels=net.binding_names, grayscale=False, colorbar=True)

#    def run(self, maxstep, tol=None):
#        if tol is not None:
#            self.rt, self.final_state, self.act_trace, self.t_trace, self.q_trace, self.ema_speed_trace = run(self, maxstep=maxstep, tol=tol)
#        else:
#            self.rt, self.final_state, self.act_trace, self.t_trace, self.q_trace, self.ema_speed_trace = run(self, maxstep=maxstep)
#
#    def plot_act_trace(self):
#        if hasattr(self, 'act_trace'):
#            heatmap(self.act_trace.T, xlabel="Timestep", ylabel="Bindings", 
#                    xticklabels=None, yticklabels=self.binding_names, grayscale=False)
#        else:
#            sys.exit("There is no [act_trace] attribute in the current object.")

    ##########################################################################
    ###
    ### Optim (CHECK)
    ###        
    ##########################################################################        
 
    def gharmony_given_clamping(self, x, group_name_clamped, clampval, quant_list, q):
        x0 = np.zeros(self.nbindings)
        self.clamp_group(group_name_clamped, clampval)
        x0[self.clampind] = self.clampvec[self.clampind]
        ind = list(set(np.arange(self.nbindings)) - set(self.clampind))
        x0[ind] = x
        H = 0.5 * x0.dot(self.W.dot(x0)) + x0.dot(self.b)
        for jj, subgroup in enumerate(quant_list):
            temp_ind = [self.binding_names.index(bb) for bb in subgroup]
            qx = x0[temp_ind]
            H = H + q* self.qharmony(qx)
        return H

    def optimal_state_given_clamping(self, group_name_clamped, clampval, quant_list, q, method="Nelder-Mead"):
        group_ind = self.group_names.index(group_name_clamped)
        if method=="basinhopping":
            res = optim.basinhopping(lambda x: -self.gharmony_given_clamping(x, group_name_clamped, clampval, quant_list, q), 
                                     np.random.uniform(size = self.nbindings - self.groups[group_ind].nbindings))
        else:
            res = optim.minimize(lambda x: -self.gharmony_given_clamping(x, group_name_clamped, clampval, quant_list, q), 
                           np.random.uniform(size = self.nbindings - self.groups[group_ind].nbindings), 
                           method = method)
        print(res)
        x0 = np.zeros(self.nbindings)
        x0[self.clampind] = self.clampvec[self.clampind]
        ind = list(set(np.arange(self.nbindings)) - set(self.clampind))
        x0[ind] = res.x
        self.optim_state = x0
        print(x0)    


# CHECK: for now, consider all units' act change:
def run_GscNet0(net, maxstep, plot=False, grayscale=False, colorbar=True, 
                tol=None, ema_factor=0.1, crop=True, update_q=True):
    # To-do: needs to include conversion from distributed to local representations
    # ema factor
    ngroups = len(net.groups)
    act_trace = np.zeros((maxstep, net.nunits))
    
    if tol is None:
        ema_speed_trace = np.zeros(maxstep)
        prev_act = copy.deepcopy(net.act)
        ema_speed = 0
        ema_factor = 0.1
        for tt in np.arange(maxstep): #range(1, maxstep+1):
            net.update(update_q=update_q)
            act_trace[tt, :] = net.act
        ema_speed_trace = -1

    else:
        if tol < 0:
            sys.exit("The tolerance parameter should be a positive real number.")
        else:
            ema_speed_trace = np.zeros(maxstep)
            prev_act = copy.deepcopy(net.act)
            ema_speed = 0
            ema_factor = ema_factor
            for tt in np.arange(maxstep): #range(1, maxstep+1):
                net.update(update_q=update_q)
                #diff = sqrt(sum((net.act - prev_act)**2))/net.nbindings    # CHECK: nunits vs. nbindings
                diff = max(abs(net.act - prev_act)) # divide it by dt?
                if tt == 0:
                    ema_speed = diff
                else:
                    ema_speed = ema_factor * ema_speed + (1-ema_factor) * diff
                act_trace[tt, :] = net.act
                ema_speed_trace[tt] = ema_speed
                prev_act = copy.deepcopy(net.act)
                
                if (ema_speed < tol) and (tol > 0):
                    break
                
            ema_speed_trace = ema_speed_trace[tt]
    
    if plot:
        for ii, net in enumerate(net.groups):
            curr_act_trace = act_trace[:, net.group_unit_idx[ii]]
            actC_trace = net.S2C(curr_act_trace[:tt,:].T).T
            heatmap(actC_trace.T, xlabel="Timestep", ylabel="Bindings", yticklabels=net.binding_names, grayscale=grayscale, colorbar=colorbar)

#    return [tt+1, net.act, act_trace[:tt,], t_trace[:tt], q_trace[:tt], ema_speed_trace[:tt]]
    return [tt+1, net.act, act_trace[:tt,], ema_speed_trace]



# Harmony

# Total Harmony = H_G + q * H_Q
# H_G = H_0 + H_1
# H_0: Harmonic Grammar
# H_1: Bowl
# H_Q: Quantization

def bind(fillernames, rolenames, sep='/'):
    return [f + sep + r for r in rolenames for f in fillernames]


def H(net, act, quant_list=None):
    # Evaluate total harmony at [act] in neural coordinates.
    # For not, quant_list will be ignored.
    # act: neural coordinate (n in Paul's derivation)
    # actC: conceptual coordinate (a in Paul's derivation)
    return Hg(net, act) + net.q * Q(n=act, net=net)

def HGrad(net, act, quant_list=None):
    return HgGrad(net, act) + net.q * QGrad(n=act, net=net) 

def Hg(net, act):
    # act: neural coordinate of the current state
    # WG: weights in neural coordinates
    return H0(net, act) + H1(net, act) # + constant

def HgGrad(net, act):
    return H0Grad(net, act) + H1Grad(net, act)

def H0(net, act):
    # act: neural coordinate of the current state
    # WG: weights in neural coordinates
    return 0.5 * act.dot(net.WG).dot(act) + (net.bG + net.ext).dot(act)

def H0Grad(net, act):
    return net.WG.dot(act) + net.bG + net.ext

def H1(net, act):
    # zeta: bowl center in neural coordinates
    return -0.5 * net.beta * (act - net.zeta).dot(act - net.zeta)
    #return 0.5 * act.dot(net.WB).dot(act) + net.bB.dot(act)

def H1Grad(net, act):
    return -net.beta * (act - net.zeta)



#def Hq(net, act, quant_list):
#    # Compute quantization harmony Q. By default, it considers filler competition in each role.
#    # For now, use dynamics in the pattern coordinates. 
#    # Should be replaced with distributed dynamics.
#    actC = net.S2C(act)
#        
#    if quant_list is None:
#        # By default, assumes filler-competition in each role.
#        if net.nfillers == 1:
#            const1 = 0
#            const2 = -np.sum(actC**2 * (1-actC)**2)
#            Q = const2
#        else:
#            const1 = -np.sum((np.sum(actC.reshape(net.nfillers,net.nroles,order='F')**2, axis=0)-1)**2)
#            const2 = -np.sum(actC**2 * (1-actC)**2)
#            Q = net.c * const1 + (1-net.c) * const2
#    else:
#        Q = 0
#        for jj, group in enumerate(quant_list):
#            temp_ind = [net.binding_names.index(bb) for bb in group]
#            actC0 = actC[temp_ind]
#            Q = Q + q2(net, actC0)
#    return Q

#def HqGrad(net, act, quant_list=None):
#    # Should be replaced with distributed dynamics.
#    actC = net.S2C(act)
#    if quant_list is None:
#        # By default, assume filler competition in each role.
#        if net.nfillers == 1:
#            # One-unit network case
#            # 'Sum_of_square = 1' does not make sense.
#            q_grad = -2 * actC * (actC-1) * (2*actC-1)
#        else:
#            const1 = -4 * actC * ( np.tile( np.sum(actC.reshape(net.nfillers,net.nroles,order='F')**2, axis=0), (net.nfillers,1) ).flatten('F') - 1)
#            const2 = -2 * actC * (actC-1) * (2*actC-1)
#            q_grad = net.c * const1 + (1-net.c) * const2
#    else:
#        q_grad = np.zeros(net.nbindings)
#        for jj, subgroup in enumerate(quant_list):
#            curr_q_grad = np.zeros(net.nbindings)
#            temp_ind = [net.binding_names.index(bb) for bb in subgroup]
#            actC0 = actC[temp_ind]
#            curr_q_grad[temp_ind] = q2grad(net, actC0)
#            q_grad = q_grad + curr_q_grad
#    return net.C2S(q_grad)






#def harmony(net, act, quant_list=None):
#    '''Compute total harmony at a given activation state.'''
#    # gharmony = H_0, bharmony = H_1, q2haromny = H_Q
#    return gharmony(net, act) + bharmony(net, act) + net.q * q2harmony(net, act, quant_list)

#def oharmony(net, act):
#    '''Compute total harmony at a given activation state.'''
#    #return gharmony(net, act) + bharmony(net, act)
#    return 0.5 * act.dot(net.W).dot(act) + act.dot(net.b) + act.dot(net.ext) # distributed version
#    #return 0.5 * net.S2C(act).dot(net.WC).dot(net.S2C(act)) + net.S2C(act).dot(net.bC) + net.S2C(act).dot(net.extC)

#def gharmony(net, act):
#    '''Compute harmony (given a harmonic grammar) at a given activation state.'''
#    return 0.5 * act.dot(net.WG).dot(act) + (net.bG + net.ext).dot(act)
    
#def bharmony(net, act):
#    '''Compute bowl harmony. act should be an array object.'''
#    return -0.5 * net.beta * (net.S2C(act) - net.z).dot(net.S2C(act) - net.z)

def q2harmony(net, act=None, quant_list=None):
    '''Compute quantization harmony Q. By default, it considers filler competition in each role.'''
    if act is None:
        act = net.act
    actC = net.S2C(act)
        
    if quant_list is None:
        # By default, assumes filler-competition in each role.
        if net.nfillers == 1:
            const1 = 0
            const2 = -np.sum(actC**2 * (1-actC)**2)
            Q = const2
        else:
            const1 = -np.sum((np.sum(actC.reshape(net.nfillers,net.nroles,order='F')**2, axis=0)-1)**2)
            const2 = -np.sum(actC**2 * (1-actC)**2)
            Q = net.c * const1 + (1-net.c) * const2
    else:
        Q = 0
        for jj, group in enumerate(quant_list):
            temp_ind = [net.binding_names.index(bb) for bb in group]
            actC0 = actC[temp_ind]
            Q = Q + q2(net, actC0)
    return Q

def q2(net, actC):
    # Within-group competition
    # actC = net.TPinv.dot(act)
    const1 = -np.sum(actC**2-1)**2
    const2 = -np.sum(actC**2 * (1-actC)**2)
    if len(actC) == 1:
        return const2
    else: 
        return net.c * const1 + (1-net.c) * const2

#def hgrad(net, act, quant_list=None):
#    #return ghgrad(net, act) + bhgrad(net, act) + net.q * q2hgrad(net, act, quant_list)
#    return ohgrad(net, act) + net.q * q2hgrad(net, act, quant_list)

#def ghgrad(net, act):
#    return net.WG.dot(act) + net.bG + net.ext

#def bhgrad(net, act):
#    return -net.beta * (net.S2C(act) - net.z)
#    #return net.WB.dot(net.act) + net.bB # should be same -net.beta * (net.S2C(act) - net.z)

#def ohgrad(net, act):
#    return net.W.dot(act) + net.b + net.ext

def q2hgrad(net, act, quant_list=None):
    actC = net.S2C(act)
    if quant_list is None:
        if net.nfillers == 1:
            const1 = 0
            const2 = -2 * actC * (actC-1) * (2*actC-1)
        else:
            const1 = -4 * actC * ( np.tile( np.sum(actC.reshape(net.nfillers,net.nroles,order='F')**2, axis=0), (net.nfillers,1) ).flatten('F') - 1)
            const2 = -2 * actC * (actC-1) * (2*actC-1)
        q_grad = net.c * const1 + (1-net.c) * const2
    else:
        q_grad = np.zeros(net.nbindings)
        for jj, subgroup in enumerate(quant_list):
            curr_q_grad = np.zeros(net.nbindings)
            temp_ind = [net.binding_names.index(bb) for bb in subgroup]
            actC0 = actC[temp_ind]
            curr_q_grad[temp_ind] = q2grad(net, actC0)
            q_grad = q_grad + curr_q_grad
    return net.C2S(q_grad)

def q2grad(net, actC):
    '''q2grad for each group'''
    const1 = -4 * actC * (np.sum(actC**2) - 1)
    const2 = -2 * actC * (actC-1) * (2*actC-1)
    return net.c * const1 + (1-net.c) * const2

# Hinton diagram from ( http://wiki.scipy.org/Cookbook/Matplotlib/HintonDiagrams )
# The functions were slightly modified for consistency
def _blob(x,y,area,colour):
    """
    Draws a square-shaped blob with the given area (< 1) at
    the given coordinates. 
    http://wiki.scipy.org/Cookbook/Matplotlib/HintonDiagrams
    """
    # Replaced N with np for consistency
    hs = np.sqrt(area) / 2
    xcorners = np.array([x - hs, x + hs, x + hs, x - hs])
    ycorners = np.array([y - hs, y - hs, y + hs, y + hs])
    P.fill(xcorners, ycorners, colour, edgecolor=colour)

def hinton(W, maxWeight=None, xlabels=None, ylabels=None):
    """
    Draws a Hinton diagram for visualizing a weight matrix. 
    Temporarily disables matplotlib interactive mode if it is on, 
    otherwise this takes forever.
    http://wiki.scipy.org/Cookbook/Matplotlib/HintonDiagrams
    """
    # Replaced N with np for consistency
    # xrange was replaced with range because it is not available in Python3
    # To do: Consider adding the f/r binding information as X and YTick labels
    reenable = False
    if P.isinteractive():
        P.ioff()
    P.clf()
    height, width = W.shape
    if not maxWeight:
        maxWeight = 2**np.ceil(np.log(np.max(np.abs(W)))/np.log(2))

    P.fill(np.array([0,width,width,0]),np.array([0,0,height,height]),'gray')
    #P.axis('off')
    #P.axis('equal')
    for x in range(width):
        for y in range(height):
            _x = x+1
            _y = y+1
            w = W[y,x]
            if w > 0:
                _blob(_x - 0.5, height - _y + 0.5, min(1,w/maxWeight),'white')
            elif w < 0:
                _blob(_x - 0.5, height - _y + 0.5, min(1,-w/maxWeight),'black')
    if reenable:
        P.ion()
        
    # PWC
    label = False
    if xlabels is not None:
        P.xticks(np.arange(len(xlabels))+0.5, xlabels, rotation='vertical')
        label = True

    if ylabels is not None:
        P.yticks(np.arange(len(ylabels))+0.5, ylabels[::-1])
        label = True
    
    if label:
        P.xlim([0,width])
        P.ylim([0,height])
        P.gca().set_aspect('equal', adjustable='box')
        P.gca().tick_params(direction='out')
        #P.axis('equal')
    else:
        P.axis('off')
        P.axis('equal')
        
    P.show()

    
def heatmap(data, xlabel=None, ylabel=None, xticklabels=None, yticklabels=None, 
            grayscale=False, colorbar=True):
    # Plot the activation trace as heatmap
    if grayscale:
        cmap = plt.cm.get_cmap("gray_r")
    else:
        cmap=plt.cm.get_cmap("Reds")        
    plt.imshow(data, cmap=cmap, interpolation="nearest", aspect='auto')
    if xlabel is not None:
        plt.xlabel(xlabel, fontsize=16)
    if ylabel is not None:
        plt.ylabel(ylabel, fontsize=16)
    if xticklabels is not None:
        plt.xticks(np.arange(len(xticklabels)), xticklabels)
    if yticklabels is not None:
        plt.yticks(np.arange(len(yticklabels)), yticklabels)
    if colorbar:
        plt.colorbar()
    plt.show()
    
# create symbol representations
def encode_symbols(nsymbols, reptype='local', dp=0, ndim=None):
    if reptype == 'local':
        sym_mat = np.eye(nsymbols)
    else:
        if ndim==None:
            ndim = nsymbols

        if isinstance(dp, Number):
            # dp is a scalar that represents the expected pairwise similarity (dot product).
            sym_mat = dot_products(nsymbols, ndim, dp)
        else:
            sym_mat = dot_products2(nsymbols, ndim, dp)
    return sym_mat

def dot_products(nsymbols, ndim, s):
    dp_mat = s * np.ones((nsymbols, nsymbols)) + (1-s) * np.eye(nsymbols, nsymbols)
    sym_mat = dot_products2(nsymbols, ndim, dp_mat)
    return sym_mat
    
def dot_products2(nsymbols, ndim, dp_mat, max_iter = 100000):
    # Given square matrix dpMatrix of dimension N-by-N, find N
    # dim-dimensional unit vectors whose pairwise dot products match
    # dpMatrix. Results are returned in the columns of M. itns is the
    # number of iterations of search required, and may be ignored.
    #
    # Algorithm: Find a matrix M such that M'*M = dpMatrix. This is done
    # via gradient descent on a cost function that is the square of the
    # frobenius norm of (M'*M-dpMatrix).
    #
    # NOTE: It has trouble finding more than about 16 vectors, possibly for
    # dumb numerical reasons (like stepsize and tolerance), which might be
    # fixable if necessary.

    if not (dp_mat.T == dp_mat).all():
        # print error
        sys.exit('dot_products2: dp_mat must be symmetric')
    
    if (np.diag(dp_mat) != 1).any():
        sys.exit('dot_products2: dp_mat must have all ones on the main diagonal')
    
    sym_mat = np.random.uniform(size=ndim*nsymbols).reshape(ndim, nsymbols, order='F')
    min_step = .1
    tol = 1e-6
    #max_iter = 100000
    converged = False
    for iter_num in range(1, max_iter+1):
        inc = sym_mat.dot(sym_mat.T.dot(sym_mat) - dp_mat)
        step = min(min_step, .01/abs(inc).max())
        sym_mat = sym_mat - step * inc
        max_diff = abs(sym_mat.T.dot(sym_mat)-dp_mat).max()
        if max_diff <= tol:
            converged = True
            break
        
    if not converged:
        print("Didn't converge after %d iterations" % max_iter)
        
    return sym_mat

def compute_TPmat(R, F):
    # TP matrix that converts local to distributed representations (C-space to S-space).
    # See http://en.wikipedia.org/wiki/Vectorization_(mathematics) for justification of kronecker product.
    TP = np.kron(R,F)
    if TP.shape[0] == TP.shape[1]:
        TPinv = la.inv(TP)
    else:
        TPinv = la.pinv(TP)  # TP may be a non-square matrix. So use pseudo-inverse.
    return TP, TPinv

def compute_sspace_biases(b_local, net):
    return net.TPinv.T.dot(b_local)

def compute_sspace_weights(W_local, net):
    return net.TPinv.T.dot(W_local).dot(net.TPinv)
    
def compute_projmat(A):
    # A is an n x m matrix of basis (column) vectors of the subspace.
    # Only works if the rank of A is equal to the nunmber of columns of A.
    return A.dot(la.inv(A.T.dot(A))).dot(A.T)

#class sim_params():
#    
#    def __init__(self):
#        
#        self.testvar = 'ema_speed'
#        self.convergence_check = True
#        self.norm_ord = np.inf
#        self.dist_space = 'S'
#        self.tol = 1e-2
#        self.ema_factor = 0.1
#        self.grid_points = None
#        self.input_method = 'ext'
#        self.update_T = True
#        self.update_q = True
#        self.maxstep = np.array(10000)
#        self.nrep = 1
#        self.seed_num = time.time()
#        
#    def set_params(self, name, val):
#        
#        setattr(self, name, val)

class sim():
    # Use this class to run simulations and store the results.
    def __init__(self, net, params):
        # net: an object of the GscNet class.
        # params: a dictionary for parameter setting.
        self.net = net
        self.params = params
        
        # Before running, save all gridpoints
        if self.net.getGPset:
            self.net.all_grid_points()
            self.gpset = copy.deepcopy(self.net.gpset)
            self.gpset_gh = copy.deepcopy(self.net.gpset_gh)
        
    def set_params(self, p_name, p_val):
        self.params[p_name] = p_val
        
    def set_seed(self, num):        
        np.random.seed(num)
        
    def set_input(self, input_list, input_vals):
        
        if not isinstance(input_list, list):
            input_list = [input_list]
            
        nwords = len(input_list)
        self.inputC = np.zeros((self.net.nbindings, nwords))
        
        for ii, binding in enumerate(input_list):
            
            val = input_vals[ii]
            inputC = np.zeros(self.net.nbindings)
            curr_role = binding.split('/')[1]
            role_list = [bb.split('/')[1] for bb in self.net.binding_names]
            role_idx = [jj for jj, rr in enumerate(role_list) if curr_role == rr]
            if not isinstance(binding, list):
                binding = [binding]
            binding_idx = self.net.find_bindings(binding)
            
            if self.params['input_inhib']:
                if self.params['input_method'] == 'ext':
                    inputC[role_idx] = -val
                
            inputC[binding_idx] = val
            self.inputC[:, ii] = inputC
        
        if self.params['cumulative_input']:
            
            self.inputC = self.inputC.cumsum(axis=1)
        
    def simulate(self, params=None):
        
        if params is None:
            params = self.params
            
        nrep = self.params['nrep']
        maxstep = self.params['maxstep'].max()
        nwords = len(self.input_list) # including the end of sentence markers, if any.
        self.act_trace = np.zeros((nrep, nwords, maxstep+1, self.net.nunits))
        self.T_trace = np.zeros((nrep, nwords, maxstep+1))
        self.q_trace = np.zeros((nrep, nwords, maxstep+1))
        self.speed_trace = np.zeros((nrep, nwords, maxstep+1))
        self.ema_speed_trace = np.zeros((nrep, nwords, maxstep+1))
        self.converged = np.ones((nrep, nwords), dtype=bool)
        self.rt = np.zeros((nrep, nwords))
        #self.grid_point = np.chararray((nrep, nwords))
        self.gp = [[[ [] for _ in range(maxstep+1)] for _ in range(nwords)] for _ in range(nrep)]

        if self.params['grid_points'] is not None:
            grid_points = self.params['grid_points']
            if not any(isinstance(pp, list) for pp in grid_points): 
                grid_points = [grid_points]
            self.dist_trace = np.zeros((nrep, nwords, maxstep+1, len(grid_points)))
        
        # create a storage to save the states at the end of processing a word.
        # extC hist

        for rep_num in np.arange(self.params['nrep']):
            
            #self.net.unclamp()
            #self.net.clear_input()
            self.net.reset() # Now reset() will do unclamp and clear_input
            self.set_input(self.input_list, self.input_vals)
            
            for word_num, curr_input in enumerate(self.input_list):
                # curr word
#                if self.params['cumulative_input']:
#                    curr_input = self.input_list[0:word_num]
#                    curr_vals = self.input_vals[0:word_num]
#                else:
#                    curr_input = self.input_list[word_num]
#          act_trace          curr_vals = self.input_vals[word_num]
#                    
#                curr_val =  self.input_vals[word_num]
#                if not isinstance(curr_input, list):
#                    curr_input = [curr_input]
#                if not isinstance(curr_val, list):
#                    curr_val = [curr_val]
                    
                curr_maxstep = self.params['maxstep'][word_num]
                
                if self.params['input_method'] == 'ext':
                    self.net.extC = copy.deepcopy(self.inputC[:,word_num])
                    self.net.ext = compute_sspace_biases(self.net.extC, self.net)
                    
                if word_num > 0:
                    self.net.prev_extC = copy.deepcopy(self.inputC[:,word_num-1])
                    self.net.prev_ext = compute_sspace_biases(self.net.prev_extC, self.net)
                
                elif self.params['input_method'] == 'clamp':
                    if self.params['cumulative_input'] == True:
                        self.net.clamp(self.input_list[:word_num+1])
                    else:
                        self.net.clamp(self.input_list[word_num])
                    #sys.exit('Not yet implementd')
                    
                self.net.run(curr_maxstep, 
                             norm_ord=self.params['norm_ord'],
                             tol=self.params['tol'], 
                             update_T=self.params['update_T'], 
                             update_q=self.params['update_q'],
                             q_fun=self.params['q_fun'],
                             grid_points=self.params['grid_points'],
                             ext_overlap=self.params['ext_overlap'], 
                             ext_overlap_steps=self.params['ext_overlap_steps']) # ADDED
                
                #self.net.prev_extC = copy.deepcopy(self.net.extC)
                
                self.act_trace[rep_num, word_num, :(curr_maxstep+1), :] = self.net.act_trace
                self.rt[rep_num, word_num] = self.net.rt
                self.T_trace[rep_num, word_num, :(curr_maxstep+1)] = self.net.T_trace
                self.q_trace[rep_num, word_num, :(curr_maxstep+1)] = self.net.q_trace
                self.speed_trace[rep_num, word_num, :(curr_maxstep+1)] = self.net.speed_trace
                self.ema_speed_trace[rep_num, word_num, :(curr_maxstep+1)] = self.net.ema_speed_trace
                self.converged[rep_num, word_num] = self.net.converged
                if self.params['grid_points'] is not None:
                    self.dist_trace[rep_num, word_num, :(curr_maxstep+1), :] = self.net.dist_trace
                self.gp[rep_num][word_num][:(curr_maxstep+1)] = self.net.gp_trace

    def plot_act_trace(self, rep_num):
        
        rep_ind = rep_num - 1
        
        nwords = self.act_trace.shape[1]
        curr_act_trace = self.act_trace[rep_ind, 0, :(self.rt[rep_ind, 0]+1), :]
        if nwords > 1:
            for wind in np.arange(1, nwords):
                temp_act_trace = self.act_trace[rep_ind, wind, 1:(self.rt[rep_ind, wind]+1), :]
                curr_act_trace = np.concatenate((curr_act_trace, temp_act_trace), axis=0)
        
        actC_trace = self.net.S2C(curr_act_trace.T).T
        heatmap(actC_trace.T, xlabel="Timestep", ylabel="Bindings", 
            xticklabels=None, yticklabels=self.net.binding_names, grayscale=False)
        
    def plot_trace(self, varname, rep_num):
        
        rep_ind = rep_num - 1
        nwords = self.act_trace.shape[1]
        curr_trace = getattr(self, varname+'_trace')[rep_ind, 0, :(self.rt[rep_ind, 0]+1)]
        if nwords > 1:
            for wind in np.arange(1, nwords):
                temp_trace = getattr(self, varname+'_trace')[rep_ind, wind, 1:(self.rt[rep_ind, 0]+1)]
                curr_trace = np.concatenate((curr_trace, temp_trace), axis=0)
                    
        curr_rt_trace = self.rt[rep_ind, :]
        curr_rt_trace = curr_rt_trace.cumsum()
        
        plt.plot(curr_trace)
        for ii, rt in enumerate(curr_rt_trace):
            plt.plot([rt, rt], [min(curr_trace[1:,:].flatten()), max(curr_trace[1:,:].flatten())], 'g-')
            
        plt.xlabel('Timestep', fontsize=16)
        plt.ylabel(varname, fontsize=16)
        
    def read_state(self, rep_num, word_num):
        
        act = self.act_trace[rep_num-1, word_num-1, self.rt[rep_num-1, word_num-1], :]
        self.net.read_state(act)
        #actC = self.net.S2C(self.act_trace[rep_num-1, word_num-1, self.rt[rep_num-1, word_num-1], :])
        #print(pd.DataFrame(self.net.vec2mat(actC), index=self.net.filler_names, columns=self.net.role_names))
    
    def read_grid_point(self, rep_num, word_num):
        act = self.act_trace[rep_num-1, word_num-1, self.rt[rep_num-1, word_num-1], :]
        self.net.read_grid_point(act)
        
    def compute_gpdist(self):
        # create gpset here after unclamping and clearing ext input.
        nrep = self.params['nrep']
        nwords = len(self.input_list)
        maxstep = self.params['maxstep']
        gp_count = np.zeros((nrep, nwords, maxstep+1, len(self.gpset)))
        for rep_ind in range(nrep):
            for word_ind in range(nwords):
                curr_rt = self.rt[rep_ind][word_ind].astype(int)
                for step in range(curr_rt):
                    idx = self.gpset.index(self.gp[rep_ind][word_ind][step])
                    gp_count[rep_ind, word_ind, step, idx] += 1
        gp_prob = gp_count.sum(axis=0)/nrep
        return gp_prob
        
    def plot_gp_prob_trace(self, word_num, gp_prob_trace, ngp=None, yticklab=True, grayscale=False, colorbar=True):
        if ngp is None:
            ngp = gp_prob_trace.shape[-1]
        
        minrt = np.min(self.rt[:][word_num-1]).astype(int)
        curr_trace = gp_prob_trace[word_num-1,:minrt,:(ngp+1)]
        if yticklab:
            yticklabels = [','.join(gp)+('(%.2f)' % self.gpset_gh[ii]) for ii, gp in enumerate(self.gpset[:(ngp+1)])]
            heatmap(curr_trace.T, xlabel="Timestep", ylabel="", \
                xticklabels=None, yticklabels=yticklabels, 
                grayscale=grayscale, colorbar=colorbar)            
        else:
            heatmap(curr_trace.T, xlabel="Timestep", ylabel="", \
                xticklabels=None, yticklabels=None, 
                grayscale=grayscale, colorbar=colorbar)
                


def b_ind(f, r, net):
    return f + r * net.nfillers
    
def u_ind(phi, rho, net):
    return phi + rho * net.ndim_f
    
def w(f, r, phi, rho, net):
    # A = net.TPinv
    return net.TPinv[b_ind(f, r, net), u_ind(phi, rho, net)]

def get_a(n, net, f, r):
    act = 0
    for phi in range(net.ndim_f):
        for rho in range(net.ndim_r):
            act += w(f, r, phi, rho, net) * n[u_ind(phi, rho, net)]
    return act

def n2a(n, net, f=None, r=None):
    # quant_list
    if (f is None) and (r is None):
        avec = np.zeros(net.nbindings)
        for f in range(net.nfillers):
            for r in range(net.nroles):
                avec[b_ind(f, r, net)] = get_a(n, net, f, r)
        return avec
    elif (f is None) and (r is not None):
        avec = np.zeros(net.nfillers)
        for f in range(net.nfillers):
            avec[f] = get_a(n, net, f, r)
        return avec
    elif (f is not None) and (r is None):
        avec = np.zeros(net.nroles)
        for r in range(net.nroles):
            avec[r] = get_a(n, net, f, r)
        return avec    
    else:
        return get_a(n, net, f, r)

def Q0(net, n):
#    q0 = 0.0
#    for f in range(net.nfillers):
#        for r in range(net.nroles):
#            q0 += n2a(n, net, f=f, r=r)**2 * (1-n2a(n, net, f=f, r=r))**2
#    return q0
#    return -np.sum(n2a(n,net)**2 * (1-n2a(n,net))**2)
    a = net.S2C(n)
    return -np.sum(a**2 * (1-a)**2)
    
def Q0GradE(net, n):
    # Elementwise computation. Very slow.
    # Based on the first derivation
    q0grad = np.zeros(net.nunits)
    for phi in range(net.ndim_f):
        for rho in range(net.ndim_r):
            q0grad[u_ind(phi, rho, net)] = 0.0
            for f in range(net.nfillers):
                for r in range(net.nroles):
                    a_fr = n2a(n, net, f, r)
                    g_fr = 2 * a_fr * (1 - a_fr) * (1 - 2 * a_fr)
                    q0grad[u_ind(phi, rho, net)] += w(f, r, phi, rho, net) * g_fr
    return -q0grad

def Q0Grad(net, n):
    # Based on the first derivation
    a = net.S2C(n)
    g = 2 * a * (1-a) * (1-2*a)   # a vectorized version of g_{fr}
    #A_fr_phirho = np.sum(net.TPinv[:, phirho] * g)
    gmat = np.tile(g, (net.nunits, 1)).T
    q0grad = np.sum(net.TPinv * gmat, axis=0)
    return -q0grad


def Q1(net, n):
    # quant_list (for now, ignore this)
#    q1 = 0.0
#    for r in range(net.nroles):
#        q1 += (np.sum(n2a(n, net, r=r)**2)-1)**2
#    return -np.sum((np.sum(net.vec2mat(n2a(n, net))**2, axis=0) - 1)**2)
    return -np.sum(np.sum(net.vec2mat(net.S2C(n))**2, axis=0)-1)**2

def Q1GradE(net, n):  # Elementwise computation
    # Not yet tested.
    q1grad = np.zeros(net.nunits)
    for phi in range(net.ndim_f):
        for rho in range(net.ndim_r):
            unit_grad = 0.0
            for r in range(net.nroles):
                var1 = np.sum(n2a(n, net, r=r)**2) - 1
                var2 = 0.0
                for f in range(net.nfillers):
                    var2 += n2a(n, net, f, r) * w(f, r, phi, rho, net)
                unit_grad += 4 * var1 * var2
            q1grad[u_ind(phi, rho, net)] = unit_grad
    return -q1grad

def Q1Grad(net, n):
    # Not yet tested.
    # term 1
    a = net.S2C(n)
    q1grad = 0.0
    for r_ind, rr in enumerate(net.role_names):
        curr_binding_ind = net.find_roles(rr)
        amat = np.tile(a[curr_binding_ind], (net.nunits, 1)).T
        term2 = np.sum(net.TPinv[curr_binding_ind, :] * amat, axis=0)
        term1 = np.sum(a[curr_binding_ind] ** 2) - 1
        q1grad += term1 * term2
    q1grad = 4 * q1grad
    return -q1grad


def QGrad(n, net):
    # So far, the c parameter indicates the relative strength of Q1. 
    # In this function, it indicates the relative strength of Q0.
    return net.c * Q0Grad(net, n) + (1-net.c) * Q1Grad(net, n)

def Q(net, n): # Q
    return net.c * Q0(net, n) + (1-net.c) * Q1(net, n)