Fourvec_Constrainer.cc

Go to the documentation of this file.

//
//
// File: src/Fourvec_Constrainer.cc
// Purpose: Do a kinematic fit for a set of 4-vectors, given a set
//          of mass constraints.
// Created: Jul, 2000, sss, based on run 1 mass analysis code.
//
// CMSSW File      : src/Fourvec_Constrainer.cc
// Original Author : Scott Stuart Snyder <snyder@bnl.gov> for D0
// Imported to CMSSW by Haryo Sumowidagdo <Suharyo.Sumowidagdo@cern.ch>
//

#include "TopQuarkAnalysis/TopHitFit/interface/Fourvec_Constrainer.h"
#include "TopQuarkAnalysis/TopHitFit/interface/Fourvec_Event.h"
#include "TopQuarkAnalysis/TopHitFit/interface/Pair_Table.h"
#include "TopQuarkAnalysis/TopHitFit/interface/Chisq_Constrainer.h"
#include "TopQuarkAnalysis/TopHitFit/interface/matutil.h"
#include "TopQuarkAnalysis/TopHitFit/interface/Defaults.h"
#include <cmath>
#include <iostream>

using std::cos;
using std::exp;
using std::ostream;
using std::sin;
using std::sqrt;
using std::vector;

namespace hitfit {

  //*************************************************************************
  // Argument handling.
  //

  Fourvec_Constrainer_Args::Fourvec_Constrainer_Args(const Defaults& defs)
      //
      // Purpose: Constructor.
      //
      // Inputs:
      //   defs -        The Defaults instance from which to initialize.
      //
      : _use_e(defs.get_bool("use_e")),
        _e_com(defs.get_float("e_com")),
        _ignore_met(defs.get_bool("ignore_met")),
        _chisq_constrainer_args(defs) {}

  bool Fourvec_Constrainer_Args::use_e() const
  //
  // Purpose: Return the use_e parameter.
  //          See the header for documentation.
  //
  {
    return _use_e;
  }

  double Fourvec_Constrainer_Args::e_com() const
  //
  // Purpose: Return the e_com parameter.
  //          See the header for documentation.
  //
  {
    return _e_com;
  }

  bool Fourvec_Constrainer_Args::ignore_met() const
  //
  // Purpose: Return the ignore_met parameter.
  //          See the header for documentation.
  //
  {
    return _ignore_met;
  }

  const Chisq_Constrainer_Args& Fourvec_Constrainer_Args::chisq_constrainer_args() const
  //
  // Purpose: Return the contained subobject parameters.
  //
  {
    return _chisq_constrainer_args;
  }

  //*************************************************************************
  // Variable layout.
  //
  // We need to map the quantities we fit onto the vectors of well- and
  // poorly-measured quantities.
  //

  //
  // The well-measured variables consist of three variables for each
  // object.  If we are using transverse momentum constraints,
  // these fill be followed by the two cartesian components of kt.
  //
  // Each object is represented by three variables: the momentum (or 1/p
  // if the muon flag was set), and the two spherical angles, phi and eta.
  // Here is how they're ordered.
  //
  typedef enum { p_offs = 0, phi_offs = 1, eta_offs = 2 } Offsets;

  typedef enum { x_offs = 0, y_offs = 1 } Kt_Offsets;

  //
  // If there is a neutrino, then it is at index 1 of the poorly-measured
  // set (otherwise, that set is empty).
  //
  typedef enum { nu_z = 1 } Unmeasured_Variables;

  namespace {

    int obj_index(int i)
    //
    // Purpose: Return the starting variable index for object I.
    //
    // Inputs:
    //   i -           The object index.
    //
    // Returns:
    //   The index in the well-measured set of the first variable
    //   for object I.
    //
    {
      return i * 3 + 1;
    }

  }  // unnamed namespace

  //*************************************************************************
  // Object management.
  //

  Fourvec_Constrainer::Fourvec_Constrainer(const Fourvec_Constrainer_Args& args)
      //
      // Purpose: Constructor.
      //
      // Inputs:
      //   args -        The parameter settings for this instance.
      //
      : _args(args) {}

  void Fourvec_Constrainer::add_constraint(std::string s)
  //
  // Purpose: Specify an additional constraint S for the problem.
  //          The format for S is described in the header.
  //
  // Inputs:
  //   s -           The constraint to add.
  //
  {
    _constraints.push_back(Constraint(s));
  }

  void Fourvec_Constrainer::mass_constraint(std::string s)
  //
  // Purpose: Specify the combination of objects that will be returned by
  //          constrain() as the mass.  The format of S is the same as for
  //          normal constraints.  The LHS specifies the mass to calculate;
  //          the RHS should be zero.
  //          This should only be called once.
  //
  // Inputs:
  //   s -           The constraint defining the mass.
  //
  {
    assert(_mass_constraint.empty());
    _mass_constraint.push_back(Constraint(s));
  }

  std::ostream& operator<<(std::ostream& s, const Fourvec_Constrainer& c)
  //
  // Purpose: Print the object to S.
  //
  // Inputs:
  //   s -           The stream to which to write.
  //   c -           The object to write.
  //
  // Returns:
  //   The stream S.
  //
  {
    s << "Constraints: (e_com = " << c._args.e_com() << ") ";
    if (c._args.use_e())
      s << "(E)";
    s << "\n";

    for (std::vector<Constraint>::size_type i = 0; i < c._constraints.size(); i++)
      s << "  " << c._constraints[i] << "\n";

    if (!c._mass_constraint.empty()) {
      s << "Mass constraint:\n";
      s << c._mass_constraint[0] << "\n";
    }
    return s;
  }

  //*************************************************************************
  // Event packing and unpacking.
  //

  namespace {

    void adjust_fourvecs(Fourvec_Event& ev, bool use_e_flag)
    //
    // Purpose: For all objects in EV, adjust their 4-momenta
    //          to have their requested masses.
    //
    // Inputs:
    //   ev -          The event on which to operate.
    //   use_e_flag -  If true, keep E and scale 3-momentum.
    //                 If false, keep the 3-momentum and scale E.
    //
    {
      int nobjs = ev.nobjs();
      for (int i = 0; i < nobjs; i++) {
        const FE_Obj& obj = ev.obj(i);
        Fourvec p = obj.p;
        if (use_e_flag)
          adjust_p_for_mass(p, obj.mass);
        else
          adjust_e_for_mass(p, obj.mass);
        ev.set_obj_p(i, p);
      }
    }

    Fourvec get_p_eta_phi_vec(const Column_Vector& c, int ndx, const FE_Obj& obj, bool use_e_flag)
    //
    // Purpose: Convert object NDX from its representation in the set
    //          of well-measured variables C to a 4-vector.
    //
    // Inputs:
    //   c -           The vector of well-measured variables.
    //   ndx -         The index of the object in which we're interested.
    //   obj -         The object from the Fourvec_Event.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //
    // Returns:
    //   The object's 4-momentum.
    //
    {
      // Get the energy and momentum of the object.
      double e, p;

      if (use_e_flag) {
        // We're using E as a fit variable.  Get it directly.
        e = c(ndx + p_offs);

        // Take into account the muon case.
        if (obj.muon_p)
          e = 1 / e;

        // Find the momentum given the energy.
        if (obj.mass == 0)
          p = e;
        else {
          double xx = e * e - obj.mass * obj.mass;
          if (xx >= 0)
            p = sqrt(xx);
          else
            p = 0;
        }
      } else {
        // We're using P as a fit variable.  Fetch it.
        p = c(ndx + p_offs);

        // Take into account the muon case.
        if (obj.muon_p)
          p = 1 / p;

        // Find the energy given the momentum.
        e = (obj.mass == 0 ? p : sqrt(obj.mass * obj.mass + p * p));
      }

      // Get angular variables.
      double phi = c(ndx + phi_offs);
      double eta = c(ndx + eta_offs);
      if (fabs(eta) > 50) {
        // Protect against ridiculously large etas
        eta = eta > 0 ? 50 : -50;
      }
      double exp_eta = exp(eta);
      double iexp_eta = 1 / exp_eta;
      double sin_theta = 2 / (exp_eta + iexp_eta);
      double cos_theta = (exp_eta - iexp_eta) / (exp_eta + iexp_eta);

      // Form the 4-momentum.
      return Fourvec(p * sin_theta * cos(phi), p * sin_theta * sin(phi), p * cos_theta, e);
    }

    void set_p_eta_phi_vec(const FE_Obj& obj, Column_Vector& c, int ndx, bool use_e_flag)
    //
    // Purpose: Initialize the variables in the well-measured set C describing
    //          object NDX from its Fourvec_Event representation OBJ.
    //
    // Inputs:
    //   obj -         The object from the Fourvec_Event.
    //   c -           The vector of well-measured variables.
    //   ndx -         The index of the object in which we're interested.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //
    //
    {
      if (use_e_flag)
        c(ndx + p_offs) = obj.p.e();
      else
        c(ndx + p_offs) = obj.p.vect().mag();

      if (obj.muon_p)
        c(ndx + p_offs) = 1 / c(ndx + p_offs);
      c(ndx + phi_offs) = obj.p.phi();
      c(ndx + eta_offs) = obj.p.pseudoRapidity();
    }

    Matrix error_matrix(double p_sig, double phi_sig, double eta_sig)
    //
    // Purpose: Set up the 3x3 error matrix for an object.
    //
    // Inputs:
    //   p_sig -       The momentum uncertainty.
    //   phi_sig -     The phi uncertainty.
    //   eta_sig -     The eta uncertainty.
    //
    // Returns:
    //   The object's error matrix.
    //
    {
      Matrix err(3, 3, 0);
      err(1 + p_offs, 1 + p_offs) = p_sig * p_sig;
      err(1 + phi_offs, 1 + phi_offs) = phi_sig * phi_sig;
      err(1 + eta_offs, 1 + eta_offs) = eta_sig * eta_sig;
      return err;
    }

    void pack_event(const Fourvec_Event& ev,
                    bool use_e_flag,
                    bool use_kt_flag,
                    Column_Vector& xm,
                    Column_Vector& ym,
                    Matrix& G_i,
                    Diagonal_Matrix& Y)
    //
    // Purpose: Take the information from a Fourvec_Event EV and pack
    //          it into vectors of well- and poorly-measured variables.
    //          Also set up the error matrices.
    //
    // Inputs:
    //   ev -          The event to pack.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //   use_kt_flag - True if we're to pack kt variables.
    //
    // Outputs:
    //   xm -          Vector of well-measured variables.
    //   ym -          Vector of poorly-measured variables.
    //   G_i -         Error matrix for well-measured variables.
    //   Y -           Inverse error matrix for poorly-measured variables.
    //
    {
      // Number of objects in the event.
      int nobjs = ev.nobjs();

      int n_measured_vars = nobjs * 3;
      if (use_kt_flag)
        n_measured_vars += 2;

      // Clear the error matrix.
      G_i = Matrix(n_measured_vars, n_measured_vars, 0);

      // Loop over objects.
      for (int i = 0; i < nobjs; i++) {
        const FE_Obj& obj = ev.obj(i);
        int this_index = obj_index(i);
        set_p_eta_phi_vec(obj, xm, this_index, use_e_flag);
        G_i.sub(this_index, this_index, error_matrix(obj.p_error, obj.phi_error, obj.eta_error));
      }

      if (use_kt_flag) {
        // Set up kt.
        int kt_ndx = obj_index(nobjs);
        xm(kt_ndx + x_offs) = ev.kt().x();
        xm(kt_ndx + y_offs) = ev.kt().y();

        // And its error matrix.
        G_i(kt_ndx + x_offs, kt_ndx + x_offs) = ev.kt_x_error() * ev.kt_x_error();
        G_i(kt_ndx + y_offs, kt_ndx + y_offs) = ev.kt_y_error() * ev.kt_y_error();
        G_i(kt_ndx + x_offs, kt_ndx + y_offs) = ev.kt_xy_covar();
        G_i(kt_ndx + y_offs, kt_ndx + x_offs) = ev.kt_xy_covar();
      }

      // Handle a neutrino.
      if (ev.has_neutrino()) {
        ym(nu_z) = ev.nu().z();
        Y = Diagonal_Matrix(1, 0);
      }
    }

    void unpack_event(
        Fourvec_Event& ev, const Column_Vector& x, const Column_Vector& y, bool use_e_flag, bool use_kt_flag)
    //
    // Purpose: Update the contents of EV from the variable sets X and Y.
    //
    // Inputs:
    //   ev -          The event.
    //   x -           Vector of well-measured variables.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //   use_kt_flag - True if we're too unpack kt variables.
    //
    // Outputs:
    //   ev -          The event after updating.
    //
    {
      // Do all the objects.
      Fourvec sum = Fourvec(0);
      for (int j = 0; j < ev.nobjs(); j++) {
        const FE_Obj& obj = ev.obj(j);
        ev.set_obj_p(j, get_p_eta_phi_vec(x, obj_index(j), obj, use_e_flag));
        sum += obj.p;
      }

      if (use_kt_flag) {
        int kt_ndx = obj_index(ev.nobjs());
        Fourvec kt = Fourvec(x(kt_ndx + x_offs), x(kt_ndx + y_offs), 0, 0);
        Fourvec nu = kt - sum;
        if (ev.has_neutrino()) {
          nu.setPz(y(nu_z));
          adjust_e_for_mass(nu, 0);
          ev.set_nu_p(nu);
        } else {
          adjust_e_for_mass(nu, 0);
          ev.set_x_p(nu);
        }
      }
    }

  }  // unnamed namespace

  //*************************************************************************
  // Constraint evaluation.
  //

  namespace {

    double dot_and_gradient(
        const Fourvec& v1, const Fourvec& v2, bool use_e_flag, double v1_x[3], double v2_x[3], bool& badflag)
    //
    // Purpose: Compute the dot product v1.v2 and its gradients wrt
    //          p, phi, and theta of each 4-vector.
    //
    // Inputs:
    //   v1 -          The first 4-vector in the dot product.
    //   v2 -          The second 4-vector in the dot product.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //
    // Outputs:
    //   v1_x -        Gradients of the dot product wrt v1's p, phi, theta.
    //   v2_x -        Gradients of the dot product wrt v2's p, phi, theta.
    //   badflag -     Set to true for the singular case (vectors vanish).
    //
    // Returns:
    //   The dot product.
    //
    {
      // Calculate the dot product.
      double dot = v1 * v2;

      double p1 = v1.vect().mag();
      double p2 = v2.vect().mag();
      double e1 = v1.e();
      double e2 = v2.e();
      double pt1 = v1.vect().perp();
      double pt2 = v2.vect().perp();

      // Protect against the singular case.
      badflag = false;
      if (p1 == 0 || p2 == 0 || e1 == 0 || e2 == 0 || pt1 == 0 || pt2 == 0) {
        badflag = true;
        v1_x[p_offs] = v1_x[phi_offs] = v1_x[eta_offs] = 0;
        v2_x[p_offs] = v2_x[phi_offs] = v2_x[eta_offs] = 0;
        return false;
      }

      // Calculate the gradients.
      v1_x[p_offs] = (dot - v1.m2() * e2 / e1) / p1;
      v2_x[p_offs] = (dot - v2.m2() * e1 / e2) / p2;

      if (use_e_flag) {
        v1_x[p_offs] *= e1 / p1;
        v2_x[p_offs] *= e2 / p2;
      }

      v1_x[phi_offs] = v1(1) * v2(0) - v1(0) * v2(1);
      v2_x[phi_offs] = -v1_x[phi_offs];

      double fac = v1(0) * v2(0) + v1(1) * v2(1);
      v1_x[eta_offs] = pt1 * v2(2) - v1(2) / pt1 * fac;
      v2_x[eta_offs] = pt2 * v1(2) - v2(2) / pt2 * fac;

      return dot;
    }

    void addin_obj_gradient(int constraint_no, int sign, int base_index, const double grad[], Matrix& Bx)
    //
    // Purpose: Tally up the dot product gradients for an object
    //          into Bx.
    //
    // Inputs:
    //   constraint_no-The number of the constraint.
    //   base_index -  The index in the well-measured variable list
    //                 of the first variable for this object.
    //   sign -        The sign with which these gradients should be
    //                 added into Bx, either +1 or -1.  (I.e., which
    //                 side of the constraint equation.)
    //   grad -        The gradients for this object, vs. p, phi, theta.
    //   Bx -          The well-measured variable gradients.
    //
    // Outputs:
    //   Bx -          The well-measured variable gradients, updated.
    //
    {
      Bx(base_index + p_offs, constraint_no) += sign * grad[p_offs];
      Bx(base_index + phi_offs, constraint_no) += sign * grad[phi_offs];
      Bx(base_index + eta_offs, constraint_no) += sign * grad[eta_offs];
    }

    void addin_nu_gradient(int constraint_no, int sign, int kt_ndx, const double grad[], Matrix& Bx, Matrix& By)
    //
    // Purpose: Tally up the dot product gradients for a neutrino
    //          into Bx and By.
    //
    // Inputs:
    //   constraint_no-The number of the constraint.
    //   sign -        The sign with which these gradients should be
    //                 added into Bx, either +1 or -1.  (I.e., which
    //                 side of the constraint equation.)
    //   kt_ndx -      The index of the kt variables in the variables array.
    //   grad -        The gradients for this object, vs. p, phi, theta.
    //   Bx -          The well-measured variable gradients.
    //   By -          The poorly-measured variable gradients.
    //
    // Outputs:
    //   Bx -          The well-measured variable gradients, updated.
    //   By -          The poorly-measured variable gradients, updated.
    //
    {
      Bx(kt_ndx + x_offs, constraint_no) += sign * grad[p_offs];    // Really p for now.
      Bx(kt_ndx + y_offs, constraint_no) += sign * grad[phi_offs];  // Really phi for now.
      By(nu_z, constraint_no) += sign * grad[eta_offs];             // Really theta ...
    }

    void addin_gradient(
        const Fourvec_Event& ev, int constraint_no, int sign, int obj_no, const double grad[], Matrix& Bx, Matrix& By)
    //
    // Purpose: Tally up the dot product gradients for an object (which may
    //          or may not be a neutrino) into Bx and By.
    //
    // Inputs:
    //   ev -          The event we're fitting.
    //   constraint_no-The number of the constraint.
    //   sign -        The sign with which these gradients should be
    //                 added into Bx, either +1 or -1.  (I.e., which
    //                 side of the constraint equation.)
    //   obj_no -      The number of the object.
    //   grad -        The gradients for this object, vs. p, phi, theta.
    //   Bx -          The well-measured variable gradients.
    //   By -          The poorly-measured variable gradients.
    //
    // Outputs:
    //   Bx -          The well-measured variable gradients, updated.
    //   By -          The poorly-measured variable gradients, updated.
    //
    {
      if (obj_no >= ev.nobjs()) {
        assert(obj_no == ev.nobjs());
        addin_nu_gradient(constraint_no, sign, obj_index(obj_no), grad, Bx, By);
      } else
        addin_obj_gradient(constraint_no, sign, obj_index(obj_no), grad, Bx);
    }

    void addin_gradients(const Fourvec_Event& ev,
                         int constraint_no,
                         int sign,
                         int i,
                         const double igrad[],
                         int j,
                         const double jgrad[],
                         Matrix& Bx,
                         Matrix& By)
    //
    // Purpose: Tally up the gradients from a single dot product into Bx and By.
    //
    // Inputs:
    //   ev -          The event we're fitting.
    //   constraint_no-The number of the constraint.
    //   sign -        The sign with which these gradients should be
    //                 added into Bx, either +1 or -1.  (I.e., which
    //                 side of the constraint equation.)
    //   i -           The number of the first object.
    //   igrad -       The gradients for the first object, vs. p, phi, theta.
    //   j -           The number of the second object.
    //   jgrad -       The gradients for the second object, vs. p, phi, theta.
    //   Bx -          The well-measured variable gradients.
    //   By -          The poorly-measured variable gradients.
    //
    // Outputs:
    //   Bx -          The well-measured variable gradients, updated.
    //   By -          The poorly-measured variable gradients, updated.
    //
    {
      addin_gradient(ev, constraint_no, sign, i, igrad, Bx, By);
      addin_gradient(ev, constraint_no, sign, j, jgrad, Bx, By);
    }

    void add_mass_terms(const Fourvec_Event& ev, const vector<Constraint>& constraints, Row_Vector& F)
    //
    // Purpose: Add the m^2 terms into the constraint values.
    //
    // Inputs:
    //   ev -          The event we're fitting.
    //   constraints - The list of constraints.
    //   F -           Vector of constraint values.
    //
    // Outputs:
    //   F -           Vector of constraint values, updated.
    //
    {
      for (std::vector<Constraint>::size_type i = 0; i < constraints.size(); i++)
        F(i + 1) += constraints[i].sum_mass_terms(ev);
    }

    bool calculate_mass_constraints(const Fourvec_Event& ev,
                                    const Pair_Table& pt,
                                    const vector<Constraint>& constraints,
                                    bool use_e_flag,
                                    Row_Vector& F,
                                    Matrix& Bx,
                                    Matrix& By)
    //
    // Purpose: Calculate the mass constraints and gradients.
    //          Note: At this stage, the gradients are calculated not
    //                quite with respect to the fit variables; instead, for
    //                all objects (including the neutrino) we calculate
    //                the gradients with respect to p, phi, theta.  They'll
    //                be converted via appropriate Jacobian transformations
    //                later.
    //
    // Inputs:
    //   ev -          The event we're fitting.
    //   pt -          Table of cached pair assignments.
    //   constraints - The list of constraints.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //
    // Outputs:
    //   (nb. these should be passed in correctly dimensioned.)
    //   F -           Vector of constraint values.
    //   Bx -          The well-measured variable gradients.
    //   By -          The poorly-measured variable gradients.
    //
    {
      int npairs = pt.npairs();
      for (int p = 0; p < npairs; p++) {
        const Objpair& objpair = pt.get_pair(p);
        int i = objpair.i();
        int j = objpair.j();
        double igrad[3], jgrad[3];
        bool badflag = false;
        double dot = dot_and_gradient(ev.obj(i).p, ev.obj(j).p, use_e_flag, igrad, jgrad, badflag);
        if (badflag)
          return false;

        for (std::vector<Constraint>::size_type k = 0; k < constraints.size(); k++)
          if (objpair.for_constraint(k)) {
            F(k + 1) += objpair.for_constraint(k) * dot;
            addin_gradients(ev, k + 1, objpair.for_constraint(k), i, igrad, j, jgrad, Bx, By);
          }
      }

      add_mass_terms(ev, constraints, F);
      return true;
    }

    void add_nuterm(unsigned ndx, const Fourvec& v, Matrix& Bx, bool use_e_flag, int kt_ndx)
    //
    // Purpose: Carry out the Jacobian transformation for
    //          (p_nu^x,p_nu_y) -> (kt^x, kt_y) for a given object.
    //
    // Inputs:
    //   ndx -         Index of the object for which to transform gradients.
    //   v -           The object's 4-momentum.
    //   Bx -          The well-measured variable gradients.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //   kt_ndx -      The index of the kt variables in the variables array.
    //
    // Outputs:
    //   Bx -          The well-measured variable gradients, updated.
    //
    {
      double px = v.px();
      double py = v.py();
      double cot_theta = v.pz() / v.vect().perp();

      for (int j = 1; j <= Bx.num_col(); j++) {
        double dxnu = Bx(kt_ndx + x_offs, j);
        double dynu = Bx(kt_ndx + y_offs, j);

        if (dxnu != 0 || dynu != 0) {
          double fac = 1 / v.vect().mag();
          if (use_e_flag)
            fac = v.e() * fac * fac;
          Bx(ndx + p_offs, j) -= (px * dxnu + py * dynu) * fac;
          Bx(ndx + phi_offs, j) += py * dxnu - px * dynu;
          Bx(ndx + eta_offs, j) -= (px * dxnu + py * dynu) * cot_theta;
        }
      }
    }

    void convert_neutrino(const Fourvec_Event& ev, bool use_e_flag, Matrix& Bx, Matrix& By)
    //
    // Purpose: Carry out the Jacobian transformations the neutrino.
    //          First, convert from spherical (p, phi, theta) coordinates
    //          to rectangular (x, y, z).  Then convert from neutrino pt
    //          components to kt components.
    //
    // Inputs:
    //   ev -          The event we're fitting.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //   Bx -          The well-measured variable gradients.
    //   By -          The poorly-measured variable gradients.
    //
    // Outputs:
    //   Bx -          The well-measured variable gradients, updated.
    //   By -          The poorly-measured variable gradients, updated.
    //
    {
      int nconstraints = Bx.num_col();

      const Fourvec& nu = ev.nu();

      // convert neutrino from polar coordinates to rectangular.
      double pnu2 = nu.vect().mag2();
      double pnu = sqrt(pnu2);
      double ptnu2 = nu.vect().perp2();
      double ptnu = sqrt(ptnu2);

      // Doesn't matter whether we use E or P here, since nu is massless.

      double thfac = nu.z() / pnu2 / ptnu;
      double fac[3][3];
      fac[0][0] = nu(0) / pnu;
      fac[0][1] = nu(1) / pnu;
      fac[0][2] = nu(2) / pnu;
      fac[1][0] = -nu(1) / ptnu2;
      fac[1][1] = nu(0) / ptnu2;
      fac[1][2] = 0;
      fac[2][0] = nu(0) * thfac;
      fac[2][1] = nu(1) * thfac;
      fac[2][2] = -ptnu / pnu2;

      int kt_ndx = obj_index(ev.nobjs());
      for (int j = 1; j <= nconstraints; j++) {
        double tmp1 = fac[0][0] * Bx(kt_ndx + x_offs, j) + fac[1][0] * Bx(kt_ndx + y_offs, j) + fac[2][0] * By(nu_z, j);
        Bx(kt_ndx + y_offs, j) =
            fac[0][1] * Bx(kt_ndx + x_offs, j) + fac[1][1] * Bx(kt_ndx + y_offs, j) + fac[2][1] * By(nu_z, j);
        By(nu_z, j) = fac[0][2] * Bx(kt_ndx + x_offs, j) + fac[2][2] * By(nu_z, j);

        Bx(kt_ndx + x_offs, j) = tmp1;
      }

      // Add nu terms.
      for (int j = 0; j < ev.nobjs(); j++) {
        add_nuterm(obj_index(j), ev.obj(j).p, Bx, use_e_flag, kt_ndx);
      }
    }

    void calculate_kt_constraints(const Fourvec_Event& ev, bool use_e_flag, Row_Vector& F, Matrix& Bx)
    //
    // Purpose: Calculate the overall kt constraints (kt = 0) for the case
    //          where there is no neutrino.
    //
    // Inputs:
    //   ev -          The event we're fitting.
    //   use_e_flag -  If true, we're using E as the fit variable, otherwise p.
    //   F -           Vector of constraint values.
    //   Bx -          The well-measured variable gradients.
    //
    // Outputs:
    //   F -           Vector of constraint values, updated.
    //   Bx -          The well-measured variable gradients, updated.
    //
    {
      Fourvec tmp = Fourvec(0);
      int base = F.num_col() - 2;
      int nobjs = ev.nobjs();
      for (int j = 0; j < nobjs; j++) {
        const Fourvec& obj = ev.obj(j).p;
        tmp += obj;

        int ndx = obj_index(j);
        double p = obj.vect().mag();
        double cot_theta = obj.z() / obj.vect().perp();

        Bx(ndx + p_offs, base + 1) = obj(0) / p;
        Bx(ndx + phi_offs, base + 1) = -obj(1);
        Bx(ndx + eta_offs, base + 1) = obj(0) * cot_theta;

        Bx(ndx + p_offs, base + 2) = obj(1) / p;
        Bx(ndx + phi_offs, base + 2) = obj(0);
        Bx(ndx + eta_offs, base + 2) = obj(1) * cot_theta;

        if (use_e_flag) {
          Bx(ndx + p_offs, base + 1) *= obj.e() / p;
          Bx(ndx + p_offs, base + 2) *= obj.e() / p;
        }
      }

      int kt_ndx = obj_index(nobjs);
      Bx(kt_ndx + x_offs, base + 1) = -1;
      Bx(kt_ndx + y_offs, base + 2) = -1;

      F(base + 1) = tmp(0) - ev.kt().x();
      F(base + 2) = tmp(1) - ev.kt().y();
    }

    void ddtheta_to_ddeta(int i, double cos_theta, Matrix& Bx)
    //
    // Purpose: Do the Jacobian transformation from theta -> eta
    //          for a single object.
    //
    // Inputs:
    //   i -           The index of the object.
    //   cos_theta -   cos(theta) for the object.
    //   Bx -          The well-measured variable gradients.
    //
    // Outputs:
    //   Bx -          The well-measured variable gradients, updated.
    //
    {
      double sin_theta = sqrt(1 - cos_theta * cos_theta);
      for (int j = 1; j <= Bx.num_col(); j++)
        Bx(i, j) *= -sin_theta; /* \sin\theta = 1 / \cosh\eta */
    }

  }  // unnamed namespace

  class Fourvec_Constraint_Calculator : public Constraint_Calculator
  //
  // Purpose: Constraint evaluator.
  //
  {
  public:
    // Constructor, destructor.
    Fourvec_Constraint_Calculator(Fourvec_Event& ev,
                                  const vector<Constraint>& constraints,
                                  const Fourvec_Constrainer_Args& args);
    ~Fourvec_Constraint_Calculator() override {}

    // Evaluate constraints at the point described by X and Y (well-measured
    // and poorly-measured variables, respectively).  The results should
    // be stored in F.  BX and BY should be set to the gradients of F with
    // respect to X and Y, respectively.
    //
    // Return true if the point X, Y is accepted.
    // Return false if it is rejected (i.e., in an unphysical region).
    // The constraints need not be evaluated in that case.
    bool eval(const Column_Vector& x, const Column_Vector& y, Row_Vector& F, Matrix& Bx, Matrix& By) override;

    // Calculate the constraint functions and gradients.
    bool calculate_constraints(Row_Vector& F, Matrix& Bx, Matrix& By) const;

  private:
    // The event we're fitting.
    Fourvec_Event& _ev;

    // Vector of constraints.
    const vector<Constraint>& _constraints;

    // Argument values.
    const Fourvec_Constrainer_Args& _args;

    // The pair table.
    Pair_Table _pt;
  };

  Fourvec_Constraint_Calculator::Fourvec_Constraint_Calculator(Fourvec_Event& ev,
                                                               const vector<Constraint>& constraints,
                                                               const Fourvec_Constrainer_Args& args)
      //
      // Purpose: Constructor.
      //
      // Inputs:
      //   ev -          The event we're fitting.
      //   constraints - The list of constraints.
      //   args -        The parameter settings.
      //
      //
      : Constraint_Calculator(constraints.size() + ((ev.has_neutrino() || args.ignore_met()) ? 0 : 2)),
        _ev(ev),
        _constraints(constraints),
        _args(args),
        _pt(constraints, ev)

  {}

  bool Fourvec_Constraint_Calculator::calculate_constraints(Row_Vector& F, Matrix& Bx, Matrix& By) const
  //
  // Purpose: Calculate the constraint functions and gradients.
  //
  // Outputs:
  //   F -           Vector of constraint values.
  //   Bx -          The well-measured variable gradients.
  //   By -          The poorly-measured variable gradients.
  //
  {
    // Clear the matrices.
    Bx = Matrix(Bx.num_row(), Bx.num_col(), 0);
    By = Matrix(By.num_row(), By.num_col(), 0);
    F = Row_Vector(F.num_col(), 0);

    const double p_eps = 1e-10;

    if (_ev.has_neutrino() && _ev.nu().z() > _args.e_com()) {
      return false;
    }

    int nobjs = _ev.nobjs();

    // Reject the point if any of the momenta get too small.
    for (int j = 0; j < nobjs; j++) {
      if (_ev.obj(j).p.perp() <= p_eps || _ev.obj(j).p.e() <= p_eps) {
        return false;
      }
    }

    if (!calculate_mass_constraints(_ev, _pt, _constraints, _args.use_e(), F, Bx, By))
      return false;

    if (_ev.has_neutrino())
      convert_neutrino(_ev, _args.use_e(), Bx, By);
    else if (!_args.ignore_met()) {
      /* kt constraints */
      calculate_kt_constraints(_ev, _args.use_e(), F, Bx);
    }

    /* convert d/dtheta to d/deta */
    for (int j = 0; j < nobjs; j++) {
      ddtheta_to_ddeta(obj_index(j) + eta_offs, _ev.obj(j).p.cosTheta(), Bx);
    }

    /* handle muons */
    for (int j = 0; j < nobjs; j++) {
      const FE_Obj& obj = _ev.obj(j);
      if (obj.muon_p) {
        // Again, E vs. P doesn't matter here since we assume muons to be massless.
        double pmu2 = obj.p.vect().mag2();
        int ndx = obj_index(j) + p_offs;
        for (int k = 1; k <= Bx.num_col(); k++)
          Bx(ndx, k) = -Bx(ndx, k) * pmu2;
      }
    }

    return true;
  }

  bool Fourvec_Constraint_Calculator::eval(
      const Column_Vector& x, const Column_Vector& y, Row_Vector& F, Matrix& Bx, Matrix& By)
  //
  // Purpose: Evaluate constraints at the point described by X and Y
  //          (well-measured and poorly-measured variables, respectively).
  //          The results should be stored in F.  BX and BY should be set
  //          to the gradients of F with respect to X and Y, respectively.
  //
  // Inputs:
  //   x -           Vector of well-measured variables.
  //   y -           Vector of poorly-measured variables.
  //
  // Outputs:
  //   F -           Vector of constraint values.
  //   Bx -          The well-measured variable gradients.
  //   By -          The poorly-measured variable gradients.
  //
  {
    int nobjs = _ev.nobjs();

    const double p_eps = 1e-10;
    const double eta_max = 10;

    // Give up if we've gone into an obviously unphysical region.
    for (int j = 0; j < nobjs; j++)
      if (x(obj_index(j) + p_offs) < p_eps || fabs(x(obj_index(j) + eta_offs)) > eta_max) {
        return false;
      }

    unpack_event(_ev, x, y, _args.use_e(), _ev.has_neutrino() || !_args.ignore_met());

    return calculate_constraints(F, Bx, By);
  }

  //*************************************************************************
  // Mass calculation.
  //

  namespace {

    double calculate_sigm(
        const Matrix& Q, const Matrix& R, const Matrix& S, const Column_Vector& Bx, const Column_Vector& By)
    //
    // Purpose: Do error propagation to find the uncertainty in the final mass.
    //
    // Inputs:
    //   Q -           The final error matrix for the well-measured variables.
    //   R -           The final error matrix for the poorly-measured variables.
    //   S -           The final cross error matrix for the two sets of variables.
    //   Bx -          The well-measured variable gradients.
    //   By -          The poorly-measured variable gradients.
    //
    {
      double sig2 = scalar(Bx.T() * Q * Bx);

      if (By.num_row() > 0) {
        sig2 += scalar(By.T() * R * By);
        sig2 += 2 * scalar(Bx.T() * S * By);
      }

      assert(sig2 >= 0);
      return sqrt(sig2);
    }

    void calculate_mass(Fourvec_Event& ev,
                        const vector<Constraint>& mass_constraint,
                        const Fourvec_Constrainer_Args& args,
                        double chisq,
                        const Matrix& Q,
                        const Matrix& R,
                        const Matrix& S,
                        double& m,
                        double& sigm)
    //
    // Purpose: Calculate the final requested mass and its uncertainty.
    //
    // Inputs:
    //   ev -          The event we're fitting.
    //   mass_constraint- The description of the mass we're to calculate.
    //   args -        Parameter settings.
    //   chisq -       The chisq from the fit.
    //   Q -           The final error matrix for the well-measured variables.
    //   R -           The final error matrix for the poorly-measured variables.
    //   S -           The final cross error matrix for the two sets of variables.
    //
    // Outputs:
    //   m -           The mass.
    //   sigm -        Its uncertainty.
    //
    {
      // Don't do anything if the mass wasn't specified.
      if (mass_constraint.empty()) {
        m = 0;
        sigm = 0;
        return;
      }

      // Do the constraint calculation.
      int n_measured_vars = ev.nobjs() * 3 + 2;
      int n_unmeasured_vars = 0;
      if (ev.has_neutrino())
        n_unmeasured_vars = 1;

      Row_Vector F(1);
      Matrix Bx(n_measured_vars, 1);
      Matrix By(n_unmeasured_vars, 1);

      Fourvec_Constraint_Calculator cc(ev, mass_constraint, args);
      cc.calculate_constraints(F, Bx, By);

      // Calculate the mass.
      //assert (F(1) >= 0);
      if (F(1) >= 0.)
        m = sqrt(F(1) * 2);
      else {
        m = 0.;
        chisq = -100.;
      }

      // And the uncertainty.
      // We can only do this if the fit converged.
      if (chisq < 0)
        sigm = 0;
      else {
        //assert (F(1) > 0);
        Bx = Bx / (m);
        By = By / (m);

        sigm = calculate_sigm(Q, R, S, Bx, By);
      }
    }

  }  // unnamed namespace

  double Fourvec_Constrainer::constrain(
      Fourvec_Event& ev, double& m, double& sigm, Column_Vector& pullx, Column_Vector& pully)
  //
  // Purpose: Do a constrained fit for EV.  Returns the requested mass and
  //          its error in M and SIGM, and the pull quantities in PULLX and
  //          PULLY.  Returns the chisq; this will be < 0 if the fit failed
  //          to converge.
  //
  // Inputs:
  //   ev -          The event we're fitting.
  //
  // Outputs:
  //   ev -          The fitted event.
  //   m -           Requested invariant mass.
  //   sigm -        Uncertainty on m.
  //   pullx -       Pull quantities for well-measured variables.
  //   pully -       Pull quantities for poorly-measured variables.
  //
  // Returns:
  //   The fit chisq, or < 0 if the fit didn't converge.
  //
  {
    adjust_fourvecs(ev, _args.use_e());

    bool use_kt = ev.has_neutrino() || !_args.ignore_met();
    int nobjs = ev.nobjs();
    int n_measured_vars = nobjs * 3;
    int n_unmeasured_vars = 0;
    int n_constraints = _constraints.size();

    if (use_kt) {
      n_measured_vars += 2;

      if (ev.has_neutrino())
        n_unmeasured_vars = 1;
      else
        n_constraints += 2;
    }

    Matrix G_i(n_measured_vars, n_measured_vars);
    Diagonal_Matrix Y(n_unmeasured_vars);
    Column_Vector x(n_measured_vars);
    Column_Vector y(n_unmeasured_vars);
    pack_event(ev, _args.use_e(), use_kt, x, y, G_i, Y);

    Column_Vector xm = x;
    Column_Vector ym = y;

    // ??? Should allow for using a different underlying fitter here.
    Chisq_Constrainer fitter(_args.chisq_constrainer_args());

    Fourvec_Constraint_Calculator cc(ev, _constraints, _args);

    Matrix Q;
    Matrix R;
    Matrix S;
    double chisq = fitter.fit(cc, xm, x, ym, y, G_i, Y, pullx, pully, Q, R, S);

    unpack_event(ev, x, y, _args.use_e(), use_kt);

    calculate_mass(ev, _mass_constraint, _args, chisq, Q, R, S, m, sigm);

    return chisq;
  }

}  // namespace hitfit