View Javadoc
1   //******************************************************************************
2   //
3   // Title:       Force Field X.
4   // Description: Force Field X - Software for Molecular Biophysics.
5   // Copyright:   Copyright (c) Michael J. Schnieders 2001-2025.
6   //
7   // This file is part of Force Field X.
8   //
9   // Force Field X is free software; you can redistribute it and/or modify it
10  // under the terms of the GNU General Public License version 3 as published by
11  // the Free Software Foundation.
12  //
13  // Force Field X is distributed in the hope that it will be useful, but WITHOUT
14  // ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
15  // FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
16  // details.
17  //
18  // You should have received a copy of the GNU General Public License along with
19  // Force Field X; if not, write to the Free Software Foundation, Inc., 59 Temple
20  // Place, Suite 330, Boston, MA 02111-1307 USA
21  //
22  // Linking this library statically or dynamically with other modules is making a
23  // combined work based on this library. Thus, the terms and conditions of the
24  // GNU General Public License cover the whole combination.
25  //
26  // As a special exception, the copyright holders of this library give you
27  // permission to link this library with independent modules to produce an
28  // executable, regardless of the license terms of these independent modules, and
29  // to copy and distribute the resulting executable under terms of your choice,
30  // provided that you also meet, for each linked independent module, the terms
31  // and conditions of the license of that module. An independent module is a
32  // module which is not derived from or based on this library. If you modify this
33  // library, you may extend this exception to your version of the library, but
34  // you are not obligated to do so. If you do not wish to do so, delete this
35  // exception statement from your version.
36  //
37  //******************************************************************************
38  package ffx.potential.constraint;
39  
40  import ffx.numerics.Constraint;
41  import ffx.potential.bonded.Angle;
42  import ffx.potential.bonded.Atom;
43  import ffx.potential.bonded.Bond;
44  import org.apache.commons.math3.linear.OpenMapRealMatrix;
45  import org.apache.commons.math3.linear.QRDecomposition;
46  import org.apache.commons.math3.linear.RealMatrix;
47  import org.apache.commons.math3.linear.RealVectorPreservingVisitor;
48  import org.apache.commons.math3.util.FastMath;
49  
50  import java.util.ArrayList;
51  import java.util.Arrays;
52  import java.util.HashSet;
53  import java.util.List;
54  import java.util.Set;
55  import java.util.logging.Logger;
56  import java.util.stream.IntStream;
57  
58  public class CcmaConstraint implements Constraint {
59  
60    public static final double DEFAULT_CCMA_NONZERO_CUTOFF = 0.01;
61    private static final Logger logger = Logger.getLogger(CcmaConstraint.class.getName());
62    private static final int DEFAULT_MAX_ITERS = 150;
63    // Might be slightly more elegant to have "component constraint" objects.
64    // This is how OpenMM does it, though, and arrays are likely more performative.
65    private final int[] atoms1;
66    private final int[] atoms2;
67    private final double[] lengths;
68    private final int nConstraints;
69    private final int[] uniqueIndices;
70    private final int maxIters = DEFAULT_MAX_ITERS;
71    private final double elementCutoff = DEFAULT_CCMA_NONZERO_CUTOFF;
72    private final double[] reducedMasses;
73    private final RealMatrix kInvSparse;
74  
75    /**
76     * Constructs a set of bond length Constraints to be satisfied using the Constaint Constraint
77     * Matrix Approximation, a parallelizable stable numeric method.
78     *
79     * <p>The nonzeroCutoff field specifies how large an element of K-1 needs to be to be kept;
80     * smaller elements (indicating weak constraint-constraint couplings) are set to zero to make K-1
81     * sparse and thus keep CCMA tractable.
82     *
83     * <p>Due to leaking-this issues, a thin Factory method is provided to build this.
84     *
85     * <p>The constrainedAngles will constrain both component bonds and the triangle-closing distance.
86     * The constrainedBonds list should not include any bonds included in the constrained angles.
87     *
88     * @param constrainedBonds Bonds to be constrained, not included in constrainedAngles
89     * @param constrainedAngles Angles to be constrained indirectly (by constructing a rigid
90     *     triangle).
91     * @param allAtoms All Atoms of the system, including unconstrained Atoms.
92     * @param masses All masses of the system, including unconstrained atom masses.
93     * @param nonzeroCutoff CCMA parameter defining how sparse/dense K-1 should be.
94     */
95    private CcmaConstraint(
96        List<Bond> constrainedBonds,
97        List<Angle> constrainedAngles,
98        final Atom[] allAtoms,
99        final double[] masses,
100       double nonzeroCutoff) {
101 
102     long time = -System.nanoTime();
103     int nBonds = constrainedBonds.size();
104     int nAngles = constrainedAngles.size();
105     assert constrainedAngles.stream()
106         .flatMap((Angle a) -> a.getBondList().stream())
107         .noneMatch(constrainedBonds::contains);
108     nConstraints = nBonds + (3 * nAngles);
109     atoms1 = new int[nConstraints];
110     atoms2 = new int[nConstraints];
111     lengths = new double[nConstraints];
112 
113     for (int i = 0; i < nBonds; i++) {
114       Bond bi = constrainedBonds.get(i);
115       atoms1[i] = bi.getAtom(0).getXyzIndex() - 1;
116       atoms2[i] = bi.getAtom(1).getXyzIndex() - 1;
117       lengths[i] = bi.bondType.distance;
118     }
119 
120     /*logger.info(" Atoms 1: " + Arrays.toString(atoms1));
121     logger.info(" Atoms 2: " + Arrays.toString(atoms2));
122     logger.info(" Lengths: " + Arrays.toString(lengths));*/
123 
124     for (int i = 0; i < nAngles; i++) {
125       int iAng = nBonds + (3 * i);
126       Angle ai = constrainedAngles.get(i);
127       Atom center = ai.getCentralAtom();
128       Bond b1 = ai.getBond(0);
129       Bond b2 = ai.getBond(1);
130       Atom at0 = b1.get1_2(center);
131       Atom at2 = b2.get1_2(center);
132 
133       double angVal = ai.angleType.angle[ai.nh];
134       double dist1 = b1.bondType.distance;
135       double dist2 = b2.bondType.distance;
136       double dist3 = SettleConstraint.lawOfCosines(dist1, dist2, angVal);
137 
138       int index0 = at0.getXyzIndex() - 1;
139       int index1 = center.getXyzIndex() - 1;
140       int index2 = at2.getXyzIndex() - 1;
141 
142       atoms1[iAng] = index0;
143       atoms2[iAng] = index1;
144       lengths[iAng] = dist1;
145 
146       ++iAng;
147       atoms1[iAng] = index1;
148       atoms2[iAng] = index2;
149       lengths[iAng] = dist2;
150 
151       ++iAng;
152       atoms1[iAng] = index0;
153       atoms2[iAng] = index2;
154       lengths[iAng] = dist3;
155     }
156 
157     uniqueIndices =
158         IntStream.concat(Arrays.stream(atoms1), Arrays.stream(atoms2))
159             .sorted()
160             .distinct()
161             .toArray();
162 
163     OpenMapRealMatrix kSparse = new OpenMapRealMatrix(nConstraints, nConstraints);
164 
165     int nAtoms = allAtoms.length;
166 
167     // Points from an Atom index to the Set of all Constraint indices it's involved in.
168     List<Set<Integer>> atomsToConstraints = new ArrayList<>(nAtoms);
169     for (int i = 0; i < nAtoms; i++) {
170       // Can assume 4 constraints max, since CHONPS is rarely bonded to 5+ atoms.
171       atomsToConstraints.add(new HashSet<>(4));
172     }
173 
174     for (int i = 0; i < nConstraints; i++) {
175       atomsToConstraints.get(atoms1[i]).add(i);
176       atomsToConstraints.get(atoms2[i]).add(i);
177     }
178 
179     logger.info(
180         String.format(" Initial CCMA setup: %10.6g sec", 1.0E-9 * (time + System.nanoTime())));
181     long subTime = -System.nanoTime();
182 
183     for (int i = 0; i < nConstraints; i++) {
184       int atomi0 = atoms1[i];
185       int atomi1 = atoms2[i];
186       // DO YOU HAVE ANY IDEA HOW LONG IT TOOK FOR ME TO REALIZE MASSES WERE XYZ-INDEXED?
187       double invMassI0 = 1.0 / masses[atomi0 * 3];
188       double invMassI1 = 1.0 / masses[atomi1 * 3];
189       double sumInv = invMassI0 + invMassI1;
190 
191       // Indices of all Constraints that involve either atomi0 or atomi1.
192       Set<Integer> coupledConstraints = new HashSet<>(atomsToConstraints.get(atomi0));
193       coupledConstraints.addAll(atomsToConstraints.get(atomi1));
194 
195       // Iterate over all coupled Constraints, sharing at least one common Atom with constraint i.
196       for (int j : coupledConstraints) {
197         // Diagonal element, coupling is obviously 1.0.
198         if (i == j) {
199           kSparse.setEntry(i, j, 1.0);
200           continue;
201         }
202 
203         double scale;
204         int atomj0 = atoms1[j];
205         int atomj1 = atoms2[j];
206         int atoma; // Atom unique to constraint i.
207         int atomb; // Atom shared between both constraints.
208         int atomc; // Atom unique to constraint j.
209 
210         if (atomi0 == atomj0) {
211           assert atomi1 != atomj1;
212           atoma = atomi1;
213           atomb = atomi0;
214           atomc = atomj1;
215           scale = invMassI0 / sumInv;
216         } else if (atomi0 == atomj1) {
217           assert atomi1 != atomj0;
218           atoma = atomi1;
219           atomb = atomi0;
220           atomc = atomj0;
221           scale = invMassI0 / sumInv;
222         } else if (atomi1 == atomj0) {
223           // Yes IntelliJ, it should always be true. That's why I'm asserting it.
224           assert atomi0 != atomj1;
225           atoma = atomi0;
226           atomb = atomi1;
227           atomc = atomj1;
228           scale = invMassI1 / sumInv;
229         } else if (atomi1 == atomj1) {
230           assert atomi0 != atomj0;
231           atoma = atomi1;
232           atomb = atomi0;
233           atomc = atomj1;
234           scale = invMassI1 / sumInv;
235         } else {
236           throw new IllegalArgumentException(
237               " Despite by necessity sharing an atom, these constraints don't share an atom.");
238         }
239 
240         // We now have a pair of constraints a-b b-c. Find the a-b-c angle.
241 
242         // Search for a constraint a-c that closes the triangle.
243         boolean foundAngle = false;
244         for (int constraintK : atomsToConstraints.get(atoma)) {
245           if (atoms1[constraintK] == atomc || atoms2[constraintK] == atomc) {
246             double dab = lengths[i];
247             double dbc = lengths[j];
248             double dac = lengths[constraintK];
249 
250             // Apply the Law of Cosines to find the angle defined by the a-b b-c a-c lengths.
251             double angle = (dab * dab) + (dbc * dbc) - (dac * dac);
252             angle /= (2 * dab * dbc);
253             // The angle is formally the cosine of its current value, but all we need is that
254             // cosine.
255             double coupling = scale * angle;
256             kSparse.setEntry(i, j, coupling);
257             foundAngle = true;
258             break;
259           }
260         }
261 
262         // Otherwise, look for a force field term. In this case, the approximate matrix K
263         // will not quite match the true Jacobian matrix J.
264         if (!foundAngle) {
265           Atom atA = allAtoms[atoma];
266           Atom atB = allAtoms[atomb];
267           Atom atC = allAtoms[atomc];
268           Angle angleB = atA.getAngle(atB, atC);
269           double angVal = angleB.angleType.angle[angleB.nh];
270           double coupling = scale * FastMath.cos(FastMath.toRadians(angVal));
271           kSparse.setEntry(i, j, coupling);
272           foundAngle = true;
273         }
274 
275         if (!foundAngle) {
276           logger.severe(
277               String.format(" Could not find the angle between coupled constraints %d, %d", i, j));
278         }
279       }
280     }
281 
282     subTime += System.nanoTime();
283     logger.info(
284         String.format(" Time to construct K as a sparse matrix: %10.6g sec", 1.0E-9 * subTime));
285 
286     // K is constructed. Invert it.
287     // TODO: Find a fast way to invert sparse matrices, since this is at least O(n^2).
288     subTime = -System.nanoTime();
289     QRDecomposition qrd = new QRDecomposition(kSparse);
290     RealMatrix kInvDense = qrd.getSolver().getInverse();
291     subTime += System.nanoTime();
292     logger.info(String.format(" Time to invert K: %10.6g sec", 1.0E-9 * subTime));
293     subTime = -System.nanoTime();
294 
295     kInvSparse = new OpenMapRealMatrix(nConstraints, nConstraints);
296     IntStream.range(0, nConstraints)
297         .
298         // parallel(). // I have no idea if OpenMapRealMatrix is thread-safe.
299             forEach(
300             (int i) -> {
301               double[] rowI = kInvDense.getRow(i);
302               for (int j = 0; j < nConstraints; j++) {
303                 double val = rowI[j];
304                 if (Math.abs(val) > elementCutoff) {
305                   kInvSparse.setEntry(i, j, val);
306                 }
307               }
308             });
309 
310     // TODO: Delete this block.
311     double[][] dataDense = kInvDense.getData();
312     double[][] dataSparse = kInvSparse.getData();
313     logger.fine(" Dense array:");
314     for (int i = 0; i < nConstraints; i++) {
315       logger.fine(Arrays.toString(dataDense[i]));
316     }
317     logger.fine(" Sparse array:");
318     for (int i = 0; i < nConstraints; i++) {
319       logger.fine(Arrays.toString(dataSparse[i]));
320     }
321 
322     // TODO: Actually do this.
323     reducedMasses = new double[nConstraints];
324     for (int i = 0; i < nConstraints; i++) {
325       int atI = atoms1[i];
326       atI *= 3; // The mass array is XYZ-indexed, not atom-indexed.
327       int atJ = atoms2[i];
328       atJ *= 3;
329       double invMassI = 1.0 / masses[atI];
330       double invMassJ = 1.0 / masses[atJ];
331       reducedMasses[i] = 0.5 / (invMassI + invMassJ);
332     }
333   }
334 
335   /**
336    * Constructs a set of bond length Constraints to be satisfied using the Constaint Constraint
337    * Matrix Approximation, a parallelizable stable numeric method.
338    *
339    * <p>The nonzeroCutoff field specifies how large an element of K-1 needs to be to be kept;
340    * smaller elements (indicating weak constraint-constraint couplings) are set to zero to make K-1
341    * sparse and thus keep CCMA tractable.
342    *
343    * <p>The constrainedAngles will constrain both component bonds and the triangle-closing distance.
344    * The constrainedBonds list should not include any bonds included in the constrained angles.
345    *
346    * @param constrainedBonds Bonds to be constrained, not included in constrainedAngles
347    * @param constrainedAngles Angles to be constrained indirectly (by constructing a rigid
348    *     triangle).
349    * @param allAtoms All Atoms of the system, including unconstrained Atoms.
350    * @param masses All masses of the system, including unconstrained atom masses.
351    * @param nonzeroCutoff CCMA parameter defining how sparse/dense K-1 should be.
352    * @return Returns a new CcmaConstraint instance.
353    */
354   public static CcmaConstraint ccmaFactory(
355       List<Bond> constrainedBonds,
356       List<Angle> constrainedAngles,
357       final Atom[] allAtoms,
358       final double[] masses,
359       double nonzeroCutoff) {
360     CcmaConstraint newC =
361         new CcmaConstraint(constrainedBonds, constrainedAngles, allAtoms, masses, nonzeroCutoff);
362     constrainedBonds.forEach((Bond b) -> b.setConstraint(newC));
363     constrainedAngles.forEach((Angle a) -> a.setConstraint(newC));
364     return newC;
365   }
366 
367   /** {@inheritDoc} */
368   @Override
369   public void applyConstraintToStep(double[] xPrior, double[] xNew, double[] masses, double tol) {
370     applyConstraints(xPrior, xNew, masses, false, tol);
371   }
372 
373   /** {@inheritDoc} */
374   @Override
375   public void applyConstraintToVelocities(double[] x, double[] v, double[] masses, double tol) {
376     applyConstraints(x, v, masses, true, tol);
377   }
378 
379   /** {@inheritDoc} */
380   @Override
381   public int[] constrainedAtomIndices() {
382     return Arrays.copyOf(uniqueIndices, uniqueIndices.length);
383   }
384 
385   /** {@inheritDoc} */
386   @Override
387   public boolean constraintSatisfied(double[] x, double tol) {
388     return false;
389   }
390 
391   /** {@inheritDoc} */
392   @Override
393   public boolean constraintSatisfied(double[] x, double[] v, double xTol, double vTol) {
394     return false;
395   }
396 
397   /** {@inheritDoc} */
398   @Override
399   public int getNumDegreesFrozen() {
400     return nConstraints;
401   }
402 
403   /**
404    * Much of the math for applying to coordinates/velocities is the same. As such, OpenMM just uses a
405    * single driver method with a flag to indicate velocities or positions.
406    *
407    * @param xPrior Either pre-step coordinates (positions) or current coordinates (velocities)
408    * @param output Output vector, either constrained coordinates or velocities
409    * @param masses Masses of all particles
410    * @param constrainV Whether to apply constraints to velocities, or to coordinates.
411    * @param tol Numerical tolerance.
412    */
413   private void applyConstraints(
414       double[] xPrior, double[] output, double[] masses, boolean constrainV, double tol) {
415     if (xPrior == output) {
416       throw new IllegalArgumentException(" xPrior and output must be different arrays!");
417     }
418     long time = -System.nanoTime();
419 
420     /*// Distance vector for all constraints. TODO: Flatten to 1D array.
421     double[][] r_ij = new double[nConstraints][3];
422     // Distance magnitude for all constraints.
423     double[] d_ij = new double[nConstraints];
424     // Squared distance magnitude for all constraints.
425     double[] d_ij2 = new double[nConstraints];
426 
427     for (int i = 0; i < nConstraints; i++) {
428         int atom1 = 3 * atoms1[i];
429         int atom2 = 3 * atoms2[i];
430         for (int j = 0; j < 3; j++) {
431             r_ij[i][j] = xPrior[atom1 + j] - xPrior[atom2 + j];
432         }
433         d_ij2[i] = VectorMath.dot(r_ij[i], r_ij[i]);
434     }
435 
436     double lowerTol = 1 - 2*tol + tol*tol;
437     // TODO: Determine if adding tol+tol to upperTol is a typo in OpenMM. Also why they do it in the first place.
438     double upperTol = 1 + 2*tol + tol*tol;
439 
440     int nConverged = 0;
441     double[] constraintDelta = new double[nConstraints];
442     double[] tempDelta = new double[nConstraints];
443 
444     // Main CCMA loop.
445     for (int constraintIter = 0; constraintIter <= maxIters; constraintIter++) {
446         if (constraintIter >= maxIters) {
447             throw new IllegalArgumentException(String.format(" CCMA constraint failed to converge in %d iterations!", maxIters));
448         }
449         nConverged = 0;
450         double rmsd = 0;
451         for (int i = 0; i < nConstraints; i++) {
452             int atom1 = atoms1[i] * 3;
453             int atom2 = atoms2[i] * 3;
454 
455             // Separation vector I-J at this iteration.
456             double[] rp_ij = new double[3];
457             for (int j = 0; j < 3; j++) {
458                 rp_ij[j] = output[atom1 + j] - output[atom2 + j];
459             }
460             if (constrainV) {
461                 double rrpr = VectorMath.dot(rp_ij, r_ij[i]);
462                 constraintDelta[i] = -2 * reducedMasses[i] * rrpr / d_ij2[i];
463                 if (Math.abs(constraintDelta[i]) <= tol) {
464                     ++nConverged;
465                 }
466             } else {
467                 double rp2 = VectorMath.dot(rp_ij, rp_ij);
468                 double dist2 = lengths[i] * lengths[i];
469                 double diff = dist2 - rp2;
470                 double rrpr = VectorMath.dot(rp_ij, r_ij[i]);
471                 constraintDelta[i] = reducedMasses[i] * diff / rrpr;
472                 if (rp2 >= lowerTol * dist2 && rp2 <= upperTol * dist2) {
473                     ++nConverged;
474                 }
475             }
476         }
477 
478         // Test if the last iteration satisfied all constraints.
479         if (nConverged == nConstraints) {
480             break;
481         }
482 
483         if (kInvSparse != null) {
484             for (int i = 0; i < nConstraints; i++) {
485                 RealVector row = kInvSparse.getRowVector(i);
486                 RealVectorPreservingVisitor neo = new MatrixWalker(constraintDelta);
487                 double sum = row.walkInOptimizedOrder(neo);
488                 tempDelta[i] = sum;
489             }
490             System.arraycopy(tempDelta, 0, constraintDelta, 0, nConstraints);
491         }
492         for (int i = 0; i < nConstraints; i++) {
493             int atom1 = atoms1[i];
494             int atom2 = atoms2[i];
495             int at13 = 3 * atom1;
496             int at23 = 3 * atom2;
497 
498             double[] dr = new double[3];
499             VectorMath.scalar(r_ij[i], constraintDelta[i], dr);
500             for (int j = 0; j < 3; j++) {
501                 output[at13 + j] = dr[j] / masses[at13];
502                 output[at23 + j] = dr[j] / masses[at23];
503             }
504         }
505         logger.info(String.format(" %d converged at iteration %d", nConverged, constraintIter));
506     }
507     time += System.nanoTime();
508     logger.info(String.format(" Application of CCMA constraint: %10.4g sec", (time * 1.0E-9)));*/
509   }
510 
511   private static class MatrixWalker implements RealVectorPreservingVisitor {
512 
513     private final double[] constraintDelta; // DO NOT MODIFY.
514     private double sum = 0.0;
515 
516     MatrixWalker(final double[] cDelta) {
517       constraintDelta = cDelta;
518     }
519 
520     /** {@inheritDoc} */
521     @Override
522     public double end() {
523       return sum;
524     }
525 
526     /** {@inheritDoc} */
527     @Override
528     public void start(int dimension, int start, int end) {
529     }
530 
531     /** {@inheritDoc} */
532     @Override
533     public void visit(int index, double value) {
534       sum += (value * constraintDelta[index]);
535     }
536   }
537 }