In this tutorial, we show how to generate a CG mapping matrix for a molecule given a bead mapping. The trajectory and topology file come from a AA simulation done in gromacs (see Molecules_CG_Mapping
folder). The protein is FF (diphenylalanine) and the solvent is a mixture of water and methanol.
# disable GPU
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import hoomd, hoomd.htf as htf, hoomd.md
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import MDAnalysis as mda
from os import path
import matplotlib.pyplot as plt,matplotlib
# Loading inputs
TPR = 'Molecules_CG_Mapping/nvt_prod.tpr'
tpr = mda.Universe(TPR)
TRAJECTORY = 'Molecules_CG_Mapping/traj.trr'
u = mda.Universe(TPR, TRAJECTORY)
# Generating Mapping Matrix for FF
protein_FF = u.select_atoms("resname PHE and resid 0:1")
beads_mapping_FF = [['N','H1','H2','H3'],
['CA','HA','CB','HB1','HB2'],
['CG','CD1','HD1','CD2','HD2','CE1','HE1','CE2','HE2','CZ','HZ'],
['C','O'],
['N','H'],
['CA','HA','CB','HB1','HB2'],
['CG','CD1','HD1','CD2','HD2','CE1','HE1','CE2','HE2','CZ','HZ'],
['C','O1','O2']]
mapping_FF = htf.matrix_mapping(protein_FF,beads_mapping_FF)
print (mapping_FF)
/home/mgholiza/.conda/envs/hoomd-tf2/lib/python3.7/site-packages/MDAnalysis/core/universe.py:171: UserWarning: No coordinate reader found for Molecules_CG_Mapping/nvt_prod.tpr. Skipping this file. 'this file.'.format(filename))
[[0.8224383 0.05918723 0.05918723 0.05918723 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0.44409524 0.03726984 0.44409524 0.03726984 0.03726984 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.15577257 0.15577257 0.01307291 0.15577257 0.01307291 0.15577257 0.01307291 0.15577257 0.01307291 0.15577257 0.01307291 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.42880501 0.57119499 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.93286579 0.06713421 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.44409524 0.03726984 0.44409524 0.03726984 0.03726984 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.15577257 0.15577257 0.01307291 0.15577257 0.01307291 0.15577257 0.01307291 0.15577257 0.01307291 0.15577257 0.01307291 0. 0. 0. ] [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.27291648 0.36354176 0.36354176]]
Since the bonded ineractions were removed from the GROMACS traj, we need to consider and exclusion list to only account for non-bonded interactions between the beads.
mapped_exclusion_list = htf.gen_mapped_exclusion_list(u, protein_FF, beads_mapping_FF, selection="resname PHE")
print(mapped_exclusion_list[0][0:8])
print(mapped_exclusion_list[5][0:8])
[False True False False False False False False] [False False False False True False True True]
This means that bead 1
is bonded with bead 2
and bead 6
is bonded with beads 5
, 7
, 8
.
Let's look at the number of bead types:
def bead_type_finder(beads_mapping):
# beads_mapping input is a list of lists of mappings in a molecule
unique_beads = []
bead_types = []
for m in beads_mapping:
if m not in unique_beads:
unique_beads.append(m)
bead_types.append(unique_beads.index(m))
n_beads = len(unique_beads)
# bead_types is an array for the type of the beads in the molecule
# n_types is the total number of unique beads in the molecule
return n_beads, np.array(bead_types, dtype=np.float32)[np.newaxis,...]
n_bead_types, bead_types = bead_type_finder(beads_mapping_FF)
print (n_bead_types, bead_types)
6 [[0. 1. 2. 3. 4. 1. 2. 5.]]
The molecular-level mapping operator can be generated for the solvents too, but we don't really need them.
# Generating Mapping Matrix for Water
water = u.select_atoms("resname SOL and resid 500")
bead_mapping_water = [['OW','HW1','HW2']]
mapping_water = htf.matrix_mapping(water,bead_mapping_water)
print (mapping_water)
[[0.88809574 0.05595213 0.05595213]]
# Generating Mapping Matrix for Methanol
methanol = u.select_atoms("resname MET and resid 11665 ")
beads_mapping_methanol = [['C','H','H','H','OA','HO']]
mapping_methanol = htf.matrix_mapping(methanol,beads_mapping_methanol)
print (mapping_methanol)
[[0.37484707 0.03145832 0.03145832 0.03145832 0.49931966 0.03145832]]
# Getting the segment id of each molecule in topology
_,idx = np.unique(u.select_atoms('all').segids,return_index=True)
seg_id_list = u.select_atoms('all').segids[np.sort(idx)].tolist()
# Getting the list of every molecule type name in topology
_,idx = np.unique(u.atoms.resnames,return_index=True)
resname_list = u.atoms.resnames[np.sort(idx)].tolist()
# Getting list of atoms in each type of molecule
atoms_in_molecule_list = [protein_FF.names,
water.names,
methanol.names]
atoms_in_molecule_list = [protein_FF.names]
molecule_mapping_index = htf.find_molecules_from_topology(u,atoms_in_molecule_list, selection = "resname PHE")
molecule_mapping = mapping_FF
# get number of atoms
N = sum([len(m) for m in molecule_mapping_index])
# get number of molecules
M = len(molecule_mapping_index)
# get number of beads
B = molecule_mapping.shape[0]
#create a mass-weighted (M * bead_number) x N mapping operator
cg_mapping = htf.sparse_mapping([molecule_mapping for _ in molecule_mapping_index],
molecule_mapping_index)
assert cg_mapping.shape == (M * B, N)
print (M*B)
160
l1 = 0.1
lr = 1e-3
class PotentialLayer(tf.keras.layers.Layer):
def __init__(self, N_hidden, CG_NN, rcut, n_total_bead_type_interactions):
super(PotentialLayer, self).__init__(self, name='pf')
w_avg = (rcut / N_hidden / 2) ** 2
self.rcut = rcut
self.N_hidden = N_hidden
self.CG_NN = CG_NN
self.n_total_bead_type_interactions = n_total_bead_type_interactions
def build(self,input_shape):
self.widths = self.add_weight(
shape = (self.N_hidden,), initializer = tf.constant_initializer(value=0.1),
trainable = False, name='widths')
self.heights = self.add_weight(
shape = (self.n_total_bead_type_interactions, self.N_hidden),
initializer = tf.random_uniform_initializer(maxval=0.001, minval=-0.001),
regularizer=tf.keras.regularizers.l1(l1=l1),
trainable = True, name='heights')
self.lj_loc = self.add_weight(
shape = (self.CG_NN,), initializer = "ones",
regularizer=tf.keras.regularizers.l2(),
trainable = True, name='lj_loc')
def call(self, m_r, interactions):
flat_r = tf.tile(tf.reshape(m_r, [-1, 1]), [1, self.N_hidden])
offsets = tf.linspace(2.5, self.rcut, self.N_hidden, name='offsets')[tf.newaxis,...]
w_avg = (self.rcut / self.N_hidden / 2) ** 2
activation = tf.math.exp(-(flat_r - offsets) ** 2 / self.widths)
activation = tf.reshape(activation, [-1,self.CG_NN, self.N_hidden])
activation_g = activation[:, :, tf.newaxis,:] * self.heights
interactions = tf.reshape(interactions,[-1, self.CG_NN, 1, self.n_total_bead_type_interactions])
output_layer = tf.reduce_sum(interactions @ activation_g, axis=3)
nn_energies = tf.reshape(output_layer, shape=[-1, self.CG_NN]) # + lj_force
mask = tf.reshape(tf.cast((m_r > 1e-1), tf.float32), [-1, self.CG_NN])
neighs = tf.reduce_mean(tf.reduce_sum(mask, axis=1), axis=0)
predicted_CG_energy = tf.reduce_sum(nn_energies * mask, axis=1, name='predicted_energies')
return predicted_CG_energy
def get_config(self):
config = super(PotentialLayer, self).get_config()
config.update({
'widths': self.widths.numpy(),
'heights': self.widths.numpy(),
'lj_loc': self.lj_loc.numpy(),
'rcut': self.rcut,
'N_hidden': self.N_hidden,
'CG_NN': self.CG_NN,
'n_total_bead_type_interactions': self.n_total_bead_type_interactions})
return config
class TrainablePF(htf.SimModel):
def setup(self, cg_mapping, rcut, CG_NN, N_hidden, n_bead_types):
self.rcut = rcut
self.cg_mapping = cg_mapping
self.CG_NN = CG_NN
self.N_hidden = N_hidden
I = n_bead_types*(n_bead_types-1)//2 + n_bead_types
self.pf = PotentialLayer(N_hidden, CG_NN, rcut, I)
self.wca = htf.WCARepulsion(1e-3)
self.avg_mapped_rdf = tf.keras.metrics.MeanTensor('avg-mapped-rdf')
self.n_bead_types = n_bead_types
def compute(self, mapped_nlist, mapped_pos_with_type):
m_r = htf.safe_norm(tensor=mapped_nlist[:, :, :3], axis=2)
nlist_btype = tf.cast(mapped_nlist[..., -1], dtype = tf.int32)
pos_btype = tf.cast(mapped_pos_with_type[...,-1], dtype = tf.int32)
interactions = htf.compute_ohe_bead_type_interactions(pos_btype,nlist_btype, self.n_bead_types)
predicted_m_energy = self.pf(m_r, interactions) + tf.reduce_sum(self.wca(mapped_nlist), axis=1) #adding repulsion from WCA
predicted_m_forces = htf.compute_nlist_forces(mapped_nlist, predicted_m_energy)
predicted_m_rdf = []
for i in range(self.n_bead_types):
for j in range(self.n_bead_types):
p_index_m_rdf = htf.compute_rdf(mapped_nlist, [0.01, self.rcut],
type_tensor = mapped_pos_with_type[:,3], type_i=i, type_j=j)
predicted_m_rdf.append(p_index_m_rdf)
self.avg_mapped_rdf.update_state(predicted_m_rdf)
return predicted_m_forces[:,:3], predicted_m_energy
# defining bead types
n_bead_types = 6
system_bead_types = tf.repeat(bead_types, M, axis=0)
r_cut = 25
CG_NN = 64
N_hidden =48
model = TrainablePF(64, cg_mapping = cg_mapping, rcut = r_cut, CG_NN = CG_NN, N_hidden = N_hidden,
n_bead_types=n_bead_types)
from tensorflow import keras
optimizer=keras.optimizers.Adam(lr=lr, clipnorm=1.0)
model.compile(optimizer=optimizer, loss=['MeanSquaredError', None])
N_CG_particles = M*B
losses_train = []
L_Train = []
Here we redefine the TRAJECTORY (forces included) and TPR so that they only contain protein atoms. This was done using trjconv
command in GROMACS. The training was done for 20
epochs on a 90 ns trajectory. The code is available here but results are just loaded from the Molecules_CG_Mapping
folder.
run_training = False
epochs = 20
if run_training:
for i in range(epochs):
for inputs, ts in htf.iter_from_trajectory(512, u, selection='resname PHE', r_cut=r_cut, period =1):
aa_forces = ts.forces
true_mapped_force = tf.sparse.sparse_dense_matmul(cg_mapping, aa_forces, name='true-mapped-forces')
nlist = inputs[0]
positions = inputs[1]
box_size = tf.cast(htf.box_size(inputs[2]), tf.float32)
# calculate the center of mass of a CG bead
mapped_pos = htf.center_of_mass(positions[:,:3], cg_mapping, box_size)
system_bead_types = tf.reshape(system_bead_types,[-1,1])
mapped_pos_with_type = tf.concat([mapped_pos, system_bead_types], axis =1)
# Getting mapped_nlist based on non-bonded bead type interactions
mapped_nlist = htf.compute_nlist(mapped_pos_with_type, r_cut, CG_NN, box_size, sorted=False,
return_types=True, exclusion_matrix=mapped_exclusion_list)
predicted_mapped_forces, predicted_energy = model([mapped_nlist, mapped_pos_with_type])
model_inputs = [mapped_nlist, mapped_pos_with_type, box_size] + inputs[3:]
labels = true_mapped_force
loss_train = model.train_on_batch(model_inputs, labels, reset_metrics=False)
losses_train.append(loss_train)
loss_values_train = np.mean(np.array(losses_train)[:,0]) # get the mean for time-step
L_Train.append(loss_values_train/3/N_CG_particles)
print(f'epoch: {i+1}', f' Train Loss = {loss_values_train/3/N_CG_particles}')
model.save_weights(f'Molecules_CG_Mapping/weights/weights_epoch_{i+1}_lr_{lr}_l1_{l1}')
dat = {'L_Train':L_Train, 'mapped_possition_t0':mapped_possition_t0.numpy(),
'mapped_nlist_t0':mapped_nlist_t0.numpy(), 'box_size':box_size.numpy()}
with open(f'Molecules_CG_Mapping/training_lr_{lr}_l1_{l1}.npz', 'wb') as f:
np.savez(f,**dat)
# model.save(f'Molecules_CG_Mapping/model_lr_{lr}_l1_{l1}', overwrite=True)
dat = np.load(f'Molecules_CG_Mapping/training_lr_{lr}_l1_{l1}.npz')
L_Train = dat['L_Train']
mapped_possition_t0 = dat['mapped_possition_t0']
mapped_nlist_t0 = dat['mapped_nlist_t0']
box_size = dat['box_size']
plt.figure(figsize=(5,3), dpi = 150)
plt.plot(range(epochs), L_Train, label='Training loss')
plt.ylabel('Loss')
plt.xlabel('Epochs')
plt.legend()
plt.title(f'lr:{lr} l1:{l1}')
plt.tight_layout()
# plt.savefig(f'Loss_mixed_lr_{lr}_l1_{l1}.png', dpi=300)
Let's load the weights into a new model:
new_model = TrainablePF(64, cg_mapping = cg_mapping, rcut = r_cut, CG_NN = CG_NN, N_hidden = N_hidden,
n_bead_types=n_bead_types)
new_model.load_weights(f'Molecules_CG_Mapping/weights/weights_epoch_{epochs}_lr_{lr}_l1_{l1}')
<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f21d3805bd0>
Calculating mapped potentials per bead type interaction:
r = np.linspace(0.1, r_cut, 500)
p_list = []
plt.figure(figsize=(5,3), dpi = 150)
for i in range(n_bead_types):
output_forces, potential_energy = htf.compute_pairwise(new_model, r, type_i=i, type_j=i)
plt.plot(r, potential_energy[:,0] - potential_energy[-1,0], label=f'{i,i}', color=f'C{i}')
p_list.append(potential_energy)
plt.axhline(y=0, color='k', linestyle='--')
plt.xlabel(r'r')
plt.ylabel(r'U(r)')
plt.legend(bbox_to_anchor=(1, 1),title='Pair types:')
plt.title(f'Pair Mapped Potentials. lr:{lr} l1:{l1}')
plt.tight_layout()
# plt.savefig(f'Mapped_Potential_mixed_lr_{lr}_l1_{l1}.png', dpi=300)
plt.figure(figsize=(5,3), dpi = 150)
for i in range(n_bead_types):
output_forces, potential_energy = htf.compute_pairwise(new_model, r, type_i=1, type_j=i)
plt.plot(r, potential_energy[:,0] - potential_energy[-1,0], label=f'{1,i}', color=f'C{i}')
p_list.append(potential_energy)
plt.axhline(y=0, color='k', linestyle='--')
plt.xlabel(r'r')
plt.ylabel(r'U(r)')
plt.legend(bbox_to_anchor=(1, 1),title='Pair types:')
plt.title('Pair Mapped Potentials')
Text(0.5, 1.0, 'Pair Mapped Potentials')
plt.figure(figsize=(5,3), dpi = 150)
for i in range(n_bead_types):
output_forces, potential_energy = htf.compute_pairwise(new_model, r, type_i=2, type_j=i)
plt.plot(r, potential_energy[:,0] - potential_energy[-1,0], label=f'{2,i}', color=f'C{i}')
p_list.append(potential_energy)
plt.axhline(y=0, color='k', linestyle='--')
plt.xlabel(r'r')
plt.ylabel(r'U(r)')
plt.legend(bbox_to_anchor=(1, 1),title='Pair types:')
plt.title('Pair Mapped Potentials')
Text(0.5, 1.0, 'Pair Mapped Potentials')
Plotting mapped rdfs:
mapped_RDF = new_model.avg_mapped_rdf.result().numpy().reshape([n_bead_types,n_bead_types,2,-1])
fig, axs = plt.subplots(nrows=n_bead_types, ncols= n_bead_types, figsize=(18,18), dpi = 300)
for i in range(n_bead_types):
for j in range(n_bead_types):
axs[i,j].plot(mapped_RDF[i,j,1, :], mapped_RDF[i,j,0, :], label='Mapped RDF', color='C0')
axs[i,j].set_title(f'{i,j}')
axs[i,0].set_ylabel(r'g(r)')
axs[-1,i].set_xlabel(r'r')
fig.suptitle('Mapped RDF between bead types', y=1.00, fontsize=25)
plt.tight_layout()
########### Hoomd-Sim Code ################
hoomd.context.initialize('--mode=cpu')
tfcompute = htf.tfcompute(model)
snapshot = hoomd.data.make_snapshot(N=N_CG_particles,
box=hoomd.data.boxdim(*box_size),
particle_types=[str(i) for i in range(n_bead_types)],
bond_types=['bonded-interaction'],
angle_types=[str(i) for i in range(n_bead_types)],
dihedral_types=[str(i) for i in range(n_bead_types)])
snapshot.particles.position[:] = mapped_possition_t0[:,0:3]
snapshot.particles.typeid[:] = mapped_possition_t0[:,-1]
bonds_group = htf.gen_bonds_group(mapped_exclusion_list)
snapshot.bonds.resize(bonds_group.shape[0])
snapshot.bonds.group[:] = np.array(bonds_group)
TPR = 'Molecules_CG_Mapping/nvt_prod.tpr'
TRAJECTORY = 'Molecules_CG_Mapping/traj.trr'
u = mda.Universe(TPR, TRAJECTORY)
# note that htf.gen_mapped_exclusion_list can also be used to get the mapped adgacency matrix
mapped_adjacency_matrix = htf.gen_mapped_exclusion_list(u, protein_FF, beads_mapping_FF, selection="resname PHE and resid 0:1").astype(int)
rs,ang,dihe = htf.compute_cg_graph(DSGPM=False, adj_mat=mapped_adjacency_matrix, cg_beads=8, group_atoms=False)
r_ids,a_ids,d_ids = htf.mol_features_multiple(bnd_indices=rs, ang_indices=ang, dih_indices=dihe, molecules=20, beads=8)
snapshot.angles.resize(a_ids.shape[0])
snapshot.angles.group[:] = a_ids
snapshot.dihedrals.resize(d_ids.shape[0])
snapshot.dihedrals.group[:] = d_ids
system = hoomd.init.read_snapshot(snapshot)
HOOMD-blue v2.9.0 CUDA (10.1) DOUBLE HPMC_MIXED SSE SSE2 SSE3 SSE4_1 SSE4_2 AVX AVX2 Compiled: 03/03/2021 Copyright (c) 2009-2019 The Regents of the University of Michigan. ----- You are using HOOMD-blue. Please cite the following: * J A Anderson, C D Lorenz, and A Travesset. "General purpose molecular dynamics simulations fully implemented on graphics processing units", Journal of Computational Physics 227 (2008) 5342--5359 * J Glaser, T D Nguyen, J A Anderson, P Lui, F Spiga, J A Millan, D C Morse, and S C Glotzer. "Strong scaling of general-purpose molecular dynamics simulations on GPUs", Computer Physics Communications 192 (2015) 97--107 ----- CUDA driver version is insufficient for CUDA runtime version HOOMD-blue is running on the CPU CG coordinates are not caculated. Only connectivities are calculated notice(2): Group "all" created containing 160 particles
group_all = hoomd.group.all()
kT = 1.9872/1000
#NVT Simulation in Hoomd
#dt=0.002ps
im = hoomd.md.integrate.mode_standard(dt=5.0/489.0)
nvt = hoomd.md.integrate.nvt(group=group_all, kT=298.15 * kT, tau=350 / 489.)
nvt.randomize_velocities(1234)
#simulation
nvt_dump = hoomd.dump.gsd(filename='Molecules_CG_Mapping/CG_inference.gsd', period=10, group=group_all, phase=0, overwrite=True)
nlist = hoomd.md.nlist.cell()
tfcompute.attach(nlist, r_cut=r_cut,
save_output_period=100, train=False)
hoomd.run(2000)
notice(2): Force mode is FORCE_MODE.tf2hoomd notice(2): Starting TensorflowCompute notice(2): completed reallocate notice(2): Setting flag indicating virial modification will occur notice(2): -- Neighborlist exclusion statistics -- : notice(2): Particles with 1 exclusions : 80 notice(2): Particles with 2 exclusions : 40 notice(2): Particles with 3 exclusions : 40 notice(2): Neighbors included by diameter : no notice(2): Neighbors excluded when in the same body: no ** starting run ** Time 00:00:10 | Step 479 / 2000 | TPS 47.8893 | ETA 00:00:31 Time 00:00:20 | Step 1481 / 2000 | TPS 100.135 | ETA 00:00:05 Time 00:00:25 | Step 2000 / 2000 | TPS 97.3588 | ETA 00:00:00 Average TPS: 78.9228 --------- -- Neighborlist stats: 257 normal updates / 20 forced updates / 0 dangerous updates n_neigh_min: 6 / n_neigh_max: 55 / n_neigh_avg: 28.9625 shortest rebuild period: 7 -- Cell list stats: Dimension: 3, 3, 3 n_min : 0 / n_max: 16 / n_avg: 5.92593 ** run complete **
output_forces, potential_energy = htf.compute_pairwise(new_model, r, type_i=1, type_j=1)
output_forces_hoomd, potential_energy_hoomd= hoomd.htf.compute_pairwise(model, r, type_i=1, type_j=1)
plt.plot(r, potential_energy[:,0] - potential_energy[-1,0], label=f'pair 1-1 mapped', color=f'C0')
plt.plot(r, potential_energy_hoomd[:,0] - potential_energy_hoomd[-1,0], label=f'pair 1-1 CG', color=f'C3')
# plt.figure(figsize=(5,3), dpi = 150)
# plt.plot(r, potential_energy[:,0] - potential_energy[-1,0], label='Mapped Potential Energy', color='C0')
plt.axhline(y=0, color='k', linestyle='--')
plt.xlabel(r'r')
plt.ylabel(r'U(r)')
plt.legend(bbox_to_anchor=(1, 1),title='Pair types:')
plt.title(f'Pair Mapped Potentials. lr:{lr}, l1:{l1}')
plt.tight_layout()
e = tf.reduce_sum(tfcompute.outputs[0], axis=1)
plt.plot(e)
plt.show()