Add classes for adjusting score for parameter changes

This is primarily used when an unintuitive transformation of a
constrained parameter exists in an unconstrained space. To sample in the
unconstrained space but represent the parameter in the constrained space
requires an adjustment to the score and gradient, handled by
JacobianScoreAdjuster.

Relates #1007

modules/core/include/JacobianAdjuster.h

@@ -0,0 +1,161 @@
+/**
+ *  \file IMP/core/JacobianAdjuster.h
+ *  \brief Adjust score and gradients with Jacobian of parameter transformation.
+ *
+ *  Copyright 2007-2018 IMP Inventors. All rights reserved.
+ */
+
+#ifndef IMPCORE_JACOBIAN_ADJUSTER_H
+#define IMPCORE_JACOBIAN_ADJUSTER_H
+
+#include <IMP/core/core_config.h>
+#include <IMP/ModelObject.h>
+#include <IMP/tuple_macros.h>
+#include <IMP/set_map_macros.h>
+
+IMPCORE_BEGIN_NAMESPACE
+
+//! Store Jacobian factors of a univariate transformation
+/** For a univariate transformation, \f$y = f(x)\f$, it stores the Jacobian
+    of the inverse transformation \f$f^{-1}(y)\f$ and the necessary
+    Jacobian-derived adjustment to the score and its gradient.
+
+    \param[in] jacobian Jacobian value \f$J(y) = \frac{d}{dy}f^{-1}(y)\f$.
+    \param[in] score_adjustment Adjustment to score due to transformation,
+                                \f$-\log |J(y)|\f$.
+    \param[in] gradient_adjustment Adjustment to gradient of score due to
+                                   transformation,
+                                   \f$-\frac{d}{dy} \log |J(y)|\f$.
+ */
+IMP_NAMED_TUPLE_3(UnivariateJacobian, UnivariateJacobians, Float, jacobian,
+                  Float, score_adjustment, Float, gradient_adjustment,);
+
+//! Store Jacobian factors of a multivariate transformation
+/** For a multivariate transformation, \f$Y = f(X)\f$, it stores the Jacobian
+    of the inverse transformation \f$f^{-1}(X)\f$ and the necessary
+    Jacobian-derived adjustment to the score and its gradient.
+
+    \param[in] jacobian Jacobian matrix \f$J(Y)\f$.
+    \param[in] score_adjustment Adjustment to score due to transformation,
+                                \f$-\log |\det J(y)|\f$.
+    \param[in] gradient_adjustment Adjustment to gradient of score due to
+                                   transformation,
+                                   \f$-\nabla_Y \log |\det J(y)|\f$.
+ */
+IMP_NAMED_TUPLE_3(MultivariateJacobian, MultivariateJacobians,
+                  FloatsList, jacobian, Float, score_adjustment, Floats,
+                  gradient_adjustment,{
+  IMP_IF_CHECK(USAGE_AND_INTERNAL) {
+    int N = gradient_adjustment_.size();
+    IMP_USAGE_CHECK(N > 0, "Number of variables must be positive.");
+    IMP_USAGE_CHECK(jacobian_.size() == N,
+      "Jacobian number of rows must be equal to gradient adjustment size.");
+    for (unsigned int n = 0; n < N; ++n) {
+      IMP_USAGE_CHECK(jacobian_[n].size() == N, "Jacobian must be square.");
+    }
+  }
+});
+
+//! Adjust score and gradients with Jacobian of parameter transformation.
+/** It stores UnivariateJacobians and MultivariateJacobians as well as the
+    particle and float attributes to which they are related.
+
+    It computes the total necessary Jacobian adjustment to the score due to
+    transformation with get_score_adjustment(). It also adds the necessary
+    Jacobian gradient modification to all derivatives with
+    apply_gradient_adjustment().
+
+    \note This class should not be created directly. Instead,
+          get_jacobian_adjuster() should be used so that only one is
+          associated with each module.
+ */
+class IMPCOREEXPORT JacobianAdjuster : public IMP::ModelObject {
+  typedef IMP::FloatIndex FloatIndex;
+  IMP_NAMED_TUPLE_2(FloatsIndex, FloatsIndexes, ParticleIndex, particle,
+                    FloatKeys, keys, );
+  typedef std::pair<FloatIndex, UnivariateJacobian> UP;
+  typedef std::pair<FloatsIndex, MultivariateJacobian> MP;
+  typedef IMP_KERNEL_SMALL_UNORDERED_MAP<FloatIndex, UnivariateJacobian>
+      UnivariateJacobianMap;
+  typedef IMP_KERNEL_SMALL_UNORDERED_MAP<FloatsIndex, MultivariateJacobian>
+      MultivariateJacobianMap;
+  UnivariateJacobianMap uni_map_;
+  MultivariateJacobianMap multi_map_;
+
+ public:
+  JacobianAdjuster(Model* m, const std::string name = "JacobianAdjuster%1%");
+
+  //! Set Jacobian for univariate transformation corresponding to attribute.
+  /** \param[in] ks Attribute of transformed parameter.
+      \param[in] pi Index of decorated particle.
+      \param[in] j Container of Jacobian, score adjustment, and gradient
+                   adjustment due to transformation.
+   */
+  void set_jacobian(FloatKey k, ParticleIndex pi, const UnivariateJacobian& j);
+
+  //! Set Jacobian for multivariate transformation corresponding to attribute.
+  /** \param[in] ks Attributes of transformed parameter.
+      \param[in] pi Index of decorated particle.
+      \param[in] j Container of Jacobian matrix, score adjustment, and gradient
+                   adjustment due to transformation.
+   */
+  void set_jacobian(FloatKeys ks, ParticleIndex pi,
+                    const MultivariateJacobian& j);
+
+  //! Get stored Jacobian for univariate transformation.
+  const UnivariateJacobian& get_jacobian(FloatKey k, ParticleIndex pi) const;
+
+  //! Get stored Jacobian for multivariate transformation.
+  const MultivariateJacobian& get_jacobian(FloatKeys ks,
+                                           ParticleIndex pi) const;
+
+#ifndef SWIG
+  UnivariateJacobian& access_jacobian(FloatKey k, ParticleIndex pi);
+
+  MultivariateJacobian& access_jacobian(FloatKeys ks, ParticleIndex pi);
+#endif
+
+  //! Get total adjustment to score from stored Jacobians.
+  /** Only optimized attributes are considered. */
+  double get_score_adjustment() const;
+
+  //! Adjust derivatives of particles using stored Jacobians.
+  /** Only derivatives of optimized attributes are adjusted. */
+  void apply_gradient_adjustment();
+
+  virtual ModelObjectsTemp do_get_inputs() const IMP_OVERRIDE;
+
+  virtual ModelObjectsTemp do_get_outputs() const IMP_OVERRIDE;
+
+ private:
+  static ModelKey get_score_state_key();
+  void create_score_state();
+
+  IMP_OBJECT_METHODS(JacobianAdjuster);
+};
+
+IMP_OBJECTS(JacobianAdjuster, JacobianAdjusters);
+
+//! Get adjuster of Model's Jacobian
+/** If one does not yet exist, it is created. */
+IMPCOREEXPORT JacobianAdjuster* get_jacobian_adjuster(Model* m);
+
+
+
+#if !defined(IMP_DOXYGEN) && !defined(SWIG)
+//! Adjust all transformed parameter gradients with JacobianAdjuster.
+class IMPCOREEXPORT JacobianAdjustGradient : public ScoreState {
+ public:
+  JacobianAdjustGradient(Model *m)
+      : ScoreState(m, "JacobianAdjustGradient%1%") {}
+  virtual void do_before_evaluate() IMP_OVERRIDE;
+  virtual void do_after_evaluate(DerivativeAccumulator *da) IMP_OVERRIDE;
+  virtual ModelObjectsTemp do_get_inputs() const IMP_OVERRIDE;
+  virtual ModelObjectsTemp do_get_outputs() const IMP_OVERRIDE;
+  IMP_OBJECT_METHODS(JacobianAdjustGradient);
+};
+#endif
+
+IMPCORE_END_NAMESPACE
+
+#endif /* IMPCORE_JACOBIAN_ADJUSTER_H */

modules/core/pyext/swig.i-in

@@ -198,6 +198,9 @@ IMP_SWIG_OBJECT(IMP::core, IsCollisionPairPredicate, IsCollisionPairPredicates);
 IMP_SWIG_VALUE(IMP::core, BinormalTerm, BinormalTermList);
 IMP_SWIG_OBJECT(IMP::core, MultipleBinormalRestraint, MultipleBinormalRestraints);
 
+IMP_SWIG_OBJECT(IMP::core, JacobianAdjuster, JacobianAdjusters);
+IMP_SWIG_VALUE(IMP::core, UnivariateJacobian, UnivariateJacobians);
+IMP_SWIG_VALUE(IMP::core, MultivariateJacobian, MultivariateJacobians);
 
 /* Don't wrap internal classes */
 %ignore IMP::core::internal::ChildArrayTraits;
@@ -330,6 +333,7 @@ void visit_depth_first(Hierarchy d, HierarchyVisitor *f)
 %include "IMP/core/MCCGSampler.h"
 %include "IMP/core/MinimumRestraint.h"
 %include "IMP/core/Gaussian.h"
+%include "IMP/core/JacobianAdjuster.h"
 
 %include "IMP/core/PairRestraint.h"
 %include "IMP/core/SingletonRestraint.h"

modules/core/src/JacobianAdjuster.cpp

@@ -0,0 +1,218 @@
+/**
+ *  \file JacobianAdjuster.cpp
+ *  \brief Classes for dealing with transformations of parameters.
+ *
+ *  Copyright 2007-2018 IMP Inventors. All rights reserved.
+ *
+ */
+
+#include <IMP/core/JacobianAdjuster.h>
+#include <IMP/object_cast.h>
+
+IMPCORE_BEGIN_NAMESPACE
+
+JacobianAdjuster::JacobianAdjuster(Model* m, const std::string name) 
+  : ModelObject(m, name) {}
+
+void JacobianAdjuster::set_jacobian(FloatKey k, ParticleIndex pi,
+                                    const UnivariateJacobian& j) {
+  IMP_IF_CHECK(USAGE_AND_INTERNAL) {
+    IMP_USAGE_CHECK(get_model()->get_has_attribute(k, pi),
+                    "Particle does not have Float attribute.");
+  }
+
+  FloatIndex fi(pi, k);
+  UnivariateJacobianMap::iterator lb = uni_map_.lower_bound(fi);
+  if(lb != uni_map_.end() && !(uni_map_.key_comp()(fi, lb->first))) {
+    lb->second = j;
+  } else {
+    uni_map_.insert(lb, UnivariateJacobianMap::value_type(fi, j));
+    set_has_dependencies(false);
+    create_score_state();
+  }
+}
+
+void JacobianAdjuster::set_jacobian(
+    FloatKeys ks, ParticleIndex pi, const MultivariateJacobian& j) {
+  IMP_IF_CHECK(USAGE_AND_INTERNAL) {
+    int N = ks.size();
+    IMP_USAGE_CHECK(
+        N > 0,
+        "Multivariate transformed parameter must at least one key.");
+    IMP_USAGE_CHECK(
+      j.get_jacobian().size() == N && j.get_gradient_adjustment().size() == N,
+      "Jacobian and gradient adjustment dimensions must equal number of keys.");
+    IMP_FOREACH(FloatKey k, ks) {
+      IMP_USAGE_CHECK(get_model()->get_has_attribute(k, pi),
+                      "Particle does not have Float attribute.");
+    }
+  }
+
+  FloatsIndex fi(pi, ks);
+  MultivariateJacobianMap::iterator lb = multi_map_.lower_bound(fi);
+  if(lb != multi_map_.end() && !(multi_map_.key_comp()(fi, lb->first))) {
+    lb->second = j;
+  } else {
+    multi_map_.insert(lb, MultivariateJacobianMap::value_type(fi, j));
+    set_has_dependencies(false);
+    create_score_state();
+  }
+}
+
+const UnivariateJacobian& JacobianAdjuster::get_jacobian(FloatKey k, ParticleIndex pi) const {
+  FloatIndex fi(pi, k);
+  UnivariateJacobianMap::const_iterator it;
+  it = uni_map_.find(fi);
+  IMP_USAGE_CHECK(it != uni_map_.end(), "Jacobian has not been added.");
+  return it->second;
+}
+
+UnivariateJacobian& JacobianAdjuster::access_jacobian(FloatKey k, ParticleIndex pi) {
+  FloatIndex fi(pi, k);
+  UnivariateJacobianMap::iterator it;
+  it = uni_map_.find(fi);
+  IMP_USAGE_CHECK(it != uni_map_.end(), "Jacobian has not been added.");
+  return it->second;
+}
+
+const MultivariateJacobian& JacobianAdjuster::get_jacobian(FloatKeys ks, ParticleIndex pi) const {
+  FloatsIndex fi(pi, ks);
+  MultivariateJacobianMap::const_iterator it;
+  it = multi_map_.find(fi);
+  IMP_USAGE_CHECK(it != multi_map_.end(), "Jacobian has not been added.");
+  return it->second;
+}
+
+MultivariateJacobian& JacobianAdjuster::access_jacobian(FloatKeys ks, ParticleIndex pi) {
+  FloatsIndex fi(pi, ks);
+  MultivariateJacobianMap::iterator it;
+  it = multi_map_.find(fi);
+  IMP_USAGE_CHECK(it != multi_map_.end(), "Jacobian has not been added.");
+  return it->second;
+}
+
+double JacobianAdjuster::get_score_adjustment() const {
+  double adj = 0;
+  Model *m = get_model();
+
+  FloatIndex fi;
+  IMP_FOREACH(UP up, uni_map_) {
+    fi = up.first;
+    if (m->get_is_optimized(fi.get_key(), fi.get_particle())) {
+      adj += up.second.get_score_adjustment();
+    }
+  }
+
+  FloatsIndex fsi;
+  IMP_FOREACH(MP mp, multi_map_) {
+    fsi = mp.first;
+    if (m->get_is_optimized(fsi.get_keys()[0], fsi.get_particle())) {
+      adj += mp.second.get_score_adjustment();
+    }
+  }
+
+  return adj;
+}
+
+void JacobianAdjuster::apply_gradient_adjustment() {
+  Model *m = get_model();
+  DerivativeAccumulator da = DerivativeAccumulator();
+
+  FloatIndex fi;
+  IMP_FOREACH(UP up, uni_map_) {
+    fi = up.first;
+    if (m->get_is_optimized(fi.get_key(), fi.get_particle())) {
+        m->add_to_derivative(fi.get_key(), fi.get_particle(),
+                             up.second.get_gradient_adjustment(), da);
+    }
+  }
+
+  FloatsIndex fsi;
+  ParticleIndex pi;
+  Floats grad_adj;
+  FloatKeys ks;
+  unsigned int n, N;
+  IMP_FOREACH(MP mp, multi_map_) {
+    fsi = mp.first;
+    FloatKeys ks = fsi.get_keys();
+    pi = fsi.get_particle();
+    if (m->get_is_optimized(ks[0], pi)) {
+      grad_adj = mp.second.get_gradient_adjustment();
+      N = ks.size();
+      IMP_INTERNAL_CHECK(
+        grad_adj.size() == N,
+        "Size of gradient doesn't match number of attributes.");
+      for (n = 0; n < N; ++n) {
+        m->add_to_derivative(ks[n], pi, grad_adj[n], da);
+      }
+    }
+  }
+}
+
+ModelKey JacobianAdjuster::get_score_state_key() {
+  static ModelKey k("jacobian adjust gradient score state");
+  return k;
+}
+
+void JacobianAdjuster::create_score_state() {
+  ModelKey k = get_score_state_key();
+  if (!get_model()->get_has_data(k)) {
+    IMP_NEW(JacobianAdjustGradient, ss, (get_model()));
+    get_model()->add_data(k, ss);
+    get_model()->add_score_state(ss.release());
+  } else {
+    IMP::Pointer<JacobianAdjustGradient> ss =
+        IMP::object_cast<JacobianAdjustGradient>(get_model()->get_data(k));
+    ss->set_has_dependencies(false);
+  }
+}
+
+ModelObjectsTemp JacobianAdjuster::do_get_inputs() const {
+  ModelObjectsTemp ret;
+  Model *m = get_model();
+
+  ret.reserve(uni_map_.size() + multi_map_.size());
+
+  IMP_FOREACH(UP up, uni_map_) {
+    ret.push_back(m->get_particle(up.first.get_particle()));
+  }
+
+  IMP_FOREACH(MP mp, multi_map_) {
+    ret.push_back(m->get_particle(mp.first.get_particle()));
+  }
+
+  return ret;
+}
+
+ModelObjectsTemp JacobianAdjuster::do_get_outputs() const {
+  return get_inputs();
+}
+
+JacobianAdjuster *get_jacobian_adjuster(Model *m) {
+  static ModelKey k("jacobian adjuster");
+  if (!m->get_has_data(k)) {
+    IMP_NEW(JacobianAdjuster, ja, (m));
+    m->add_data(k, ja);
+    return ja.release();
+  } else {
+    IMP::Pointer<JacobianAdjuster> ja =
+        IMP::object_cast<JacobianAdjuster>(m->get_data(k));
+    return ja.release();
+  }
+}
+
+void JacobianAdjustGradient::do_before_evaluate() {}
+
+void JacobianAdjustGradient::do_after_evaluate(DerivativeAccumulator *da) {
+  get_jacobian_adjuster(get_model())->apply_gradient_adjustment();
+}
+
+ModelObjectsTemp JacobianAdjustGradient::do_get_inputs() const {
+  return get_jacobian_adjuster(get_model())->get_inputs();
+}
+
+ModelObjectsTemp JacobianAdjustGradient::do_get_outputs() const {
+  return get_jacobian_adjuster(get_model())->get_outputs();
+}
+
+IMPCORE_END_NAMESPACE