From 593925c6e6d973c4760dd6f634f218166fac4655 Mon Sep 17 00:00:00 2001 From: mac/cli Date: Fri, 15 Aug 2025 12:01:36 -0400 Subject: [PATCH 1/4] fix restart --- cmake/examples/minichem.cmake | 2 +- examples/2024-XZhang-minichem/mini.inp | 1 + examples/2024-XZhang-minichem/minichem.cpp | 7 ++++++- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/cmake/examples/minichem.cmake b/cmake/examples/minichem.cmake index f4895240..552505ed 100644 --- a/cmake/examples/minichem.cmake +++ b/cmake/examples/minichem.cmake @@ -17,5 +17,5 @@ set(PNETCDF 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) diff --git a/examples/2024-XZhang-minichem/mini.inp b/examples/2024-XZhang-minichem/mini.inp index 3b4f92ed..65fc92c6 100644 --- a/examples/2024-XZhang-minichem/mini.inp +++ b/examples/2024-XZhang-minichem/mini.inp @@ -118,6 +118,7 @@ hrate = -0.1 #K/day cooling rate prad = 2.e5 Ttop_tau = 100. Tbot_tau = 100. +Tbot = 1244.91 init_Ttol = 0.1 diagnostics = div, curl, b, mean, anomaly, radflux, hydroflux diff --git a/examples/2024-XZhang-minichem/minichem.cpp b/examples/2024-XZhang-minichem/minichem.cpp index 444731fb..78a091b3 100644 --- a/examples/2024-XZhang-minichem/minichem.cpp +++ b/examples/2024-XZhang-minichem/minichem.cpp @@ -206,6 +206,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); @@ -352,7 +353,11 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) { auto temp = get_temp(pimpl->peos, phydro->w); auto temp_a = temp.accessor(); - 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; From 8fc37eeec749ff99de9f26ccaf2d4d4145ce628d Mon Sep 17 00:00:00 2001 From: mac/cli Date: Fri, 5 Sep 2025 11:09:45 -0400 Subject: [PATCH 2/4] wip --- src/snap/setup_moist_adiabatic_profile.cpp | 90 ++++++++++++++++++++++ src/snap/setup_moist_adiabatic_profile.hpp | 7 ++ 2 files changed, 97 insertions(+) create mode 100644 src/snap/setup_moist_adiabatic_profile.cpp create mode 100644 src/snap/setup_moist_adiabatic_profile.hpp diff --git a/src/snap/setup_moist_adiabatic_profile.cpp b/src/snap/setup_moist_adiabatic_profile.cpp new file mode 100644 index 00000000..0b5d4d33 --- /dev/null +++ b/src/snap/setup_moist_adiabatic_profile.cpp @@ -0,0 +1,90 @@ +// yaml +#include + +// torch +#include + +torch::Tensor setup_moist_adiabatic_profile(std::string infile) { + auto config = YAML::LoadFile(infile); + auto Ps = config["problem"]["Ps"].as(1.e5); + auto Ts = config["problem"]["Ts"].as(300.); + auto Tmin = config["problem"]["Tmin"].as(200.); + auto grav = -config["forcing"]["const-gravity"]["grav1"].as(); + + // 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(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(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(); + 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()) break; + dz = pcoord->dx1f[i].item(); + 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(); + 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; +} diff --git a/src/snap/setup_moist_adiabatic_profile.hpp b/src/snap/setup_moist_adiabatic_profile.hpp new file mode 100644 index 00000000..27eccbe1 --- /dev/null +++ b/src/snap/setup_moist_adiabatic_profile.hpp @@ -0,0 +1,7 @@ +// C/C++ +#include + +// torch +#include + +torch::Tensor setup_moist_adiabatic_profile(std::string infile); From 943bfe2ea87ade40475392b28ae871af2fac6dc5 Mon Sep 17 00:00:00 2001 From: mac/cli Date: Fri, 5 Sep 2025 15:49:05 -0400 Subject: [PATCH 3/4] wip --- cmake/athena.cmake | 3 +- cmake/examples/minichem.cmake | 1 + cmake/parameters.cmake | 6 +++ examples/2024-XZhang-minichem/mini.inp | 8 ++- examples/2024-XZhang-minichem/minichem.cpp | 36 ++++++++++++- examples/2024-XZhang-minichem/read_tensor.py | 9 ++++ src/outputs/CMakeLists.txt | 1 + src/outputs/pt.cpp | 55 ++++++++++++++++++++ 8 files changed, 115 insertions(+), 4 deletions(-) create mode 100644 examples/2024-XZhang-minichem/read_tensor.py create mode 100644 src/outputs/pt.cpp diff --git a/cmake/athena.cmake b/cmake/athena.cmake index f63cf0cd..3893e65c 100644 --- a/cmake/athena.cmake +++ b/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) diff --git a/cmake/examples/minichem.cmake b/cmake/examples/minichem.cmake index 552505ed..942c5829 100644 --- a/cmake/examples/minichem.cmake +++ b/cmake/examples/minichem.cmake @@ -14,6 +14,7 @@ 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) diff --git a/cmake/parameters.cmake b/cmake/parameters.cmake index a2bfdfb1..7f8f0df0 100644 --- a/cmake/parameters.cmake +++ b/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() diff --git a/examples/2024-XZhang-minichem/mini.inp b/examples/2024-XZhang-minichem/mini.inp index 65fc92c6..5076a33c 100644 --- a/examples/2024-XZhang-minichem/mini.inp +++ b/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 + +file_type = pt # torch's jit save +dt = 1.E4 # time increment + # #file_type = netcdf # Netcdf format output #variable = diag # diagnostic variables @@ -64,7 +68,7 @@ ox3_bc = periodic # Outer-X3 boundary condition flag nx1 = 150 -nx2 = 16 +nx2 = 32 nx3 = 1 @@ -120,5 +124,7 @@ 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 diff --git a/examples/2024-XZhang-minichem/minichem.cpp b/examples/2024-XZhang-minichem/minichem.cpp index 78a091b3..12641655 100644 --- a/examples/2024-XZhang-minichem/minichem.cpp +++ b/examples/2024-XZhang-minichem/minichem.cpp @@ -8,6 +8,9 @@ // yaml #include +// torch +#include + // athena #include #include @@ -25,6 +28,9 @@ #include #include +// kintera +#include + // snap #include #include @@ -314,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 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) @@ -401,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); + } } diff --git a/examples/2024-XZhang-minichem/read_tensor.py b/examples/2024-XZhang-minichem/read_tensor.py new file mode 100644 index 00000000..09d34097 --- /dev/null +++ b/examples/2024-XZhang-minichem/read_tensor.py @@ -0,0 +1,9 @@ +import torch + +# 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) diff --git a/src/outputs/CMakeLists.txt b/src/outputs/CMakeLists.txt index cfb9dc67..e9eb2b19 100644 --- a/src/outputs/CMakeLists.txt +++ b/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) diff --git a/src/outputs/pt.cpp b/src/outputs/pt.cpp new file mode 100644 index 00000000..fbc60a6c --- /dev/null +++ b/src/outputs/pt.cpp @@ -0,0 +1,55 @@ +// athena +#include +#include +#include +#include +#include +#include + +// kintera +#include + +// canoe +#include + +#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 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); +} From e7c18d3d8fb6ffce72e4be041459f76ade34b670 Mon Sep 17 00:00:00 2001 From: mac/cli Date: Fri, 5 Sep 2025 15:55:05 -0400 Subject: [PATCH 4/4] wip --- .../2024-XZhang-minichem/read_save_tensor.py | 29 +++++++++++++++++++ examples/2024-XZhang-minichem/read_tensor.py | 9 ------ 2 files changed, 29 insertions(+), 9 deletions(-) create mode 100644 examples/2024-XZhang-minichem/read_save_tensor.py delete mode 100644 examples/2024-XZhang-minichem/read_tensor.py diff --git a/examples/2024-XZhang-minichem/read_save_tensor.py b/examples/2024-XZhang-minichem/read_save_tensor.py new file mode 100644 index 00000000..f0572b9f --- /dev/null +++ b/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) diff --git a/examples/2024-XZhang-minichem/read_tensor.py b/examples/2024-XZhang-minichem/read_tensor.py deleted file mode 100644 index 09d34097..00000000 --- a/examples/2024-XZhang-minichem/read_tensor.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch - -# 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)