# * **************************************************************************
# *
# * CDTK, Chemical Dynamics Toolkit
# * A modular system for chemical dynamics applications and more
# *
# * Copyright (C) 2018
# * Ralph Welsch, DESY, <ralph.welsch@desy.de>
# *
#* Copyright (C) 2020, 2021, 2022, 2023
#* Ludger Inhester, DESY, ludger.inhester@cfel.de
#*
# * This file is part of CDTK.
# *
# * CDTK is free software: you can redistribute it and/or modify
# * it under the terms of the GNU General Public License as published by
# * the Free Software Foundation, either version 3 of the License, or
# * (at your option) any later version.
# *
# * This program is distributed in the hope that it will be useful,
# * but WITHOUT ANY WARRANTY; without even the implied warranty of
# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# * GNU General Public License for more details.
# *
# * You should have received a copy of the GNU General Public License
# * along with this program. If not, see <http://www.gnu.org/licenses/>.
# *
# * **************************************************************************
import numpy as np
import warnings
import CDTK.Tools.Conversion as conv
[docs]
class Constraints(object):
def __init__(self, f_constraint):
self.f_constraint = f_constraint
self.lam_r = None
[docs]
def get_dg0dg1_matrix(self, r1, dgdt0overM, lambda_r):
"""
Calculates the matrix resulting from multipying the constraints Jacobian
calculated at r(t) times the Jacobian calculated at r_c(t+dt)
Arguments:
----------
r1 : (natoms, 3) ndarray of float
Positions of atoms at t+dt, r(t+dt).
dgdt0overM : (nconstraints, natoms, 3) ndarray of float
Derivative of the constraints at r(t) times dt over the mass of each atom.
lambda_r : (nconstraints) nndarray of float
Lagrange multiplier for the constrins of the position (\lambda_R) for each constraint.
Return Values:
--------------
g: (nconstraints) ndarray of float
Constraints calculated at the new positions r_c(t+dt) after applying constraints.
dg0dg1_mat: (nconstraints, nconstraints) ndarray of float.
Derivative of the constrains at r(t) [\nabla g(r(t))] times the derivative of the constraints at r_c(t+dt).
"""
# Apply constrains to position, r(t+dt)->r_c(t+dt)
r_new = r1 - np.einsum(
"ijk,i->jk", dgdt0overM, lambda_r
) # sum over [constrains] axis
g1, dg1 = self.f_constraint(r_new, der=True) # Calculate new constraints
dg0dg1_mat = -np.einsum(
"ijk,ljk->il", dg1, dgdt0overM
) # sums over [natoms], [xyz] axes
return g1, dg0dg1_mat
[docs]
def get_correction_r(self, r0, r1, dt, m, tolerance=1e-5):
"""
Calculates the correction to the position according to the RATTLE algorithm
to a geometry with constraints.
todo: Put reference to Andersen
todo: Molecular simulations of fluids: Theory, Algorithms, Object-orientation
Parameters
----------
r0 : (natoms, 3) ndarray of float
Positions of atoms at t, r(t).
r1 : (natoms, 3) ndarray of float
Positions of atoms at t+dt, r(t+dt).
dt : float
Timestep
m : (natoms)
Array with atomic masses
tolerance : float, optional
Tolearance of the constrain function for the iterative procedure, by default 1e-5
Returns
-------
correction: (natom, 3) ndarray of float
"""
# Calculate constrains of r(t)
g0, dg0 = self.f_constraint(r0, der=True)
lambda_r = np.zeros(g0.shape[0])
m = np.reshape(m, (1, len(m), 1))
dg0overM = dg0 / m
g1, dg0dg1_mat = self.get_dg0dg1_matrix(r1, dg0overM * dt, lambda_r)
# Start iterative loop to refine lambda_r
it = 0
while np.max(np.abs(g1)) > tolerance and it < 40:
l = np.linalg.solve(dg0dg1_mat, g1)
lambda_r = lambda_r - l
g1, dg0dg1_mat = self.get_dg0dg1_matrix(r1, dg0overM * dt, lambda_r)
it += 1
if np.max(np.abs(g1)) > tolerance:
warnings.warn("Warning constraints failed in {:d} iterations".format(it))
return -np.einsum("ijk,i->jk", dg0overM, lambda_r) # sum over [constrains] axis
[docs]
def get_correction_v(self, r1, v1, dt, m):
"""
Calculates the correction to the velocity according to the RATTLE algorithm
to a geometry with constraints. This variant takes advantage of
r"\frac{d\sigma(X(t))}{dt} = 0" and r"\frac{d\sigma}{dX} \cdot \frac{dX}{dt} = 0".
Some information can be found in [Ref] Andersen JOURNAL OF COMPUTATIONAL PHYSICS 52, 24-34
Arguments:
r1 ((natoms, 3) ndarray of float): Positions of atoms at t+dt with constraints applied, r(t+dt).
v1 ((natoms, 3) ndarray of float): Velocity of atoms at t+dt without constraints, r(t+dt).
dt (float): Timestep
m ((natoms) ndarray of float): Array with atomic masses
Returns:
((natom, 3) - ndarray of float): correction
"""
_, dg1 = self.f_constraint(r1, der=True)
dgOverM = dg1 / m[:, np.newaxis]
mOverdg1dg1 = 1 / np.einsum("ijk,ijk->i", dg1, dgOverM)
dg1v1 = np.einsum("ijk,jk->i", dg1, v1)
lambda_v = mOverdg1dg1 * dg1v1
return -0.5 * dt * (np.einsum("ijk,i->jk",dg1, lambda_v) / m[:, np.newaxis])
[docs]
def check_constraints(self, x1):
tol = 1e-4
g = self.f_constraint(x1, der=False)
av_g = np.sum(abs(g)) / float(len(g))
max_g = np.max(abs(g))
if av_g > tol or max_g > tol:
print("WARNING: constraints not fulfilled")
print("av_g: ", av_g)
print("max_g: ", max_g)
print("g: ", g)
return
[docs]
def enforceConstraints(self, x1, m, tol = 1e-6, nsteps=100):
"""Enforces constraints by making Newton steps for the coordinates
Keeps the center of mass of the molecule the same
TODO: address the velocities
TODO: prevent Rotation
Args:
x1 (numpy array shape (natoms,3)): coordinates
m (numpy array shape (natoms)): atom masses
tol (float, optional): tolerance. Defaults to 1e-6.
nsteps (int, optional): maximum number of steps. Defaults to 100.
Returns:
x1: (numpy array shape (natoms,3)): new coordinates
"""
# go to center of mass of the molecule
com0 = (x1[:,:] * m[:,np.newaxis]).sum(axis=0)/m.sum()
x1[:,:] = x1[:,:] - com0[np.newaxis,:]
# make Newton steps to enforce constraints
for i in range(nsteps):
g, dg = self.f_constraint(x1, der=True)
av_g = abs(np.mean(g))
max_g = np.max(abs(g))
#print(max_g)
if(av_g < tol and max_g < tol):
break
step = (g[:,np.newaxis,np.newaxis] * dg[:,:,:]).sum(axis=0)
x1 = x1 - 0.5 * step
# move to the old center of mass
com1 = (x1[:,:] * m[:,np.newaxis]).sum(axis=0)/m.sum()
x1 = x1 - com1[np.newaxis,:] + com0[np.newaxis,:]
return x1