Set internal coordinates of linked rigid members

When linking an RMF to an existing IMP hierarchy,
make sure that we set the internal coordinates
of any IMP rigid members to match that in the RMF,
when we read the first frame (we only need to do it
once since they are static in the RMF).

modules/rmf/include/internal/atom_links_xyzs.h

@@ -22,7 +22,8 @@ class IMPRMFEXPORT HierarchyLoadXYZs {
   typedef std::pair<RMF::NodeID, ParticleIndex> Pair;
   typedef Vector<Pair> NodeParticlePairs;
 
-  NodeParticlePairs global_, local_;
+  NodeParticlePairs global_, local_, local_first_;
+  bool first_load_after_link_;
 
   // backwards compat
   RMF::IntKey rb_index_key_;
@@ -34,7 +35,8 @@ class IMPRMFEXPORT HierarchyLoadXYZs {
                       const ParticleIndexes &rigid_bodies);
   void link_particle(RMF::NodeConstHandle n, Model *m,
                      ParticleIndex p,
-                     const ParticleIndexes &rigid_bodies);
+                     const ParticleIndexes &rigid_bodies,
+                     bool link_to_existing);
   void load(RMF::FileConstHandle fh, Model *m);
 };
 

modules/rmf/src/HierarchyLoadLink.cpp

@@ -253,7 +253,7 @@ void HierarchyLoadLink::add_link_recursive(Model *m,
   } // alternative factory
   data.load_static.link_particle(node, m, cur, rigid_bodies);
   data.load_rigid_bodies.link_particle(node, m, cur, rigid_bodies);
-  data.load_xyzs.link_particle(node, m, cur, rigid_bodies);
+  data.load_xyzs.link_particle(node, m, cur, rigid_bodies, true);
   data.load_gaussians.link_particle(node, m, cur, rigid_bodies);
 
   do_link_particle(m, root, cur, node);

modules/rmf/src/internal/atom_links_xyzs.cpp

@@ -13,7 +13,7 @@
 IMPRMF_BEGIN_INTERNAL_NAMESPACE
 
 HierarchyLoadXYZs::HierarchyLoadXYZs(RMF::FileConstHandle f)
-  : reference_frame_factory_(f), ip_factory_(f) {
+  : reference_frame_factory_(f), ip_factory_(f), first_load_after_link_(false) {
   // backwards compat
   RMF::Category cat = f.get_category("IMP");
   rb_index_key_ = f.get_key(cat, "rigid body", RMF::IntTraits());
@@ -59,12 +59,13 @@ void HierarchyLoadXYZs::setup_particle(
           .set_internal_coordinates(get_coordinates(n, ip_factory_));
     }
   }
-  link_particle(n, m, p, rigid_bodies);
+  link_particle(n, m, p, rigid_bodies, false);
 }
 
 void HierarchyLoadXYZs::link_particle(
     RMF::NodeConstHandle n, Model *m, ParticleIndex p,
-    const ParticleIndexes &rigid_bodies) {
+    const ParticleIndexes &rigid_bodies,
+    bool link_to_existing) {
   if (!ip_factory_.get_is(n)) return;
   if (rigid_bodies.empty()) {
     global_.push_back(std::make_pair(n.get_id(), p));
@@ -77,6 +78,12 @@ void HierarchyLoadXYZs::link_particle(
       // backwards compat: need to read coordinates of rigid members in order
       // to set the reference frame of the rigid body later
       global_.push_back(std::make_pair(n.get_id(), p));
+    } else if (link_to_existing && core::RigidMember::get_is_setup(m, p)) {
+      // if we are linking to an existing IMP rigid member, we need to update
+      // its internal coordinates to match those in the RMF, but only on the
+      // load of the first frame (since they are static in the RMF)
+      first_load_after_link_ = true;
+      local_first_.push_back(std::make_pair(n.get_id(), p));
     }
   }
 }
@@ -94,6 +101,15 @@ void HierarchyLoadXYZs::load(RMF::FileConstHandle fh, Model *m) {
     algebra::Vector3D rf(get_coordinates(fh.get_node(pp.first), ip_factory_));
     core::RigidBodyMember(m, pp.second).set_internal_coordinates(rf);
   }
+  if (first_load_after_link_) {
+    for(Pair pp : local_first_) {
+      IMP_LOG_VERBOSE("Loading local coordinates on first frame read for "
+                      << m->get_particle_name(pp.second) << std::endl);
+      algebra::Vector3D rf(get_coordinates(fh.get_node(pp.first), ip_factory_));
+      core::RigidMember(m, pp.second).set_internal_coordinates(rf);
+    }
+    first_load_after_link_ = false;
+  }
 }
 
 HierarchySaveXYZs::HierarchySaveXYZs(RMF::FileHandle f) : ip_factory_(f) {}

modules/rmf/test/test_rigid_bodies.py

@@ -175,6 +175,48 @@ def test_parent_non_rigid(self):
             new_hs.append(new_hs[0].get_child(0))
             self.compare_coords(hs, new_hs, m)
 
+    def test_linked_internal_coordinates(self):
+        """Check that internal coordinates of linked hierarchies are updated"""
+        for suffix in IMP.rmf.suffixes:
+            # Create rigid body containing one particle
+            m = IMP.Model()
+            h = IMP.atom.Hierarchy.setup_particle(IMP.Particle(m))
+            h.set_name("rb")
+            tr = IMP.algebra.Transformation3D(
+                         IMP.algebra.get_identity_rotation_3d(),
+                         IMP.algebra.Vector3D(1,2,3))
+            rbd = IMP.core.RigidBody.setup_particle(IMP.Particle(m),
+                                          IMP.algebra.ReferenceFrame3D(tr))
+            p = IMP.Particle(m)
+            d = IMP.core.XYZR.setup_particle(p)
+            d.set_coordinates(IMP.algebra.Vector3D(0., 2., 0.))
+            d.set_radius(.5)
+            IMP.atom.Mass.setup_particle(p, .1)
+            rbd.add_member(p)
+            orig = d.get_coordinates()
+            h.add_child(IMP.atom.Hierarchy.setup_particle(p))
+            # Write to RMF
+            fn = self.get_tmp_file_name("rigid" + suffix)
+            f = RMF.create_rmf_file(fn)
+            IMP.rmf.add_hierarchies(f, [h])
+            IMP.rmf.save_frame(f, str(0))
+            del f
+            # Change rigid body reference frame and particle internal coords
+            tr = IMP.algebra.Transformation3D(
+                         IMP.algebra.get_identity_rotation_3d(),
+                         IMP.algebra.Vector3D(3,2,1))
+            rbd.set_reference_frame(IMP.algebra.ReferenceFrame3D(tr))
+            self.assertTrue(IMP.core.RigidMember.get_is_setup(p))
+            IMP.core.RigidMember(p).set_internal_coordinates(
+                IMP.algebra.Vector3D(2., 0., 0.))
+            m.update()
+            # Reread RMF and make sure coordinates match
+            f = RMF.open_rmf_file_read_only(fn)
+            IMP.rmf.link_hierarchies(f, [h])
+            IMP.rmf.load_frame(f, RMF.FrameID(0))
+            new = d.get_coordinates()
+            self.assertLess(IMP.algebra.get_squared_distance(orig, new), 0.01)
+
     def compare_coords(self, hs, new_hs, m):
         coords = [IMP.core.XYZ(x).get_coordinates() for x in hs]
         new_coords = [IMP.core.XYZ(x).get_coordinates() for x in new_hs]