diff --git a/python/csrc/pymesh.cpp b/python/csrc/pymesh.cpp index 0d5c8b0..1304253 100644 --- a/python/csrc/pymesh.cpp +++ b/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::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; }); } diff --git a/python/csrc/pyoutput.cpp b/python/csrc/pyoutput.cpp index 795be13..2cfc8f4 100644 --- a/python/csrc/pyoutput.cpp +++ b/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, 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_(m, "OutputType"); + auto pyOutputType = + py::class_>( + m, "OutputType"); pyOutputType.def(py::init<>()) .def(py::init()) @@ -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_(m, "NetcdfOutput"); + py::class_>(m, "NetcdfOutput"); pyNetcdfOutput.def(py::init()) .def("__repr__", diff --git a/src/bc/internal_boundary.hpp b/src/bc/internal_boundary.hpp index cbe4834..abea1b6 100644 --- a/src/bc/internal_boundary.hpp +++ b/src/bc/internal_boundary.hpp @@ -35,6 +35,9 @@ struct InternalBoundaryOptionsImpl { const YAML::Node &node); InternalBoundaryOptionsImpl() = default; + std::shared_ptr clone() const { + return std::make_shared(*this); + } void report(std::ostream &os) const { os << "-- internal boundary options --\n"; os << "* MAXRUN = " << MAXRUN << "\n" diff --git a/src/coord/coordinate.hpp b/src/coord/coordinate.hpp index 20080df..d6d2374 100644 --- a/src/coord/coordinate.hpp +++ b/src/coord/coordinate.hpp @@ -34,6 +34,10 @@ struct CoordinateOptionsImpl { std::string const &filename); CoordinateOptionsImpl() = default; + std::shared_ptr clone() const { + return std::make_shared(*this); + } + void report(std::ostream &os) const { os << "-- coordinate options --\n"; os << "* type = " << type() << "\n" diff --git a/src/forcing/forcing.hpp b/src/forcing/forcing.hpp index e1b0e9d..a09fbfd 100644 --- a/src/forcing/forcing.hpp +++ b/src/forcing/forcing.hpp @@ -30,6 +30,10 @@ struct ConstGravityOptionsImpl { YAML::Node const& forcing); ConstGravityOptionsImpl() = default; + std::shared_ptr clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*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 clone() const { + return std::make_shared(*this); + } void report(std::ostream& os) const { os << "-- plume forcing options --\n"; os << "* entrainment = " << entrainment() << "\n" diff --git a/src/hydro/hydro.hpp b/src/hydro/hydro.hpp index 8d30c4a..b7d7c0b 100644 --- a/src/hydro/hydro.hpp +++ b/src/hydro/hydro.hpp @@ -31,6 +31,7 @@ struct HydroOptionsImpl { std::string const& filename, bool verbose = false); HydroOptionsImpl() = default; + std::shared_ptr clone() const; void report(std::ostream& os) const { os << "-- hydro options --\n"; os << "* verbose = " << verbose() << "\n" diff --git a/src/hydro/hydro_options.cpp b/src/hydro/hydro_options.cpp index d7be324..14782af 100644 --- a/src/hydro/hydro_options.cpp +++ b/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 diff --git a/src/mesh/meshblock.cpp b/src/mesh/meshblock.cpp index bde508e..49c370d 100644 --- a/src/mesh/meshblock.cpp +++ b/src/mesh/meshblock.cpp @@ -257,12 +257,16 @@ std::vector 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."); } } diff --git a/src/output/netcdf.cpp b/src/output/netcdf.cpp index 29b7203..d2265d7 100644 --- a/src/output/netcdf.cpp +++ b/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: ....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 diff --git a/src/output/output_type.cpp b/src/output/output_type.cpp index 506851e..98cd35e 100644 --- a/src/output/output_type.cpp +++ b/src/output/output_type.cpp @@ -3,6 +3,9 @@ #include // snap +#include +#include + #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(false); + options->super_resolution() = node["super-resolution"].as(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(*(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) { diff --git a/src/output/output_type.hpp b/src/output/output_type.hpp index 0d58ddd..117f7ef 100644 --- a/src/output/output_type.hpp +++ b/src/output/output_type.hpp @@ -82,6 +82,7 @@ struct OutputOptionsImpl { ADD_ARG(bool, combine) = true; ADD_ARG(bool, verbose) = false; + ADD_ARG(bool, super_resolution) = false; }; using OutputOptions = std::shared_ptr; @@ -138,7 +139,7 @@ class OutputType { // functions //! \brief Create doubly linked list of OutputData's containing requested //! variables - void LoadOutputData(MeshBlockImpl *pmb, Variables const &vars); + MeshBlockImpl *LoadOutputData(MeshBlockImpl *pmb, Variables const &vars); void AppendOutputDataNode(OutputData *pdata); void ReplaceOutputDataNode(OutputData *pold, OutputData *pnew); diff --git a/src/scalar/scalar.hpp b/src/scalar/scalar.hpp index ebc3212..24ed050 100644 --- a/src/scalar/scalar.hpp +++ b/src/scalar/scalar.hpp @@ -24,6 +24,7 @@ struct ScalarOptionsImpl { } static std::shared_ptr from_yaml( std::string const& filename, bool verbose = false); + ScalarOptionsImpl() = default; void report(std::ostream& os) const { os << "-- scalar options --\n"; diff --git a/src/utils/refine.cpp b/src/utils/refine.cpp index 19d6617..08450d2 100644 --- a/src/utils/refine.cpp +++ b/src/utils/refine.cpp @@ -3,18 +3,29 @@ namespace snap { torch::Tensor conservative_refine(torch::Tensor x) { - auto opts_y = torch::nn::functional::InterpolateFuncOptions() - .scale_factor(std::vector({2.0, 2.0, 1.0})) - .mode(torch::kTrilinear) - .align_corners(false); + double scale_x3_y = x.size(-3) > 1 ? 2.0 : 1.0; + double scale_x2_y = x.size(-2) > 1 ? 2.0 : 1.0; + double scale_x3_x = x.size(-3) > 1 ? 0.5 : 1.0; + double scale_x2_x = x.size(-2) > 1 ? 0.5 : 1.0; - auto opts_x = torch::nn::functional::InterpolateFuncOptions() - .scale_factor(std::vector({0.5, 0.5, 1.0})) - .mode(torch::kArea); + auto opts_y = + torch::nn::functional::InterpolateFuncOptions() + .scale_factor(std::vector({scale_x3_y, scale_x2_y, 1.0})) + .mode(torch::kTrilinear) + .align_corners(false) + .recompute_scale_factor(false); - auto opts_dy = torch::nn::functional::InterpolateFuncOptions() - .scale_factor(std::vector({2.0, 2.0, 1.0})) - .mode(torch::kArea); + auto opts_x = + torch::nn::functional::InterpolateFuncOptions() + .scale_factor(std::vector({scale_x3_x, scale_x2_x, 1.0})) + .mode(torch::kArea) + .recompute_scale_factor(false); + + auto opts_dy = + torch::nn::functional::InterpolateFuncOptions() + .scale_factor(std::vector({scale_x3_y, scale_x2_y, 1.0})) + .mode(torch::kArea) + .recompute_scale_factor(false); // bilinear refine int dim = 0; @@ -36,9 +47,13 @@ torch::Tensor conservative_refine(torch::Tensor x) { } torch::Tensor conservative_coarsen(torch::Tensor x) { + double scale_x3 = x.size(-3) > 1 ? 0.5 : 1.0; + double scale_x2 = x.size(-2) > 1 ? 0.5 : 1.0; + auto opts = torch::nn::functional::InterpolateFuncOptions() - .scale_factor(std::vector({0.5, 0.5, 1.0})) - .mode(torch::kArea); + .scale_factor(std::vector({scale_x3, scale_x2, 1.0})) + .mode(torch::kArea) + .recompute_scale_factor(false); int dim = 0; while (x.dim() < 5) { ++dim;