diff --git a/CMakeLists.txt b/CMakeLists.txt index 079fcac2..bb7b0a76 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -64,7 +64,7 @@ include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/parameters.cmake) message(STATUS "${PROJECT_NAME}-3. Setting up system libraries") find_package(Eigen3 REQUIRED) find_package(Cantera) -#find_package(Torch) +find_package(Torch) include(${CMAKE_SOURCE_DIR}/cmake/yamlpp.cmake) include(${CMAKE_SOURCE_DIR}/cmake/gtest.cmake) diff --git a/cmake/examples/jcrm.cmake b/cmake/examples/jcrm.cmake index b705863c..554f8712 100644 --- a/cmake/examples/jcrm.cmake +++ b/cmake/examples/jcrm.cmake @@ -4,11 +4,12 @@ set(NUMBER_GHOST_CELLS 3) # canoe configure -set(NVAPOR 2) -set(NCLOUD 3) +set(NVAPOR 1) +set(NCLOUD 1) +set(NPRECIP 1) set(NETCDF ON) set(PNETCDF ON) set(MPI ON) set(EQUATION_OF_STATE ideal_moist) -# set(RSOLVER lmars) -set(RSOLVER hllc_transform) +set(RSOLVER lmars) +# set(RSOLVER hllc_transform) diff --git a/cmake/examples/jupiter_h2o.cmake b/cmake/examples/jupiter_h2o.cmake index b0fd9320..9937f314 100644 --- a/cmake/examples/jupiter_h2o.cmake +++ b/cmake/examples/jupiter_h2o.cmake @@ -1,13 +1,7 @@ # configure file for test jupiter crm -macro(SET_IF_EMPTY _variable) - if("${${_variable}}" STREQUAL "") - set(${_variable} ${ARGN}) - endif() -endmacro() - # athena variables -set_if_empty(NUMBER_GHOST_CELLS 3) +set(NUMBER_GHOST_CELLS 3) # canoe configure set(NVAPOR 1) @@ -15,6 +9,7 @@ set(NCLOUD 2) set(NPHASE_LEGACY 3) set(PNETCDF ON) set(MPI ON) +set(EQUATION_OF_STATE ideal_moist) set(TASKLIST ImplicitHydroTasks) -#set_if_empty(RSOLVER hllc_transform) +# set_if_empty(RSOLVER hllc_transform) set(RSOLVER lmars) diff --git a/cmake/parameters.cmake b/cmake/parameters.cmake index 9cc737a5..d18f0d48 100644 --- a/cmake/parameters.cmake +++ b/cmake/parameters.cmake @@ -14,6 +14,9 @@ set_if_empty(NSTATIC 0) set_if_empty(NINT_PARTICLE_DATA 0) set_if_empty(NREAL_PARTICLE_DATA 0) +# dirty +set_if_empty(NBLOCKS 0) + if(NOT AFFINE OR NOT DEFINED AFFINE) set(AFFINE_OPTION "NOT_AFFINE") else() diff --git a/configure.hpp.in b/configure.hpp.in index 4ccbba6a..61287efa 100644 --- a/configure.hpp.in +++ b/configure.hpp.in @@ -1,8 +1,7 @@ -#ifndef CONFIGURE_HPP_ -#define CONFIGURE_HPP_ +#pragma once enum { - //cubesphere nblocks data + // dirty cubesphere nblocks data NBLOCKS = @NBLOCKS@, // thermodynamics data @@ -57,6 +56,19 @@ enum { // nblocks for exo3 option (USE_NBLOCKS or NOT_USE_NBLOCKS) #define @USE_NBLOCKS@ -#define ENABLE_FIX - -#endif // CONFIGURE_HPP_ +//#define ENABLE_FIX + +// logger +#define SPDLOG_ACTIVE_LEVEL SPDLOG_LEVEL_@LOGGER_LEVEL@ +#define LOG_TRACE(...) SPDLOG_LOGGER_TRACE(__VA_ARGS__) +#define LOG_DEBUG(...) SPDLOG_LOGGER_DEBUG(__VA_ARGS__) +#define LOG_INFO(...) SPDLOG_LOGGER_INFO(__VA_ARGS__) +#define LOG_WARN(...) SPDLOG_LOGGER_WARN(__VA_ARGS__) +#define LOG_ERROR(...) SPDLOG_LOGGER_ERROR(__VA_ARGS__) +#define LOG_CRITICAL(...) SPDLOG_LOGGER_CRITICAL(__VA_ARGS__) + +#ifdef __CUDACC__ + #define DISPATCH_MACRO __host__ __device__ +#else + #define DISPATCH_MACRO +#endif diff --git a/examples/2019-Li-snap/jcrm.cpp b/examples/2019-Li-snap/jcrm.cpp index 2a3497ca..f57de5df 100644 --- a/examples/2019-Li-snap/jcrm.cpp +++ b/examples/2019-Li-snap/jcrm.cpp @@ -118,7 +118,7 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) { Real Ps = P0 * pow(Ts / T0, cp / Rd); // read deep composition - std::vector yfrac(1 + NVAPOR, 0.); + std::vector yfrac(IVX, 0.); Real qdry = 1.; for (int n = 1; n <= NVAPOR; ++n) { std::string name = "q" + pthermo->SpeciesName(n); @@ -132,6 +132,7 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) { while (iter++ < max_iter) { pthermo->SetMassFractions(yfrac.data()); pthermo->EquilibrateTP(Ts, Ps); + std::cout << "Ts = " << Ts << ", Ps = " << Ps << std::endl; // stop at just above Z = 0 int i = is; diff --git a/examples/2019-Li-snap/jcrm.inp b/examples/2019-Li-snap/jcrm.inp index 70218658..d7e3b57c 100644 --- a/examples/2019-Li-snap/jcrm.inp +++ b/examples/2019-Li-snap/jcrm.inp @@ -64,14 +64,17 @@ grav_acc1 = -24.79 gamma = 1.4 # override by thermodynamics implicit_flag = 1 + +particle_radius = 0., 1.e-5 +particle_density = 1.e3, 1.e3 + packages = bot_heating, top_cooling bot_heating.flux = 50. top_cooling.flux = -50. -thermodynamics_config = jup-thermo-H2O-NH3.yaml -#microphysics_config = water_ammonia.yaml +thermodynamics_config = jup-precip.yaml qH2O = 0.0500 qNH3 = 0.0028 diff --git a/examples/2019-Li-snap/jup-precip.yaml b/examples/2019-Li-snap/jup-precip.yaml new file mode 100644 index 00000000..9a549e5d --- /dev/null +++ b/examples/2019-Li-snap/jup-precip.yaml @@ -0,0 +1,55 @@ + +phases: + - name: atm + thermo: ideal-moist + species: [dry, H2O, H2O(l), H2O(p)] + kinetics: condensation + +species: + - name: dry + composition: {H: 1.5, He: 0.15} + thermo: + model: constant-cp + T0: 273.16 + cp0: 29.1 J/mol/K + + - name: H2O + composition: {H: 2, O: 1} + thermo: + model: constant-cp + T0: 273.16 + cp0: 21.1 J/mol/K + + - name: H2O(l) + composition: {H: 2, O: 1} + thermo: + model: constant-cp + T0: 273.16 + cp0: 62.556 J/mol/K + h0: -45.103 kJ/mol # h0 = - (beta - delta) * R * T0 + equation-of-state: + model: constant-volume + molar-volume: 0.0 cm^3/mol + + - name: H2O(p) + composition: {H: 2, O: 1} + thermo: + model: constant-cp + T0: 273.16 + cp0: 62.556 J/mol/K + h0: -45.103 kJ/mol # h0 = - (beta - delta) * R * T0 + equation-of-state: + model: constant-volume + molar-volume: 0.0 cm^3/mol + +reactions: + - equation: H2O <=> H2O(l) + type: nucleation + rate-constant: {formula: ideal, T3: 273.16, P3: 611.7, beta: 24.845, delta: 4.986, minT: 100.} + + - equation: H2O(l) => H2O(p) + rate-constant: {A: 0.001, b: 0} + + - equation: H2O(p) => H2O + type: evaporation + rate-constant: {A: 0.01, b: 0} diff --git a/examples/2024-YZhou-jupiter/jupiter_crm.cpp b/examples/2024-YZhou-jupiter/jupiter_crm.cpp index 46f9ea13..b3001b13 100644 --- a/examples/2024-YZhou-jupiter/jupiter_crm.cpp +++ b/examples/2024-YZhou-jupiter/jupiter_crm.cpp @@ -150,7 +150,7 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) { Real Ts = T0 - grav / cp * x1min; Real Ps = P0 * pow(Ts / T0, cp / Rd); Real xH2O = pin->GetReal("problem", "qH2O.ppmv") / 1.E6; - Real xNH3 = pin->GetReal("problem", "qNH3.ppmv") / 1.E6; + // Real xNH3 = pin->GetReal("problem", "qNH3.ppmv") / 1.E6; while (iter++ < max_iter) { // read in vapors diff --git a/examples/2024-YZhou-jupiter/jupiter_h2o.cpp b/examples/2024-YZhou-jupiter/jupiter_h2o.cpp index 7744decc..5b3ea1b5 100644 --- a/examples/2024-YZhou-jupiter/jupiter_h2o.cpp +++ b/examples/2024-YZhou-jupiter/jupiter_h2o.cpp @@ -37,7 +37,7 @@ #include // special includes -#include +// #include Real grav, P0, T0, Tmin, prad, hrate; int iH2O; @@ -53,20 +53,21 @@ void MeshBlock::InitUserMeshBlockData(ParameterInput *pin) { void MeshBlock::UserWorkBeforeOutput(ParameterInput *pin) { auto pthermo = Thermodynamics::GetInstance(); + auto &w = phydro->w; for (int k = ks; k <= ke; ++k) for (int j = js; j <= je; ++j) for (int i = is; i <= ie; ++i) { user_out_var(0, k, j, i) = pthermo->GetTemp(w.at(k, j, i)); - user_out_var(1, k, j, i) = pthermo->PotentialTemp(w.at(k, j, i), P0); + user_out_var(1, k, j, i) = potential_temp(pthermo, w.at(k, j, i), P0); // theta_v user_out_var(2, k, j, i) = - user_out_var(1, k, j, i) * pthermo->RovRd(this, k, j, i); + user_out_var(1, k, j, i) * pthermo->RovRd(w.at(k, j, i)); // mse user_out_var(3, k, j, i) = - pthermo->MoistStaticEnergy(this, grav * pcoord->x1v(i), k, j, i); + moist_static_energy(pthermo, w.at(k, j, i), grav * pcoord->x1v(i)); // relative humidity - user_out_var(4, k, j, i) = pthermo->RelativeHumidity(this, 1, k, j, i); + user_out_var(4, k, j, i) = relative_humidity(pthermo, w.at(k, j, i))[1]; } } @@ -81,12 +82,7 @@ void Forcing(MeshBlock *pmb, Real const time, Real const dt, for (int k = ks; k <= ke; ++k) for (int j = js; j <= je; ++j) for (int i = is; i <= ie; ++i) { - auto &&air = AirParcelHelper::gather_from_primitive(pmb, k, j, i); - - air.ToMoleFraction(); - Real cv = - get_cv_mass(air, 0, pthermo->GetRd(), pthermo->GetCvRatioMass()); - + Real cv = pthermo->GetCv(w.at(k, j, i)); if (w(IPR, k, j, i) < prad) { du(IEN, k, j, i) += dt * hrate * w(IDN, k, j, i) * cv * (1. + 1.E-4 * sin(2. * M_PI * rand() / RAND_MAX)); @@ -95,6 +91,7 @@ void Forcing(MeshBlock *pmb, Real const time, Real const dt, } void Mesh::InitUserMeshData(ParameterInput *pin) { + auto pthermo = Thermodynamics::GetInstance(); grav = -pin->GetReal("hydro", "grav_acc1"); P0 = pin->GetReal("problem", "P0"); @@ -105,8 +102,7 @@ void Mesh::InitUserMeshData(ParameterInput *pin) { hrate = pin->GetReal("problem", "hrate") / 86400.; // index - auto pindex = IndexMap::GetInstance(); - iH2O = pindex->GetVaporId("H2O"); + iH2O = pthermo->SpeciesIndex("H2O"); EnrollUserExplicitSourceFunction(Forcing); } @@ -117,6 +113,7 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) { app->Log("ProblemGenerator: jupiter_crm"); auto pthermo = Thermodynamics::GetInstance(); + auto &w = phydro->w; // mesh limits Real x1min = pmy_mesh->mesh_size.x1min; @@ -137,20 +134,21 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) { AirParcel air(AirParcel::Type::MoleFrac); - // estimate surface temperature and pressure Real Ts = T0 - grav / cp * x1min; Real Ps = P0 * pow(Ts / T0, cp / Rd); - Real xH2O = pin->GetReal("problem", "qH2O.ppmv") / 1.E6; + Real yH2O = pin->GetReal("problem", "qH2O.gkg") / 1.E3; + + std::vector yfrac(IVX, 0.); + yfrac[iH2O] = yH2O; + yfrac[0] = 1. - yH2O; while (iter++ < max_iter) { - // read in vapors - air.w[iH2O] = xH2O; - air.w[IPR] = Ps; - air.w[IDN] = Ts; + pthermo->SetMassFractions(yfrac.data()); + pthermo->EquilibrateTP(Ts, Ps); // stop at just above P0 for (int i = is; i <= ie; ++i) { - pthermo->Extrapolate(&air, pcoord->dx1f(i), "pseudo", grav); + pthermo->Extrapolate_inplace(pcoord->dx1f(i), "pseudo", grav); if (air.w[IPR] < P0) break; } @@ -169,26 +167,24 @@ void MeshBlock::ProblemGenerator(ParameterInput *pin) { // construct atmosphere from bottom up for (int k = ks; k <= ke; ++k) for (int j = js; j <= je; ++j) { - air.SetZero(); - air.w[iH2O] = xH2O; - air.w[IPR] = Ps; - air.w[IDN] = Ts; + pthermo->SetMassFractions(yfrac.data()); + pthermo->EquilibrateTP(Ts, Ps); // half a grid to cell center - pthermo->Extrapolate(&air, pcoord->dx1f(is) / 2., "reversible", grav); + pthermo->Extrapolate_inplace(pcoord->dx1f(is) / 2., "reversible", grav); int i = is; for (; i <= ie; ++i) { - if (air.w[IDN] < Tmin) break; - AirParcelHelper::distribute_to_conserved(this, k, j, i, air); - pthermo->Extrapolate(&air, pcoord->dx1f(i), "pseudo", grav, 1.e-5); + pthermo->GetPrimitive(w.at(k, j, i)); + if (pthermo->GetTemp() < Tmin) break; + pthermo->Extrapolate_inplace(pcoord->dx1f(i), "pseudo", grav); } // Replace adiabatic atmosphere with isothermal atmosphere if temperature // is too low for (; i <= ie; ++i) { - AirParcelHelper::distribute_to_conserved(this, k, j, i, air); - pthermo->Extrapolate(&air, pcoord->dx1f(i), "isothermal", grav); + pthermo->GetPrimitive(w.at(k, j, i)); + pthermo->Extrapolate_inplace(pcoord->dx1f(i), "isothermal", grav); } } } diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 23d13ba6..e7f8d716 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -62,10 +62,10 @@ add_subdirectory(special) add_subdirectory(diagnostics) add_subdirectory(forcing) -if (${Torch_FOUND}) - add_subdirectory(torch) - add_subdirectory(modules) -endif() +#if (${Torch_FOUND}) +# add_subdirectory(torch) +# add_subdirectory(modules) +#endif() #add_subdirectory(nbody) if(UNIX AND NOT APPLE) diff --git a/src/add_arg.h b/src/add_arg.h new file mode 100644 index 00000000..fc87344a --- /dev/null +++ b/src/add_arg.h @@ -0,0 +1,21 @@ +#define ADD_ARG(T, name) \ + public: \ + inline auto name(const T& new_##name) -> decltype(*this) { /* NOLINT */ \ + this->name##_ = new_##name; \ + return *this; \ + } \ + inline auto name(T&& new_##name) -> decltype(*this) { /* NOLINT */ \ + this->name##_ = std::move(new_##name); \ + return *this; \ + } \ + DISPATCH_MACRO \ + inline const T& name() const noexcept { /* NOLINT */ \ + return this->name##_; \ + } \ + DISPATCH_MACRO \ + inline T& name() noexcept { /* NOLINT */ \ + return this->name##_; \ + } \ + \ + private: \ + T name##_ /* NOLINT */ diff --git a/src/checks.cpp b/src/checks.cpp index 0e4278eb..47af3196 100644 --- a/src/checks.cpp +++ b/src/checks.cpp @@ -195,7 +195,7 @@ void fix_eos_cons2prim(MeshBlock* pmb, AthenaArray& prim, } } - Real temp = pthermo->GetTemp(pmb, k, j, ifix - 1); + Real temp = pthermo->GetTemp(pmb->phydro->w.at(k, j, ifix - 1)); for (int i = ifix; i <= iu; ++i) { Real z = pcoord->x1v(i) - pcoord->x1v(ifix - 1); prim(IDN, k, j, i) = diff --git a/src/microphysics/kessler94.cpp b/src/microphysics/kessler94.cpp.old similarity index 98% rename from src/microphysics/kessler94.cpp rename to src/microphysics/kessler94.cpp.old index 1e31ddec..cb092d65 100644 --- a/src/microphysics/kessler94.cpp +++ b/src/microphysics/kessler94.cpp.old @@ -11,7 +11,7 @@ #include // microphysics -#include +//#include #include "microphysical_schemes.hpp" diff --git a/src/microphysics/microphysics.cpp b/src/microphysics/microphysics.cpp index d1c7cb36..91a0919f 100644 --- a/src/microphysics/microphysics.cpp +++ b/src/microphysics/microphysics.cpp @@ -6,6 +6,7 @@ // athena #include +#include #include #include #include @@ -19,6 +20,11 @@ #include #include #include +#include +#include + +// utils +#include // microphysics #include "microphysical_schemes.hpp" @@ -26,7 +32,7 @@ const std::string Microphysics::input_key = "microphysics_config"; -enum { NMASS = 0 }; +enum { NMASS = NCLOUD + NPRECIP }; Microphysics::Microphysics(MeshBlock *pmb, ParameterInput *pin) : pmy_block_(pmb) { @@ -58,6 +64,26 @@ Microphysics::Microphysics(MeshBlock *pmb, ParameterInput *pin) vsed_[2].ZeroClear(); systems_ = MicrophysicalSchemesFactory::Create(pmb, pin); + + // set up sedimentation options + auto str = pin->GetOrAddString("microphysics", "particle_radius", "0."); + sed_opts_.radius() = Vectorize(str.c_str(), " ,"); + + str = pin->GetOrAddString("microphysics", "particle_density", "0."); + sed_opts_.density() = Vectorize(str.c_str(), " ,"); + + auto vsed1 = pin->GetOrAddReal("microphysics", "particle_vsed1", 0.); + auto vsed2 = pin->GetOrAddReal("microphysics", "particle_vsed2", 0.); + + if (vsed1 != 0.0) { + sed_opts_.const_vsed().push_back(vsed1); + if (vsed2 != 0.0) sed_opts_.const_vsed().push_back(vsed2); + } else if (vsed2 != 0.0) { + sed_opts_.const_vsed().push_back(0.); + sed_opts_.const_vsed().push_back(vsed2); + } + + sed_opts_.gravity() = pin->GetOrAddReal("hydro", "grav_acc1", 0.0); } Microphysics::~Microphysics() { @@ -79,14 +105,44 @@ void Microphysics::EvolveSystems(AirColumn &ac, Real time, Real dt) { } void Microphysics::SetVsedFromConserved(Hydro const *phydro) { + if (NMASS == 0) return; + auto pmb = pmy_block_; int ks = pmb->ks, js = pmb->js, is = pmb->is; int ke = pmb->ke, je = pmb->je, ie = pmb->ie; - for (auto &system : systems_) { - system->SetVsedFromConserved(vsed_, phydro, ks, ke, js, je, is, ie); - } + // for (auto &system : systems_) { + // system->SetVsedFromConserved(vsed_, phydro, ks, ke, js, je, is, ie); + // } + + // Following this article + // https://stackoverflow.com/questions/59676983/in-torch-c-api-how-to-write-to-the-internal-data-of-a-tensor-fastly + + // set temperature + auto pthermo = Thermodynamics::GetInstance(); + torch::Tensor temp = + torch::zeros({pmb->ncells3, pmb->ncells2, pmb->ncells1}, torch::kDouble); + auto accessor = temp.accessor(); + + for (int k = 0; k < pmb->ncells3; ++k) + for (int j = 0; j < pmb->ncells2; ++j) + for (int i = 0; i < pmb->ncells1; ++i) { + accessor[k][j][i] = pthermo->GetTemp(phydro->w.at(k, j, i)); + } + + auto sed = Sedimentation(sed_opts_); + sed->to(torch::kCPU, torch::kDouble); + + auto diag = std::make_shared(); + (*diag)["temperature"] = + std::async(std::launch::async, [&]() { return temp; }).share(); + sed->set_shared_data(diag); + + // calculate sedimentation velocity + auto vel = sed->forward(to_torch(phydro->w)); + + vsed_[X1DIR] = to_athena(vel); // interpolation to cell interface for (int n = 0; n < NMASS; ++n) diff --git a/src/microphysics/microphysics.hpp b/src/microphysics/microphysics.hpp index 4702c35d..04977c04 100644 --- a/src/microphysics/microphysics.hpp +++ b/src/microphysics/microphysics.hpp @@ -13,6 +13,9 @@ #include #include +// microphysics +#include "sedimentation.hpp" + class MeshBlock; class ParameterInput; class MicrophysicalSchemeBase; @@ -72,6 +75,9 @@ class Microphysics { private: //! meshblock pointer MeshBlock const *pmy_block_; + + //! sedimentation options + SedimentationOptions sed_opts_; }; using MicrophysicsPtr = std::shared_ptr; diff --git a/src/microphysics/sedimentation.cpp b/src/microphysics/sedimentation.cpp new file mode 100644 index 00000000..872b8a35 --- /dev/null +++ b/src/microphysics/sedimentation.cpp @@ -0,0 +1,66 @@ +// athena +#include + +// shared +#include + +// microphysics +#include "sedimentation.hpp" + +#define sqr(x) ((x) * (x)) + +void SedimentationImpl::reset() { + if (options.radius().size() != options.density().size()) { + throw std::runtime_error( + "Sedimentation: radius and density must have the same size"); + } + + radius = register_parameter( + "radius", + torch::clamp(torch::tensor(options.radius()), options.min_radius())); + + density = register_parameter( + "density", torch::clamp(torch::tensor(options.density()), 0.)); +} + +torch::Tensor SedimentationImpl::forward(torch::Tensor hydro_w) { + const auto d = options.a_diameter(); + const auto epsilon_LJ = options.a_epsilon_LJ(); + const auto m = options.a_mass(); + + auto tem = shared_->at("temperature").get(); + + // cope with float precision + auto eta = (5.0 / 16.0) * std::sqrt(M_PI * Constants::kBoltz) * std::sqrt(m) * + torch::sqrt(tem) * + torch::pow(Constants::kBoltz / epsilon_LJ * tem, 0.16) / + (M_PI * d * d * 1.22); + + // Calculate mean free path, lambda + auto lambda = (eta * std::sqrt(M_PI * sqr(Constants::kBoltz))) / + (hydro_w[IPR] * std::sqrt(2.0 * m)); + + // Calculate Knudsen number, Kn + auto Kn = lambda / radius.view({-1, 1, 1, 1}); + + // Calculate Cunningham slip factor, beta + auto beta = 1.0 + Kn * (1.256 + 0.4 * torch::exp(-1.1 / Kn)); + + // Calculate vsed + auto vel = beta / (9.0 * eta) * + (2.0 * sqr(radius.view({-1, 1, 1, 1})) * options.gravity() * + (density.view({-1, 1, 1, 1}) - hydro_w[IDN])); + + // Set velocity to zero for particles with radius less than min_radius + for (int i = 0; i < options.radius().size(); ++i) { + if (options.radius()[i] < options.min_radius()) { + vel[i] = torch::zeros_like(vel[i]); + } + + if (options.const_vsed().size() > i) { + vel[i] += options.const_vsed()[i]; + } + } + + return torch::clamp(vel, -options.upper_limit(), options.upper_limit()); +} diff --git a/src/microphysics/sedimentation.hpp b/src/microphysics/sedimentation.hpp new file mode 100644 index 00000000..58dfad9f --- /dev/null +++ b/src/microphysics/sedimentation.hpp @@ -0,0 +1,84 @@ +#pragma once + +// C/C++ +#include + +// torch +#include +#include +#include + +// share +// clang-format off +#include +#include +// clang-format on + +using SharedData = std::shared_ptr< + std::unordered_map>>; + +struct SedimentationOptions { + //! radius and density of particles + ADD_ARG(std::vector, radius) = {10.0e-6}; + ADD_ARG(std::vector, density) = {1.0e3}; + + //! additional constant sedimentation velocity + ADD_ARG(std::vector, const_vsed) = {}; + + ADD_ARG(double, gravity) = -10.0; + + //! default H2-atmosphere properties + //! diameter of molecule [m] + ADD_ARG(double, a_diameter) = 2.827e-10; + + //! Lennard-Jones potential in J [J] + ADD_ARG(double, a_epsilon_LJ) = 59.7e-7; + + //! molecular mass of H2 [kg] + ADD_ARG(double, a_mass) = 3.34e-27; + + //! minimum radius of particles subject to sedimentation [m] + ADD_ARG(double, min_radius) = 1.e-6; + + //! upper limit of sedimentation velocity [m/s] + ADD_ARG(double, upper_limit) = 5.e3; +}; + +class SedimentationImpl : public torch::nn::Cloneable { + public: + //! particle radius and density + //! 1D tensor of number of particles + //! radius and density must have the same size + torch::Tensor radius, density; + + //! options with which this `Sedimentation` was constructed + SedimentationOptions options; + + //! Constructor to initialize the layers + explicit SedimentationImpl(SedimentationOptions const& options_) + : options(options_) { + reset(); + } + void reset() override; + + //! Calculate sedimentation velocites + /*! + * Calling this function requires the shared future diagnostic data + * `temperature` to be set. Otherwise, the function will throw an error. + * + * \param hydro_w 4D tensor of hydro variables + * \return 4D tensor of sedimentation velocities. The first dimension is the + * number of particles. + */ + torch::Tensor forward(torch::Tensor hydro_w); + + //! Set shared future diagnostic data + /*! + * \param data shared future diagnostic data + */ + void set_shared_data(SharedData data) { shared_ = data; } + + private: + SharedData shared_; +}; +TORCH_MODULE(Sedimentation); diff --git a/src/snap/athena_arrays.hpp b/src/snap/athena_arrays.hpp deleted file mode 100644 index ee76655d..00000000 --- a/src/snap/athena_arrays.hpp +++ /dev/null @@ -1,52 +0,0 @@ -#pragma once -// C/C++ -#include - -// athena -#include - -// torch -#include - -template -void AthenaArray::toDevice(c10::DeviceType device) { - if (nx5_ != 1 || nx6_ != 1) { - throw std::runtime_error("AthenaArray::toDevice: nx5 and nx6 must be 1"); - } - - int64_t str1 = 1; - int64_t str2 = nx1_; - int64_t str3 = nx2_ * nx1_; - int64_t str4 = nx3_ * nx2_ * nx1_; - - ptensor_ = std::make_shared( - torch::from_blob(pdata_, {nx4_, nx3_, nx2_, nx1_}, - {str4, str3, str2, str1}, nullptr, - torch::dtype(torch::kFloat32)) - .to(device)); -} - -template -void AthenaArray::fromDevice() { - int64_t str1 = 1; - int64_t str2 = nx1_; - int64_t str3 = nx2_ * nx1_; - int64_t str4 = nx3_ * nx2_ * nx1_; - - // create a temporary tensor holder - torch::Tensor tmp = torch::from_blob(pdata_, {nx4_, nx3_, nx2_, nx1_}, - {str4, str3, str2, str1}, nullptr, - torch::dtype(torch::kFloat32)); - - tmp.copy_(*ptensor_); -} - -template -torch::Tensor& AthenaArray::tensor() { - return *ptensor_; -} - -template -torch::Tensor const& AthenaArray::tensor() const { - return *ptensor_; -} diff --git a/src/snap/athena_torch.hpp b/src/snap/athena_torch.hpp new file mode 100644 index 00000000..03a3fbad --- /dev/null +++ b/src/snap/athena_torch.hpp @@ -0,0 +1,54 @@ +#pragma once + +// C/C++ +#include + +// athena +#include + +// torch +#include + +torch::Tensor to_torch(AthenaArray const& w) { + auto nx1 = w.GetDim1(); + auto nx2 = w.GetDim2(); + auto nx3 = w.GetDim3(); + auto nx4 = w.GetDim4(); + auto nx5 = w.GetDim5(); + auto nx6 = w.GetDim6(); + + if (nx5 != 1 || nx6 != 1) { + throw std::runtime_error("to_torch: nx5 and nx6 must be 1"); + } + + int64_t str1 = 1; + int64_t str2 = nx1; + int64_t str3 = nx2 * nx1; + int64_t str4 = nx3 * nx2 * nx1; + + return torch::from_blob(const_cast(w.data()), {nx4, nx3, nx2, nx1}, + {str4, str3, str2, str1}, nullptr, + torch::dtype(torch::kDouble)); +} + +AthenaArray to_athena(torch::Tensor tensor) { + auto nx1 = tensor.size(3); + auto nx2 = tensor.size(2); + auto nx3 = tensor.size(1); + auto nx4 = tensor.size(0); + + int64_t str1 = 1; + int64_t str2 = nx1; + int64_t str3 = nx2 * nx1; + int64_t str4 = nx3 * nx2 * nx1; + + AthenaArray p(nx4, nx3, nx2, nx1); + + // create a temporary tensor holder + torch::Tensor tmp = + torch::from_blob(p.data(), {nx4, nx3, nx2, nx1}, {str4, str3, str2, str1}, + nullptr, torch::dtype(torch::kDouble)); + tmp.copy_(tensor); + + return p; +} diff --git a/src/snap/riemann/hllc_transform.cpp b/src/snap/riemann/hllc_transform.cpp index 5710ea86..e4a5053d 100644 --- a/src/snap/riemann/hllc_transform.cpp +++ b/src/snap/riemann/hllc_transform.cpp @@ -34,8 +34,7 @@ void Hydro::RiemannSolver(const int k, const int j, const int il, const int iu, AthenaArray &wr, AthenaArray &flx, const AthenaArray &dxw) { auto pthermo = Thermodynamics::GetInstance(); - // FIXME - // auto pmicro = pmy_block->pimpl->pmicro; + auto pmicro = pmy_block->pimpl->pmicro; int ivy = IVX + ((ivx - IVX) + 1) % 3; int ivz = IVX + ((ivx - IVX) + 2) % 3; @@ -240,17 +239,18 @@ void Hydro::RiemannSolver(const int k, const int j, const int il, const int iu, flx(ivz, k, j, i) = flxi[IVZ]; flx(IEN, k, j, i) = flxi[IEN]; - /* tracer flux - Real tfl[(NMASS)], tfr[(NMASS)], tflxi[(NMASS)]; - for (int n = 0; n < NMASS; ++n) { - Real vsed = pmicro->vsedf[dir](n, k, j, i); - tfl[n] = wli[IDN] * rdl * vxl; - tfr[n] = wri[IDN] * rdr * vxr; - tflxi[n] = sl * tfl[n] + sr * tfr[n]; - if (dir==0){ // The direction of the sedimentation velocity - pmicro->mass_flux[dir](n, k, j, i) = tflxi[n] + vsed * wri[IDN] * rdr; - }else{ - pmicro->mass_flux[dir](n, k, j, i) = tflxi[n]; + // sedimentation flux + if (ivx == IVX) { + if (i == iu) return; + Real hr = (wri[IPR] * (kappar + 1.) + KE_r) / wri[IDN]; + for (int n = 0; n < NCLOUD + NPRECIP; ++n) { + Real rho = wri[IDN] * wri[1 + NVAPOR + n]; + auto vsed = pmicro->vsedf[0](n, k, j, i); + flx(1 + NVAPOR + n, k, j, i) += vsed * rho; + flx(ivx, k, j, i) += vsed * rho * wri[ivx]; + flx(ivy, k, j, i) += vsed * rho * wri[ivy]; + flx(ivz, k, j, i) += vsed * rho * wri[ivz]; + flx(IEN, k, j, i) += vsed * rho * hr; } } } diff --git a/src/snap/riemann/lmars.cpp b/src/snap/riemann/lmars.cpp index 6b5daa94..24a41a5c 100644 --- a/src/snap/riemann/lmars.cpp +++ b/src/snap/riemann/lmars.cpp @@ -18,6 +18,9 @@ // snap #include +// microphysics +#include + void Hydro::RiemannSolver(int const k, int const j, int const il, int const iu, int const ivx, AthenaArray &wl, AthenaArray &wr, AthenaArray &flx, @@ -27,6 +30,7 @@ void Hydro::RiemannSolver(int const k, int const j, int const il, int const iu, int dir = ivx - IVX; auto pthermo = Thermodynamics::GetInstance(); + auto pmicro = pmy_block->pimpl->pmicro; MeshBlock *pmb = pmy_block; Real rhobar, pbar, cbar, ubar, hl, hr; @@ -94,19 +98,14 @@ void Hydro::RiemannSolver(int const k, int const j, int const il, int const iu, // FIXME: remove this? if (ivx == IVX) { if (i == iu) return; - Real vsed_[NHYDRO]; - for (int n = 0; n <= NVAPOR + NCLOUD; ++n) vsed_[n] = 0.; - for (int n = 1 + NVAPOR + NCLOUD; n <= NVAPOR + NCLOUD + NPRECIP; ++n) { - vsed_[n] = -10.; - } - - for (int n = 1 + NVAPOR; n <= NVAPOR + NCLOUD + NPRECIP; ++n) { - Real rho = wri[IDN] * wri[n]; - flx(n, k, j, i) += vsed_[n] * rho; - flx(ivx, k, j, i) += vsed_[n] * rho * wri[ivx]; - flx(ivy, k, j, i) += vsed_[n] * rho * wri[ivy]; - flx(ivz, k, j, i) += vsed_[n] * rho * wri[ivz]; - flx(IEN, k, j, i) += vsed_[n] * rho * hr; + for (int n = 0; n < NCLOUD + NPRECIP; ++n) { + Real rho = wri[IDN] * wri[1 + NVAPOR + n]; + auto vsed = pmicro->vsedf[0](n, k, j, i); + flx(1 + NVAPOR + n, k, j, i) += vsed * rho; + flx(ivx, k, j, i) += vsed * rho * wri[ivx]; + flx(ivy, k, j, i) += vsed * rho * wri[ivy]; + flx(ivz, k, j, i) += vsed * rho * wri[ivz]; + flx(IEN, k, j, i) += vsed * rho * hr; } } } diff --git a/src/tasklist/implicit_hydro_tasks.cpp b/src/tasklist/implicit_hydro_tasks.cpp index aca21464..ee643084 100644 --- a/src/tasklist/implicit_hydro_tasks.cpp +++ b/src/tasklist/implicit_hydro_tasks.cpp @@ -282,7 +282,7 @@ TaskStatus ImplicitHydroTasks::ImplicitCorrection(MeshBlock *pmb, int stage) { TaskStatus ImplicitHydroTasks::UpdateAllConserved(MeshBlock *pmb, int stage) { if (stage <= nstages) { - // pmb->pimpl->pmicro->SetVsedFromConserved(pmb->phydro); + pmb->pimpl->pmicro->SetVsedFromConserved(pmb->phydro); } else { return TaskStatus::fail; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 490f8675..3a2153f7 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -10,11 +10,12 @@ enable_testing() # torch tests if (${Torch_FOUND}) - setup_test(test_eos) - setup_test(test_weno) - setup_test(test_reconstruct) - setup_test(test_riemann) - setup_test(test_hydro) + #setup_test(test_eos) + #setup_test(test_weno) + #setup_test(test_reconstruct) + #setup_test(test_riemann) + #setup_test(test_hydro) + setup_test(test_sedimentation) endif() # athena tests diff --git a/tests/device_testing.hpp b/tests/device_testing.hpp new file mode 100644 index 00000000..f677e495 --- /dev/null +++ b/tests/device_testing.hpp @@ -0,0 +1,55 @@ +#pragma once + +// external +#include + +// torch +#include + +struct Parameters { + torch::DeviceType device_type; + torch::Dtype dtype; +}; + +void PrintTo(const Parameters& param, std::ostream* os) { + std::string device_str = torch::Device(param.device_type).str(); + std::string dtype_str = torch::toString(param.dtype); + *os << "Device: " << device_str << ", Dtype: " << dtype_str; +} + +class DeviceTest : public testing::TestWithParam { + protected: + torch::Device device = torch::kCPU; + torch::Dtype dtype = torch::kFloat32; + + void SetUp() override { + // Get the current parameters + auto param = GetParam(); + device = torch::Device(param.device_type); + dtype = param.dtype; + + // Check if the device is available, and skip the test if not + if (device.type() == torch::kCUDA && !torch::cuda::is_available()) { + GTEST_SKIP() << "CUDA is not available, skipping test."; + } + + if (device.type() == torch::kMPS && !torch::hasMPS()) { + GTEST_SKIP() << "MPS is not available, skipping test."; + } + } +}; + +INSTANTIATE_TEST_SUITE_P( + DeviceAndDtype, DeviceTest, + testing::Values(Parameters{torch::kCPU, torch::kFloat32}, + Parameters{torch::kCPU, torch::kFloat64}, + Parameters{torch::kMPS, torch::kFloat32}, + Parameters{torch::kCUDA, torch::kFloat32}, + Parameters{torch::kCUDA, torch::kFloat64}), + [](const testing::TestParamInfo& info) { + std::string name = torch::Device(info.param.device_type).str(); + name += "_"; + name += torch::toString(info.param.dtype); + std::replace(name.begin(), name.end(), '.', '_'); + return name; + }); diff --git a/tests/test_sedimentation.cpp b/tests/test_sedimentation.cpp new file mode 100644 index 00000000..5ae396d5 --- /dev/null +++ b/tests/test_sedimentation.cpp @@ -0,0 +1,79 @@ +// external +#include + +// athena +#include + +// microphysics +#include + +// tests +#include "device_testing.hpp" + +TEST_P(DeviceTest, forward) { + auto options = SedimentationOptions(); + auto sed = Sedimentation(options); + + int nhydro = 5; + int nx3 = 1; + int nx2 = 1; + int nx1 = 9; + + sed->to(device, dtype); + + auto w = + torch::randn({nhydro, nx3, nx2, nx1}, torch::device(device).dtype(dtype)); + w[IDN] = w[IDN].abs(); + w[IPR] = w[IPR].abs(); + + auto diag = std::make_shared(); + (*diag)["temperature"] = + std::async(std::launch::async, [&]() { + return 300. * + torch::randn({nx3, nx2, nx1}, torch::device(device).dtype(dtype)) + .abs(); + }).share(); + sed->set_shared_data(diag); + + auto vel = sed->forward(w); + std::cout << "vel = " << vel << std::endl; +} + +TEST_P(DeviceTest, two_species_forward) { + auto options = SedimentationOptions(); + options.radius({0.0e-6, 10.e-6}); + options.density({1.0e3, 1.0e3}); + + auto sed = Sedimentation(options); + + int nhydro = 5; + int nx3 = 1; + int nx2 = 1; + int nx1 = 9; + + sed->to(device, dtype); + + auto w = + torch::randn({nhydro, nx3, nx2, nx1}, torch::device(device).dtype(dtype)); + w[IDN] = w[IDN].abs(); + w[IPR] = w[IPR].abs(); + + auto diag = std::make_shared(); + (*diag)["temperature"] = + std::async(std::launch::async, [&]() { + return 300. * + torch::randn({nx3, nx2, nx1}, torch::device(device).dtype(dtype)) + .abs(); + }).share(); + sed->set_shared_data(diag); + + auto vel = sed->forward(w); + std::cout << "vel = " << vel << std::endl; +} + +int main(int argc, char **argv) { + testing::InitGoogleTest(&argc, argv); + // start_logging(argc, argv); + + return RUN_ALL_TESTS(); +}