Source code for pysisyphus.dynamics.lincs

import numpy as np


[docs] def lincs_closure(geom, constraints, order=4): """ Drop conn_nums and keep conns_ in a list of list, instead of an array. We could do everything with for i, conn_constr in enumerate(conns) pass instead of the pseudo-code way as given in the paper. """ # Number of constraints K = len(constraints) # Inverse atom masses inv_masses = 1 / geom.masses # List of constraint sets constraint_sets = [set(c) for c in constraints] constraints = np.array(constraints, dtype=int) # Indices of the atoms that make up the bond constraints atom_1, atom_2 = constraints.T # Original bond lengths lengths = np.linalg.norm(geom.coords3d[atom_1] - geom.coords3d[atom_2], axis=1) # Determine number of constraints that are connected to every constraint conn_nums = np.zeros(K, dtype=int) conns_ = [list() for _ in range(K)] for i, con_1 in enumerate(constraint_sets): for j, con_2 in enumerate(constraint_sets): if i == j: continue if con_1 & con_2: conn_nums[i] += 1 conns_[i].append(j) # Maximum number of connected constraints c_max = conn_nums.max() S_diag = 1 / np.sqrt(inv_masses[atom_1] + inv_masses[atom_2]) # Array containing the indices of the connected constraints for every constraint conns = np.zeros((K, c_max), dtype=int) for i, c in enumerate(conns_): conns[i, :len(c)] = c coefs = np.zeros((K, c_max)) signs = { True: -1, False: 1, } for i, con_1 in enumerate(constraint_sets): for j in range(c_max): conns_ij = conns[i, j] sign = signs[(atom_1[i] == atom_1[conns_ij]) or (atom_2[i] == atom_2[conns_ij])] # Connected atom c = tuple(con_1 & constraint_sets[conns_ij]) assert len(c) == 1 c = int(c[0]) coefs[i, j] = sign * inv_masses[c] * S_diag[i] * S_diag[conns_ij] def solve(rhs, sol, new_coords3d, A, B): w = 1 for _ in range(order): for i, _ in enumerate(constraints): rhs[w, i] = 0. for j in range(conn_nums[i]): # rhs[w, i] += A[i, j] * rhs[2-w, conns_[i][j]] rhs[w, i] += A[i, j] * rhs[2-w, conns[i, j]] sol[i] += rhs[w, i] w = 2 - w for i, (a1, a2) in enumerate(constraints): new_coords3d[a1] -= inv_masses[a1] * S_diag[i] * sol[i] * B[i].sum() new_coords3d[a2] += inv_masses[a2] * S_diag[i] * sol[i] * B[i].sum() def lincs(prev_coords3d, new_coords3d): # Calculate constraint directions B = prev_coords3d[atom_1] - prev_coords3d[atom_2] B /= np.linalg.norm(B, axis=1)[:, None] print("B", B) A = np.zeros((K, c_max)) rhs = np.zeros((2, K)) sol = np.zeros(K) for i, (a1, a2) in enumerate(constraints): print(i, a1, a2) for j, k in enumerate(conns[i]): A[i, j] = coefs[i, j] * (B[i] * B[k]).sum() rhs[0, i] = S_diag[i] * ((B[i] * (new_coords3d[a1] - new_coords3d[a2])).sum() - lengths[i]) sol[i] = rhs[0, i] solve(rhs, sol, new_coords3d, A, B) # Correction for rotational lenghthening for i, (a1, a2) in enumerate(constraints): length_i = lengths[i] p = (2 * length_i**2 - (-(new_coords3d[a1] - new_coords3d[a2])**2).sum())**0.5 rhs[0, i] = S_diag[i] * (length_i - p) sol[i] = rhs[0, i] solve(rhs, sol, new_coords3d, A, B) # import pdb; pdb.set_trace() return new_coords3d return lincs