<jupyter_start><jupyter_text>let's plot the distributions at t=0,1<jupyter_code>xs = t.tensor(np.linspace(-8, 8, 1000))
t0s = t.tensor(np.zeros(1000))
t1s = t.tensor(np.ones(1000))
ys_t0 = np.array([(-potential(x, t)).exp() for x, t in zip(xs, t0s)])
ys_t1 = np.array([(-potential(x, t)).exp() for x, t in zip(xs, t1s)])
plt.plot(xs, ys_t0, label = f"prior distribution")
plt.plot(xs, ys_t1, label = f"posterior distribution")
plt.legend()<jupyter_output><empty_output><jupyter_text>ok, now can we do AIS?<jupyter_code>#define an initial position.
from torch.distributions.multivariate_normal import MultivariateNormal
mvn = MultivariateNormal(t.tensor([[0.]]), 0.5 * t.eye(1))
num_particles = 1000
iterations = 10
incremental_works = t.zeros(num_particles, iterations)
positions = t.zeros(num_particles, iterations)
lambda_schedule = t.tensor(np.linspace(0., 1., iterations))
mvn = MultivariateNormal(t.tensor([[0.]]), 0.5 * t.eye(1))
dt=1e-2
import tqdm
for particle_idx in tqdm.trange(num_particles):
    x = mvn.sample()
    x.requires_grad=True
    incremental_work_array, trajectory = AIS(x=x, potential=potential, lambda_schedule=lambda_schedule, propagator=MALA_propagator, dt=dt)
    incremental_works[particle_idx,:] = incremental_work_array
    positions[particle_idx] = t.flatten(trajectory)
    
    
    
    
    
sns.distplot(positions[:,0])
sns.distplot(positions[:,-1])
incremental_works
cum_works = t.cumsum(incremental_works, 1)
plt.plot(incremental_works[0,:])
last_cum_works = cum_works[:,-1]
from pymbar import EXP
print(EXP(last_cum_works.numpy()))
sns.distplot(weights)
t.log(t.mean(weights))
plt.xlim(-1, 3)<jupyter_output><empty_output><jupyter_text>ok, so it looks like AIS works, more or less...can we make an ULA uncontrolled SMC proposal and assert that the logw returned is the same as a compute_ULA_logw?<jupyter_code>x = MultivariateNormal(t.tensor([[0.]]), 0.5 * t.eye(1)).sample()
x.requires_grad=True
x_new, logw = ULA_uncontrolled_SMC(x, potential, t.tensor(0.), t.tensor(1.), dt=1e-2)
logw
x_new.requires_grad=True
old_potential_old_x = potential(x, t.tensor(0.))
new_potential_old_x = potential(x, t.tensor(1.))

new_potential_new_x = potential(x_new, t.tensor(1.))
new_potential_old_x.backward()
new_potential_new_x.backward()
compute_ULA_logw(log_gamma_old = -old_potential_old_x, 
                 log_gamma_new = -new_potential_new_x, 
                 x_old = x, 
                 x_new = x_new, 
                 grad_log_gamma_x_new = -x_new.grad, 
                 grad_log_gamma_x_old = -x.grad, 
                 dt = 1e-2)<jupyter_output><empty_output><jupyter_text>great, so we are recovering the correct log weight in the uncontrolled regime at this moment, i have to talk about the form of the controlled twisted initial distribution...<jupyter_code>x = MultivariateNormal(t.tensor([0., 1.]), 0.5 * t.eye(2)).sample()
uncontrolled_alphas = t.tensor([0.5, 0.1, 0.7])
uncontrolled_mus = t.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.2]])
uncontrolled_sigmas = t.stack([t.eye(2) for _ in range(3)])
A_0 = t.eye(2)
b_0 = t.tensor([0., 0.])
x_0 = t.tensor([4.])
x.size()
#x = MultivariateNormal(t.tensor([0., 0.]), 0.5 * t.eye(2)).sample()
uncontrolled_alphas = t.tensor([0.5, 0.0, 0.5])
uncontrolled_mus = t.tensor([[-5.0], [0.0], [5.0]])
uncontrolled_sigmas = t.stack([t.eye(1) for _ in range(3)])
A_0 = t.zeros(1)
b_0 = t.tensor([0.])
c_0 = t.tensor(0.)

twisted_gmm_proposal(
                         uncontrolled_alphas,
                         uncontrolled_mus,
                         uncontrolled_sigmas,
                         A_0,
                         b_0,
                         c_0)
num_proposals = 1000
all_proposals = t.zeros(num_proposals)
lognorms = t.zeros(num_proposals)
for _ in range(num_proposals):
    proposal, lognorm = twisted_gmm_proposal(
                         uncontrolled_alphas,
                         uncontrolled_mus,
                         uncontrolled_sigmas,
                         A_0,
                         b_0,
                         c_0)
    all_proposals[_] = proposal
    lognorms[_] = lognorm
sns.distplot(all_proposals, bins=50)
a = t.randn(3,3)
b = t.randn(3,3)
a>b
x= t.tensor(0.)
x.item()
list(x.size()) == []<jupyter_output><empty_output><jupyter_text>maybe we can build a twisted forward kernel now? can we play around with this a bit and see if it does what we want it to do?as a first sanity check, we can make sure that if the twisting is off, (i.e. uncontrolled), then there should be no twisting of a potential...right? if we define a potential, we can make a kernel density plot for uncontrolled smc, right?<jupyter_code>def uncontrolled_kernel_density(potential, parameter, x, y, dt):
    """
    report the log probability of transition from x to y under an Euler Maruyama discretized Unadjusted Langevin Algorithm
    """
    mu, sigma = EL_mu_sigma(x, 
                            potential(x, parameter),
                           dt)
    mvg = MultivariateNormal(mu, sigma)
    return mvg.log_prob(y)
positionsx = [t.tensor([[i]], requires_grad=True) for i in np.linspace(-1,1,30)]
positionsy = [t.tensor([[i]], requires_grad=True) for i in np.linspace(-1,1,30)]


_data = np.zeros((30, 30))
for idx, x in enumerate(positionsx):
    for idy, y in enumerate(positionsy):
        z = t.exp(uncontrolled_kernel_density(potential, 0., x, y, 1e-2)).detach()
        _data[idx, idy]=z
_data = np.array(_data)
plt.contourf( np.linspace(-1,1,30),  np.linspace(-1,1,30), _data)
plt.legend()
def A(x, time):
    return t.eye(1)*0.
def b(x, time):
    return t.tensor([0.])
def c(x, time):
    return t.tensor(0.)
twisted_forward_kernel(x = t.tensor([0.], requires_grad=True), 
                       potential = potential, 
                       potential_parameters = t.tensor(0.), 
                       dt = 1e-2, 
                       A = A, 
                       A_parameters = 0., 
                       b = b, 
                       b_parameters = 0., 
                       c = c, 
                       c_parameters = 0., 
                       d_t = t.tensor([0.]))
Theta = t.tensor([[1.]])
t.log((Theta.det().sqrt()) * t.exp(0.5 / t.tensor(1e-2)) * t.exp(-0.5 / t.tensor(1e-2) ))
twisted_gmm_proposal(uncontrolled_alphas = t.tensor([0.1, 0.2, 0.7]),
                         uncontrolled_mus = t.tensor([[0.], [1.], [2.]]),
                         uncontrolled_Sigmas = t.stack([t.eye(1), t.eye(1), t.eye(1)]),
                         A_0 = t.zeros(1,1),
                         b_0 = t.zeros(1),
                         c_0 = t.tensor(0.))<jupyter_output><empty_output><jupyter_text>ok, so it would seem at this point that we have an algorithm with components that are functioninig properly. CSMC says that the twisted weights look like:$$w_0^{\psi^{(i)}} \equiv \frac{\pi_0(\psi_0^{(i)})}{\psi_0^{(i)}(x_0)}$$and $$w_t^{\psi^{(i)}} \equiv \frac{w_t(x_{t-1}, x_t) K_t(\psi_t^{(i)})(x_{t-1})} {\psi_t^{(i)}(x_{t-1}, x_t)}$$<jupyter_code>def TwistedSMCSampler(potential,
                      uncontrolled_alphas,
                      uncontrolled_mus, 
                      uncontrolled_sigmas,
                      A0, 
                      b0,
                      c0,
                      A_t,
                      b_t,
                      c_t,
                      potential_parameter_sequence,
                      A_t_parameter_sequence,
                      b_t_parameter_sequence,
                      c_t_parameter_sequence,
                      dt,
                      uncontrolled=False):
    """
    given a potential, A0, b0, c0, as well as functions A_t, b_t, c_t (coupled with a sequence of appropriate parameters), conduct a
    forward pass of twisted SMC. 
    
    """
    #first thing to do is to make sure that the first dimension of all the sequences are the same length
    sequence_first_dimensions = [entry.size()[0] for entry in [A_t_parameter_sequence, b_t_parameter_sequence, c_t_parameter_sequence]]
    print(sequence_first_dimensions)
    sequence_length = sequence_first_dimensions[0]
    assert all(i ==sequence_first_dimensions[0] for i in sequence_first_dimensions)
    
    #make an object that will store the trajectory and the twisted works...
    trajectory = []
    twisted_log_weights = []
    
    #the potential parameter sequence is 1 greater than the rest since there is a parameterized _prior_ potential (index 0)
    assert sequence_first_dimensions[0] + 1 == potential_parameter_sequence.size()[0]
    
    #now we can conduct the pass of the 0th potential with a gaussian mixture model
    x, log_pi0_psi0 = twisted_gmm_proposal(uncontrolled_alphas,
                                           uncontrolled_mus,
                                           uncontrolled_sigmas,
                                           A0, b0, c0)
    trajectory.append(x.numpy())
    x.requires_grad=True #so we can compute potential gradient
    
    #compute log_twisted_w0
    log_phi0_x0 = quad_phi_0(x.detach(),A0.detach(), b0.detach(), c0.detach())
    if uncontrolled:
        assert np.isclose(log_phi0_x0, 0.)
        assert np.isclose(log_pi0_psi0, 0.)
    #print(log_psi0_x0)
    #print(log_pi0_psi0)
    initial_log_weight = log_pi0_psi0 + log_phi0_x0
    #print(initial_log_weight)
    twisted_log_weights.append(initial_log_weight)
                              
    #print(twisted_log_weights)
    
    #now we can iteration from t=1 to T
    for iteration in range(sequence_length):
        """
        we have to make a proposal and then compute the twisted weight;
        the twisted weight is the uncontrolled weight * K(\psi)(x_{t-1}) / \psi(x_{t-1}, x_t), 
        so we have to compute an uncontrolled weight first
        """
        old_potential_parameters = potential_parameter_sequence[iteration] #previous (is current by indexing)
        new_potential_parameters = potential_parameter_sequence[iteration+1] #current (is current+1 by indexing)
          
        #make a proposal
        d_t = t.tensor(0.)
#         print(b_t)
#         print(b_t_parameter_sequence[iteration])
        x_new, logK_psi = twisted_forward_kernel(x = x,
                                                 potential = potential,
                                                 potential_parameters = new_potential_parameters,
                                                 dt = dt,
                                                 A = A_t,
                                                 A_parameters=A_t_parameter_sequence[iteration],
                                                 b = b_t,
                                                 b_parameters=b_t_parameter_sequence[iteration],
                                                 c = c_t,
                                                 c_parameters=c_t_parameter_sequence[iteration],
                                                 d_t = d_t
                                                 )
        #print(f"logK_psi: {logK_psi}")
        assert not x_new.requires_grad
        x_new.requires_grad=True
        log_gamma_old = -potential(x.detach(), old_potential_parameters) # \log \gamma_{t-1}(x_{t-1})
        log_gamma_new = -potential(x_new, new_potential_parameters) # \log \gamma_{t}(x_t)
        x_auxiliary = x.clone().detach().requires_grad_(True) # get a gradable detached clone of the auxiliary x variable
        log_gamma_new_oldx = -potential(x_auxiliary, new_potential_parameters) # \log \gamma_{t}(x_{t-1})
        
        log_gamma_new.backward()
        log_gamma_new_oldx.backward()
        
        grad_log_gamma_x_old = x_auxiliary.grad
        grad_log_gamma_x_new = x_new.grad
        uncontrolled_log_weight = compute_ULA_logw(log_gamma_old = log_gamma_old,
                                                  log_gamma_new = log_gamma_new,
                                                  x_old = x_auxiliary,
                                                  x_new = x_new,
                                                  grad_log_gamma_x_new = grad_log_gamma_x_new,
                                                  grad_log_gamma_x_old = grad_log_gamma_x_old,
                                                  dt=dt)
        
        if uncontrolled:
            """
            we are going to make a uncontrolled smc proposal and validate that the log weight is the same as the `uncontrolled
            log weight`
            """
            uncontrolled_log_weight_check = log_gamma_new - log_gamma_old
            
            #forward
            x_old_validate = x_auxiliary.clone().detach().requires_grad_(True)
            #print(x_old_validate)
            mu_forward, sigma_forward = EL_mu_sigma(x_old_validate, potential(x_old_validate, new_potential_parameters), dt)
            forward_MVN = MultivariateNormal(mu_forward, sigma_forward)
            logp_forward = forward_MVN.log_prob(x_new.detach())
            
            #backward
            x_new_validate = x_new.clone().detach().requires_grad_(True)
            mu_backward, sigma_backward =  EL_mu_sigma(x_new_validate, potential(x_new_validate, new_potential_parameters), dt)
            backward_MVN = MultivariateNormal(mu_backward, sigma_backward)
            logp_backward = backward_MVN.log_prob(x_old_validate.detach())
            
            uncontrolled_log_weight_check += (logp_backward - logp_forward)
            
            assert np.isclose(uncontrolled_log_weight_check.item(), uncontrolled_log_weight.item(), atol=1e-3), f"check: {uncontrolled_log_weight_check.item()}, computed: {uncontrolled_log_weight.item()} "
            
        
        
        #print(old_potential_parameters)
        #print(new_potential_parameters)
        
        #compute phi_i
        phi_i = (square_mahalanobis_distance(x_new.detach(), A_t(x.detach(), A_t_parameter_sequence[iteration]))
                                 +x_new.detach().dot(b_t(x.detach(), b_t_parameter_sequence[iteration]))
                                 + c_t(c_t_parameter_sequence[iteration])
                                 + d_t
                )
        phi_t = quad_phi_t(x_tm1 = x.detach(), 
                           x_t = x_new.detach(), 
                           A_t = A_t, 
                           b_t = b_t, 
                           c_t = c_t, 
                           d_t = d_t, 
                           A_parameters = A_t_parameter_sequence[iteration], 
                           b_parameters = b_t_parameter_sequence[iteration], 
                           c_parameters = c_t_parameter_sequence[iteration]
                          )
        
#         print(f"phi_i: {phi_i}")
#         print(f"logK_psi: {logK_psi}")
#         print(f"uncontrolled_log_weight: {uncontrolled_log_weight}")
        twisted_log_weight = (uncontrolled_log_weight 
                              + logK_psi 
                              + phi_i
                             )
        if uncontrolled:
            assert np.isclose(logK_psi + phi_i, 0.)
        
        twisted_log_weights.append(twisted_log_weight) #log the twisted log weights
        trajectory.append(x_new.detach().numpy()) #log the new positions
        x = x_new.clone().detach().requires_grad_(True) #reinitialize the new position
    
    return trajectory, twisted_log_weights 
def dummy_A_t(x, param):
    return t.tensor([[0.]])
def dummy_b_t(x, param):
    return param[0]*(x-param[1])
    
def dummy_c_t(param):
    return t.tensor(0.)
import tqdm
#given the potential above, we are just shifting a gaussian...

iterations=100
sequence_length=20
trajs = []
twisted_log_weights = []
for iteration in tqdm.trange(iterations):
    traj, twisted_log_weight = TwistedSMCSampler(potential = potential,
                          uncontrolled_alphas = t.tensor([1.]),
                          uncontrolled_mus = t.tensor([[0.]]), 
                          uncontrolled_sigmas = t.tensor([[[0.5]]]), #this is the variance corresponding to the prior annealing protocol
                          A0 = t.tensor([[0.]]), #controlled gmm parameters are 0
                          b0 = t.tensor([0.]), #controlled gmm parameters are 0
                          c0 = t.tensor(0.), #controlled gmm parameters are 0
                          A_t = dummy_A_t,
                          b_t = dummy_b_t,
                          c_t = dummy_c_t,
                          potential_parameter_sequence = t.tensor(np.linspace(0,1,sequence_length)),
                          A_t_parameter_sequence = t.tensor([0. for _ in range(sequence_length-1)]), #this will always return 0
                          b_t_parameter_sequence = t.tensor([[0., 0.] for _ in range(sequence_length-1)]),
                          c_t_parameter_sequence = t.tensor(np.linspace(0,1,sequence_length)[1:]),
                          dt = 1e-2,
                          uncontrolled=True)
    trajs.append(traj)
    twisted_log_weights.append(twisted_log_weight)
init_posits = [i[0] for i in trajs]
final_posits = [i[-1] for i in trajs]
sns.distplot(init_posits)
sns.distplot(final_posits)
twisted_log_weights = np.array(twisted_log_weights)
cumulatie_log_weights = np.array([np.cumsum(entry) for entry in twisted_log_weights])
final_log_weights = [entry[-1] for entry in cumulatie_log_weights]
sns.distplot(final_log_weights)
free_energy = EXP(-np.array(final_log_weights))
free_energy
np.exp(-free_energy[0])<jupyter_output><empty_output><jupyter_text>can we compare this to an ULA uncontrolled SMC?<jupyter_code># trajs = []
# logws = []
# sequence_length=5
# iterations=1000
# for iteration in tqdm.trange(iterations):
#     mvg = MultivariateNormal(t.tensor([[0.]]), t.tensor([[0.5]]))
#     x = mvg.sample()
#     x.requires_grad=True
#     traj, logw = full_ULA_uncontrolled_SMC_sequence(x, potential, t.tensor(np.linspace(0,1,sequence_length)), dt=1e-3)
#     #print(f"blaws: ", traj, logw)
#     trajs.append(traj)
#     logws.append(logw)
np_logws = np.array(logws)
cum_logws = np.array([np.cumsum(entry) for entry in np_logws])
sns.distplot([entry[-1] for entry in cum_logws])
def ADP_value_training(x_sequences, 
                       twisted_log_weights, 
                       potential,
                       twisting_functions, 
                       twisting_parameters, 
                       minimization_iterations, 
                       epsilon, 
                       dt):
    """
    given a set of twisting functions \psi \in \Psi(Q), and a set of i.i.d. samples {X_{0:T}^n}_{n-1}^N from Q^{\psi},
    conduct approximate dynamic programming (ADP) for learning optimal sequences of potentials
    
    arguments
        x_sequences : t.tensor(R, T+1, M)
            trajectory sequence where R is the number of samples, T is the number of annealing steps, and M is the 
            dimension of x
        twisted_log_weights : t.tensor(R, T+1)
            log incremental weights of the R annealing samples
        twisting_functions : dict
            dictionary of functions corresponding to A, b, c
        twisting_initializers : dict
            dictionary of A0, b0, c0
        twisting_parameters : dict
            dictionary of gradable parameters 
            contains:
                A_t, b_t, c_t
        minimization_iterations : int
            number of iteration in minimization
        epsilon : float
            learning rate
        dt : t.tensor(float)
            timestep
    """
    #initialize by setting K_{T+1}^{\psi}(\exp{-V_{T+1}})(X_T^n) = 1 for n = 1, ..., N
    """
    iterate recursively from t=T, T-1, ..., 1
        a. define V_{comp, t} = -log(w_t^{\psi})(X_{t-1}, X_t) - log(K_{t+1}^{\psi})
    """
    num_samples, sequence_length = twisted_log_weights.size()[:2]
    print(f"num_samples, sequence length: {num_samples, sequence_length}")
    log_twisted_K = t.zeros(num_samples) #first log twisted Ks are 0.
    print(f"the twisting parameters are: {twisting_parameters}")
    
    #make holding variables for modified parameters:
    returnable_twisting_parameters = {key: [] for key, val in twisting_parameters.items()}
    
    #make a logger for the loss function 
    loss_functions = {}
    
    for time in tqdm.tqdm(range(sequence_length)[::-1]): #iterate backward from T to 1
        loss_logger = []
        print(f"time: {time}")
        #perform adp
        A_params = twisting_parameters['A_t'][time]
        b_params = twisting_parameters['b_t'][time]
        c_params = twisting_parameters['c_t'][time]
        #d_params = twisting_parameters['d_t'][t]
        
        for iteration in range(minimization_iterations):
            V_t_bar = -twisted_log_weights[:,time] - log_twisted_K # a.
            if time==0: #then we do not compute quad_phi_t, but rather the 0th iteration equivalent
                xis = [quad_phi_0(x_sequences[sample_idx, time], 
                                  A_params,
                                  b_params,
                                  c_params) for sample_idx in range(num_samples)
                      ]
            else:
                xis = [quad_phi_t(x_tm1 = x_sequences[sample_idx, time-1], 
                                       x_t = x_sequences[sample_idx, time], 
                                       A_t = twisting_functions['A_t'], 
                                       b_t = twisting_functions['b_t'], 
                                       c_t = twisting_functions['c_t'], 
                                       d_t = twisting_functions['d_t'], 
                                       A_parameters = A_params, 
                                       b_parameters = b_params, 
                                       c_parameters = c_params) for sample_idx in range(num_samples)
                            ]
            loss_function = sum([(i-j)**2 for i, j in zip(xis, V_t_bar)])
            loss_function.backward()
            loss_logger.append(loss_function.item())
#             print(f"loss: {loss_function.item()}")
#             print(f"b params: {b_params}")
            with t.no_grad():
                try:
                    A_params -= epsilon * A_params.grad
                    A_params.grad.zero_()
                except Exception as e:
                    #print(f"A_param optimizer: {e}")
                    pass
                try:
                    b_params -= epsilon * b_params.grad
                    b_params.grad.zero_()
                except Exception as e:
                    #print(f"b_params optimizer: {e}")
                    pass
                try:
                    c_params -= epsilon * c_params.grad
                    c_params.grad.zero_()
                except Exception as e:
                    #print(f"c_params optimizer: {e}")
                    pass
#                 try:
#                     d_params -= epsilon * d_params.grad
#                 except Exception as e:
#                     print(f"d_params optimizer: {e}")
        #recompute the log_twisted_K
        if time > 0:
            log_twisted_K = t.tensor(
                                    [twisted_forward_kernel(x = _x.clone().detach().requires_grad_(True),
                                     potential = potential, 
                                     potential_parameters=twisting_parameters['potential'][time],
                                     dt = dt,
                                     A = twisting_functions['A_t'],
                                     A_parameters = A_params.detach(),
                                     b = twisting_functions['b_t'],
                                     b_parameters = b_params.detach(),
                                     c = twisting_functions['c_t'],
                                     c_parameters = c_params.detach(),
                                     d_t = t.tensor(0.))[1] for _x in x_sequences[:, time-1]
                                    ]
                                    )
        #for each of the parameters, update the returnable_twisting_parameters_dict...
        for key, param in zip(['A_t', 'b_t', 'c_t'], [A_params, b_params, c_params]):
            returnable_twisting_parameters[key].append(param.detach())
        
        #add the loss logger to the loss dictionary
        loss_functions[time] = np.array(loss_logger)
    
    
    print(f"returnable twisting parameters: {returnable_twisting_parameters}")
    flipped_returnable_twisting_parameters = {key: lst[::-1] for key, lst in returnable_twisting_parameters.items()}
    return flipped_returnable_twisting_parameters, loss_functions
      
twisting_parameters, loss_functions = ADP_value_training(x_sequences = t.tensor(np.array(trajs)), 
                       twisted_log_weights = t.tensor(twisted_log_weights), 
                       potential = potential,
                       twisting_functions = {'A_t': dummy_A_t, 'b_t': dummy_b_t, 'c_t': dummy_c_t, 'd_t': t.tensor(0.)}, 
                       twisting_parameters = {'A_t': [t.tensor([[0.]]) for _ in range(sequence_length)],
                                              'b_t': [t.tensor([0.], requires_grad=True)] +  [t.tensor([0., 0.], requires_grad=True) for _  in range(sequence_length-1)],
                                              'c_t': [t.tensor(0.) for _ in range(sequence_length)],
                                              'potential': t.tensor(np.linspace(0,1,sequence_length))}, 
                       minimization_iterations=20, 
                       epsilon = 1e-3, 
                       dt = 1e-2)
len(loss_functions)
for i in range(100, 200):
    plt.plot(loss_functions[i])
twisting_parameters
t.stack(twisting_parameters['b_t'][1:]).squeeze(0).size()
sequence_length
t.tensor([0. for _ in range(sequence_length-1)])
import tqdm
#given the potential above, we are just shifting a gaussian...

iterations=50
sequence_length=200
trajs = []
twisted_log_weights = []
for iteration in tqdm.trange(iterations):
    traj, twisted_log_weight = TwistedSMCSampler(potential = potential,
                          uncontrolled_alphas = t.tensor([1.]),
                          uncontrolled_mus = t.tensor([[0.]]), 
                          uncontrolled_sigmas = t.tensor([[[0.5]]]), #this is the variance corresponding to the prior annealing protocol
                          A0 = t.tensor([[0.]]), #controlled gmm parameters are 0
                          b0 = twisting_parameters['b_t'][0], #controlled gmm parameters are 0
                          c0 = t.tensor(0.), #controlled gmm parameters are 0
                          A_t = dummy_A_t,
                          b_t = dummy_b_t,
                          c_t = dummy_c_t,
                          potential_parameter_sequence = t.tensor(np.linspace(0,1,sequence_length)),
                          A_t_parameter_sequence = t.tensor([0. for _ in range(sequence_length-1)]), #this will always return 0
                          c_t_parameter_sequence = t.tensor(np.linspace(0,1,sequence_length)[1:]),
                          b_t_parameter_sequence = t.stack(twisting_parameters['b_t'][1:]),
                          dt = 1e-2,
                          uncontrolled=False)
    trajs.append(traj)
    twisted_log_weights.append(twisted_log_weight)
twisted_log_weights = np.array(twisted_log_weights)
cumulatie_log_weights = np.array([np.cumsum(entry) for entry in twisted_log_weights])
log_weights = np.array([entry[-1] for entry in cumulatie_log_weights])
sns.distplot(log_weights)
-np.log(np.sqrt(2))
from pymbar import EXP
free_energy = EXP(-log_weights)
free_energy
np.exp(-free_energy[0])
initial_posts = [tra[0] for tra in trajs]
final_posts = [tra[-1] for tra in trajs]
sns.distplot(initial_posts)

sns.distplot(final_posts)<jupyter_output><empty_output><jupyter_text>perhaps we should like to rewrite everything in numpy and use pytorch _only_ for optimization<jupyter_code>def log_probability(potential):
    """
    generate a log unnormalized probability
    
    arguments
        potential : np.float
            -log unnormalized probability density
    """
    return np.exp(-potential)
def EL_mu_sigma(x, func, dt, parameters, is_force_function=False,  **kwargs):
    """
    create mean vector and covariance marix for a multivariate gaussian proposal
    """
    tau = dt/2.
    if not is_force_function:
        #we have to compute a numerical approximation of the gradient
        force = compute_force(x, func, parameters, **kwargs)
    else: #it is a force function, in which case, we just plug in the x, potential_params
        force = func(x, parameters, **kwargs)
    
    #now we can compute mu and Sigma
    mu = x + tau * force
    Sigma = 2*tau * np.eye(len(x))
    
    return mu, Sigma
def compute_force(x, potential_function, potential_parameters, eps=None):
    """
    given a potential function, compute an approximation of the Force (i.e. -grad(potential_function(x, parameters)))
    """
    from scipy.optimize import approx_fprime

    if not eps:
        eps = np.sqrt(np.finfo(float).eps)
    epsilons = [eps]*len(x)
    grad_potential = approx_fprime(x, potential_function, epsilons, potential_parameters)
    return -grad_potential
def potential(x, potential_parameters):
    return np.sum(potential_parameters * x**2)
def compute_ULA_logw():
    """
    compute the unadjusted langevin algorithm log_weight
    """
    logw = (
            -potentialt_xt
            + 0.5 * x_tm1.dot(forcet_xt)
            - 0.5 * xt.dot(forcet_xt)
            - (dt/8.)*np.sum(forcet_xt**2)
            + potentialtm1_xtm1
            - 0.5 * x_t.dot(forcet_xtm1)
            + 0.5 * xtm1.dot(forcet_xtm1)
            + (dt/8.)*np.sum(forcet_xtm1**2)
            )
    return logw

def compute_generalized_logw(log_gamma_old, log_gamma_new, log_forward_kernelt, log_backward_kernel_tm1):
    """
    compute a generalized log incremental (unnormalized) weight
    """
    logw = log_gamma_new + log_backward_kernel_tm1 - log_gamma_old - log_forward_kernelt
    return logw

def twisted_gmm_components(uncontrolled_alphas, uncontrolled_mus, uncontrolled_Sigmas, A0, b0, c0):
    """
    compute a twisted gaussian mixture model mixing components, twisted_mu, and twisted_Sigma
    """
    assert len(uncontrolled_alphas) == len(uncontrolled_mus)
    components, dimensions = uncontrolled_mus.shape
    assert uncontrolled_Sigmas.shape == (components, dimension, dimension)
    
    #compute mixture components
    Sigma_tilde_js = np.linalg.inv(np.linalg.inv(uncontrolled_Sigmas) + 2.0*A0)
    log_zetas = np.array([
                            gmm_log_zetas(sigma_tilde_j, sigma_j, mu_j, b0, c0) 
                            for sigma_tilde_j, sigma_j in 
                            zip(Sigma_tilde_js, uncontrolled_Sigmas)
                            ])
    log_alpha_tildes = np.log(uncontrolled_alphas) + log_zetas
    
    return log_alpha_tildes, Sigma_tilde_js

def twisted_gmm_proposal(log_alpha_tildes, Sigma_tilde_js, b0):
    """
    make a twisted gaussian mixture model proposal
    """
    from scipy.special import logsumexp
    from scipy.stats import multivariate_normal
    normalized_alpha_tildes = np.exp(log_alpha_tildes - logsumexp(log_alpha_tildes))
    #choose a component
    component_index = np.random.choice(range(len(normalized_alpha_tildes)), p = normalized_alpha_tildes)
    
    #then choose a position based on that gaussian
    Sigma_tilde_j = Sigma_tilde_js[component_index]
    mu_j = uncontrolled_mus[component_index]
    Sigma_j = uncontrolled_Sigmas[component_index]
    twisted_mean = np.matmul(Sigma_tilde_j, np.matmul(np.linalg.inv(Sigma_j), mu_j) - b0)
    twisted_Sigma = Sigma_tilde_j
    x = multivariate_normal.rvs(mean=twisted_mean, cov = twisted_Sigma)
    logpdf = multivariate_normal.logpdf(x, mean=twisted_mean, cov = twisted_Sigma)
    return x, logpdf

def compute_twisted_gmm_lognormalizer(log_alpha_tildes):
    """
    compute the twisted gaussian mixture model log normalization constant with unnormalized log_alpha_tildes
    """
    from scipy.special import logsumexp
    return logsumexp(log_alpha_tildes)

def gmm_log_zetas(Sigma_tilde_j, Sigma_j, mu_j, b0, c0):
    """
    compute the logzeta_js mixture components
    """
    from scipy.spatial.distance import mahalanobis
    comp1 = -0.5 * np.log(np.linalg.det(Sigma_j))
    comp2 = 0.5 * np.log(np.linalg.det(Sigma_tilde_j))
    comp3 = 0.5 * mahalanobis(np.matmul(np.linalg.inv(Sigma_j), mu_j), b0, np.linalg.inv(Sigma_tilde_j))**2
    comp4 = -0.5 * mahalanobis(mu_j, np.zeros(len(mu_j)), Sigma_j)**2
    comp5 = -c0
    return comp1 + comp2 + comp3 + comp4 + comp5

def Theta_t(x_tm1, A_t, dt):
    """
    compute Theta_t = (I_d _ 2*dt*A_t(x_tm1))^-1
    """
    theta = np.linalg.inv(np.eye(len(x_tm1)) + 2. * dt * A_t)
    return theta

def f_t(x_tm1, potential_function, parameters, dt, **kwargs):
    """
    NOTE : parameters should reflect potential function at time t (not t-1)
    compute f_t(x_tm1) = x_tm1 + 0.5 * dt * forcet(x_tm1) = mu
    """
    mu, cov = EL_mu_sigma(x_tm1, potential_function, dt, parameters, is_force_function=False,  **kwargs)
    return mu
    

def twisted_forward_tools(x, potential_function, potential_parameters, A_function, A_parameters, dt, **kwargs):
    """
    make a twisted forward 
    """
    dimension = len(x) #grab the dimension of x  
    f_t = f_t(x, potential_function, potential_parameters, dt)
    A_t = A_function(x, A_parameters)
    
    #theta and f
    theta = Theta_t(x, A_t, dt)
    f = f_t(x, potential, potential, dt) 
    
    return theta, f

def twisted_forward_proposal(theta, f, dt, b_t, **kwargs):
    """
    make a forward twisted proposal
    """
    twisted_mean = np.matmul(theta, f - dt*b_t)
    twisted_covariance = dt * theta
    x = multivariate_normal.rvs(mean=twisted_mean, cov = twisted_covariance)
    return x

def twisted_forward_log_normalizer(theta, f, b_t, dt, c, d):
    """
    perform one line computation to calculate the log normalization constant of the twisted forward proposal
    """
    from scipy.spatial.distance import mahalanobis
    
    #there are 3 components that will be computed separately and added.
    comp1 = 0.5 * np.linalg.det(theta)
    comp2 = (1./(2.*dt)) * mahalanobis(f, dt*b_t, np.linalg.inv(theta))**2
    comp3 = -(1./(2.*dt)) * f.dot(f) - c - d
    
    log_forward_normalizer = comp1 + comp2 + comp3
    return log_forward_normalizer


    
a = np.array([1, 2, 3, 4]).transpose()
b = np.array([0, 2, 3, 4])
a.dot(b)
a**2
a = np.array([[0., 1., 2.], [0., 1., 4.]]).shape
a
q = np.random.rand(2, 2)
np.matmul(q, np.array([1, 2]))
A0 = np.random.rand(2,2)
np.linalg.inv(np.linalg.inv(q) - A0)
np.array([np.linalg.inv(it) - A0 for it in q])
from scipy.stats import multivariate_normal
mv = multivariate_normal.rvs(mean = np.array([0., 1.]), cov = np.eye(2))
%%time
mv = multivariate_normal.rvs(mean = np.random.randn(2000), cov = np.eye(2000))<jupyter_output>CPU times: user 27 s, sys: 3.64 s, total: 30.7 s
Wall time: 4.19 s
{'fraction_non_alphanumeric': 0.10508416449944272, 'fraction_numerical': 0.015261067131548112, 'mean_word_length': 2.113167259786477, 'pattern_counts': {'":': 0, '<': 32, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 86, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '494604', 'license': 'MIT', 'n_tokens_mistral': 11495, 'n_tokens_neox': 10823, 'n_words': 2366, 'path': 'troubleshooting/csmc_stage.ipynb', 'repo_name': 'choderalab/coddiwomple', 'size': '34991'}