Source code for moldesign.geom.shake

# Copyright 2016 Autodesk Inc.
#
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#
# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
import itertools
import numpy as np
import moldesign as mdt
from moldesign import units as u

from .constraints import FixedCoordinate, FixedPosition

# TODO: create dynamics wrapper that uses timestep to explicitly calculate constraint forces

[docs]def shake_positions(mol, prev_positions, max_cycles=100, use_masses=True):
""" Satisfy all molecular constraints using the SHAKE algorithm

Args:
mol (moldesign.Molecule): molecule with unsatisfied constraints
prev_positions (mdt.units.Array[length]): positions prior to move
max_cycles (int): halt and raise an exception after this many cycles

Note:

References:
R. Elber, A. Ruymgaart, B. Hess: SHAKE Parallelization.
Eur Phys J Spec Top. 2011 Nov 1; 200(1): 211.
doi:10.1140/epjst/e2011-01525-9
"""
constraints = []
for c in mol.constraints:  # Replace FixedPosition with 3 FixedCoordinates - it's better behaved
if isinstance(c, FixedPosition):
for i in xrange(3):
vec = np.zeros(3)
vec[i] = 1.0
constraints.append(FixedCoordinate(c.atom,  vec, value=c.value[i]))
else:
constraints.append(c)

# ---array shapes---
# values: (num_constr,)
# A: (num_constr, num_constr)
# multipliers: (num_constr, )
# delta: (3N,)
curr = mol.positions.copy()
prev = prev_positions.copy()
mol.positions = prev
mol.positions = curr

if use_masses:
dim_masses = mol.dim_masses
else:
dim_masses = np.ones((mol.num_atoms, 3)) * u.default.mass
flat_masses = dim_masses.defunits_value().flatten()

for i in xrange(max_cycles):
for c in mol.constraints:
if not c.satisfied():
break
else:
return  # e.g., we're done

# Get constraint derivatives
# Note: we remove units here because pint does not handle arrays with heterogeneous units
values = np.array([c.error().defunits_value() for c in constraints])

multipliers = np.linalg.solve(A, values)

# reapply units and adjust positions
delta = multipliers.dot(prevgrad).reshape(mol.num_atoms, 3) * (
u.default.mass * u.default.length)

mol.positions -= delta/dim_masses

else:
raise mdt.ConvergenceFailure('SHAKE did not converge after %d iterations'%
max_cycles)