diff --git a/modules/core/include/JacobianAdjuster.h b/modules/core/include/JacobianAdjuster.h index c8eab8d2fc..af001e9482 100644 --- a/modules/core/include/JacobianAdjuster.h +++ b/modules/core/include/JacobianAdjuster.h @@ -68,6 +68,11 @@ IMP_NAMED_TUPLE_3(MultivariateJacobian, MultivariateJacobians, \note This class should not be created directly. Instead, get_jacobian_adjuster() should be used so that only one is associated with each module. + \note Optimizers that divide the score and its gradient by a + pseudo-temperature to smooth the distribution should also set the + temperature here. This ensures that at high temperatures the + sampling distribution of the untransformed variable approaches + the default uniform prior. */ class IMPCOREEXPORT JacobianAdjuster : public IMP::ModelObject { typedef IMP::FloatIndex FloatIndex; @@ -81,6 +86,7 @@ class IMPCOREEXPORT JacobianAdjuster : public IMP::ModelObject { MultivariateJacobianMap; UnivariateJacobianMap uni_map_; MultivariateJacobianMap multi_map_; + double temp_; public: JacobianAdjuster(Model* m, const std::string name = "JacobianAdjuster%1%"); @@ -109,6 +115,12 @@ class IMPCOREEXPORT JacobianAdjuster : public IMP::ModelObject { const MultivariateJacobian& get_jacobian(FloatKeys ks, ParticleIndex pi) const; + //! Set temperature applied to score and its gradient. + void set_temperature(double temperature); + + //! Get temperature applied to score and its gradient. + double get_temperature() const; + #ifndef SWIG UnivariateJacobian& access_jacobian(FloatKey k, ParticleIndex pi); diff --git a/modules/core/src/JacobianAdjuster.cpp b/modules/core/src/JacobianAdjuster.cpp index ec104490a5..e5f38c174f 100644 --- a/modules/core/src/JacobianAdjuster.cpp +++ b/modules/core/src/JacobianAdjuster.cpp @@ -12,7 +12,7 @@ IMPCORE_BEGIN_NAMESPACE JacobianAdjuster::JacobianAdjuster(Model* m, const std::string name) - : ModelObject(m, name) {} + : ModelObject(m, name), temp_(1) {} void JacobianAdjuster::set_jacobian(FloatKey k, ParticleIndex pi, const UnivariateJacobian& j) { @@ -91,6 +91,13 @@ MultivariateJacobian& JacobianAdjuster::access_jacobian(FloatKeys ks, ParticleIn return it->second; } +void JacobianAdjuster::set_temperature(double temperature) { + IMP_INTERNAL_CHECK(temperature > 0, "Temperature must be positive."); + temp_ = temperature; +} + +double JacobianAdjuster::get_temperature() const { return temp_; } + double JacobianAdjuster::get_score_adjustment() const { double adj = 0; Model *m = get_model(); @@ -111,12 +118,12 @@ double JacobianAdjuster::get_score_adjustment() const { } } - return adj; + return adj * temp_; } void JacobianAdjuster::apply_gradient_adjustment() { Model *m = get_model(); - DerivativeAccumulator da = DerivativeAccumulator(); + DerivativeAccumulator da = DerivativeAccumulator(temp_); FloatIndex fi; IMP_FOREACH(UP up, uni_map_) { diff --git a/modules/core/test/test_jacobian_adjustment.py b/modules/core/test/test_jacobian_adjustment.py index abe35a42d8..80f36041c7 100644 --- a/modules/core/test/test_jacobian_adjustment.py +++ b/modules/core/test/test_jacobian_adjustment.py @@ -111,6 +111,13 @@ def test_get_jacobian_adjuster(self): ja2 = IMP.core.get_jacobian_adjuster(m) self.assertEqual(ja1, ja2) + def test_get_set_temperature(self): + m = IMP.Model() + ja = IMP.core.get_jacobian_adjuster(m) + self.assertAlmostEqual(ja.get_temperature(), 1) + ja.set_temperature(10) + self.assertAlmostEqual(ja.get_temperature(), 10) + def test_univariate_jacobian_roundtrip(self): m = IMP.Model() p = IMP.Particle(m) @@ -371,6 +378,44 @@ def test_multivariate_score_state_applies_gradient_adjustment(self): ) self.assertAlmostEqual(ps[1].get_derivative(k), 0, delta=1e-6) + def test_get_tempered_score_adjustment(self): + m = IMP.Model() + ps = [IMP.Particle(m) for i in range(2)] + vs = [np.random.normal(size=3) for _ in ps] + ja = IMP.core.get_jacobian_adjuster(m) + ja.set_temperature(100) + for p, v in zip(ps, vs): + k = IMP.FloatKey("dummy") + p.add_attribute(k, 10.) + j = IMP.core.UnivariateJacobian(*v) + ja.set_jacobian(k, p.get_index(), j) + + self.assertAlmostEqual(ja.get_score_adjustment(), 0, delta=1e-6) + + ps[0].set_is_optimized(k, True) + self.assertAlmostEqual( + ja.get_score_adjustment(), 100 * vs[0][1], delta=1e-6) + + def test_apply_tempered_gradient_adjustment(self): + m = IMP.Model() + p = IMP.Particle(m) + k = IMP.FloatKey("dummy") + p.add_attribute(k, 10.) + vs = np.random.normal(size=3) + j = IMP.core.UnivariateJacobian(*vs) + ja = IMP.core.get_jacobian_adjuster(m) + ja.set_temperature(100) + ja.set_jacobian(k, p.get_index(), j) + + self.assertAlmostEqual(p.get_derivative(k), 0., delta=1e-6) + + ja.apply_gradient_adjustment() + self.assertAlmostEqual(p.get_derivative(k), 0., delta=1e-6) + + p.set_is_optimized(k, True) + ja.apply_gradient_adjustment() + self.assertAlmostEqual(p.get_derivative(k), 100 * vs[2], delta=1e-6) + if __name__ == "__main__": IMP.test.main()