Add output APIs

python/csrc/pymesh.cpp

@@ -77,6 +77,10 @@ void bind_mesh(py::module &m) {
             }
           },
           py::arg("dx3"), py::arg("dx2"), py::arg("dx1"), py::arg("func"))
+      .ADD_OPTION(bool, snap::MeshBlockOptionsImpl, verbose)
+      .ADD_OPTION(std::string, snap::MeshBlockOptionsImpl, basename)
+      .ADD_OPTION(std::vector<snap::OutputOptions>, snap::MeshBlockOptionsImpl,
+                  outputs)
       .ADD_OPTION(harp::IntegratorOptions, snap::MeshBlockOptionsImpl, intg)
       .ADD_OPTION(snap::CoordinateOptions, snap::MeshBlockOptionsImpl, coord)
       .ADD_OPTION(snap::HydroOptions, snap::MeshBlockOptionsImpl, hydro)
@@ -146,5 +150,7 @@ void bind_mesh(py::module &m) {
       .def("print_cycle_info", &snap::MeshBlockImpl::print_cycle_info)
       .def("finalize", &snap::MeshBlockImpl::finalize)
       .def("device", &snap::MeshBlockImpl::device)
-      .def("check_redo", &snap::MeshBlockImpl::check_redo);
+      .def("check_redo", &snap::MeshBlockImpl::check_redo)
+      .def("get_outputs",
+           [](snap::MeshBlockImpl &self) { return self.output_types; });
 }

python/csrc/pyoutput.cpp

@@ -39,9 +39,14 @@ void bind_output(py::module &m) {
       .ADD_OPTION(double, snap::OutputOptionsImpl, x3_slice)
       .ADD_OPTION(std::vector<std::string>, snap::OutputOptionsImpl, variables)
       .ADD_OPTION(std::string, snap::OutputOptionsImpl, file_type)
-      .ADD_OPTION(std::string, snap::OutputOptionsImpl, data_format);
+      .ADD_OPTION(std::string, snap::OutputOptionsImpl, data_format)
+      .ADD_OPTION(bool, snap::OutputOptionsImpl, combine)
+      .ADD_OPTION(bool, snap::OutputOptionsImpl, verbose)
+      .ADD_OPTION(bool, snap::OutputOptionsImpl, super_resolution);
 
-  auto pyOutputType = py::class_<snap::OutputType>(m, "OutputType");
+  auto pyOutputType =
+      py::class_<snap::OutputType, std::shared_ptr<snap::OutputType>>(
+          m, "OutputType");
 
   pyOutputType.def(py::init<>())
       .def(py::init<snap::OutputOptions>())
@@ -50,11 +55,12 @@ void bind_output(py::module &m) {
              return fmt::format("OutputType(file_number = {}; next_time = {})",
                                 a.file_number, a.next_time);
            })
-      .def("increment_file_number",
-           [](snap::OutputType &a) { return ++a.file_number; });
+      .def_readwrite("file_number", &snap::OutputType::file_number)
+      .def_readwrite("next_time", &snap::OutputType::next_time);
 
   auto pyNetcdfOutput =
-      py::class_<snap::NetcdfOutput, snap::OutputType>(m, "NetcdfOutput");
+      py::class_<snap::NetcdfOutput, snap::OutputType,
+                 std::shared_ptr<snap::NetcdfOutput>>(m, "NetcdfOutput");
 
   pyNetcdfOutput.def(py::init<snap::OutputOptions>())
       .def("__repr__",

src/bc/internal_boundary.hpp

@@ -35,6 +35,9 @@ struct InternalBoundaryOptionsImpl {
       const YAML::Node &node);
 
   InternalBoundaryOptionsImpl() = default;
+  std::shared_ptr<InternalBoundaryOptionsImpl> clone() const {
+    return std::make_shared<InternalBoundaryOptionsImpl>(*this);
+  }
   void report(std::ostream &os) const {
     os << "-- internal boundary options --\n";
     os << "* MAXRUN = " << MAXRUN << "\n"

src/coord/coordinate.hpp

@@ -34,6 +34,10 @@ struct CoordinateOptionsImpl {
       std::string const &filename);
 
   CoordinateOptionsImpl() = default;
+  std::shared_ptr<CoordinateOptionsImpl> clone() const {
+    return std::make_shared<CoordinateOptionsImpl>(*this);
+  }
+
   void report(std::ostream &os) const {
     os << "-- coordinate options --\n";
     os << "* type = " << type() << "\n"

src/forcing/forcing.hpp

@@ -30,6 +30,10 @@ struct ConstGravityOptionsImpl {
       YAML::Node const& forcing);
 
   ConstGravityOptionsImpl() = default;
+  std::shared_ptr<ConstGravityOptionsImpl> clone() const {
+    return std::make_shared<ConstGravityOptionsImpl>(*this);
+  }
+
   void report(std::ostream& os) const {
     os << "-- constant gravity options --\n";
     os << "* grav1 = " << grav1() << "\n"
@@ -71,6 +75,9 @@ struct CoriolisOptionsImpl {
       YAML::Node const& forcing);
 
   CoriolisOptionsImpl() = default;
+  std::shared_ptr<CoriolisOptionsImpl> clone() const {
+    return std::make_shared<CoriolisOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- coriolis options --\n";
     os << "* omega1 = " << omega1() << "\n"
@@ -138,6 +145,9 @@ struct DiffusionOptionsImpl {
       YAML::Node const& forcing);
 
   DiffusionOptionsImpl() = default;
+  std::shared_ptr<DiffusionOptionsImpl> clone() const {
+    return std::make_shared<DiffusionOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- diffusion options --\n";
     os << "* K = " << K() << "\n"
@@ -176,6 +186,9 @@ struct FricHeatOptionsImpl {
       YAML::Node const& forcing);
 
   FricHeatOptionsImpl() = default;
+  std::shared_ptr<FricHeatOptionsImpl> clone() const {
+    return std::make_shared<FricHeatOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- frictional heating options --\n";
   }
@@ -211,6 +224,9 @@ struct BodyHeatOptionsImpl {
       YAML::Node const& forcing);
 
   BodyHeatOptionsImpl() = default;
+  std::shared_ptr<BodyHeatOptionsImpl> clone() const {
+    return std::make_shared<BodyHeatOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- body heating options --\n";
     os << "* dTdt = " << dTdt() << "\n"
@@ -253,6 +269,9 @@ struct TopCoolOptionsImpl {
       YAML::Node const& forcing);
 
   TopCoolOptionsImpl() = default;
+  std::shared_ptr<TopCoolOptionsImpl> clone() const {
+    return std::make_shared<TopCoolOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- top cooling options --\n";
     os << "* flux = " << flux() << "\n"
@@ -293,6 +312,9 @@ struct BotHeatOptionsImpl {
       YAML::Node const& forcing);
 
   BotHeatOptionsImpl() = default;
+  std::shared_ptr<BotHeatOptionsImpl> clone() const {
+    return std::make_shared<BotHeatOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- bottom heating options --\n";
     os << "* flux = " << flux() << "\n"
@@ -333,6 +355,9 @@ struct RelaxBotCompOptionsImpl {
       YAML::Node const& forcing);
 
   RelaxBotCompOptionsImpl() = default;
+  std::shared_ptr<RelaxBotCompOptionsImpl> clone() const {
+    return std::make_shared<RelaxBotCompOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- relax bottom composition options --\n";
     os << "* tau = " << tau() << "\n"
@@ -374,6 +399,9 @@ struct RelaxBotTempOptionsImpl {
       YAML::Node const& forcing);
 
   RelaxBotTempOptionsImpl() = default;
+  std::shared_ptr<RelaxBotTempOptionsImpl> clone() const {
+    return std::make_shared<RelaxBotTempOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- relax bottom temperature options --\n";
     os << "* tau = " << tau() << "\n"
@@ -413,6 +441,9 @@ struct RelaxBotVeloOptionsImpl {
       YAML::Node const& forcing);
 
   RelaxBotVeloOptionsImpl() = default;
+  std::shared_ptr<RelaxBotVeloOptionsImpl> clone() const {
+    return std::make_shared<RelaxBotVeloOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- relax bottom velocity options --\n";
     os << "* tau = " << tau() << "\n"
@@ -456,6 +487,9 @@ struct TopSpongeLyrOptionsImpl {
       YAML::Node const& forcing);
 
   TopSpongeLyrOptionsImpl() = default;
+  std::shared_ptr<TopSpongeLyrOptionsImpl> clone() const {
+    return std::make_shared<TopSpongeLyrOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- top sponge layer options --\n";
     os << "* tau = " << tau() << "\n"
@@ -496,6 +530,9 @@ struct BotSpongeLyrOptionsImpl {
       YAML::Node const& forcing);
 
   BotSpongeLyrOptionsImpl() = default;
+  std::shared_ptr<BotSpongeLyrOptionsImpl> clone() const {
+    return std::make_shared<BotSpongeLyrOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- bottom sponge layer options --\n";
     os << "* tau = " << tau() << "\n"
@@ -536,6 +573,9 @@ struct PlumeForcingOptionsImpl {
       YAML::Node const& forcing);
 
   PlumeForcingOptionsImpl() = default;
+  std::shared_ptr<PlumeForcingOptionsImpl> clone() const {
+    return std::make_shared<PlumeForcingOptionsImpl>(*this);
+  }
   void report(std::ostream& os) const {
     os << "-- plume forcing options --\n";
     os << "* entrainment = " << entrainment() << "\n"

src/hydro/hydro.hpp

@@ -31,6 +31,7 @@ struct HydroOptionsImpl {
       std::string const& filename, bool verbose = false);
 
   HydroOptionsImpl() = default;
+  std::shared_ptr<HydroOptionsImpl> clone() const;
   void report(std::ostream& os) const {
     os << "-- hydro options --\n";
     os << "* verbose = " << verbose() << "\n"

src/hydro/hydro_options.cpp

@@ -103,4 +103,38 @@ HydroOptions HydroOptionsImpl::from_yaml(std::string const& filename,
   return op;
 }
 
+HydroOptions HydroOptionsImpl::clone() const {
+  auto op = HydroOptionsImpl::create();
+
+  op->verbose() = verbose();
+  op->disable_flux_x1() = disable_flux_x1();
+  op->disable_flux_x2() = disable_flux_x2();
+  op->disable_flux_x3() = disable_flux_x3();
+
+  if (grav()) op->grav() = grav()->clone();
+  if (coriolis()) op->coriolis() = coriolis()->clone();
+  if (visc()) op->visc() = visc()->clone();
+  if (fricHeat()) op->fricHeat() = fricHeat()->clone();
+  if (bodyHeat()) op->bodyHeat() = bodyHeat()->clone();
+  if (topCool()) op->topCool() = topCool()->clone();
+  if (botHeat()) op->botHeat() = botHeat()->clone();
+  if (relaxBotComp()) op->relaxBotComp() = relaxBotComp()->clone();
+  if (relaxBotTemp()) op->relaxBotTemp() = relaxBotTemp()->clone();
+  if (relaxBotVelo()) op->relaxBotVelo() = relaxBotVelo()->clone();
+  if (topSpongeLyr()) op->topSpongeLyr() = topSpongeLyr()->clone();
+  if (botSpongeLyr()) op->botSpongeLyr() = botSpongeLyr()->clone();
+  if (plumeForcing()) op->plumeForcing() = plumeForcing()->clone();
+
+  // TODO(cli)
+  /*if (eos()) op->eos() = eos()->clone();
+  if (proj()) op->proj() = proj()->clone();
+  if (recon1()) op->recon1() = recon1()->clone();
+  if (recon23()) op->recon23() = recon23()->clone();
+  if (riemann()) op->riemann() = riemann()->clone();
+  if (icorr()) op->icorr() = icorr()->clone();
+  if (sed()) op->sed() = sed()->clone();*/
+
+  return op;
+}
+
 }  // namespace snap

src/mesh/meshblock.cpp

@@ -257,12 +257,16 @@ std::vector<torch::indexing::TensorIndex> MeshBlockImpl::part(
   auto slice3 = torch::indexing::Slice(start3, start3 + len3);
   auto slice4 = torch::indexing::Slice();
 
-  if (opts.ndim() == 3) {
+  if (opts.ndim() == 1) {
+    return {slice1};
+  } else if (opts.ndim() == 2) {
+    return {slice2, slice1};
+  } else if (opts.ndim() == 3) {
     return {slice3, slice2, slice1};
   } else if (opts.ndim() == 4) {
     return {slice4, slice3, slice2, slice1};
   } else {
-    throw std::runtime_error("part: ndim must be 3 or 4.");
+    throw std::runtime_error("part: ndim must be 1-4.");
   }
 }
 

src/output/netcdf.cpp

@@ -30,12 +30,16 @@ namespace snap {
 NetcdfOutput::NetcdfOutput(OutputOptions const &options_)
     : OutputType(options_) {}
 
-void NetcdfOutput::write_output_file(MeshBlockImpl *pmb, Variables const &vars,
-                                     double current_time, bool final_write) {
+void NetcdfOutput::write_output_file(MeshBlockImpl *pmb_in,
+                                     Variables const &vars, double current_time,
+                                     bool final_write) {
   // skip final write if specified
   if (final_write) return;
 
 #ifdef NETCDFOUTPUT
+  auto pmb = LoadOutputData(pmb_in, vars);
+  int rank = pmb->options->layout()->rank();
+
   auto pmeta = MetadataTable::GetInstance();
   auto phydro = pmb->phydro;
 
@@ -71,10 +75,6 @@ void NetcdfOutput::write_output_file(MeshBlockImpl *pmb, Variables const &vars,
     }
   }
 
-  // set ptrs to data in OutputData linked list, then slice/sum as needed
-  LoadOutputData(pmb, vars);
-  int rank = pmb->options->layout()->rank();
-
   // create filename: <basename>.<blockid>.<fileid>.<XXXXX>.nc
   // file_number
   std::string fname;
@@ -447,6 +447,8 @@ void NetcdfOutput::write_output_file(MeshBlockImpl *pmb, Variables const &vars,
   if (options->combine()) {
     combine_blocks(pmb, final_write);
   }
+
+  if (pmb != pmb_in) delete pmb;
 #endif  // NETCDFOUTPUT
 }
 }  // namespace snap

src/output/output_type.cpp

@@ -3,6 +3,9 @@
 #include <stdexcept>
 
 // snap
+#include <snap/mesh/meshblock.hpp>
+#include <snap/utils/refine.hpp>
+
 #include "output_formats.hpp"
 #include "output_type.hpp"
 
@@ -49,6 +52,7 @@ OutputOptions OutputOptionsImpl::from_yaml(YAML::Node const &node, int fid) {
   }
 
   options->verbose() = node["verbose"].as<bool>(false);
+  options->super_resolution() = node["super-resolution"].as<bool>(false);
 
   return options;
 }
@@ -61,16 +65,64 @@ OutputType::OutputType(OutputOptions const &options_)
       plast_data_() {  // Initialize tail node to nullptr
 }
 
-void OutputType::LoadOutputData(MeshBlockImpl *pmb, Variables const &vars) {
+MeshBlockImpl *OutputType::LoadOutputData(MeshBlockImpl *pmb_in,
+                                          Variables const &vars_in) {
   num_vars_ = 0;
   OutputData *pod;
+  MeshBlockImpl *pmb;
+  Variables vars;
+  // set ptrs to data in OutputData linked list, then slice/sum as needed
+
+  // create a refined output meshblock if super resolution is requested
+  if (options->super_resolution()) {
+    auto op = std::make_shared<MeshBlockOptionsImpl>(*(pmb_in->options));
+    op->coord() = pmb_in->options->coord()->clone();
+    if (op->coord()->nx2() > 1) op->coord()->nx2() *= 2;
+    if (op->coord()->nx3() > 1) op->coord()->nx3() *= 2;
+
+    pmb = new MeshBlockImpl(op);
+    // shall be deleted by caller of LoadOutputData
+
+    auto peos = pmb->phydro->peos;
+    auto pscalar = pmb->pscalar;
+
+    int nghost = pmb->options->coord()->nghost();
+
+    for (auto &[name, var] : vars_in) {
+      auto interior_in = pmb_in->part(
+          {0, 0, 0}, PartOptions().exterior(false).ndim(var.dim()));
+      auto interior_out =
+          pmb->part({0, 0, 0}, PartOptions().exterior(false).ndim(var.dim()));
+      auto vec = var.sizes().vec();
+
+      // dim 2
+      if (vec.size() > 1 && vec[vec.size() - 2] > 1) {
+        vec[vec.size() - 2] = 2 * (vec[vec.size() - 2] - nghost);
+      }
+
+      // dim 3
+      if (vec.size() > 2 && vec[vec.size() - 3] > 1) {
+        vec[vec.size() - 3] = 2 * (vec[vec.size() - 3] - nghost);
+      }
+
+      vars[name] = torch::zeros(vec, var.options());
+      vars[name]
+          .index(interior_out)
+          .copy_(conservative_refine(var.index(interior_in)));
+    }
+    // vars["hydro_w"] = peos->compute("U->W", {vars["hydro_u"]});
+    // scalar eos?
+  } else {
+    pmb = pmb_in;
+    vars = vars_in;
+  }
 
   loadHydroOutputData(pmb, vars);
   loadDiagOutputData(pmb, vars);
   loadScalarOutputData(pmb, vars);
-  loadUserOutputData(pmb, vars);
+  loadUserOutputData(pmb_in, vars);
 
-  return;
+  return pmb;
 }
 
 void OutputType::AppendOutputDataNode(OutputData *pnew_data) {