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 f4895240..942c5829 100644 --- a/cmake/examples/minichem.cmake +++ b/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) 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 3b4f92ed..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 @@ -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 diff --git a/examples/2024-XZhang-minichem/minichem.cpp b/examples/2024-XZhang-minichem/minichem.cpp index 444731fb..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 @@ -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 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(); - 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); + } } 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/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); +} 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);