Merge pull request #1696 from glotzerlab/fix-shapespace-retreat

Fix shapespace retreat

hoomd/hpmc/ShapeMoves.h

@@ -40,7 +40,8 @@ template<typename Shape> class ShapeMoveBase
     virtual void update_shape(uint64_t,
                               const unsigned int&,
                               typename Shape::param_type&,
-                              hoomd::RandomGenerator&)
+                              hoomd::RandomGenerator&,
+                              bool managed)
         {
         }
 
@@ -141,7 +142,8 @@ template<typename Shape> class PythonShapeMove : public ShapeMoveBase<Shape>
     void update_shape(uint64_t timestep,
                       const unsigned int& type_id,
                       typename Shape::param_type& shape,
-                      hoomd::RandomGenerator& rng)
+                      hoomd::RandomGenerator& rng,
+                      bool managed)
         {
         for (unsigned int i = 0; i < m_params[type_id].size(); i++)
             {
@@ -165,7 +167,7 @@ template<typename Shape> class PythonShapeMove : public ShapeMoveBase<Shape>
             }
         pybind11::object d = m_python_callback(type_id, m_params[type_id]);
         pybind11::dict shape_dict = pybind11::cast<pybind11::dict>(d);
-        shape = typename Shape::param_type(shape_dict);
+        shape = typename Shape::param_type(shape_dict, managed);
         }
 
     void retreat(uint64_t timestep, unsigned int type)
@@ -212,8 +214,7 @@ template<typename Shape> class PythonShapeMove : public ShapeMoveBase<Shape>
     std::vector<std::vector<Scalar>>
         m_params_backup;                       // tunable shape parameters to perform trial moves on
     std::vector<std::vector<Scalar>> m_params; // tunable shape parameters to perform trial moves on
-    // callback that takes m_params as an argiment and returns a Python dictionary with the shape
-    // params.
+    // callback that takes m_params as an argument and returns a Python dictionary of shape params.
     pybind11::object m_python_callback;
     };
 
@@ -279,7 +280,8 @@ class ConvexPolyhedronVertexShapeMove : public ShapeMoveBase<ShapeConvexPolyhedr
     void update_shape(uint64_t timestep,
                       const unsigned int& type_id,
                       param_type& shape,
-                      hoomd::RandomGenerator& rng)
+                      hoomd::RandomGenerator& rng,
+                      bool managed)
         {
         // perturb the shape.
         for (unsigned int i = 0; i < shape.N; i++)
@@ -375,7 +377,8 @@ class ElasticShapeMove<ShapeConvexPolyhedron> : public ElasticShapeMoveBase<Shap
     void update_shape(uint64_t timestep,
                       const unsigned int& type_id,
                       param_type& param,
-                      hoomd::RandomGenerator& rng)
+                      hoomd::RandomGenerator& rng,
+                      bool managed)
         {
         Matrix3S F_curr;
         // perform a scaling move
@@ -586,7 +589,8 @@ template<> class ElasticShapeMove<ShapeEllipsoid> : public ElasticShapeMoveBase<
     void update_shape(uint64_t timestep,
                       const unsigned int& type_id,
                       param_type& param,
-                      hoomd::RandomGenerator& rng)
+                      hoomd::RandomGenerator& rng,
+                      bool managed)
         {
         Scalar lnx = log(param.x / param.y);
         Scalar stepsize = this->m_step_size[type_id];

hoomd/hpmc/UpdaterShape.h

@@ -288,7 +288,11 @@ template<class Shape> void UpdaterShape<Shape>::update(uint64_t timestep)
                                          hoomd::Counter(typ_i, 0, i_sweep));
 
             // perform an in-place shape update on shape_param_new
-            m_move_function->update_shape(timestep, typ_i, shape_param_new, rng_i);
+            m_move_function->update_shape(timestep,
+                                          typ_i,
+                                          shape_param_new,
+                                          rng_i,
+                                          m_exec_conf->isCUDAEnabled());
 
             // update det(I)
             detail::MassProperties<Shape> mp(shape_param_new);
@@ -322,6 +326,7 @@ template<class Shape> void UpdaterShape<Shape>::update(uint64_t timestep)
                 m_exec_conf->msg->notice(5)
                     << "UpdaterShape move rejected -- overlaps found" << std::endl;
                 // revert shape parameter changes
+                m_move_function->retreat(timestep, typ_i);
                 h_det.data[typ_i] = h_det_old.data[typ_i];
                 m_mc->setParam(typ_i, shape_param_old);
                 }

hoomd/hpmc/pytest/test_shape_updater.py

@@ -173,8 +173,8 @@ def test_vertex_shape_move(simulation_factory, two_particle_snapshot_factory):
     assert np.isclose(updater.particle_volumes[0], 1)
 
 
-def test_python_callback_shape_move(simulation_factory,
-                                    two_particle_snapshot_factory):
+def test_python_callback_shape_move_ellipsoid(simulation_factory,
+                                              lattice_snapshot_factory):
     """Test ShapeSpace with a toy class that randomly squashes spheres \
            into oblate ellipsoids with constant volume."""
 
@@ -206,7 +206,7 @@ def __call__(self, type_id, param_list):
     mc.shape["A"] = ellipsoid
 
     # create simulation & attach objects
-    sim = simulation_factory(two_particle_snapshot_factory(d=10))
+    sim = simulation_factory(lattice_snapshot_factory(a=2.75, n=(3, 3, 5)))
     sim.operations.integrator = mc
     sim.operations += updater
 
@@ -222,7 +222,7 @@ def __call__(self, type_id, param_list):
 
     # run with 0 probability of performing a move:
     #  - shape and params should remain unchanged
-    #  - all moves accepted
+    #  - no shape moves proposed
     move.param_move_probability = 0
     sim.run(10)
     assert np.allclose(mc.shape["A"]["a"], ellipsoid["a"])
@@ -234,15 +234,109 @@ def __call__(self, type_id, param_list):
     #  - shape and params should change
     #  - volume should remain unchanged
     move.param_move_probability = 1
-    sim.run(10)
-    assert np.sum(updater.shape_moves) == 20
+    sim.run(50)
+    assert np.sum(updater.shape_moves) == 100
+
+    # Check that the shape parameters have changed
     assert not np.allclose(mc.shape["A"]["a"], ellipsoid["a"])
     assert not np.allclose(mc.shape["A"]["b"], ellipsoid["b"])
     assert not np.allclose(mc.shape["A"]["c"], ellipsoid["c"])
     assert not np.allclose(move.params["A"], [1])
+
+    # Check that the shape parameters map back to the correct geometry
+    assert np.allclose(move.params["A"], [
+        mc.shape["A"]["a"] / mc.shape["A"]["b"],
+        mc.shape["A"]["a"] / mc.shape["A"]["c"]
+    ])
+
+    # Check that the callback is conserving volume properly
     assert np.allclose(updater.particle_volumes, 4 * np.pi / 3)
 
 
+def test_python_callback_shape_move_pyramid(simulation_factory,
+                                            two_particle_snapshot_factory):
+    """Test ShapeSpace with a toy class that randomly stretches square \
+        pyramids."""
+
+    def square_pyramid_factory(h):
+        """Generate a square pyramid with unit volume."""
+        theta = np.arange(0, 2 * np.pi, np.pi / 2)
+        base_vertices = np.array(
+            [np.cos(theta), np.sin(theta),
+             np.zeros_like(theta)]).T * np.sqrt(3 / 2)
+        vertices = np.vstack([base_vertices, [0, 0, h]])
+        return vertices / np.cbrt(h), base_vertices
+
+    class ScalePyramid:
+
+        def __init__(self, h=1.1):
+            _, self.base_vertices = square_pyramid_factory(h=1.1)
+            self.default_dict = dict(sweep_radius=0, ignore_statistics=True)
+
+        def __call__(self, type_id, param_list):
+            h = param_list[0] + 0.1  # Prevent a 0-height pyramid
+            new_vertices = np.vstack([self.base_vertices, [0, 0, h]])
+            new_vertices /= np.cbrt(h)  # Rescale to unit volume
+            ret = dict(vertices=new_vertices, **self.default_dict)
+            return ret
+
+    initial_pyramid, _ = square_pyramid_factory(h=1.1)
+
+    move = ShapeSpace(callback=ScalePyramid(), default_step_size=0.2)
+    move.params["A"] = [1]
+
+    updater = hpmc.update.Shape(trigger=1, shape_move=move, nsweeps=2)
+    updater.shape_move = move
+
+    mc = hoomd.hpmc.integrate.ConvexPolyhedron()
+    mc.d["A"] = 0
+    mc.a["A"] = 0
+    mc.shape["A"] = dict(vertices=initial_pyramid)
+
+    # create simulation & attach objects
+    sim = simulation_factory(two_particle_snapshot_factory(d=2.5))
+    sim.operations.integrator = mc
+    sim.operations += updater
+
+    # test attachmet before first run
+    assert not move._attached
+    assert not updater._attached
+
+    sim.run(0)
+
+    # test attachmet after first run
+    assert move._attached
+    assert updater._attached
+
+    # run with 0 probability of performing a move:
+    #  - shape and params should remain unchanged
+    #  - all moves accepted
+    move.param_move_probability = 0
+    sim.run(10)
+    assert np.allclose(mc.shape["A"]["vertices"], initial_pyramid)
+    assert np.allclose(move.params["A"], [1])
+    assert np.allclose(updater.particle_volumes, 1)
+
+    # always attempt a shape move:
+    #  - shape and params should change
+    #  - volume should remain unchanged
+    move.param_move_probability = 1
+    sim.run(50)
+    assert np.sum(updater.shape_moves) == 100
+
+    # Check that the shape parameters have changed
+    current_h = move.params["A"][0]
+    assert not np.allclose(mc.shape["A"]["vertices"], initial_pyramid)
+    assert not np.isclose(current_h, 1)
+
+    # Check that the shape parameters map back to the correct geometry
+    assert np.allclose(
+        square_pyramid_factory(current_h + 0.1)[0], mc.shape["A"]["vertices"])
+
+    # Check that the callback is conserving volume properly
+    assert np.allclose(updater.particle_volumes, 1)
+
+
 def test_elastic_shape_move(simulation_factory, two_particle_snapshot_factory):
 
     mc = hoomd.hpmc.integrate.ConvexPolyhedron()