<jupyter_start><jupyter_text>Task 1 A template code for training an RBM on H$_2$ data for $r = 1.2$ is shown here. Modify this!Imports and loading in data:<jupyter_code>import numpy as np
import torch
import matplotlib.pyplot as plt
import os
from RBM_helper import RBM
import H2_energy_calculator
import typing
import collections

# training_data = torch.from_numpy(np.loadtxt("H2_data/R_1.2_samples.txt"))
training_data = torch.from_numpy(np.loadtxt("H2_data/R_1.2_samples.txt"))
coeff = np.loadtxt("H2_data/H2_coefficients.txt")[20,:]

true_energy = H2_energy_calculator.energy_from_freq(training_data, coeff)
print("H2 energy for r = 1.2: ",true_energy)<jupyter_output>H2 energy for r = 1.2:  -1.0642022250418146<jupyter_text>Define the RBM:<jupyter_code>n_vis = 2
n_hin = 10

rbm = RBM(n_vis, n_hin)<jupyter_output><empty_output><jupyter_text>Train the RBM:<jupyter_code>def train_RBM(training_data:torch.Tensor) ->float:
    """
        Train Restricted Bolzman Machiene.
        Args:
            training_data: A tensor of binary data of radius.
        Returns:
            Energy value.
        
    """
    
    epochs = 500 # number of training steps
    num_samples = 1000 # number of samples to generate from the RBM to calculate the H2 energy

    true_energy = H2_energy_calculator.energy_from_freq(training_data, coeff)
    print("True energy: ",true_energy)

    for e in range(1, epochs+1):
        # do one epoch of training
        rbm.train(training_data)   

        # now generate samples and calculate the energy
        if e % 100 == 0:
#             print("\nEpoch: ", e)
#             print("Sampling the RBM...")

            # For sampling the RBM, we need to do Gibbs sampling.
            # Initialize the Gibbs sampling chain with init_state as defined below.
            init_state = torch.zeros(num_samples, n_vis)
            RBM_samples = rbm.draw_samples(15, init_state)

#             print("Done sampling. Calculating energy...")       

            energies = H2_energy_calculator.energy(RBM_samples, coeff, rbm.wavefunction) 
            print("Energy from RBM samples: ", energies.item())
    return energies.item()
# 
# load R_<r_value_samples.txt from H2_data folder
# 
# evaluate true_energy on each training data, use it to plot E_bond(r)

Ebond = {}
directory = 'H2_data'
for entry in os.scandir(directory):
    if (entry.path.endswith("samples.txt")):
        training_data = torch.from_numpy(np.loadtxt(entry))
        radius = entry.name.split("_")[1]
        Ebond[radius] = train_RBM(training_data)
#         print(Ebond)

Ebond_sorted = collections.OrderedDict(sorted(Ebond.items()))
x_values = [] # radii
y_values = [] # Energy

for key, val in Ebond_sorted.items():
    x_values.append(float(key))
    y_values.append(val)

# plot Ebont(r)
plt.plot(x_values, y_values)
plt.xlabel('r')
plt.ylabel('Ebond')
plt.title('Ebond(r)')
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10742387957577831, 'fraction_numerical': 0.021553198768388643, 'mean_word_length': 3.43030303030303, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '31492', 'license': 'MIT', 'n_tokens_mistral': 1057, 'n_tokens_neox': 997, 'n_words': 241, 'path': 'Project_1_RBM_and_Tomography/Task1_Olga_V.ipynb', 'repo_name': 'olgOk/CDL_Cohort_Project_Week1', 'size': '2923'}