StanFitter.h
Go to the documentation of this file.
1 #pragma once
2 
3 #include <iostream>
4 #include <map>
5 #include <vector>
6 
7 #include "boost/random.hpp"
10 #include "CAFAna/Core/IFitVar.h"
11 #include "CAFAna/Core/Progress.h"
12 #include "CAFAna/Fit/IFitter.h"
13 #include "CAFAna/Fit/MCMCSamples.h"
14 #include "CAFAna/Fit/StanConfig.h"
16 #include "CAFAna/Core/Utilities.h"
17 
18 #include "OscLib/IOscCalc.h"
19 
21 
22 
23 // note, Stan depends on Eigen and one method inherited from the interface
24 // has an Eigen object in its signature
25 // (also, note, we're here hiding some warnings that Stan & dependencies trigger)
26 #pragma GCC diagnostic push
27 #if __GNUC__ >= 6
28 #pragma GCC diagnostic ignored "-Wmisleading-indentation"
29 #endif
30 #pragma GCC diagnostic ignored "-Wunused-function"
31 #include "Eigen/Core"
35 #include "stan/model/prob_grad.hpp"
36 
37 #pragma GCC diagnostic pop
38 
39 namespace stan
40 {
41  namespace callbacks
42  {
43  class logger;
44  class stream_logger;
45  }
46  namespace io
47  {
48  class var_context;
49 
50  template<typename T>
51  class writer;
52  }
53  namespace mcmc
54  {
55  template <typename Model, typename RNG>
57  }
58 }
59 
60 namespace ana
61 {
62  class ISyst;
63 
64  /// Stan 'writer' callback that allows us to get ahold of the values from the sampling.
65  /// Note: doesn't own its own Samples.
67  {
68  public:
69  enum class WhichSamples
70  {
71  kWarmup,
72  kPostWarmup,
73  };
74 
75  MemoryTupleWriter(MCMCSamples * samples, MCMCSamples * warmup = nullptr)
76  : fSamples(samples), fWarmup(warmup), fWhichSamples(warmup ? WhichSamples::kWarmup : WhichSamples::kPostWarmup)
77  {
78  if (!samples && !warmup)
79  {
80  std::cerr << "MemoryTupleWriter passed null MCMCSamples for both warmup and post-warmup. Abort" << std::endl;
81  abort();
82  }
83  }
84 
85  // note: there are other overloads from the base class for writing other things like strings.
86  // we're retaining those do-nothing versions here (since anything written that way is not useful here)
87  // and only overloading the one we care about
88  using stan::callbacks::writer::operator();
89 
90  void operator()(const std::vector<double>& state) override
91  {
92  (fWhichSamples == WhichSamples::kWarmup ? fWarmup : fSamples)->AddSample(state);
93  }
94  void operator()(const std::vector<std::string>& names) override
95  {
96  (fWhichSamples == WhichSamples::kWarmup ? fWarmup : fSamples)->SetNames(names);
97  }
98 
99  template <typename Sampler>
100  void SaveSamplerState(Sampler &sampler, double samplingTime)
101  {
102  // this is from Stan
103  sampler.write_sampler_state(*this);
104 
105  auto mcmcsamples = fWhichSamples == WhichSamples::kWarmup ? fWarmup : fSamples;
106  mcmcsamples->SetHyperparams(sampler.get_nominal_stepsize(), // the 'nominal' stepsize is updated at the end of warmup
107  TMatrixDFromEigenMatrixXd(sampler.z().inv_e_metric_));
108  mcmcsamples->SetSamplingTime(samplingTime);
109  }
110 
111  void SetActiveSamples(WhichSamples s) { fWhichSamples = s; }
112  WhichSamples ActiveSamples() const { return fWhichSamples; }
113 
114 
115  private:
118 
119  enum WhichSamples fWhichSamples;
120  };
121 
122  /// \brief Fitter type that bolts the Stan fitting tools onto CAFAna.
123  ///
124  /// Unfortunately CAFAna and Stan have different case conventions
125  /// so the capitalization of the methods here is a haphazard mix :(
127  {
128  public:
129  StanFitter(const IExperiment *expt,
130  std::vector<const IFitVar *> vars,
131  std::vector<const ISyst *> systs = {});
132 
133  /// \brief Return names of parameters . (Required by Stan interface.)
134  ///
135  /// \param param_names Vector of parameter names to be filled.
136  /// \param include_tparams Unused in this implementation. (Other types of Stan models need it.)
137  /// \param include_gqs Unused in this implementation. (Other types of Stan models need it.)
138  void constrained_param_names(std::vector<std::string>& param_names,
139  bool include_tparams = true,
140  bool include_gqs = true) const;
141 
142  /// Override (and call) the base class version, doing some caching of the oscillation parameters
143  /// to avoid the problem discussed in the doc for fParamCache
144  virtual std::unique_ptr<IFitSummary> Fit(osc::IOscCalcAdjustable *seed,
145  SystShifts &bestSysts = junkShifts,
146  const SeedList& seedPts = SeedList(),
147  const std::vector<SystShifts> &systSeedPts = {},
148  Verbosity verb = kVerbose) const override;
149 
150  // un-hide the other versions of Fit() inherited from IFitter
151  using IFitter::Fit;
152 
153  /// Return vector of dimensions of each parameter. (Required by Stan interface.)
154  ///
155  /// \param dims Return vector of dimensions. Always 1 for each parameter in our application.
156  void get_dims(std::vector<std::vector<size_t> >& dims) const
157  {
158  dims = std::vector<std::vector<std::size_t>>(fVars.size() + fSysts.size(), std::vector<size_t>{1});
159  }
160 
161  /// Return names of parameters. (Required by Stan interface.)
162  void get_param_names(std::vector<std::string>& names) const;
163 
164  /// Implement workhorse method that actually calculates the posterior probability from Stan interface
165  ///
166  /// \tparam propto__
167  /// \tparam jacobian__
168  /// \tparam T__ Template type for "real" parameters. (Stan uses multiple internal types to allow for matrices and other things.)
169  /// \param params_real Real-valued parameters in Stan's unconstrained parameter space.
170  /// \param params_i__ Integer-valued parameters in Stan's unconstrained parameter space.
171  /// \param pstream__
172  /// \return
173  template <bool propto__, bool jacobian__, typename T__>
174  T__ log_prob(std::vector<T__>& params_real,
175  std::vector<int>& params_i__,
176  std::ostream* pstream__ = 0) const;
177 
178  /// Version with Matrix is req'd by Stan interface, but we won't use it
179  template <bool propto, bool jacobian, typename T_>
180  T_ log_prob(Eigen::Matrix<T_,Eigen::Dynamic,1>& params_r,
181  std::ostream* pstream = 0) const;
182 
183  /// Take the seed points and convert them into Stan's unconstrained internal space
184  void transform_inits(const stan::io::var_context& context,
185  std::vector<int>& params_int,
186  std::vector<double>& params_real,
187  std::ostream* pstream__) const;
188 
189  /// Method for writing out data inherited from Stan interface
190  template <typename RNG>
191  void write_array(RNG& base_rng__,
192  std::vector<double>& params_real,
193  std::vector<int>& params_int,
194  std::vector<double>& vars,
195  bool include_tparams__ = true,
196  bool include_gqs__ = true,
197  std::ostream* pstream__ = 0) const;
198 
199  /// Peruse the samples the MCMC generated. If you want to take ownership of them instead, see ExtractSamples()
200  const MCMCSamples& GetSamples(MemoryTupleWriter::WhichSamples ws = MemoryTupleWriter::WhichSamples::kPostWarmup) const
201  {
202  return ws == MemoryTupleWriter::WhichSamples::kWarmup ? fMCMCWarmup : fMCMCSamples;
203  }
204 
205  /// Extract the MCMC samples from the fitter (you own them now).
207  MemoryTupleWriter::WhichSamples ws = MemoryTupleWriter::WhichSamples::kPostWarmup)
208  {
209  samples = std::move(ws == MemoryTupleWriter::WhichSamples::kWarmup ? fMCMCWarmup : fMCMCSamples);
210  };
211 
212  /// Another way of extracting the MCMC samples (so you own them now) if you prefer unique_ptrs.
213  std::unique_ptr<MCMCSamples> ExtractSamples(MemoryTupleWriter::WhichSamples ws = MemoryTupleWriter::WhichSamples::kPostWarmup)
214  {
215  return std::make_unique<MCMCSamples>(std::move(ws == MemoryTupleWriter::WhichSamples::kWarmup ? fMCMCWarmup : fMCMCSamples));
216  }
217 
218  /// Supply samples and hyperparameters from a previously run warmup of MCMC.
219  /// Burden is on the user to ensure they are compatible with the sampling required.
220  /// (Call this prior to calling Fit() in order to use it.)
221  void ReuseWarmup(MCMCSamples && warmup);
222 
223  /// Change the config used for Stan. See the StanConfig struct documentation for ideas
224  void SetStanConfig(const StanConfig& cfg) { fStanConfig = cfg; }
225 
226  /// Run Stan's test of its auto-differentiation (comparing to a finite-differences calculation)
227  void TestGradients(osc::IOscCalcAdjustable *seed,
228  SystShifts &systSeed) const;
229 
230  /// \brief Return names of parameters in Stan's unconstrained space. (Required by Stan interface.)
231  ///
232  /// \param param_names Vector of parameter names to be filled.
233  /// \param include_tparams Unused in this implementation. (Other types of Stan models need it.)
234  /// \param include_gqs Unused in this implementation. (Other types of Stan models need it.)
235  void unconstrained_param_names(std::vector<std::string>& param_names,
236  bool include_tparams = true,
237  bool include_gqs = true) const
238  {
239  // other kinds of Stan models need more here, but we don't.
240  std::vector<std::string> names;
241  get_param_names(names);
242  param_names.insert(param_names.end(), names.begin(), names.end());
243  }
244 
245  private:
247  {
248  public:
249  samplecounter_callback(std::size_t warmupSampleCount, std::size_t regularSampleCount)
250  : fWarmupCounts(warmupSampleCount),
251  fRegularCounts(regularSampleCount),
252  fCurrCount(0)
253  {
254  }
255 
256  void operator()()
257  {
258  if (fWarmupCounts + fRegularCounts == 0)
259  return;
260 
261  if (fCurrCount < fWarmupCounts && !fProgressWarmup)
262  fProgressWarmup = std::make_unique<ana::Progress>(
263  "Warming up MCMC sampler (" + std::to_string(fWarmupCounts) + " samples)");
264  if (fCurrCount >= fWarmupCounts)
265  {
266  if (fProgressWarmup)
267  fProgressWarmup->Done();
268 
269  if (!fProgressNormal)
270  fProgressNormal = std::make_unique<ana::Progress>(
271  "Sampling using MCMC (" + std::to_string(fRegularCounts) + " samples)");
272 
273  if (fCurrCount - fWarmupCounts >= fRegularCounts)
274  fProgressNormal->Done();
275  else
276  fProgressNormal->SetProgress(double(++fCurrCount - fWarmupCounts) / fRegularCounts);
277  }
278  else if (fWarmupCounts > 0)
279  fProgressWarmup->SetProgress(double(++fCurrCount) / fWarmupCounts);
280  }
281 
282  private:
283  // unique_ptrs so that we can initialize as needed later
284  std::unique_ptr<ana::Progress> fProgressWarmup;
285  std::unique_ptr<ana::Progress> fProgressNormal;
286  std::size_t fWarmupCounts;
287  std::size_t fRegularCounts;
288  std::size_t fCurrCount;
289  };
290 
291  /// Helper function to build the initial seed point context for Stan
293  BuildInitContext(osc::IOscCalcAdjustable *seed,
294  const SystShifts &systSeed) const;
295 
296  /// Convert a 'normal' calculator into the Stan-aware variant used internally
297  void CreateCalculator(osc::IOscCalcAdjustable * seed) const;
298 
299  /// Unroll the parameters and stuff them into the calculator and/or syst shifts obj.
300  /// Needs to be templated because Stan wants to instantiate it with <double>,
301  /// even though that'll never actually be used (sigh)...
302  template <typename T>
303  void DecodePars(const std::vector<T>& pars, osc::IOscCalcAdjustableStan * calc) const
304  {
306  "DecodePars() can only be used with double or stan::math::var");
307  assert(pars.size() == fVars.size()+fSysts.size());
308 
309  if (fVars.size() > 0)
310  {
311  assert(calc);
312  for(unsigned int i = 0; i < fVars.size(); ++i){
313  auto val = pars[i];
314  // sadly StanFitSupport<IConstrainedFitVar> is not a derived class StanFitSupport<IFitVar>
315  auto var = dynamic_cast<const StanFitSupport<IFitVar>*>(fVars[i]);
316  auto constrVar = dynamic_cast<const StanFitSupport<IConstrainedFitVar>*>(fVars[i]);
317  assert(var || constrVar);
318  if (constrVar)
319  constrVar->SetValue(calc, val);
320  else
321  var->SetValue(calc, val);
322  }
323  }
324 
325  fShifts->ResetToNominal();
326  for(unsigned int j = 0; j < fSysts.size(); ++j){
327  auto val = pars[fVars.size()+j];
328  fShifts->SetShift(fSysts[j], val);
329  }
330  }
332  {
333  public:
334  StanFitSummary(double bestLL)
335  : fBestLL(bestLL)
336  {}
337 
338  bool IsBetterThan(const IFitSummary* other) const override
339  {
340  if (!other)
341  return true;
342 
343  if (auto otherSummary = dynamic_cast<const StanFitSummary*>(other))
344  return EvalMetricVal() > otherSummary->EvalMetricVal();
345  else
346  {
347  assert (false && "Can't compare a StanFitSummary to another kind of IFitSummary!");
348  return false; // prevent compiler warnings
349  }
350  }
351 
352  /// What's the value of the thing being minimized or maximized? (LL, chi2, etc.)
353  /// Interpretation depends on the derived class...
354  double EvalMetricVal() const override
355  {
356  return fBestLL;
357  }
358 
359  private:
360  double fBestLL;
361  };
362  /// Implement workhorse method from IFitter interface
363  std::unique_ptr<IFitSummary> FitHelperSeeded(osc::IOscCalcAdjustable *seed,
364  SystShifts &systSeed,
365  Verbosity verb) const override;
366 
367  /// Run Stan with HMC.
368  /// Lots of copy-paste from stan::services::sample::hmc_nuts_diag_e_adapt()
369  /// (in stan/services/sample/hmc_nuts_diag_e_adapt.hpp)
370  /// so that we can customize in order to save the sampler state after warmup is finished.
371  ///
372  /// \param init_writer Stream to write initialization info to
373  /// \param interrupt Object whose operator()() will be called after every sample
374  /// \param diagStream Stream to write diagnostic info to
375  /// \param logger General object to write log to
376  /// \param init_context MCMC initialization info
377  /// \param procId Process ID
378  /// \return Stan return code
379  template <typename Sampler>
380  int RunHMC(stan::callbacks::writer &init_writer,
382  std::ostream &diagStream,
383  const std::unique_ptr<stan::callbacks::stream_logger> &logger,
384  stan::io::array_var_context &init_context,
385  unsigned int procId) const;
386 
387  /// Run the HMC sampler explicitly.
388  /// Cribbed from stan::services::run_adaptive_sampler() (in stan/services/util/run_adaptive_sampler.hpp)
389  /// in order to insert save/restore state code after warmup is done.
390  ///
391  /// \param sampler The MCMC sampler object to use. (In Stan this is a templated type, but we'll only ever use this version)
392  /// \param cont_vector Starting parameter values (transformed into Stan's internal working space)
393  /// \param rng The random number generator being used
394  /// \param interrupt Object whose operator()() will be called after every sample
395  /// \param logger General object to write log to
396  /// \param diagnostic_writer Object to write diagnostics to
397  template <typename Sampler>
398  void RunSampler(Sampler& sampler,
399  std::vector<double>& cont_vector,
400  boost::ecuyer1988& rng,
401  stan::callbacks::interrupt& interrupt,
402  stan::callbacks::logger& logger,
403  stan::callbacks::writer& diagnostic_writer) const;
404  /// Helper function that actually does the unconstraining Stan needs
405  template <typename T>
406  void transform_helper(const stan::io::var_context& context,
407  stan::io::writer<double>& writer,
408  const T& var) const;
409 
410  // members below
411  mutable std::unique_ptr<osc::IOscCalcAdjustableStan> fCalc;
413  StanConfig fStanConfig; ///< Configuration passed to Stan for fitting. See the StanConfig struct documentation for ideas
416 
417  /// See MemoryTupleWriter class documentation for more info.
418  /// Pointer so it can be initialized lazily
419  mutable std::unique_ptr<MemoryTupleWriter> fValueWriter;
420 
421  /// stan::math::var objects have one 'gotcha' associated with them:
422  /// after the gradient of the log-prob is calculated, Stan internally
423  /// 'recovers' the memory associated with every stan::math::var's value
424  /// by calling recover_memory() (inside log_prob_grad()).
425  /// This is no problem for any variables that are computed every fit iteration,
426  /// like all the ones we're fitting. But any unfitted parameters that
427  /// happen to have stan::math::var type (i.e: unfitted oscillation params)
428  /// get the rug yanked from under them.
429  /// This cache is made at the beginning from all those values and any that aren't
430  /// part of the values passed by the fitter are rewritten at the beginning of
431  /// each iteration.
432  mutable std::unique_ptr<osc::IOscCalcAdjustable> fOscCalcCache;
433  };
434 
435 }
samplecounter_callback(std::size_t warmupSampleCount, std::size_t regularSampleCount)
Definition: StanFitter.h:249
Cuts and Vars for the 2020 FD DiF Study.
Definition: vars.h:6
std::vector< SystGroupDef > systs
Definition: syst_header.h:385
TMatrixD TMatrixDFromEigenMatrixXd(const Eigen::MatrixXd &mat)
Definition: EigenUtils.cxx:23
void ExtractSamples(MCMCSamples &samples, MemoryTupleWriter::WhichSamples ws=MemoryTupleWriter::WhichSamples::kPostWarmup)
Extract the MCMC samples from the fitter (you own them now).
Definition: StanFitter.h:206
Simple record of shifts applied to systematic parameters.
Definition: SystShifts.h:20
fVtxDx Fit("f")
MCMCSamples * fWarmup
Definition: StanFitter.h:117
Configuration parameters for the Stan MCMC fitter, bundled up together for easy storage and parameter...
Definition: StanConfig.h:6
OStream cerr
Definition: OStream.cxx:7
MCMCSamples fMCMCWarmup
Definition: StanFitter.h:415
void DecodePars(const std::vector< T > &pars, osc::IOscCalcAdjustableStan *calc) const
Definition: StanFitter.h:303
virtual void SetValue(osc::IOscCalcAdjustableStan *osc, stan::math::var val) const =0
void get_dims(std::vector< std::vector< size_t > > &dims) const
Definition: StanFitter.h:156
osc::OscCalcDumb calc
void operator()(const std::vector< double > &state) override
Definition: StanFitter.h:90
MCMCSamples fMCMCSamples
Definition: StanFitter.h:414
void SetStanConfig(const StanConfig &cfg)
Change the config used for Stan. See the StanConfig struct documentation for ideas.
Definition: StanFitter.h:224
WhichSamples ActiveSamples() const
Definition: StanFitter.h:112
const XML_Char * s
Definition: expat.h:262
expt
Definition: demo5.py:34
void SaveSamplerState(Sampler &sampler, double samplingTime)
Definition: StanFitter.h:100
std::unique_ptr< ana::Progress > fProgressWarmup
Definition: StanFitter.h:284
const XML_Char int const XML_Char * value
Definition: expat.h:331
void SetActiveSamples(WhichSamples s)
Definition: StanFitter.h:111
unsigned int seed
Definition: runWimpSim.h:102
const MCMCSamples & GetSamples(MemoryTupleWriter::WhichSamples ws=MemoryTupleWriter::WhichSamples::kPostWarmup) const
Peruse the samples the MCMC generated. If you want to take ownership of them instead, see ExtractSamples()
Definition: StanFitter.h:200
def callbacks(model_name, group, tensorboard=True)
Definition: regression.py:123
MCMCSamples * fSamples
Definition: StanFitter.h:116
MemoryTupleWriter(MCMCSamples *samples, MCMCSamples *warmup=nullptr)
Definition: StanFitter.h:75
const std::map< std::pair< std::string, std::string >, Variable > vars
Base class for fitters.
Definition: IFitter.h:16
const double j
Definition: BetheBloch.cxx:29
void operator()(const std::vector< std::string > &names) override
Definition: StanFitter.h:94
const IExperiment * fExpt
Definition: StanFitter.h:412
const XML_Char * context
Definition: expat.h:434
StanConfig fStanConfig
Configuration passed to Stan for fitting. See the StanConfig struct documentation for ideas...
Definition: StanFitter.h:413
std::unique_ptr< osc::IOscCalcAdjustableStan > fCalc
Definition: StanFitter.h:411
std::string pars("Th23Dmsq32")
double EvalMetricVal() const override
Definition: StanFitter.h:354
Base class defining interface for experiments.
Definition: IExperiment.h:14
std::unique_ptr< MemoryTupleWriter > fValueWriter
Definition: StanFitter.h:419
std::unique_ptr< osc::IOscCalcAdjustable > fOscCalcCache
Definition: StanFitter.h:432
assert(nhit_max >=nhit_nbins)
void unconstrained_param_names(std::vector< std::string > &param_names, bool include_tparams=true, bool include_gqs=true) const
Return names of parameters in Stan&#39;s unconstrained space. (Required by Stan interface.)
Definition: StanFitter.h:235
double T
Definition: Xdiff_gwt.C:5
std::unique_ptr< MCMCSamples > ExtractSamples(MemoryTupleWriter::WhichSamples ws=MemoryTupleWriter::WhichSamples::kPostWarmup)
Another way of extracting the MCMC samples (so you own them now) if you prefer unique_ptrs.
Definition: StanFitter.h:213
std::string to_string(ModuleType mt)
Definition: ModuleType.h:32
std::unique_ptr< ana::Progress > fProgressNormal
Definition: StanFitter.h:285
bool IsBetterThan(const IFitSummary *other) const override
Definition: StanFitter.h:338
Fitter type that bolts the Stan fitting tools onto CAFAna.
Definition: StanFitter.h:126