Add mini restart

cmake/athena.cmake

@@ -19,5 +19,6 @@ set(patch_command
 
 set(PACKAGE_NAME athenapp)
 set(REPO_URL "https://github.com/chengcli/athenapp")
-set(REPO_TAG "f9d3ae8bb84cac6954e6a224cb84bc917450e43d")
+# set(REPO_TAG "f9d3ae8bb84cac6954e6a224cb84bc917450e43d")
+set(REPO_TAG "8c221aa23155f2a1692664557b9ddf4c9fbe8262")
 add_package(${PACKAGE_NAME} ${REPO_URL} ${REPO_TAG} "${patch_command}" ON)

cmake/examples/minichem.cmake

@@ -14,8 +14,9 @@ set(NCHEM 13)
 # set(NVAPOR 1) set(NCLOUD 2) set(NPHASE_LEGACY 3)
 set(NETCDF ON)
 set(PNETCDF ON)
+set(PT ON)
 set(MPI ON)
 set(DISORT OFF)
 set(PYTHON_BINDINGS ON)
-set(RSOLVER lmars)
+set(RSOLVER hllc)
 # set(CUBED_SPHERE ON) set_if_empty(RSOLVER hllc_transform) set(GLOG ON)

cmake/parameters.cmake

@@ -42,6 +42,12 @@ else()
   find_package(cfitsio REQUIRED)
 endif()
 
+if(NOT PT OR NOT DEFINED PT)
+  set(PT_OPTION "NO_PTOUTPUT")
+else()
+  set(PT_OPTION "PTOUTPUT")
+endif()
+
 if(NOT HYDROSTATIC OR NOT DEFINED HYDROSTATIC)
   set(HYDROSTATIC_OPTION "NOT_HYDROSTATIC")
 else()

examples/2024-XZhang-minichem/mini.inp

@@ -22,6 +22,10 @@ file_type   = pnetcdf        # Netcdf format output
 variable    = uov           # diagnostic variables
 dt          = 1.E6          # time increment
 
+<output4>
+file_type   = pt            # torch's jit save
+dt          = 1.E4          # time increment
+
 #<output4>
 #file_type   = netcdf        # Netcdf format output
 #variable    = diag           # diagnostic variables
@@ -64,7 +68,7 @@ ox3_bc      = periodic    # Outer-X3 boundary condition flag
 
 <meshblock>
 nx1         = 150
-nx2         = 16
+nx2         = 32
 nx3         = 1
 
 <hydro>
@@ -118,6 +122,9 @@ hrate       = -0.1  #K/day cooling rate
 prad        = 2.e5
 Ttop_tau    = 100.
 Tbot_tau    = 100.
+Tbot        = 1244.91
+
+init_cond   = 2dmini.out4.00001
 
 init_Ttol   = 0.1
 diagnostics = div, curl, b, mean, anomaly, radflux, hydroflux

examples/2024-XZhang-minichem/minichem.cpp

@@ -8,6 +8,9 @@
 // yaml
 #include <yaml-cpp/yaml.h>
 
+// torch
+#include <torch/serialize.h>
+
 // athena
 #include <athena/athena.hpp>
 #include <athena/athena_arrays.hpp>
@@ -25,6 +28,9 @@
 #include <application/application.hpp>
 #include <application/exceptions.hpp>
 
+// kintera
+#include <kintera/utils/serialize.hpp>
+
 // snap
 #include <snap/eos/ideal_moist.hpp>
 #include <snap/mesh/meshblock.hpp>
@@ -206,6 +212,7 @@ void Mesh::InitUserMeshData(ParameterInput *pin) {
   hrate = pin->GetReal("problem", "hrate") / 86400.;
   Ttop_tau = pin->GetReal("problem", "Ttop_tau");
   Tbot_tau = pin->GetReal("problem", "Tbot_tau");
+  Tbot = pin->GetOrAddReal("problem", "Tbot", 0.);
   use_mini = pin->GetOrAddBoolean("problem", "use_mini", true);
   use_mini_ic = pin->GetOrAddBoolean("problem", "use_mini_ic", true);
   use_tra_ic = pin->GetOrAddBoolean("problem", "use_tra_ic", true);
@@ -313,6 +320,26 @@ torch::Tensor setup_moist_adiabatic_profile(std::string infile) {
   return w;
 }
 
+void load_torch_tensor(MeshBlock *pmb, std::string base_name) {
+  auto fname = base_name + "." + std::to_string(Globals::my_rank) + ".pt";
+  std::map<std::string, torch::Tensor> data;
+
+  data["hydro_u/0"] = torch::Tensor();
+  data["scalar_s/0"] = torch::Tensor();
+
+  kintera::load_tensors(data, fname);
+
+  // std::cout << data["hydro_u/0"].sizes() << std::endl;
+
+  auto hydro_u = get_all(pmb->phydro->u);
+  auto scalar_s = get_all(pmb->pscalars->s);
+
+  hydro_u.copy_(data["hydro_u/0"]);
+  scalar_s.copy_(data["scalar_s/0"]);
+
+  // std::cout << pmb->phydro->u(IDN, 0, 5, 5) << std::endl;
+}
+
 void MeshBlock::ProblemGenerator(ParameterInput *pin) {
   // auto infile = pin->GetString("problem", "config_file");
   // auto w = setup_moist_adiabatic_profile(infile)
@@ -352,7 +379,11 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) {
   auto temp = get_temp(pimpl->peos, phydro->w);
   auto temp_a = temp.accessor<Real, 3>();
 
-  Tbot = temp_a[ks][js][is];
+  if (std::abs(Tbot - temp_a[ks][js][is]) > 1.e-2) {
+    std::cout << "Please set Tbot = " << temp_a[ks][js][is]
+              << " K in the input file" << std::endl;
+    throw std::runtime_error("Wrong Tbot");
+  }
 
   std::cout << "ptra = " << ptra << " Pa" << std::endl;
 
@@ -396,6 +427,12 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) {
     }
   }
 
-  peos->PrimitiveToConserved(phydro->w, pfield->bcc, phydro->u, pcoord, is, ie,
-                             js, je, ks, ke);
+  auto init_cond = pin->GetOrAddString("problem", "init_cond", "");
+  if (init_cond != "") {
+    std::cout << "Loading initial condition from " << init_cond << std::endl;
+    load_torch_tensor(this, init_cond);
+  } else {
+    peos->PrimitiveToConserved(phydro->w, pfield->bcc, phydro->u, pcoord, is,
+                               ie, js, je, ks, ke);
+  }
 }

examples/2024-XZhang-minichem/read_save_tensor.py

@@ -0,0 +1,29 @@
+import torch
+
+
+def save_tensors(tensor_map: dict[str, torch.Tensor], filename: str):
+    class TensorModule(torch.nn.Module):
+        def __init__(self, tensors):
+            super().__init__()
+            for name, tensor in tensors.items():
+                self.register_buffer(name, tensor)
+
+    module = TensorModule(tensor_map)
+    scripted = torch.jit.script(module)  # Needed for LibTorch compatibility
+    scripted.save(filename)
+
+
+if __name__ == "__main__":
+    tensors = {
+        "foo": torch.randn(3, 4),
+        "bar": torch.randn(5, 6),
+    }
+    save_tensors(tensors, "foo_bar.pt")
+
+    # load the first file to infer shapes
+    module = torch.jit.load(pt_files[0][1])
+    data = {name: param for name, param in module.named_parameters()}
+
+    print(data.keys())
+    print(data["hydro_u/0"].shape)
+    print(data["scalar_s/1"].shape)

src/outputs/CMakeLists.txt

@@ -10,6 +10,7 @@ file(GLOB src_files
   netcdf.cpp
   pnetcdf.cpp
   fits.cpp
+  pt.cpp
   )
 
 string(TOLOWER ${CMAKE_BUILD_TYPE} buildl)

src/outputs/pt.cpp

@@ -0,0 +1,55 @@
+// athena
+#include <athena/athena.hpp>
+#include <athena/globals.hpp>
+#include <athena/hydro/hydro.hpp>
+#include <athena/mesh/mesh.hpp>
+#include <athena/outputs/user_outputs.hpp>
+#include <athena/scalars/scalars.hpp>
+
+// kintera
+#include <kintera/utils/serialize.hpp>
+
+// canoe
+#include <interface/hydro.hpp>
+
+#include "output_utils.hpp"
+
+PTOutput::PTOutput(OutputParameters oparams) : OutputType(oparams) {}
+
+void PTOutput::WriteOutputFile(Mesh *pm, ParameterInput *pin, bool flag) {
+  // create filename: "file_basename"+"."+"file_id"+"."+XXXXX+".pt",
+  // where XXXXX = 5-digit file_number
+  auto pmeta = MetadataTable::GetInstance();
+
+  std::string fname;
+  char number[6];
+  int err;
+  snprintf(number, sizeof(number), "%05d", output_params.file_number);
+
+  fname.assign(output_params.file_basename);
+  fname.append(".");
+  fname.append(output_params.file_id);
+  fname.append(".");
+  fname.append(number);
+  fname.append(".");
+  fname.append(std::to_string(Globals::my_rank));
+  fname.append(".pt");
+
+  // Loop over MeshBlocks
+  std::map<std::string, torch::Tensor> data;
+  for (int b = 0; b < pm->nblocal; ++b) {
+    MeshBlock *pmb = pm->my_blocks(b);
+
+    data["hydro_u/" + std::to_string(b)] = get_all(pmb->phydro->u);
+    data["scalar_s/" + std::to_string(b)] = get_all(pmb->pscalars->s);
+  }  // end loop over MeshBlocks
+
+  kintera::save_tensors(data, fname);
+
+  // increment counters
+  output_params.file_number++;
+  output_params.next_time += output_params.dt;
+  pin->SetInteger(output_params.block_name, "file_number",
+                  output_params.file_number);
+  pin->SetReal(output_params.block_name, "next_time", output_params.next_time);
+}

src/snap/setup_moist_adiabatic_profile.cpp

@@ -0,0 +1,90 @@
+// yaml
+#include <yaml-cpp/yaml.h>
+
+// torch
+#include <torch/torch.h>
+
+torch::Tensor setup_moist_adiabatic_profile(std::string infile) {
+  auto config = YAML::LoadFile(infile);
+  auto Ps = config["problem"]["Ps"].as<double>(1.e5);
+  auto Ts = config["problem"]["Ts"].as<double>(300.);
+  auto Tmin = config["problem"]["Tmin"].as<double>(200.);
+  auto grav = -config["forcing"]["const-gravity"]["grav1"].as<double>();
+
+  // initialize the block
+  auto block = snap::MeshBlock(snap::MeshBlockOptions::from_yaml(infile));
+
+  // useful modules
+  auto phydro = block->phydro;
+  auto pcoord = phydro->pcoord;
+  auto peos = phydro->peos;
+  auto m = block->named_modules()["hydro.eos.thermo"];
+  auto thermo_y = std::dynamic_pointer_cast<kintera::ThermoYImpl>(m);
+
+  // dimensions and indices
+  int nc3 = pcoord->x3v.size(0);
+  int nc2 = pcoord->x2v.size(0);
+  int nc1 = pcoord->x1v.size(0);
+  int ny = thermo_y->options.species().size() - 1;
+  int nvar = peos->nvar();
+
+  // construct an adiabatic atmosphere
+  kintera::ThermoX thermo_x(thermo_y->options);
+
+  auto temp = Ts * torch::ones({nc3, nc2}, torch::kFloat64);
+  auto pres = Ps * torch::ones({nc3, nc2}, torch::kFloat64);
+  auto xfrac = torch::zeros({nc3, nc2, 1 + ny}, torch::kFloat64);
+  auto w = torch::zeros({nvar, nc3, nc2, nc1}, torch::kFloat64);
+
+  // read in compositions
+  for (int i = 1; i <= ny; ++i) {
+    auto name = thermo_y->options.species()[i];
+    auto xmixr = config["problem"]["x" + name].as<double>(0.);
+    xfrac.select(2, i) = xmixr;
+  }
+
+  // dry air mole fraction
+  xfrac.select(2, 0) = 1. - xfrac.narrow(-1, 1, ny).sum(-1);
+
+  // adiabatic extrapolate half a grid to cell center
+  int is = pcoord->is();
+  int ie = pcoord->ie();
+  auto dz = pcoord->dx1f[is].item<double>();
+  thermo_x->extrapolate_ad(temp, pres, xfrac, grav, dz / 2.);
+
+  int i = is;
+  int nvapor = thermo_x->options.vapor_ids().size();
+  int ncloud = thermo_x->options.cloud_ids().size();
+  for (; i <= ie; ++i) {
+    auto conc = thermo_x->compute("TPX->V", {temp, pres, xfrac});
+
+    w[IPR].select(2, i) = pres;
+    w[IDN].select(2, i) = thermo_x->compute("V->D", {conc});
+
+    auto result = thermo_x->compute("X->Y", {xfrac});
+    w.narrow(0, snap::ICY, ny).select(3, i) =
+        thermo_x->compute("X->Y", {xfrac});
+
+    if ((temp < Tmin).any().item<double>()) break;
+    dz = pcoord->dx1f[i].item<double>();
+    thermo_x->extrapolate_ad(temp, pres, xfrac, grav, dz);
+  }
+
+  // isothermal extrapolation
+  for (; i <= ie; ++i) {
+    auto mu = (thermo_x->mu * xfrac).sum(-1);
+    dz = pcoord->dx1f[i].item<double>();
+    pres *= exp(-grav * mu * dz / (kintera::constants::Rgas * temp));
+    auto conc = thermo_x->compute("TPX->V", {temp, pres, xfrac});
+    w[IPR].select(2, i) = pres;
+    w[IDN].select(2, i) = thermo_x->compute("V->D", {conc});
+    w.narrow(0, snap::ICY, ny).select(3, i) =
+        thermo_x->compute("X->Y", {xfrac});
+  }
+
+  // add noise
+  w[IVX] += 0.01 * torch::rand_like(w[IVX]);
+  w[IVY] += 0.01 * torch::rand_like(w[IVY]);
+
+  return w;
+}

src/snap/setup_moist_adiabatic_profile.hpp

@@ -0,0 +1,7 @@
+// C/C++
+#include <string>
+
+// torch
+#include <torch/torch.h>
+
+torch::Tensor setup_moist_adiabatic_profile(std::string infile);