import logging
from time import time
import numpy as np
from pymatgen.core.lattice import Lattice
from pyxtal.operations import *
logging.basicConfig(filename="test.log", level=logging.DEBUG)
eV2GPa = 160.217
GPa2eV = 1.0 / eV2GPa
"""
Scripts to compute LJ energy and force
"""
[docs]
def get_neighbors(struc, i, rcut):
"""
script to get the neighbors info for an atom in the struct
"""
lat = struc.lattice_matrix
center = struc.cart_coords[i]
struc.frac_coords[i]
fcoords = struc.frac_coords
r, dists, inds, images = Lattice(lat).get_points_in_sphere(fcoords, center=center, r=rcut, zip_results=False)
ids = np.where(dists > 0.1)
r = r[ids]
r = np.dot(r, lat) - center
return dists[ids] ** 2, r
[docs]
class LJ:
"""
LJ model for 3D crystals, maybe extended to 0D, 1D, 2D later
"""
def __init__(self, epsilon=1.0, sigma=1.0, rcut=8.0):
"""
passing the parameter to LJ model
- epsilon
- sigma
- rcut
"""
self.epsilon = epsilon
self.sigma = sigma
self.rcut = rcut
[docs]
def calc(self, struc, press=1e-4):
pos = struc.cart_coords
lat = struc.lattice_matrix
volume = np.linalg.det(lat)
# initiate the energy, force, stress
energy = 0
force = np.zeros([len(pos), 3])
stress = np.zeros([3, 3])
sigma6 = self.sigma**6
sigma12 = sigma6 * sigma6
# calculating the energy/force, needs to update sigma/epsilons
for i, _pos0 in enumerate(pos):
[r2, r] = get_neighbors(struc, i, self.rcut)
r6 = np.power(r2, 3)
r12 = np.power(r6, 2)
energy += np.sum(4.0 * self.epsilon * (sigma12 / r12 - sigma6 / r6))
f = (24 * self.epsilon * (2.0 * sigma12 / r12 - sigma6 / r6) / r2)[:, np.newaxis] * r
force[i] = f.sum(axis=0)
stress += np.dot(f.T, r)
energy = 0.5 * energy
enthalpy = energy + press * volume * GPa2eV
force = force
stress = -0.5 * stress / volume * eV2GPa
return energy, enthalpy, force, stress
[docs]
class FIRE:
"""
FIRE algorithm for structure optimization
"""
def __init__(
self,
struc,
model,
symmetrize=False,
e_tol=1e-5,
f_tol=1e-2,
dt=0.1,
maxmove=0.1,
dtmax=1.0,
Nmin=10,
finc=1.1,
fdec=0.5,
astart=0.1,
fa=0.99,
a=0.1,
):
"""
Parameters for FIRE:
"""
self.dt = dt
self.maxmove = maxmove
self.symmetrize = symmetrize
self.f_tol = f_tol
self.e_tol = e_tol
self.dtmax = dtmax
self.Nmin = Nmin
self.finc = finc
self.fdec = fdec
self.astart = astart
self.fa = fa
self.a = a
self.nsteps = 0
self.Nsteps = 0
self.time0 = time()
self.model = model
self.struc = struc
self.initialize()
[docs]
def initialize(self):
"""
initilaize the positions, energy, force for the 1st step
"""
self.energy, self.enthalpy, self.force, self.stress = self.model.calc(self.struc)
self.fmax = np.max(np.abs(np.vstack((self.stress, self.force)).flatten()))
self.v = np.zeros((len(self.force) + 3, 3))
# print('Initial Energy: {:12.4f}'.format(self.energy))
[docs]
def update(self, freq=50):
self.energy, self.enthalpy, self.force, self.stress = self.model.calc(self.struc)
self.fmax = np.max(np.abs(np.vstack((self.stress, self.force)).flatten()))
if self.nsteps % freq == 0:
logging.debug(
f"Step {self.nsteps:4d} Enthalpy: {self.enthalpy:12.4f} Fmax: {self.fmax:12.4f} Vol: {self.volume:6.2f}"
)
[docs]
def step(self):
f = np.vstack((self.stress, self.force))
pos = np.vstack((self.struc.lattice_matrix, self.struc.cart_coords))
vf = np.vdot(f, self.v)
if vf > 0.0:
self.v = (1.0 - self.a) * self.v + self.a * f / np.sqrt(np.vdot(f, f)) * np.sqrt(np.vdot(self.v, self.v))
if self.Nsteps > self.Nmin:
self.dt = min(self.dt * self.finc, self.dtmax)
self.a *= self.fa
self.Nsteps += 1
else:
self.v[:] *= 0.0
self.a = self.astart
self.dt *= self.fdec
self.Nsteps = 0
self.v += self.dt * f
dr = self.dt * self.v # needs to impose constraints
normdr = np.sqrt(np.vdot(dr, dr))
if normdr > self.maxmove:
dr = self.maxmove * dr / normdr
# print('frac0', np.dot(pos[3:,:], np.linalg.inv(pos[:3,:])))
# Symmetrize the force
if self.symmetrize:
# f[:3, :] is the gradient for force, need to symmetrize it as well
dr[:3, :] = self.symmetrized_stress(dr[:3, :])
dr[3:, :] = np.dot(self.struc.frac_coords, dr[:3, :])
f_frac = np.dot(dr[3:, :], np.linalg.inv(pos[:3, :] - dr[:3, :]))
f_frac = self.symmetrized_force(f_frac)
dr[3:, :] += np.dot(f_frac, pos[:3, :] - dr[:3, :])
# pos = pos - dr
# Symmetrize positions
# if self.symmetrize:
# pos = self.symmetrized_coords(pos)
# print(self.force)
# print(self.v)
self.struc.lattice_matrix = pos[:3, :] - dr[:3, :]
self.struc.cart_coords = pos[3:, :] - dr[3:, :]
self.struc.frac_coords = np.dot(self.struc.cart_coords, np.linalg.inv(self.struc.lattice_matrix))
self.volume = np.linalg.det(self.struc.lattice_matrix)
# sg = get_symmetry_dataset((pos[:3, :], self.struc.frac_coords, [6]*4), symprec=0.1)['number']
# print(sg)
# if self.symmetrize and sg <19:
# print(sg)
# print('dr\n', dr)
# print('pos\n', pos)
# print('frac\n', self.struc.frac_coords)
# print('force\n', f_frac)
# sys.exit()
self.update()
[docs]
def run(self, max_steps=1000):
self.max_steps = max_steps
while self.nsteps < max_steps:
if self.check_convergence():
break
self.step()
self.nsteps += 1
logging.info(
f"Finish at Step {self.nsteps:4d} Enthalpy: {self.enthalpy:12.4f} Fmax: {self.fmax:12.4f} Time: {time() - self.time0:6.2f} seconds"
)
[docs]
def check_convergence(self):
"""
There should be two options to terminate the optimization
before it reaches the maximum number of cycles
1, by energy difference compared to the previous step: self.e_tol
2, by the residual force: self.f_tol
Will implement both options later.
"""
converged = False
if self.fmax < self.f_tol and self.nsteps > 0:
converged = True
return converged
[docs]
def symmetrized_coords(self, coords):
start_index = 0
new_coords = []
for ws in self.struc.wyckoff_sites:
# Get coordinates associated with WP
original_coords = coords[start_index : start_index + ws.multiplicity]
# Re-generate the points from the first generating point
gen_coord = coords[start_index]
wp_coords0 = apply_ops(gen_coord, ws.wp.generators_m)
# Calculate the PBC translation applied to each point
translations = original_coords - wp_coords0
frac_translations = np.dot(translations, np.linalg.inv(self.struc.lattice_matrix))
frac_translations = np.round(frac_translations)
cart_translations = np.dot(frac_translations, self.struc.lattice_matrix)
# Subtract translations from original coords
translated_coords = original_coords - cart_translations
# Apply inverse WP operations to get generating_points
inverse_ops = ws.wp.inverse_generators_m
generating_points = apply_ops_diagonal(translated_coords, inverse_ops)
# Find the average of the generating points
average_point = np.sum(generating_points, axis=0) / ws.multiplicity
# Convert to fractional coordinates
average_point = np.dot(average_point, np.linalg.inv(self.struc.lattice_matrix))
# Apply the WP operations to the averaged generating point
wp_coords = apply_ops(average_point, ws.wp)
# Convert back to Cartesian coordintes
wp_coords = np.dot(wp_coords, self.struc.lattice_matrix)
# Re-apply the cartesian translations
wp_coords = wp_coords + cart_translations
new_coords = wp_coords if len(new_coords) == 0 else np.vstack([new_coords, wp_coords])
start_index += ws.multiplicity
return new_coords
[docs]
def symmetrized_force(self, coords):
start_index = 0
new_coords = []
for ws in self.struc.wyckoff_sites:
# Get coordinates associated with WP
coords[start_index : start_index + ws.multiplicity]
# Re-generate the points from the first generating point
gen_coord = coords[start_index]
wp_coords0 = apply_ops(gen_coord, ws.wp.generators_m)
# Apply inverse WP operations to get generating_points
inverse_ops = ws.wp.inverse_generators_m
generating_points = apply_ops_diagonal(wp_coords0, inverse_ops)
# Find the average of the generating points
average_point = np.sum(generating_points, axis=0) / ws.multiplicity
# Convert to fractional coordinates
average_point = np.dot(average_point, np.linalg.inv(self.struc.lattice_matrix))
# For forces, we do not apply translational WP operations, so we truncate the ops
matrices = np.array([op.affine_matrix[:3, :3] for op in ws.wp])
# Apply the truncated WP operations to the averaged generating point
wp_coords = np.dot(average_point, matrices)
# Convert back to Cartesian coordintes
wp_coords = np.dot(wp_coords, self.struc.lattice_matrix)
new_coords = wp_coords if len(new_coords) == 0 else np.vstack([new_coords, wp_coords])
start_index += ws.multiplicity
# if np.sum(np.abs(coords-new_coords))>1e-1:
# print(coords)
# print(new_coords)
# print(coords-new_coords)
# sys.exit()
return new_coords
[docs]
def symmetrized_stress(self, stress):
# Make the matrix lower-diagonal
for i, j in [(1, 0), (2, 1), (2, 2)]:
stress[i][j] = stress[i][j] + stress[j][i]
stress[j][i] = 0
# Normalize the matrix
snm = self.struc.lattice.stress_normalization_matrix
m2 = np.multiply(stress, snm)
# Normalize the on-diagonal elements
indices = self.struc.lattice.stress_indices
if len(indices) == 2:
total = 0
for index in indices:
total += stress[index]
value = np.sqrt(total)
for index in inices:
m2[index] = value
elif len(indices) == 3:
total = 0
for index in indices:
total += stress[index]
value = np.cbrt(total)
for index in inices:
m2[index] = value
return m2
if __name__ == "__main__":
from spglib import get_symmetry_dataset
from pyxtal import pyxtal
for _i in range(10):
crystal = pyxtal()
crystal.from_random(3, 11, ["C"], [4])
if crystal.valid:
crystal1 = deepcopy(crystal)
test = LJ(epsilon=0.01, sigma=3.40, rcut=8.0)
struc = (crystal1.lattice_matrix, crystal1.frac_coords, [6] * 4)
eng, enth, force, stress = test.calc(crystal1)
sg = get_symmetry_dataset(struc)["number"]
print(f"\nBefore relaxation Space group: {sg:4d} Energy: {eng:12.4} Enthalpy: {enth:12.4}")
dyn1 = FIRE(crystal1, test, f_tol=1e-5, dt=0.2, maxmove=0.2) # , symmetrize=True)
dyn1.run(500)
eng, enth, force, stress = test.calc(crystal1)
struc = (dyn1.struc.lattice_matrix, dyn1.struc.frac_coords, [6] * 4)
sg = get_symmetry_dataset(struc, symprec=0.1)["number"]
print(f"After relaxation without symm Space group: {sg:4d} Energy: {eng:12.4} Enthalpy: {enth:12.4}")
dyn1 = FIRE(crystal, test, f_tol=1e-5, dt=0.2, maxmove=0.2, symmetrize=True)
dyn1.run(500)
eng, enth, force, stress = test.calc(crystal)
struc = (dyn1.struc.lattice_matrix, dyn1.struc.frac_coords, [6] * 4)
sg = get_symmetry_dataset(struc, symprec=0.02)["number"]
if sg is None:
sg = 0
print(f"After relaxation with symm Space group: {sg:4d} Energy: {eng:12.4} Enthalpy: {enth:12.4}")
# dyn1 = FIRE(crystal, test, f_tol=1e-5, dt=0.2, maxmove=0.2)
# struc = (dyn1.struc.lattice_matrix, dyn1.struc.frac_coords, [6]*4)
# sg = get_symmetry_dataset(struc, symprec=0.02)['number']
# print('After relaxation without symm Space group: {:4d} Energy: {:12.4} Enthalpy: {:12.4}'.format(sg, eng, enth))
# right now, it seems structures goes to either HCP of FCC after relaxation, which is expected for 3D LJ system
# need to compare with other code to see if the energy is correct