from typing import Dict, List, Optional, Tuple
import numpy as np
from kliff.dataset.dataset import Configuration
from kliff.models.model import ComputeArguments, Model
from kliff.models.parameter import Parameter
from kliff.neighbor import NeighborList, assemble_forces, assemble_stress
from kliff.transforms.parameter_transforms import ParameterTransform
[docs]
class LJComputeArguments(ComputeArguments):
"""
KLIFF built-in Lennard-Jones 6-12 potential computation functions.
"""
implemented_property = ["energy", "forces", "stress"]
def __init__(
self,
conf: Configuration,
supported_species: Dict[str, int],
influence_distance: float,
compute_energy: bool = True,
compute_forces: bool = True,
compute_stress: bool = False,
):
if supported_species is None:
species = sorted(set(conf.species))
supported_species = {si: i for i, si in enumerate(species)}
# using a single parameter for all species pairs
self.specie_pairs_to_param_index = {
(si, sj): 0 for si in species for sj in species
}
else:
self.specie_pairs_to_param_index = (
self._get_specie_pairs_to_param_index_map(
list(supported_species.keys())
)
)
super(LJComputeArguments, self).__init__(
conf,
supported_species,
influence_distance,
compute_energy,
compute_forces,
compute_stress,
)
self.neigh = NeighborList(
self.conf, influence_distance, padding_need_neigh=False
)
# TODO, rewrite this function to move the tow phi functions to LennardJones
# class, and then we can pass the model to this function, instead of the params.
# With this, we can unify the calling function in the Calculator.
[docs]
def compute(self, params: Dict[str, Parameter]):
coords = self.conf.coords
species = self.conf.species
coords_including_padding = self.neigh.coords
if self.compute_forces:
forces_including_padding = np.zeros_like(coords_including_padding)
energy = 0
for i, (xyz_i, si) in enumerate(zip(coords, species)):
neighlist, _, neigh_species = self.neigh.get_neigh(i)
for j, sj in zip(neighlist, neigh_species):
idx = self.specie_pairs_to_param_index[(si, sj)]
epsilon = params["epsilon"][idx]
sigma = params["sigma"][idx]
rcut = params["cutoff"][idx]
xyz_j = coords_including_padding[j]
rij = xyz_j - xyz_i
r = np.linalg.norm(rij)
if self.compute_forces:
phi, dphi = self.calc_phi_dphi(epsilon, sigma, r, rcut)
energy += 0.5 * phi
pair = 0.5 * dphi / r * rij
forces_including_padding[i] = forces_including_padding[i] + pair
forces_including_padding[j] = forces_including_padding[j] - pair
elif self.compute_energy:
phi = self.calc_phi(epsilon, sigma, r, rcut)
energy += 0.5 * phi
if self.compute_energy:
self.results["energy"] = energy
if self.compute_forces:
forces = assemble_forces(
forces_including_padding, len(coords), self.neigh.padding_image
)
self.results["forces"] = forces
if self.compute_stress:
volume = self.conf.get_volume()
stress = assemble_stress(
coords_including_padding, forces_including_padding, volume
)
self.results["stress"] = stress
[docs]
@staticmethod
def calc_phi(epsilon, sigma, r, rcut):
if r > rcut:
phi = 0
else:
sor = sigma / r
sor6 = sor * sor * sor
sor6 = sor6 * sor6
sor12 = sor6 * sor6
phi = 4 * epsilon * (sor12 - sor6)
return phi
[docs]
@staticmethod
def calc_phi_dphi(epsilon, sigma, r, rcut):
if r > rcut:
phi = 0.0
dphi = 0.0
else:
sor = sigma / r
sor6 = sor * sor * sor
sor6 = sor6 * sor6
sor12 = sor6 * sor6
phi = 4 * epsilon * (sor12 - sor6)
dphi = 24 * epsilon * (-2 * sor12 + sor6) / r
return phi, dphi
@staticmethod
def _get_specie_pairs_to_param_index_map(
species: List[str],
) -> Dict[Tuple[str, str], int]:
"""
Return a map from a tuple of two species to the index of the corresponding
parameter in the parameter array.
For example, if the supported species are ["A", "B", "C"], then the map will be
{(A, A): 0, (A, B): 1, (B, A): 1, (A, C): 2, (C, A): 2,
(B, B): 3, (B, C): 4, (C, B): 4,
(C, C): 5}.
"""
n = len(species)
speices_to_param_index_map = {}
index = 0
for i in range(n):
si = species[i]
for j in range(i, n):
sj = species[j]
speices_to_param_index_map[(si, sj)] = index
if i != j:
speices_to_param_index_map[(sj, si)] = index
index += 1
return speices_to_param_index_map
[docs]
class LennardJones(Model):
"""
KLIFF built-in Lennard-Jones 6-12 potential model.
This model supports multiple species, where a different set of parameters is used
for each species pair. For example if species A, B, and C are provided, then there
will be 6 values for each of the epsilon and sigma parameters. The order of the
parameters is as follows: A-A, A-B, A-C, B-B, B-C, and C-C.
Args:
model_name: name of the model
species: list of species. If None, there model will create a single value for
each parameter, and all species pair will use the same parameters. If a
list of species is provided, then the model will create a different set of
parameters for each species pair.
params_transform: parameter transform object. If None, no transformation is
performed.
"""
def __init__(
self,
model_name: str = "LJ6-12",
species: List[str] = None,
):
self.species = species
super(LennardJones, self).__init__(model_name)
[docs]
def init_model_params(self):
n = self._get_num_params()
model_params = {
"epsilon": Parameter(np.ones(n)),
"sigma": Parameter(2 * np.ones(n)),
"cutoff": Parameter(5 * np.ones(n)),
}
return model_params
[docs]
def init_influence_distance(self):
return self.model_params["cutoff"][0]
[docs]
def init_supported_species(self):
if self.species is None:
return None
else:
return {s: i for i, s in enumerate(self.species)}
[docs]
def get_compute_argument_class(self):
return LJComputeArguments
def _get_num_params(self):
if self.species is None:
n = 1
else:
n = len(self.species)
return (n + 1) * n // 2