Forces with PBC

[owainbeynon]

Hi,

I'm trying to train a model using forces and energies for silicon, based on the tutorial found here: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_03_force_models.html. The training runs ok, but using the resultant model for a geometry optimisation on a Si supercell, the resultant structure is incorrect, where it looks like bonds between Si atoms aren't properly conserved (image below).

Here is my script fro training on forces and energies:

import torch
import torchmetrics
import schnetpack as spk
import schnetpack.transform as trn
import pytorch_lightning as pl
import os

forcetut = './forcetut'
if not os.path.exists(forcetut):
    os.makedirs(forcetut)

custom_data = spk.data.AtomsDataModule(
        './new_dataset.db',
    batch_size=10,
    distance_unit='Ang',
   property_units={'energy':'eV', 'forces':'eV/Ang'},
   num_train=900,
   num_test=800,
    num_val=100,
    transforms=[
        trn.ASENeighborList(cutoff=5.),
        trn.RemoveOffsets("energy", remove_mean=True, remove_atomrefs=False),
        trn.CastTo32()
    ],
    num_workers=1,
    pin_memory=True, # set to false, when not using a GPU
)
custom_data.prepare_data()
custom_data.setup() 

properties = custom_data.dataset[0]
print('Loaded properties:\n', *['{:s}\n'.format(i) for i in properties.keys()])

cutoff = 7.
n_atom_basis = 30

pairwise_distance = spk.atomistic.PairwiseDistances() # calculates pairwise distances between atoms
radial_basis = spk.nn.GaussianRBF(n_rbf=20, cutoff=cutoff)
schnet = spk.representation.SchNet(
    n_atom_basis=n_atom_basis, n_interactions=3,
    radial_basis=radial_basis,
    cutoff_fn=spk.nn.CosineCutoff(cutoff)
)

pred_energy = spk.atomistic.Atomwise(n_in=n_atom_basis, output_key='energy')
#pred_forces = spk.atomistic.Forces(energy_key='energy', force_key='forces', stress_key='stress', calc_stress=True)
pred_forces = spk.atomistic.Forces(energy_key='energy', force_key='forces')

nnpot = spk.model.NeuralNetworkPotential(
    representation=schnet,
    input_modules=[pairwise_distance],
    output_modules=[pred_energy, pred_forces],
)

output_energy = spk.task.ModelOutput(
    name='energy',
    loss_fn=torch.nn.MSELoss(),
    loss_weight=0.01,
    metrics={
        "MAE": torchmetrics.MeanAbsoluteError()
    }
)

output_forces = spk.task.ModelOutput(
    name='forces',
    loss_fn=torch.nn.MSELoss(),
    loss_weight=0.99,
    metrics={
        "MAE": torchmetrics.MeanAbsoluteError()
    }
)

task = spk.task.AtomisticTask(
   model=nnpot,
    outputs=[output_energy, output_forces],
    optimizer_cls=torch.optim.AdamW,
   optimizer_args={"lr": 1e-4}
   )

logger = pl.loggers.TensorBoardLogger(save_dir=forcetut)
callbacks = [
    spk.train.ModelCheckpoint(
        model_path=os.path.join(forcetut, "Si.model"),
        save_top_k=1,
        monitor="val_loss"
    )
]

trainer = pl.Trainer(
    callbacks=callbacks,
    logger=logger,
    default_root_dir=forcetut,
#   max_epochs=5, # for testing, we restrict the number of epochs
)
trainer.fit(task, datamodule=custom_data)

I initially thought that this could be because I hadn't trained using stress, as found in this tutorial: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_05_materials.html. So I adjusted my training script to include stress accordingly, as detailed in this issue: https://github.com/atomistic-machine-learning/schnetpack/issues/551. However, training with this script, no validation data is loaded, therefore a model is not produced. Here is my training script included stress.

import torch
import torchmetrics
import schnetpack as spk
import schnetpack.transform as trn
import pytorch_lightning as pl
import os

forcetut = './forcetut'
if not os.path.exists(forcetut):
    os.makedirs(forcetut)

custom_data = spk.data.AtomsDataModule(
        './new_dataset.db',
    batch_size=10,
    distance_unit='Ang',
    property_units={'energy':'eV', 'forces':'eV/Ang', 'stress':'eV/Ang/Ang'},
   num_train=400,
    num_val=400,
    transforms=[
        trn.ASENeighborList(cutoff=5.),
        trn.RemoveOffsets("energy", remove_mean=True, remove_atomrefs=False),
        trn.CastTo32()
    ],
    num_workers=1,
    pin_memory=True, # set to false, when not using a GPU
)
custom_data.prepare_data()
custom_data.setup() 

properties = custom_data.dataset[0]
print('Loaded properties:\n', *['{:s}\n'.format(i) for i in properties.keys()])

cutoff = 7.
n_atom_basis = 30

radial_basis = spk.nn.GaussianRBF(n_rbf=20, cutoff=cutoff)
schnet = spk.representation.SchNet(
    n_atom_basis=n_atom_basis, n_interactions=3,
    radial_basis=radial_basis,
    cutoff_fn=spk.nn.CosineCutoff(cutoff)
)

pred_energy = spk.atomistic.Atomwise(n_in=n_atom_basis, output_key='energy')
pred_forces = spk.atomistic.Forces(energy_key='energy', force_key='forces', stress_key='stress', calc_stress=True)

nnpot = spk.model.NeuralNetworkPotential(
    representation=schnet,
    input_modules=[spk.atomistic.Strain(), spk.atomistic.PairwiseDistances()],
    output_modules=[pred_energy, pred_forces],
)

output_energy = spk.task.ModelOutput(
    name='energy',
    loss_fn=torch.nn.MSELoss(),
    loss_weight=0.01,
    metrics={"MAE": torchmetrics.MeanAbsoluteError()}  # Provide an empty dictionary as default
)

output_forces = spk.task.ModelOutput(
    name='forces',
    loss_fn=torch.nn.MSELoss(),
    loss_weight=0.98,
   metrics={"MAE": torchmetrics.MeanAbsoluteError()} # Provide an empty dictionary as default
)

output_stress = spk.task.ModelOutput(
    name='stress',
    loss_fn=torch.nn.MSELoss(),
    loss_weight=0.01,
    metrics={"MAE": torchmetrics.MeanAbsoluteError()}  # Provide an empty dictionary as default
)

task = spk.task.AtomisticTask(
    model=nnpot,
    outputs=[output_energy, output_forces],  # Only include energy output for now
    optimizer_cls=torch.optim.AdamW,
    optimizer_args={"lr": 1e-4}
)

logger = pl.loggers.TensorBoardLogger(save_dir=forcetut)
callbacks = [
    spk.train.ModelCheckpoint(
        model_path=os.path.join(forcetut, "best_inference_model"),
        save_top_k=1,
        monitor="val_loss"
    )
]

trainer = pl.Trainer(
    callbacks=callbacks,
    logger=logger,
    default_root_dir=forcetut,
#   max_epochs=5, # for testing, we restrict the number of epochs
)
trainer.fit(task, datamodule=custom_data)

Do you have any advice for proceeding or there any obvious/ immediate solutions to my outlined issues?

Thanks!

[jnsLs]

Dear @owainbeynon ,

often those weird structures appear when the units are not consistent or when something goes wrong with the PBC. I can see in the code you sent that your cutoff radius is not consistent and you atom_basis (feature size) seems to small.

cutoff = 7.
n_atom_basis = 30

i would suggest

cutoff = 5. # should be significantly smaller than size of the simulation box
n_atom_basis = 128

Best, Jonas

[owainbeynon]

Hi @jnsLs,

Thanks for the suggestions, now the optimisation for Si system looks OK.

My next question is now for running MD on the system. Currently following this tutorial (https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_04_molecular_dynamics.html), and exchanging the example molecule and model for my own Si system and trained model, I get a similar issue where the atoms move as if they weren't in a lattice (move outside the cell and behave as if they were not bonded).

What alterations and considerations are needed when adjusting the MD tutorial to a periodic system? As I'm aware the example demonstrates MD for a molecule.

Any advice would be appreciated.

Thanks, Owain

[jnsLs]

Hi Owain,

you need to make sure that you are using a neighborlist that supports pbc and your input xyz file needs to contain information about the pbc. This usually means, pbc should be True along all three cartesian directions and the proper lattice (simulation box) axes are required.

Best, Jonas

[owainbeynon]

Hi @jnsLs,

Thanks for the advice.

I'm running an MD using the Schnetpack MD driver as per this tutorial: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_04_molecular_dynamics.html and here is my script:

from schnetpack.md import System
from ase.io import read
import os

md_workdir = 'mdtut'

# Gnerate a directory of not present
if not os.path.exists(md_workdir):
    os.mkdir(md_workdir)
test_path = "./model/"

model_path = os.path.join(test_path, 'si.model')
molecule_path = "si.xyz"

# Load atoms with ASE
molecule = read(molecule_path)
print(molecule)
# Number of molecular replicas

# Create system instance and load molecule
md_system = System()
md_system.load_molecules(
    molecule,
    position_unit_input="Angstrom"
)

from schnetpack.md import UniformInit

system_temperature = 300 # Kelvin

# Set up the initializer
md_initializer = UniformInit(
    system_temperature,
    remove_center_of_mass=True,
    remove_translation=True,
    remove_rotation=True,
)

# Initialize the system momenta
md_initializer.initialize_system(md_system)

from schnetpack.md.integrators import VelocityVerlet

time_step = 0.5  # fs

# Set up the integrator
md_integrator = VelocityVerlet(time_step)

from schnetpack.md.neighborlist_md import NeighborListMD
from schnetpack.transform import ASENeighborList

# set cutoff and buffer region
cutoff = 5.0  # Angstrom (units used in model) (same as training) 

cutoff_shell = 2 # Angstrom

# initialize neighbor list for MD using the ASENeighborlist as basis
md_neighborlist = NeighborListMD(
    cutoff_shell,
    cutoff,
    ASENeighborList,
)

from schnetpack.md.calculators import SchNetPackCalculator
from schnetpack import properties

md_calculator = SchNetPackCalculator(
    model_path,  # path to stored model
    "forces", #force key
    "eV", #Energy units
    "Angstrom", # lenth units 
    md_neighborlist, # neighbor list
    energy_key="energy",
    required_properties=[],  # additional properties extracted from the model
)

from schnetpack.md import Simulator

md_simulator = Simulator(
    md_system,
    md_integrator,
    md_calculator
)

import torch

# check if a GPU is available and use a CPU otherwise
if torch.cuda.is_available():
    md_device = "cuda"
else:
    md_device = "cpu"

# use single precision
md_precision = torch.float32

# set precision
md_simulator = md_simulator.to(md_precision)
# move everything to target device
md_simulator = md_simulator.to(md_device)

from schnetpack.md.simulation_hooks import LangevinThermostat

# Set temperature and thermostat constant
bath_temperature = 300  # K
time_constant = 100  # fs

# Initialize the thermostat
langevin = LangevinThermostat(bath_temperature, time_constant)

simulation_hooks = [
    langevin
]

from schnetpack.md.simulation_hooks import callback_hooks

# Path to database
log_file = os.path.join(md_workdir, "simulation.hdf5")

# Size of the buffer
buffer_size = 100

# Set up data streams to store positions, momenta and the energy
data_streams = [
    callback_hooks.MoleculeStream(store_velocities=True),
    callback_hooks.PropertyStream(target_properties=[properties.energy]),
]

# Create the file logger
file_logger = callback_hooks.FileLogger(
    log_file,
    buffer_size,
    data_streams=data_streams,
    every_n_steps=1,  # logging frequency
    precision=32,  # floating point precision used in hdf5 database
)

# Update the simulation hooks
simulation_hooks.append(file_logger)

chk_file = os.path.join(md_workdir, 'simulation.chk')

# Create the checkpoint logger
checkpoint = callback_hooks.Checkpoint(chk_file, every_n_steps=100)

# Update the simulation hooks
simulation_hooks.append(checkpoint)

# directory where tensorboard log will be stored to
tensorboard_dir = os.path.join(md_workdir, 'logs')

tensorboard_logger = callback_hooks.TensorBoardLogger(
    tensorboard_dir,
    ["energy", "temperature"], # properties to log
)

# update simulation hooks
simulation_hooks.append(tensorboard_logger)

md_simulator = Simulator(md_system, md_integrator, md_calculator, simulator_hooks=simulation_hooks)

md_simulator = md_simulator.to(md_precision)
md_simulator = md_simulator.to(md_device)

n_steps = 2000

md_simulator.simulate(n_steps)

from schnetpack.md.data import HDF5Loader

data = HDF5Loader(log_file)

for prop in data.properties:
    print(prop)

from ase.io import write

# extract structure information from HDF5 data
md_atoms = data.convert_to_atoms()

# write list of Atoms to XYZ file
write(
    os.path.join(md_workdir, "trajectory.traj"),
    md_atoms,
    format="traj"
)

but I get an unreasonable looking MD trajectory (atypical behaviour for solids) , as seen here:

What could be going wrong with the simulation to get this behaviour?

For reference, I've attached the training script, AIMD data used for training, and a geometry optimisation output(which works) using the resultant model. Si.zip

Thanks

[jnsLs]

Dear @owainbeynon

there are many possible reasons why your structure does not look reasonable. However, we don't see a specific issue with our code there and unfortunately it is not feasible for us to evaluate your approach in detail.

Some general remarks:

• your dataset might not be representative enough, resulting in unfaithful model predictions for conformations far away from the training set
• you might have too many atoms in the simulation box (for example because you are using the wrong distance units)
• you can check at which point your MD trajectory diverges from something reasonable
• you can simulate different ensembles (NVE, NVT, NPT)
• you can check if the corresponding conserved properties actually stay constant (temperature, ...)
• you can reduce the temperature (if optimization works, MD at T=0K should also work)

Also we can recommend our CLI. You can use the spkmd script instead of writing your own code.

Best regards, Jonas

[jnsLs]

Since we do not see a software issue here, we close this issue.

[owainbeynon]

Ok, thanks for the help