# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import dgl
import numpy as np
import torch

# %load_ext autoreload
# %autoreload 2
import espaloma as esp
# -

# ! git pull

# +
g = esp.Graph('C')

forcefield = esp.graphs.legacy_force_field.LegacyForceField(
    "smirnoff99Frosst"
)

forcefield.parametrize(g)

from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
    n_samples=100,
    n_steps_per_sample=10,
)
simulation.run(g)

# +
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)

net = torch.nn.Sequential(
        representation, 
        esp.mm.geometry.GeometryInGraph(),
        esp.mm.energy.EnergyInGraph(), # predicted energy -> u
        esp.mm.energy.EnergyInGraph(suffix='_ref') # reference energy -> u_ref,
)

optimizer = torch.optim.Adam(
    net.parameters(),
    0.1,
)

# optimizer = torch.optim.LBFGS(
#     net.parameters(),
#     0.1,
#     line_search_fn='strong_wolfe',
# )
# -

states = []
losses = []

for _ in range(100):
    optimizer.zero_grad()
    
    def l():
        net(g.heterograph)
        
        loss = torch.nn.MSELoss()(
            g.nodes['n2'].data['u_ref'],
            g.nodes['n2'].data['u'],
        )

        loss = loss.sum()
        
        losses.append(loss.detach().numpy())
        
        loss.backward()
        
        # print(loss)
        return loss
    
    optimizer.step(l)
    
    states.append(
        {
            '%s_%s' % (term, param): getattr(
                net[0],
                '%s_%s' % (term, param)
            ).detach().clone().numpy()
            for term in ['n2'] for param in ['k', 'eq']
        }
    )

plt.plot(losses)

ks = np.array([state['n2_k'].flatten() for state in states])
eqs = np.array([state['n2_eq'].flatten() for state in states])

eqs.std(axis=0)

for idx in range(8):
    plt.plot(np.diff(ks[:, idx]))
    

for idx in range(8):
    plt.plot(eqs[:, idx])

g.nodes['n2'].data['eq_ref']

# +
from matplotlib import pyplot as plt
plt.scatter(
    g.nodes['n2'].data['u_ref'].detach(),
    g.nodes['n2'].data['u'].detach()
)

plt.xlabel('ref')
plt.ylabel('pred')
# -

# # init at answer

# +
g = esp.Graph('C')

forcefield = esp.graphs.legacy_force_field.LegacyForceField(
    "smirnoff99Frosst"
)

forcefield.parametrize(g)

from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
    n_samples=100,
    n_steps_per_sample=10,
)
simulation.run(g)

# +
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)


for term in ['n2', 'n3']:
    for param in ['k', 'eq']:
        setattr(
            representation, '%s_%s' % (term, param),
            torch.nn.Parameter(
                g.nodes[term].data[param + '_ref'].data
            )
        
        )
        
net = torch.nn.Sequential(
        representation, 
        esp.mm.geometry.GeometryInGraph(),
        esp.mm.energy.EnergyInGraph(), # predicted energy -> u
        esp.mm.energy.EnergyInGraph(suffix='_ref') # reference energy -> u_ref,
)

optimizer = torch.optim.Adam(
    net.parameters(),
    0.1,
)

# optimizer = torch.optim.LBFGS(
#     net.parameters(),
#     0.1,
#     line_search_fn='strong_wolfe',
# )
# -

list(net.named_parameters())

net(g.heterograph)

states = []
losses = []


# +
def l():
    net(g.heterograph)

    loss = torch.nn.MSELoss()(
        g.nodes['n2'].data['u_ref'],
        g.nodes['n2'].data['u'],
    )

    loss = loss.sum()

    losses.append(loss.detach().numpy())

    # loss.backward()

    print(loss)
    return loss

l()
# optimizer.step(l)
# -

g.nodes['n2'].data['k']

for _ in range(100):
    optimizer.zero_grad()
    
    def l():
        net(g.heterograph)
        
        loss = torch.nn.MSELoss()(
            g.nodes['n2'].data['u_ref'],
            g.nodes['n2'].data['u'],
        )

        loss = loss.sum()
        
        losses.append(loss.detach().numpy())
        
        loss.backward()
        
        print(loss)
        return loss
    
    optimizer.step(l)
    
    states.append(
        {
            '%s_%s' % (term, param): getattr(
                net[0],
                '%s_%s' % (term, param)
            ).detach().clone().numpy()
            for term in ['n2'] for param in ['k', 'eq']
        }
    )

plt.plot(losses)

ks = np.array([state['n2_k'].flatten() for state in states])
eqs = np.array([state['n2_eq'].flatten() for state in states])

eqs.std(axis=0)

for idx in range(8):
    plt.plot(np.diff(ks[:, idx]))
    

for idx in range(8):
    plt.plot(eqs[:, idx])

g.nodes['n2'].data['eq_ref']

# # param gaussian noise

# +
g = esp.Graph('C')

forcefield = esp.graphs.legacy_force_field.LegacyForceField(
    "smirnoff99Frosst"
)

forcefield.parametrize(g)

from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
    n_samples=100,
    n_steps_per_sample=10,
)
simulation.run(g)

# +
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)

epsilon = 0.1

for term in ['n2', 'n3']:
    for param in ['k', 'eq']:
        setattr(
            representation, '%s_%s' % (term, param),
            torch.nn.Parameter(
                g.nodes[term].data[param + '_ref'].data
                + torch.distributions.normal.Normal(
                    loc=torch.zeros_like(g.nodes[term].data[param + '_ref']),
                    scale=epsilon * torch.ones_like(g.nodes[term].data[param + '_ref']),
                ).sample()
            )
        
        )
        
net = torch.nn.Sequential(
        representation, 
        esp.mm.geometry.GeometryInGraph(),
        esp.mm.energy.EnergyInGraph(), # predicted energy -> u
        esp.mm.energy.EnergyInGraph(suffix='_ref') # reference energy -> u_ref,
)

optimizer = torch.optim.Adam(
    net.parameters(),
    0.1,
)

# optimizer = torch.optim.LBFGS(
#     net.parameters(),
#     0.1,
#     line_search_fn='strong_wolfe',
# )
# -

net(g.heterograph)

states = []
losses = []


# +
def l():
    net(g.heterograph)

    loss = torch.nn.MSELoss()(
        g.nodes['n2'].data['u_ref'],
        g.nodes['n2'].data['u'],
    )

    loss = loss.sum()

    losses.append(loss.detach().numpy())

    # loss.backward()

    print(loss)
    return loss

l()
# optimizer.step(l)
# -

g.nodes['n2'].data['k']

for _ in range(1000):
    optimizer.zero_grad()
    
    def l():
        net(g.heterograph)
        
        loss = torch.nn.MSELoss()(
            g.nodes['n2'].data['u_ref'],
            g.nodes['n2'].data['u'],
        )

        loss = loss.sum()
        
        losses.append(loss.detach().numpy())
        
        loss.backward()
        
        print(loss)
        return loss
    
    optimizer.step(l)
    
    states.append(
        {
            '%s_%s' % (term, param): getattr(
                net[0],
                '%s_%s' % (term, param)
            ).detach().clone().numpy()
            for term in ['n2'] for param in ['k', 'eq']
        }
    )

plt.plot(losses)

ks = np.array([state['n2_k'].flatten() for state in states])
eqs = np.array([state['n2_eq'].flatten() for state in states])

eqs.std(axis=0)

for idx in range(8):
    plt.plot(np.diff(ks[:, idx]))
    

for idx in range(8):
    plt.plot(eqs[:, idx])

plt.scatter(
    g.nodes['n2'].data['u_ref'].flatten().detach(),
    g.nodes['n2'].data['u'].flatten().detach(),
)

# # angle

# +
g = esp.Graph('C')

forcefield = esp.graphs.legacy_force_field.LegacyForceField(
    "smirnoff99Frosst"
)

forcefield.parametrize(g)

from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
    n_samples=100,
    n_steps_per_sample=10,
)
simulation.run(g)

# +
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)

epsilon = 0.1

for term in ['n2', 'n3']:
    for param in ['k', 'eq']:
        setattr(
            representation, '%s_%s' % (term, param),
            torch.nn.Parameter(
                g.nodes[term].data[param + '_ref'].data
                + torch.distributions.normal.Normal(
                    loc=torch.zeros_like(g.nodes[term].data[param + '_ref']),
                    scale=epsilon * torch.ones_like(g.nodes[term].data[param + '_ref']),
                ).sample()
            )
        
        )
        
net = torch.nn.Sequential(
        representation, 
        esp.mm.geometry.GeometryInGraph(),
        esp.mm.energy.EnergyInGraph(), # predicted energy -> u
        esp.mm.energy.EnergyInGraph(suffix='_ref') # reference energy -> u_ref,
)

optimizer = torch.optim.Adam(
    net.parameters(),
    0.1,
)

# optimizer = torch.optim.LBFGS(
#     net.parameters(),
#     0.1,
#     line_search_fn='strong_wolfe',
# )
# -

net(g.heterograph)

states = []
losses = []


# +
def l():
    net(g.heterograph)

    loss = torch.nn.MSELoss()(
        g.nodes['n3'].data['u_ref'],
        g.nodes['n3'].data['u'],
    )

    loss = loss.sum()

    losses.append(loss.detach().numpy())

    # loss.backward()

    print(loss)
    return loss

l()
# optimizer.step(l)
# -

g.nodes['n3'].data['k']

for _ in range(1000):
    optimizer.zero_grad()
    
    def l():
        net(g.heterograph)
        
        loss = torch.nn.MSELoss()(
            g.nodes['n3'].data['u_ref'],
            g.nodes['n3'].data['u'],
        )

        loss = loss.sum()
        
        losses.append(loss.detach().numpy())
        
        loss.backward()
        
        print(loss)
        return loss
    
    optimizer.step(l)
    
    states.append(
        {
            '%s_%s' % (term, param): getattr(
                net[0],
                '%s_%s' % (term, param)
            ).detach().clone().numpy()
            for term in ['n3'] for param in ['k', 'eq']
        }
    )

plt.plot(losses)

ks = np.array([state['n3_k'].flatten() for state in states])
eqs = np.array([state['n3_eq'].flatten() for state in states])

eqs.std(axis=0)

for idx in range(8):
    plt.plot(np.diff(ks[:, idx]))
    

for idx in range(8):
    plt.plot(eqs[:, idx])

plt.scatter(
    g.nodes['n3'].data['u_ref'].flatten().detach(),
    g.nodes['n3'].data['u'].flatten().detach(),
)

# # all

# +
g = esp.Graph('C')

forcefield = esp.graphs.legacy_force_field.LegacyForceField(
    "smirnoff99Frosst"
)

forcefield.parametrize(g)

from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
    n_samples=100,
    n_steps_per_sample=10,
)
simulation.run(g)

# +
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)

epsilon = 0.1

for term in ['n2', 'n3']:
    for param in ['k', 'eq']:
        setattr(
            representation, '%s_%s' % (term, param),
            torch.nn.Parameter(
                g.nodes[term].data[param + '_ref'].data
                + torch.distributions.normal.Normal(
                    loc=torch.zeros_like(g.nodes[term].data[param + '_ref']),
                    scale=epsilon * torch.ones_like(g.nodes[term].data[param + '_ref']),
                ).sample()
            )
        
        )
        
net = torch.nn.Sequential(
        representation, 
        esp.mm.geometry.GeometryInGraph(),
        esp.mm.energy.EnergyInGraph(), # predicted energy -> u
        esp.mm.energy.EnergyInGraph(suffix='_ref') # reference energy -> u_ref,
)

optimizer = torch.optim.Adam(
    net.parameters(),
    0.1,
)

# optimizer = torch.optim.LBFGS(
#     net.parameters(),
#     0.1,
#     line_search_fn='strong_wolfe',
# )
# -

net(g.heterograph)

states = []
losses = []


# +
def l():
    net(g.heterograph)

    loss = torch.nn.MSELoss()(
        g.nodes['g'].data['u_ref'],
        g.nodes['g'].data['u'],
    )

    loss = loss.sum()

    losses.append(loss.detach().numpy())

    # loss.backward()

    print(loss)
    return loss

l()
# optimizer.step(l)
# -

g.nodes['n3'].data['k']

for _ in range(1000):
    optimizer.zero_grad()
    
    def l():
        net(g.heterograph)
        
        loss = torch.nn.MSELoss()(
            g.nodes['n3'].data['u_ref'],
            g.nodes['n3'].data['u'],
        )

        loss = loss.sum()
        
        losses.append(loss.detach().numpy())
        
        loss.backward()
        
        print(loss)
        return loss
    
    optimizer.step(l)
    
    states.append(
        {
            '%s_%s' % (term, param): getattr(
                net[0],
                '%s_%s' % (term, param)
            ).detach().clone().numpy()
            for term in ['n3'] for param in ['k', 'eq']
        }
    )

plt.plot(losses)

ks = np.array([state['n3_k'].flatten() for state in states])
eqs = np.array([state['n3_eq'].flatten() for state in states])

eqs.std(axis=0)

for idx in range(8):
    plt.plot(np.diff(ks[:, idx]))
    

for idx in range(8):
    plt.plot(eqs[:, idx])

plt.scatter(
    g.nodes['n3'].data['u_ref'].flatten().detach(),
    g.nodes['n3'].data['u'].flatten().detach(),
)

# # normalize

# +
g = esp.Graph('C')

forcefield = esp.graphs.legacy_force_field.LegacyForceField(
    "smirnoff99Frosst"
)

forcefield.parametrize(g)

from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
    n_samples=100,
    n_steps_per_sample=10,
)
simulation.run(g)

# +
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)

normalize = esp.data.normalize.ESOL100LogNormalNormalize()

epsilon = 0.1

for term in ['n2', 'n3']:
    for param in ['k', 'eq']:
        setattr(
            representation, '%s_%s' % (term, param),
            torch.nn.Parameter(
                torch.zeros_like(
                    g.nodes[term].data[param + '_ref'],
                )
                
            )
        
        )
        
net = torch.nn.Sequential(
        representation, 
        esp.mm.geometry.GeometryInGraph(),
        esp.mm.energy.EnergyInGraph(), # predicted energy -> u
        esp.mm.energy.EnergyInGraph(suffix='_ref') # reference energy -> u_ref,
)

optimizer = torch.optim.Adam(
    net.parameters(),
    0.1,
)

# optimizer = torch.optim.LBFGS(
#     net.parameters(),
#     0.1,
#     line_search_fn='strong_wolfe',
# )
# -

net(g.heterograph)

states = []
losses = []


# +
def l():
    net(g.heterograph)

    loss = torch.nn.MSELoss()(
        g.nodes['g'].data['u_ref'],
        g.nodes['g'].data['u'],
    )

    loss = loss.sum()

    losses.append(loss.detach().numpy())

    # loss.backward()

    print(loss)
    return loss

l()
# optimizer.step(l)
# -

g.nodes['n3'].data['k']

for _ in range(1000):
    optimizer.zero_grad()
    
    def l():
        normalize.unnorm(net(g.heterograph))
        
        loss = torch.nn.MSELoss()(
            g.nodes['g'].data['u_ref'],
            g.nodes['g'].data['u'],
        )

        loss = loss.sum()
        
        losses.append(loss.detach().numpy())
        
        loss.backward()
        
        print(loss)
        return loss
    
    optimizer.step(l)
    
    states.append(
        {
            '%s_%s' % (term, param): getattr(
                net[0],
                '%s_%s' % (term, param)
            ).detach().clone().numpy()
            for term in ['n3'] for param in ['k', 'eq']
        }
    )

plt.plot(losses)

ks = np.array([state['n3_k'].flatten() for state in states])
eqs = np.array([state['n3_eq'].flatten() for state in states])

eqs.std(axis=0)

for idx in range(8):
    plt.plot(np.diff(ks[:, idx]))
    

eqs

for idx in range(8):
    plt.plot(eqs[:, idx])

plt.scatter(
    g.nodes['g'].data['u_ref'].flatten().detach(),
    g.nodes['g'].data['u'].flatten().detach(),
)

# # initialize clever

# +
g = esp.Graph('C')

forcefield = esp.graphs.legacy_force_field.LegacyForceField(
    "smirnoff99Frosst"
)

forcefield.parametrize(g)

from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
    n_samples=100,
    n_steps_per_sample=10,
)
simulation.run(g)

# +
for term in ['n2', 'n3']:
    for param in ['k', 'eq']:
        print(term, param)
        print(g.nodes[term].data[param + '_ref'].mean())
        print(g.nodes[term].data[param + '_ref'].shape)
        



# +
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)

representation.n2_k = torch.nn.Parameter(torch.ones(8, 1) * 200000.0)
representation.n2_eq = torch.nn.Parameter(torch.ones(8, 1) * 0.10)

representation.n3_k = torch.nn.Parameter(torch.ones(12, 1) * 200.)
representation.n3_eq = torch.nn.Parameter(torch.ones(12, 1) * 1.0)

net = torch.nn.Sequential(
        representation, 
        esp.mm.geometry.GeometryInGraph(),
        esp.mm.energy.EnergyInGraph(), # predicted energy -> u
        esp.mm.energy.EnergyInGraph(suffix='_ref') # reference energy -> u_ref,
)

# optimizer = torch.optim.Adam(
#     net.parameters(),
#     0.1,
# )

optimizer = torch.optim.LBFGS(
    net.parameters(),
    0.1,
    line_search_fn='strong_wolfe',
)
# -

print(net[0].n2_k)

net(g.heterograph)

states = []
losses = []

for _ in range(1000):
    optimizer.zero_grad()
    
    def l():
        net(g.heterograph)
        
        loss = torch.nn.MSELoss()(
            g.nodes['n2'].data['u_ref'],
            g.nodes['n2'].data['u'],
        )
        
        loss = loss.sum()
        
        losses.append(loss.detach().numpy())
        
        loss.backward()
        
        print(loss)
        return loss
    
    optimizer.step(l)
    
    states.append(
        {
            '%s_%s' % (term, param): getattr(
                net[0],
                '%s_%s' % (term, param)
            ).detach().clone().numpy()
            for term in ['n2'] for param in ['k', 'eq']
        }
    )

plt.plot(losses)
plt.yscale('log')

ks = np.array([state['n2_k'].flatten() for state in states])
eqs = np.array([state['n2_eq'].flatten() for state in states])

eqs.std(axis=0)

eqs

for idx in range(8):
    plt.plot(np.diff(ks[:, idx]))
    

eqs

for idx in range(8):
    plt.plot(eqs[:, idx])

# +

plt.rc('font', family='serif', size=10)
plt.rc('xtick', labelsize=8)
plt.rc('ytick', labelsize=8)



plt.scatter(
    g.nodes['n2'].data['u_ref'].flatten().detach(),
    g.nodes['n2'].data['u'].flatten().detach(),
)

plt.xlabel('$U_\mathtt{ref}$')
plt.ylabel('$U_\mathtt{pred}$')
# -



 

StarcoderdataJupyterScriptsDedupFiltered