Enable atom reordering in WholeMoleculeTransform
authorBerk Hess <hess@kth.se>
Wed, 25 Aug 2021 15:04:48 +0000 (15:04 +0000)
committerMark Abraham <mark.j.abraham@gmail.com>
Wed, 25 Aug 2021 15:04:48 +0000 (15:04 +0000)
src/gromacs/domdec/mdsetup.cpp
src/gromacs/mdlib/forcerec.cpp
src/gromacs/mdlib/tests/CMakeLists.txt
src/gromacs/mdlib/tests/wholemoleculetransform.cpp [new file with mode: 0644]
src/gromacs/mdlib/wholemoleculetransform.cpp
src/gromacs/mdlib/wholemoleculetransform.h

index 7ff47990f7dfd68d355cf33ce4e4fdb8af0d78c9..f6799a6c23dc81ec6b83e1f546638783bd1653ed 100644 (file)
@@ -44,6 +44,7 @@
 #include "gromacs/mdlib/constr.h"
 #include "gromacs/mdlib/mdatoms.h"
 #include "gromacs/mdlib/vsite.h"
+#include "gromacs/mdlib/wholemoleculetransform.h"
 #include "gromacs/mdtypes/commrec.h"
 #include "gromacs/mdtypes/forcebuffers.h"
 #include "gromacs/mdtypes/forcerec.h"
@@ -113,6 +114,11 @@ void mdAlgorithmsSetupAtomData(const t_commrec*     cr,
         gmx_mtop_generate_local_top(top_global, top, inputrec.efep != FreeEnergyPerturbationType::No);
     }
 
+    if (fr->wholeMoleculeTransform && usingDomDec)
+    {
+        fr->wholeMoleculeTransform->updateAtomOrder(cr->dd->globalAtomIndices, *cr->dd->ga2la);
+    }
+
     if (vsite)
     {
         vsite->setVirtualSites(top->idef.il,
index f5e784134a17bdd0b88f0bf5c450468452e9e5e3..70a2b1e1f54b2590032c94d301726d00b667667b 100644 (file)
@@ -740,28 +740,32 @@ void init_forcerec(FILE*                            fplog,
     }
     else
     {
+        forcerec->bMolPBC = (!DOMAINDECOMP(commrec) || dd_bonded_molpbc(*commrec->dd, forcerec->pbcType));
+
+        // Check and set up PBC for Ewald surface corrections or orientation restraints
         const bool useEwaldSurfaceCorrection =
                 (EEL_PME_EWALD(inputrec.coulombtype) && inputrec.epsilon_surface != 0);
         const bool haveOrientationRestraints = (gmx_mtop_ftype_count(mtop, F_ORIRES) > 0);
-        if (!DOMAINDECOMP(commrec))
+        const bool moleculesAreAlwaysWhole =
+                (DOMAINDECOMP(commrec) && dd_moleculesAreAlwaysWhole(*commrec->dd));
+        // WholeMoleculeTransform is only supported with a single PP rank
+        if (!moleculesAreAlwaysWhole && !havePPDomainDecomposition(commrec)
+            && (useEwaldSurfaceCorrection || haveOrientationRestraints))
         {
-            forcerec->bMolPBC = true;
-
-            if (useEwaldSurfaceCorrection || haveOrientationRestraints)
-            {
-                forcerec->wholeMoleculeTransform =
-                        std::make_unique<gmx::WholeMoleculeTransform>(mtop, inputrec.pbcType);
-            }
+            GMX_RELEASE_ASSERT(
+                    !havePPDomainDecomposition(commrec),
+                    "WholeMoleculesTransform only works when all atoms are in the same domain");
+            forcerec->wholeMoleculeTransform = std::make_unique<gmx::WholeMoleculeTransform>(
+                    mtop, inputrec.pbcType, DOMAINDECOMP(commrec));
         }
         else
         {
-            forcerec->bMolPBC = dd_bonded_molpbc(*commrec->dd, forcerec->pbcType);
-
             /* With Ewald surface correction it is useful to support e.g. running water
              * in parallel with update groups.
              * With orientation restraints there is no sensible use case supported with DD.
              */
-            if ((useEwaldSurfaceCorrection && !dd_moleculesAreAlwaysWhole(*commrec->dd))
+            if ((useEwaldSurfaceCorrection
+                 && !(DOMAINDECOMP(commrec) && dd_moleculesAreAlwaysWhole(*commrec->dd)))
                 || haveOrientationRestraints)
             {
                 gmx_fatal(FARGS,
@@ -774,7 +778,7 @@ void init_forcerec(FILE*                            fplog,
 
         if (useEwaldSurfaceCorrection)
         {
-            GMX_RELEASE_ASSERT(!DOMAINDECOMP(commrec) || dd_moleculesAreAlwaysWhole(*commrec->dd),
+            GMX_RELEASE_ASSERT(moleculesAreAlwaysWhole || forcerec->wholeMoleculeTransform,
                                "Molecules can not be broken by PBC with epsilon_surface > 0");
         }
     }
index 5854759fd50be71daa015fa4551b091e54d5457d..6345f48f31698ba6ec61f40d8382611073470193 100644 (file)
@@ -58,4 +58,5 @@ gmx_add_unit_test(MdlibUnitTest mdlib-test HARDWARE_DETECTION
         constrtestrunners_gpu.cpp
         leapfrogtestrunners_gpu.cpp
         settletestrunners_gpu.cpp
+       wholemoleculetransform.cpp
         )
diff --git a/src/gromacs/mdlib/tests/wholemoleculetransform.cpp b/src/gromacs/mdlib/tests/wholemoleculetransform.cpp
new file mode 100644 (file)
index 0000000..23cf0db
--- /dev/null
@@ -0,0 +1,149 @@
+/*
+ * This file is part of the GROMACS molecular simulation package.
+ *
+ * Copyright (c) 2021, by the GROMACS development team, led by
+ * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
+ * and including many others, as listed in the AUTHORS file in the
+ * top-level source directory and at http://www.gromacs.org.
+ *
+ * GROMACS is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public License
+ * as published by the Free Software Foundation; either version 2.1
+ * of the License, or (at your option) any later version.
+ *
+ * GROMACS 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
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with GROMACS; if not, see
+ * http://www.gnu.org/licenses, or write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
+ *
+ * If you want to redistribute modifications to GROMACS, please
+ * consider that scientific software is very special. Version
+ * control is crucial - bugs must be traceable. We will be happy to
+ * consider code for inclusion in the official distribution, but
+ * derived work must not be called official GROMACS. Details are found
+ * in the README & COPYING files - if they are missing, get the
+ * official version at http://www.gromacs.org.
+ *
+ * To help us fund GROMACS development, we humbly ask that you cite
+ * the research papers on the package. Check out http://www.gromacs.org.
+ */
+/*! \internal \file
+ * \brief
+ * Tests for the WholeMoleculeTransform class.
+ *
+ * \author berk Hess <hess@kth.se>
+ * \ingroup module_mdlib
+ */
+#include "gmxpre.h"
+
+#include "gromacs/mdlib/wholemoleculetransform.h"
+
+#include <gtest/gtest.h>
+
+#include "gromacs/domdec/ga2la.h"
+#include "gromacs/pbcutil/pbc.h"
+#include "gromacs/topology/topology.h"
+
+#include "testutils/testasserts.h"
+
+namespace gmx
+{
+
+namespace
+{
+
+//! The number of atoms for the test molecule
+constexpr int c_numAtoms = 4;
+//! The number of bonds, we use 2 bonds so we have one unbound atom
+constexpr int c_numBonds = 2;
+
+//! Returns a topology with a three atom molecule that is linearly connected
+std::unique_ptr<gmx_mtop_t> triAtomMoleculeSystem()
+{
+    gmx_moltype_t moltype;
+
+    moltype.atoms.nr = c_numAtoms;
+
+    for (int b = 0; b < c_numBonds; b++)
+    {
+        const int                parameterType = -1; // value is not actually used
+        const std::array<int, 2> atomIndices   = { b, b + 1 };
+        moltype.ilist[F_CONNBONDS].push_back(parameterType, atomIndices);
+    }
+
+    std::unique_ptr<gmx_mtop_t> mtop = std::make_unique<gmx_mtop_t>();
+
+    mtop->moltype.push_back(moltype);
+
+    gmx_molblock_t molblock;
+    molblock.type = 0;
+    molblock.nmol = 1;
+    mtop->molblock.push_back(molblock);
+    mtop->natoms = c_numAtoms;
+
+    return mtop;
+}
+
+
+//! Coordinates, broken over PBC
+const std::array<RVec, c_numAtoms> c_coords = {
+    { { 2.5, 0.5, 1.5 }, { 0.5, 0.5, 1.5 }, { 1.5, 0.5, 1.5 }, { 2.7, 0.5, 1.5 } }
+};
+
+//! A box that works with \p c_coords
+const matrix box = { { 3, 0, 0 }, { 0, 2, 0 }, { 0, 0, 2 } };
+
+TEST(WholeMoleculeTransform, MakesMoleculesWhole)
+{
+    const auto mtop = triAtomMoleculeSystem();
+
+    WholeMoleculeTransform wmt(*mtop, PbcType::Xyz, false);
+
+    wmt.updateForAtomPbcJumps(c_coords, box);
+
+    auto wholeCoords = wmt.wholeMoleculeCoordinates(c_coords, box);
+
+    for (int b = 0; b < c_numBonds; b++)
+    {
+        EXPECT_FLOAT_EQ(std::abs(wholeCoords[b][XX] - wholeCoords[b + 1][XX]), 1);
+    }
+}
+
+TEST(WholeMoleculeTransform, HandlesReordering)
+{
+    const auto mtop = triAtomMoleculeSystem();
+
+    WholeMoleculeTransform wmt(*mtop, PbcType::Xyz, true);
+
+    const std::array<int, c_numAtoms> reordering = { 2, 0, 3, 1 };
+
+    gmx_ga2la_t       ga2la(c_numAtoms, c_numAtoms);
+    std::vector<RVec> coords;
+    for (int a = 0; a < c_numAtoms; a++)
+    {
+        ga2la.insert(reordering[a], { a, 0 });
+
+        coords.push_back(c_coords[reordering[a]]);
+    }
+
+    wmt.updateAtomOrder(reordering, ga2la);
+
+    wmt.updateForAtomPbcJumps(coords, box);
+
+    auto wholeCoords = wmt.wholeMoleculeCoordinates(coords, box);
+
+    for (int b = 0; b < c_numBonds; b++)
+    {
+        EXPECT_FLOAT_EQ(
+                std::abs(wholeCoords[ga2la.find(b)->la][XX] - wholeCoords[ga2la.find(b + 1)->la][XX]), 1);
+    }
+}
+
+} // namespace
+
+} // namespace gmx
index 14b45802e758ae9fefba7a6d3527f21efc51d7d0..bffde3f3ef502dddeb915accd87a227e297f14ac 100644 (file)
@@ -1,7 +1,7 @@
 /*
  * This file is part of the GROMACS molecular simulation package.
  *
- * Copyright (c) 2020, by the GROMACS development team, led by
+ * Copyright (c) 2020,2021, by the GROMACS development team, led by
  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
  * and including many others, as listed in the AUTHORS file in the
  * top-level source directory and at http://www.gromacs.org.
 
 #include "wholemoleculetransform.h"
 
+#include "gromacs/domdec/ga2la.h"
 #include "gromacs/topology/mtop_util.h"
 
 namespace gmx
 {
 
-WholeMoleculeTransform::WholeMoleculeTransform(const gmx_mtop_t& mtop, const PbcType pbcType) :
+WholeMoleculeTransform::WholeMoleculeTransform(const gmx_mtop_t& mtop,
+                                               const PbcType     pbcType,
+                                               const bool        useAtomReordering) :
     pbcType_(pbcType)
 {
     gmx_localtop_t localTop(mtop.ffparams);
@@ -52,6 +55,53 @@ WholeMoleculeTransform::WholeMoleculeTransform(const gmx_mtop_t& mtop, const Pbc
     graph_ = mk_graph(localTop.idef, mtop.natoms);
 
     wholeMoleculeCoordinates_.resize(mtop.natoms);
+
+    if (useAtomReordering)
+    {
+        // Store a copy of the global edges
+        globalEdgeAtomBegin_       = graph_.edgeAtomBegin;
+        graphGlobalAtomOrderEdges_ = graph_.edges;
+        // Reset the graph range to cover the whole system
+        graph_.edgeAtomBegin = 0;
+        graph_.edgeAtomEnd   = graph_.shiftAtomEnd;
+
+        // Resize the edge color list for potential addition of non-connected atoms
+        graph_.edgeColor.resize(graph_.edgeAtomEnd);
+    }
+}
+
+void WholeMoleculeTransform::updateAtomOrder(ArrayRef<const int> globalAtomIndices, const gmx_ga2la_t& ga2la)
+{
+    GMX_RELEASE_ASSERT(!graphGlobalAtomOrderEdges_.empty(),
+                       "We need the edges in the global atom order");
+
+    GMX_RELEASE_ASSERT(globalAtomIndices.ssize() == graph_.shiftAtomEnd,
+                       "The number of indices should match the number of atoms we shift");
+
+    const int globalEdgeAtomEnd = globalEdgeAtomBegin_ + graphGlobalAtomOrderEdges_.size();
+    graph_.edges.clear();
+    for (const int globalAtomIndex : globalAtomIndices)
+    {
+        if (globalAtomIndex >= globalEdgeAtomBegin_ && globalAtomIndex < globalEdgeAtomEnd)
+        {
+            // Get the list of global atoms linked to this atom and add their local indices
+            const ArrayRef<const int> globalList =
+                    graphGlobalAtomOrderEdges_[globalAtomIndex - globalEdgeAtomBegin_];
+            graph_.edges.pushBackListOfSize(globalList.size());
+            ArrayRef<int> lastList = graph_.edges.back();
+            std::transform(globalList.begin(), globalList.end(), lastList.begin(), [&ga2la](int i) -> int {
+                return ga2la.find(i)->la;
+            });
+        }
+        else
+        {
+            // The atom has no edges, push back an empty list
+            graph_.edges.pushBackListOfSize(0);
+        }
+    }
+
+    GMX_RELEASE_ASSERT(int(graph_.edges.size()) == graph_.shiftAtomEnd,
+                       "We should have as many lists of edges as the system (shift) size");
 }
 
 void WholeMoleculeTransform::updateForAtomPbcJumps(ArrayRef<const RVec> x, const matrix box)
index b79644c64e60d1df3166ed14392b1d493a45bb57..8d42b7e0dad7fae058efc7c6624a15a38c4e8425 100644 (file)
@@ -1,7 +1,7 @@
 /*
  * This file is part of the GROMACS molecular simulation package.
  *
- * Copyright (c) 2020, by the GROMACS development team, led by
+ * Copyright (c) 2020,2021, by the GROMACS development team, led by
  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
  * and including many others, as listed in the AUTHORS file in the
  * top-level source directory and at http://www.gromacs.org.
@@ -49,6 +49,7 @@
 #include "gromacs/pbcutil/mshift.h"
 #include "gromacs/utility/arrayref.h"
 
+class gmx_ga2la_t;
 struct gmx_mtop_t;
 enum class PbcType : int;
 
@@ -66,8 +67,23 @@ namespace gmx
 class WholeMoleculeTransform
 {
 public:
-    /*! \brief Constructor */
-    WholeMoleculeTransform(const gmx_mtop_t& mtop, PbcType pbcType);
+    /*! \brief Constructor
+     *
+     * \param[in] mtop               The global topology use for getting the connections between atoms
+     * \param[in] pbcType            The type of PBC
+     * \param[in] useAtomReordering  Whether we will use atom reordering
+     */
+    WholeMoleculeTransform(const gmx_mtop_t& mtop, PbcType pbcType, bool useAtomReordering);
+
+    /*! \brief Changes the atom order to the one provided
+     *
+     * This method is called after domain repartitioning.
+     * The object should have been constructed with \p useAtomReordering set to \p true.
+     *
+     * \param[in] globalAtomIndices  The global atom indices for the local atoms, size should be the system size
+     * \param[in] ga2la              Global to local atom index lookup (the inverse of \p globalAtomIndices)
+     */
+    void updateAtomOrder(ArrayRef<const int> globalAtomIndices, const gmx_ga2la_t& ga2la);
 
     /*! \brief Updates the graph when atoms have been shifted by periodic vectors */
     void updateForAtomPbcJumps(ArrayRef<const RVec> x, const matrix box);
@@ -88,6 +104,10 @@ private:
     PbcType pbcType_;
     //! The graph
     t_graph graph_;
+    //! The atom index at which graphGlobalAtomOrderEdges_ starts
+    int globalEdgeAtomBegin_;
+    //! The edges for the global atom order
+    ListOfLists<int> graphGlobalAtomOrderEdges_;
     //! Buffer for storing coordinates for whole molecules
     std::vector<RVec> wholeMoleculeCoordinates_;
 };