Add ability to output cmdstan config as JSON

examples/bernoulli/.gitignore

@@ -1,2 +1,3 @@
 bernoulli
 bernoulli.hpp
+output_config.json

src/cmdstan/arguments/arg_output.hpp

@@ -6,6 +6,7 @@
 #include <cmdstan/arguments/arg_output_sig_figs.hpp>
 #include <cmdstan/arguments/arg_profile_file.hpp>
 #include <cmdstan/arguments/arg_refresh.hpp>
+#include <cmdstan/arguments/arg_single_bool.hpp>
 #include <cmdstan/arguments/categorical_argument.hpp>
 
 namespace cmdstan {
@@ -21,6 +22,11 @@ class arg_output : public categorical_argument {
     _subarguments.push_back(new arg_refresh());
     _subarguments.push_back(new arg_output_sig_figs());
     _subarguments.push_back(new arg_profile_file());
+    _subarguments.push_back(new arg_single_bool(
+        "save_cmdstan_config",
+        "Save the CmdStan configuration (parsed arguments + default values) as "
+        "JSON alongside the output files",
+        false));
   }
 };
 

src/cmdstan/arguments/argument.hpp

@@ -2,6 +2,7 @@
 #define CMDSTAN_ARGUMENTS_ARGUMENT_HPP
 
 #include <stan/callbacks/writer.hpp>
+#include <stan/callbacks/structured_writer.hpp>
 #include <iomanip>
 #include <sstream>
 #include <string>
@@ -26,6 +27,8 @@ class argument {
                      const std::string &prefix)
       = 0;
 
+  virtual void print(stan::callbacks::structured_writer &j) = 0;
+
   virtual void print_help(stan::callbacks::writer &w, const int depth,
                           const bool recurse)
       = 0;

src/cmdstan/arguments/argument_parser.hpp

@@ -140,6 +140,12 @@ class argument_parser {
     }
   }
 
+  void print(stan::callbacks::structured_writer &j) {
+    for (size_t i = 0; i < _arguments.size(); ++i) {
+      _arguments.at(i)->print(j);
+    }
+  }
+
   void print_help(stan::callbacks::writer &w, bool recurse) {
     for (size_t i = 0; i < _arguments.size(); ++i) {
       _arguments.at(i)->print_help(w, 1, recurse);

src/cmdstan/arguments/categorical_argument.hpp

@@ -28,6 +28,14 @@ class categorical_argument : public argument {
       (*it)->print(w, depth + 1, prefix);
   }
 
+  void print(stan::callbacks::structured_writer &j) {
+    j.begin_record(_name);
+    for (std::vector<argument *>::iterator it = _subarguments.begin();
+         it != _subarguments.end(); ++it)
+      (*it)->print(j);
+    j.end_record();
+  }
+
   void print_help(stan::callbacks::writer &w, const int depth,
                   const bool recurse) {
     std::string indent(indent_width * depth, ' ');

src/cmdstan/arguments/list_argument.hpp

@@ -28,6 +28,13 @@ class list_argument : public valued_argument {
     _values.at(_cursor)->print(w, depth + 1, prefix);
   }
 
+  virtual void print(stan::callbacks::structured_writer &j) {
+    j.begin_record(_name);
+    j.write("value", print_value());
+    _values.at(_cursor)->print(j);
+    j.end_record();
+  }
+
   void print_help(stan::callbacks::writer &w, int depth, bool recurse) {
     _default = _values.at(_default_cursor)->name();
 

src/cmdstan/arguments/singleton_argument.hpp

@@ -115,6 +115,10 @@ class singleton_argument : public valued_argument {
 
   bool is_default() { return _value == _default_value; }
 
+  virtual void print(stan::callbacks::structured_writer &j) {
+    j.write(_name, _value);
+  }
+
  protected:
   std::string _validity;
   virtual bool is_valid(T value) { return true; }

src/cmdstan/arguments/unvalued_argument.hpp

@@ -15,6 +15,8 @@ class unvalued_argument : public argument {
   void print(stan::callbacks::writer &w, const int depth,
              const std::string &prefix) {}
 
+  void print(stan::callbacks::structured_writer &j) {}
+
   void print_help(stan::callbacks::writer &w, const int depth,
                   const bool recurse = false) {
     std::string indent(indent_width * depth, ' ');

src/cmdstan/arguments/valued_argument.hpp

@@ -19,6 +19,10 @@ class valued_argument : public argument {
     w(message);
   }
 
+  virtual void print(stan::callbacks::structured_writer &j) {
+    j.write(_name, print_value());
+  }
+
   virtual void print_help(stan::callbacks::writer &w, const int depth,
                           const bool recurse = false) {
     std::string indent(indent_width * depth, ' ');

src/cmdstan/command.hpp

@@ -12,15 +12,11 @@
 #include <cmdstan/arguments/argument_parser.hpp>
 #include <cmdstan/command_helper.hpp>
 #include <cmdstan/return_codes.hpp>
-#include <cmdstan/write_chain.hpp>
-#include <cmdstan/write_datetime.hpp>
-#include <cmdstan/write_model_compile_info.hpp>
 #include <cmdstan/write_model.hpp>
-#include <cmdstan/write_opencl_device.hpp>
-#include <cmdstan/write_parallel_info.hpp>
-#include <cmdstan/write_profiling.hpp>
 #include <cmdstan/write_stan.hpp>
-#include <cmdstan/write_stan_flags.hpp>
+#include <cmdstan/write_config.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/core/init_threadpool_tbb.hpp>
 #include <stan/callbacks/interrupt.hpp>
 #include <stan/callbacks/json_writer.hpp>
 #include <stan/callbacks/logger.hpp>
@@ -32,7 +28,6 @@
 #include <stan/io/ends_with.hpp>
 #include <stan/io/stan_csv_reader.hpp>
 #include <stan/io/json/json_data.hpp>
-#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/model/model_base.hpp>
 #include <stan/services/diagnose/diagnose.hpp>
 #include <stan/services/experimental/advi/fullrank.hpp>
@@ -66,8 +61,6 @@
 #include <tuple>
 #include <vector>
 
-#include <stan/math/prim/core/init_threadpool_tbb.hpp>
-
 #ifdef STAN_MPI
 #include <stan/math/prim/functor/mpi_cluster.hpp>
 #include <stan/math/prim/functor/mpi_command.hpp>
@@ -210,16 +203,25 @@ int command(int argc, const char *argv[]) {
 
   std::vector<std::shared_ptr<stan::io::var_context>> init_contexts
       = get_vec_var_context(init, num_chains, id);
-  std::vector<std::string> model_compile_info = model.model_compile_info();
+
+  if (get_arg_val<bool_argument>(parser, "output", "save_cmdstan_config")) {
+    auto filename = get_basename_suffix(
+                        get_arg_val<string_argument>(parser, "output", "file"))
+                        .first
+                    + "_config.json";
+    auto ofs_args = std::make_unique<std::ofstream>(filename);
+    if (sig_figs > -1) {
+      ofs_args->precision(sig_figs);
+    }
+
+    stan::callbacks::json_writer<std::ostream> json_args(std::move(ofs_args));
+
+    write_config(json_args, parser, model);
+  }
 
   for (int i = 0; i < num_chains; ++i) {
-    write_stan(sample_writers[i]);
-    write_model(sample_writers[i], model.model_name());
-    write_datetime(sample_writers[i]);
-    parser.print(sample_writers[i]);
-    write_parallel_info(sample_writers[i]);
-    write_opencl_device(sample_writers[i]);
-    write_compile_info(sample_writers[i], model_compile_info);
+    write_config(sample_writers[i], parser, model);
+
     write_stan(diagnostic_csv_writers[i]);
     write_model(diagnostic_csv_writers[i], model.model_name());
     parser.print(diagnostic_csv_writers[i]);
@@ -279,10 +281,8 @@ int command(int argc, const char *argv[]) {
         ofs->precision(sig_figs);
       stan::callbacks::unique_stream_writer<std::ofstream> pathfinder_writer(
           std::move(ofs), "# ");
-      write_stan(pathfinder_writer);
-      write_model(pathfinder_writer, model.model_name());
-      write_datetime(pathfinder_writer);
-      parser.print(pathfinder_writer);
+      write_config(pathfinder_writer, parser, model);
+
       return_code = stan::services::pathfinder::pathfinder_lbfgs_multi<
           stan::model::model_base>(
           model, init_contexts, random_seed, id, init_radius, history_size,

src/cmdstan/command_helper.hpp

@@ -3,16 +3,6 @@
 
 #include <cmdstan/arguments/argument_parser.hpp>
 #include <cmdstan/arguments/arg_sample.hpp>
-#include <cmdstan/write_chain.hpp>
-#include <cmdstan/write_datetime.hpp>
-#include <cmdstan/write_model_compile_info.hpp>
-#include <cmdstan/write_model.hpp>
-#include <cmdstan/write_opencl_device.hpp>
-#include <cmdstan/write_parallel_info.hpp>
-#include <cmdstan/write_profiling.hpp>
-#include <cmdstan/write_stan.hpp>
-#include <cmdstan/write_stan_flags.hpp>
-#include <stan/callbacks/stream_writer.hpp>
 #include <stan/callbacks/unique_stream_writer.hpp>
 #include <stan/callbacks/json_writer.hpp>
 #include <stan/callbacks/writer.hpp>

src/cmdstan/write_config.hpp

@@ -0,0 +1,50 @@
+#ifndef CMDSTAN_WRITE_CONFIG_HPP
+#define CMDSTAN_WRITE_CONFIG_HPP
+
+#include <stan/callbacks/writer.hpp>
+#include <stan/callbacks/structured_writer.hpp>
+#include <stan/model/model_base.hpp>
+#include <stan/version.hpp>
+#include <cmdstan/arguments/argument_parser.hpp>
+#include <cmdstan/write_datetime.hpp>
+#include <cmdstan/write_model_compile_info.hpp>
+#include <cmdstan/write_model.hpp>
+#include <cmdstan/write_opencl_device.hpp>
+#include <cmdstan/write_parallel_info.hpp>
+#include <cmdstan/write_profiling.hpp>
+#include <cmdstan/write_stan.hpp>
+
+namespace cmdstan {
+
+inline void write_config(stan::callbacks::writer &writer,
+                         argument_parser &parser,
+                         stan::model::model_base &model) {
+  write_stan(writer);
+  write_model(writer, model.model_name());
+  write_datetime(writer);
+  parser.print(writer);
+  write_parallel_info(writer);
+  write_opencl_device(writer);
+  write_compile_info(writer, model);
+}
+
+inline void write_config(stan::callbacks::structured_writer &writer,
+                         argument_parser &parser,
+                         stan::model::model_base &model) {
+  writer.begin_record();
+  write_stan(writer);
+  writer.write("model_name", model.model_name());
+  writer.write("start_datetime", current_datetime());
+  parser.print(writer);
+#ifdef STAN_MPI
+  writer.write("mpi_enabled", true);
+#else
+  writer.write("mpi_enabled", false);
+#endif
+  write_opencl_device(writer);
+  write_compile_info(writer, model);
+  writer.end_record();
+}
+
+}  // namespace cmdstan
+#endif

src/cmdstan/write_datetime.hpp

@@ -9,18 +9,23 @@
 
 namespace cmdstan {
 
-void write_datetime(stan::callbacks::writer& writer) {
+std::string current_datetime() {
   const std::time_t current_datetime
       = std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
   std::tm* curr_tm = std::gmtime(&current_datetime);
   std::stringstream current_datetime_msg;
-  current_datetime_msg << "start_datetime = " << std::setfill('0')
-                       << (1900 + curr_tm->tm_year) << "-" << std::setw(2)
-                       << (curr_tm->tm_mon + 1) << "-" << std::setw(2)
-                       << curr_tm->tm_mday << " " << std::setw(2)
-                       << curr_tm->tm_hour << ":" << std::setw(2)
-                       << curr_tm->tm_min << ":" << std::setw(2)
+  current_datetime_msg << std::setfill('0') << (1900 + curr_tm->tm_year) << "-"
+                       << std::setw(2) << (curr_tm->tm_mon + 1) << "-"
+                       << std::setw(2) << curr_tm->tm_mday << " "
+                       << std::setw(2) << curr_tm->tm_hour << ":"
+                       << std::setw(2) << curr_tm->tm_min << ":" << std::setw(2)
                        << curr_tm->tm_sec << " UTC";
+  return current_datetime_msg.str();
+}
+
+void write_datetime(stan::callbacks::writer& writer) {
+  std::stringstream current_datetime_msg;
+  current_datetime_msg << "start_datetime = " << current_datetime();
   writer(current_datetime_msg.str());
 }
 

src/cmdstan/write_model_compile_info.hpp

@@ -2,16 +2,32 @@
 #define CMDSTAN_WRITE_COMPILE_INFO_HPP
 
 #include <stan/callbacks/writer.hpp>
-#include <stan/version.hpp>
-#include <string>
-#include <vector>
+#include <stan/callbacks/structured_writer.hpp>
+#include <stan/model/model_base.hpp>
 
 namespace cmdstan {
 void write_compile_info(stan::callbacks::writer& writer,
-                        std::vector<std::string>& compile_info) {
-  for (int i = 0; i < compile_info.size(); i++) {
-    writer(compile_info[i]);
+                        stan::model::model_base& model) {
+  auto compile_info = model.model_compile_info();
+  for (auto s : compile_info) {
+    writer(s);
   }
 }
+
+void write_compile_info(stan::callbacks::structured_writer& writer,
+                        stan::model::model_base& model) {
+  auto compile_info = model.model_compile_info();
+  for (auto s : compile_info) {
+    // split on "="
+    std::string::size_type pos = s.find(" = ");
+    if (pos == std::string::npos) {
+      continue;
+    }
+    std::string key = s.substr(0, pos);
+    std::string value = s.substr(pos + 3);
+    writer.write(key, value);
+  }
+}
+
 }  // namespace cmdstan
 #endif

src/cmdstan/write_opencl_device.hpp

@@ -2,6 +2,7 @@
 #define CMDSTAN_WRITE_OPENCL_DEVICE_HPP
 
 #include <stan/callbacks/writer.hpp>
+#include <stan/callbacks/structured_writer.hpp>
 #ifdef STAN_OPENCL
 #include <stan/math/opencl/opencl_context.hpp>
 #endif
@@ -27,5 +28,21 @@ void write_opencl_device(stan::callbacks::writer &writer) {
 #endif
 }
 
+void write_opencl_device(stan::callbacks::structured_writer &writer) {
+#ifdef STAN_OPENCL
+  if ((stan::math::opencl_context.platform().size() > 0)
+      && (stan::math::opencl_context.device().size() > 0)) {
+    std::stringstream msg_opencl_platform;
+    msg_opencl_platform
+        << stan::math::opencl_context.platform()[0].getInfo<CL_PLATFORM_NAME>();
+    writer.write("opencl_platform_name", msg_opencl_platform.str());
+    std::stringstream msg_opencl_device;
+    msg_opencl_device
+        << stan::math::opencl_context.device()[0].getInfo<CL_DEVICE_NAME>();
+    writer.write("opencl_device_name", msg_opencl_device.str());
+  }
+#endif
+}
+
 }  // namespace cmdstan
 #endif