diff --git a/examples/jupiter_1d_rt_toon.yaml b/examples/jupiter_1d_rt_toon.yaml new file mode 100644 index 0000000..ed9ac92 --- /dev/null +++ b/examples/jupiter_1d_rt_toon.yaml @@ -0,0 +1,67 @@ +species: + - name: H2 + composition: {H: 2} + - name: He + composition: {He: 1} + - name: CH4 + composition: {C: 1, H: 4} + - name: H2O + composition: {H: 2, O: 1} + - name: NH3 + composition: {N: 1, H: 3} + +geometry: + cells: + nx1: 100 + +opacities: + h2o_line: + type: molecule-line + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [H2O] + + nh3_line: + type: molecule-line + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [NH3] + + ch4_line: + type: molecule-line + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [CH4] + + h2_h2_cia: + type: molecule-cia + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [H2] + + h2_he_cia: + type: molecule-cia + data: [h2_0p9_he_0p1_ch4_0p004_h2o_0p002_nh3_0p0003_xsection_120_730K_1_100bar_20_2500cm1.nc] + species: [H2, He] + +bands: + - name: dump_lw + range: [20.0, 2500.0] + opacities: [h2o_line, nh3_line, ch4_line, h2_h2_cia, h2_he_cia] + solver: toon + flags: planck + nwave: 2481 + nstr: 4 + +forcing: + const-gravity: + grav1: -24.79 + +problem: + xH2: 0.9 + xHe: 0.1 + xCH4: 0.004 + xH2O: 0.002 + xNH3: 0.0003 + + Pbot: 100.e5 + Ptop: 1.e4 + + Tbot: 600. + Ttop: 130. diff --git a/python/csrc/pyrtsolver.cpp b/python/csrc/pyrtsolver.cpp index f8b0934..2d22134 100644 --- a/python/csrc/pyrtsolver.cpp +++ b/python/csrc/pyrtsolver.cpp @@ -26,10 +26,16 @@ void bind_rtsolver(py::module& m) { }) .ADD_OPTION(std::vector, harp::ToonMcKay89OptionsImpl, wave_lower) .ADD_OPTION(std::vector, harp::ToonMcKay89OptionsImpl, wave_upper) - .ADD_OPTION(bool, harp::ToonMcKay89OptionsImpl, zenith_correction) + .ADD_OPTION(std::string, harp::ToonMcKay89OptionsImpl, flags) .ADD_OPTION(int, harp::ToonMcKay89OptionsImpl, top_emission_flag) - .ADD_OPTION(bool, harp::ToonMcKay89OptionsImpl, hard_surface) - .ADD_OPTION(bool, harp::ToonMcKay89OptionsImpl, delta_eddington_lw); + .def("planck", &harp::ToonMcKay89OptionsImpl::planck) + .def("zenith_correction", + &harp::ToonMcKay89OptionsImpl::zenith_correction) + .def("hard_surface", &harp::ToonMcKay89OptionsImpl::hard_surface) + .def("delta_eddington_lw", + &harp::ToonMcKay89OptionsImpl::delta_eddington_lw) + .def("has_flag", &harp::ToonMcKay89OptionsImpl::has_flag, + py::arg("flag")); torch::python::bind_module(m, "ToonMcKay89") .def(py::init<>()) diff --git a/src/radiation/radiation_band.cpp b/src/radiation/radiation_band.cpp index 3179535..48fe80a 100644 --- a/src/radiation/radiation_band.cpp +++ b/src/radiation/radiation_band.cpp @@ -110,6 +110,37 @@ RadiationBandOptions RadiationBandOptionsImpl::from_yaml( } } else if (op->solver_name() == "toon") { op->toon() = ToonMcKay89OptionsImpl::create(); + if (band["flags"]) { + op->toon()->flags(trim_copy(band["flags"].as())); + } + if (band["toon"]) { + auto const toon = band["toon"]; + op->toon()->top_emission_flag( + toon["top_emission_flag"].as(op->toon()->top_emission_flag())); + if (toon["flags"]) { + op->toon()->flags(trim_copy(toon["flags"].as())); + } + std::vector legacy_flags; + if (toon["zenith_correction"].as(false)) { + legacy_flags.push_back("zenith_correction"); + } + if (toon["hard_surface"].as(false)) { + legacy_flags.push_back("hard_surface"); + } + if (toon["delta_eddington_lw"].as(false)) { + legacy_flags.push_back("delta_eddington_lw"); + } + if (!legacy_flags.empty()) { + auto flags = op->toon()->flags(); + for (auto const& flag : legacy_flags) { + if (!flags.empty()) { + flags += ","; + } + flags += flag; + } + op->toon()->flags(flags); + } + } op->set_wave_lower(std::vector(op->nwave(), wmin)); op->set_wave_upper(std::vector(op->nwave(), wmax)); } else if (op->solver_name() == "twostr") { @@ -180,6 +211,7 @@ void RadiationBandImpl::reset() { register_module("solver", rtsolver.ptr()); } else if (options->solver_name() == "toon") { rtsolver = torch::nn::AnyModule(ToonMcKay89(options->toon())); + register_module("solver", rtsolver.ptr()); } else { TORCH_CHECK(false, "Unknown solver: ", options->solver_name()); } diff --git a/src/rtsolver/toon_mckay89.cpp b/src/rtsolver/toon_mckay89.cpp index 417b60b..4dfbbb9 100644 --- a/src/rtsolver/toon_mckay89.cpp +++ b/src/rtsolver/toon_mckay89.cpp @@ -22,20 +22,20 @@ torch::Tensor ToonMcKay89Impl::forward(torch::Tensor prop, torch::optional temf) { // check dimensions TORCH_CHECK(prop.dim() == 4, "ToonMcKay89::forward: prop.dim() != 4"); - TORCH_CHECK(prop.size(3) >= 3, "ToonMcKay89::forward: prop.size(3) < 3"); + TORCH_CHECK(prop.size(3) >= 1, "ToonMcKay89::forward: prop.size(3) < 1"); int nwave = prop.size(0); int ncol = prop.size(1); int nlyr = prop.size(2); - // optical thickness - auto tau = prop.select(-1, 0).flip(-1); - - // single scattering albedo - auto w0 = prop.select(-1, 1).flip(-1); - - // scattering asymmetry parameter - auto g = prop.select(-1, 2).flip(-1); + // The low-level Toon kernels expect at least 3 optical-property channels: + // tau, single-scattering albedo, and asymmetry parameter. Pure-absorption + // inputs provide only tau, so pad missing scattering fields with zeros. + if (prop.size(3) < 3) { + auto padded = torch::zeros({nwave, ncol, nlyr, 3}, prop.options()); + padded.narrow(-1, 0, prop.size(3)).copy_(prop); + prop = padded; + } // add slash if (bname.size() > 0 && bname.back() != '/') { @@ -84,7 +84,7 @@ torch::Tensor ToonMcKay89Impl::forward(torch::Tensor prop, // would have strides that don't match the pointer arithmetic in the impl. prop = prop.contiguous(); - if (!temf.has_value()) { // shortwave + if (!options->planck()) { // shortwave auto iter = at::TensorIteratorConfig() .resize_outputs(false) .check_all_same_dtype(true) @@ -106,6 +106,8 @@ torch::Tensor ToonMcKay89Impl::forward(torch::Tensor prop, options->delta_eddington_lw()); return flx; } else { // longwave + TORCH_CHECK(temf.has_value(), + "ToonMcKay89::forward: 'planck' flag requires temf"); /*Eigen::VectorXd temp(nlay + 1); Eigen::VectorXd be(nlay + 1); for (int i = 0; i < nlay + 1; ++i) { diff --git a/src/rtsolver/toon_mckay89.hpp b/src/rtsolver/toon_mckay89.hpp index b17e0cb..fedd88f 100644 --- a/src/rtsolver/toon_mckay89.hpp +++ b/src/rtsolver/toon_mckay89.hpp @@ -7,6 +7,11 @@ #include #include +// C/C++ +#include +#include +#include + // harp #include @@ -21,7 +26,36 @@ struct ToonMcKay89OptionsImpl { std::shared_ptr clone() const { return std::make_shared(*this); } + bool has_flag(std::string flag) const { + auto normalize = [](std::string value) { + auto start = value.find_first_not_of(" \t\n\r"); + if (start == std::string::npos) { + return std::string{}; + } + auto end = value.find_last_not_of(" \t\n\r"); + value = value.substr(start, end - start + 1); + for (auto& c : value) { + c = static_cast(std::tolower(static_cast(c))); + } + return value; + }; + flag = normalize(flag); + std::stringstream ss(flags()); + std::string item; + while (std::getline(ss, item, ',')) { + if (normalize(item) == flag) { + return true; + } + } + return false; + } + bool planck() const { return has_flag("planck"); } + bool zenith_correction() const { return has_flag("zenith_correction"); } + bool hard_surface() const { return has_flag("hard_surface"); } + bool delta_eddington_lw() const { return has_flag("delta_eddington_lw"); } void report(std::ostream& os) const { + os << "* flags = " << flags() << "\n"; + os << "* planck = " << planck() << "\n"; os << "* zenith_correction = " << zenith_correction() << "\n"; os << "* top_emission_flag = " << top_emission_flag() << "\n"; os << "* hard_surface = " << hard_surface() << "\n"; @@ -42,8 +76,8 @@ struct ToonMcKay89OptionsImpl { //! set upper wavenumber(length) at each bin ADD_ARG(std::vector, wave_upper) = {}; - //! zenith correction - ADD_ARG(bool, zenith_correction) = false; + //! comma-separated solver flags + ADD_ARG(std::string, flags) = ""; //! top emission flag //! 0 = no incoming radiation at TOA (Toon 1989 default, GCM mode) @@ -51,15 +85,6 @@ struct ToonMcKay89OptionsImpl { //! -1 = auto-compute from first layer: tau_top = dtau[0]*exp(-1), //! Btop = (1-exp(-tau_top/mu)) * B(T_top) (FMS-style) ADD_ARG(int, top_emission_flag) = 0; - - //! hard surface (true = terrestrial with emissivity BC, - //! false = gas giant with Planck gradient BC) - ADD_ARG(bool, hard_surface) = false; - - //! apply delta-Eddington scaling in longwave - //! (true = rescale w0, dtau, g as in FMS; false = use raw values as in - //! PICASO) - ADD_ARG(bool, delta_eddington_lw) = false; }; using ToonMcKay89Options = std::shared_ptr; diff --git a/tests/test_toon.cpp b/tests/test_toon.cpp index 50dd5bf..1e60864 100644 --- a/tests/test_toon.cpp +++ b/tests/test_toon.cpp @@ -1,7 +1,12 @@ // external #include +// C/C++ +#include + // harp +#include +#include #include // tests @@ -9,6 +14,76 @@ using namespace harp; +TEST(ToonConfig, from_yaml_reads_toon_options) { + auto yaml_path = + std::filesystem::path(__FILE__).parent_path() / "toon_test.yaml"; + auto rad = harp::RadiationOptionsImpl::from_yaml(yaml_path.string()); + ASSERT_EQ(rad->bands().size(), 2u); + + auto const& op = rad->bands().front(); + + ASSERT_EQ(op->solver_name(), "toon"); + ASSERT_NE(op->toon(), nullptr); + EXPECT_EQ(op->toon()->flags(), + "planck,zenith_correction,hard_surface,delta_eddington_lw"); + EXPECT_TRUE(op->toon()->planck()); + EXPECT_TRUE(op->toon()->zenith_correction()); + EXPECT_EQ(op->toon()->top_emission_flag(), -1); + EXPECT_TRUE(op->toon()->hard_surface()); + EXPECT_TRUE(op->toon()->delta_eddington_lw()); + EXPECT_EQ(op->toon()->wave_lower(), + (std::vector{200.0, 200.0, 200.0})); + EXPECT_EQ(op->toon()->wave_upper(), + (std::vector{2000.0, 2000.0, 2000.0})); +} + +TEST(ToonConfig, radiation_band_registers_solver_module) { + auto op = harp::RadiationBandOptionsImpl::create(); + op->name("B_toon"); + op->solver_name("toon"); + op->toon(harp::ToonMcKay89OptionsImpl::create()); + op->toon()->flags("planck"); + op->nwave(2); + op->ncol(1); + op->nlyr(3); + op->wavenumber({300.0, 900.0}); + op->weight({600.0, 600.0}); + op->set_wave_lower({0.0, 600.0}); + op->set_wave_upper({600.0, 1200.0}); + + harp::RadiationBand band(op); + + EXPECT_NO_THROW({ (void)band->named_modules()["solver"]; }); +} + +TEST(ToonConfig, planck_flag_controls_thermal_emission) { + auto wave_lower = std::vector{200.0, 500.0}; + auto wave_upper = std::vector{500.0, 1000.0}; + + auto sw_op = harp::ToonMcKay89OptionsImpl::create(); + sw_op->wave_lower(wave_lower); + sw_op->wave_upper(wave_upper); + harp::ToonMcKay89 sw_toon(sw_op); + + auto lw_op = harp::ToonMcKay89OptionsImpl::create(); + lw_op->wave_lower(wave_lower); + lw_op->wave_upper(wave_upper); + lw_op->flags("planck"); + harp::ToonMcKay89 lw_toon(lw_op); + + auto prop = torch::zeros({2, 1, 3, 3}, torch::kFloat64); + prop.select(-1, 0).fill_(0.2); + auto temf = torch::ones({1, 4}, torch::kFloat64) * 300.0; + std::map sw_bc; + std::map lw_bc; + + auto sw_result = sw_toon(prop, &sw_bc, /*band=*/"", temf); + auto lw_result = lw_toon(prop, &lw_bc, /*band=*/"", temf); + + EXPECT_TRUE(torch::allclose(sw_result, torch::zeros_like(sw_result))); + EXPECT_GT(torch::max(torch::abs(lw_result)).item(), 0.0); +} + TEST_P(DeviceTest, simple_toon_mckay89) { auto op = harp::ToonMcKay89OptionsImpl::create(); op->wave_lower({200., 500., 1000.}); @@ -50,6 +125,7 @@ TEST_P(DeviceTest, simple_toon_mckay89) { } auto temf = torch::ones({ncol, nlyr + 1}, prop.options()) * tem_K; + op->flags("planck"); for (auto [w0, g] : {std::make_pair(0.1, 0.5), std::make_pair(0.5, 0.5), std::make_pair(0.9, 0.5)}) { diff --git a/tests/test_toon.py b/tests/test_toon.py index 5cc7eaf..86e6008 100644 --- a/tests/test_toon.py +++ b/tests/test_toon.py @@ -227,6 +227,7 @@ def test_toon_lw_isothermal(): opt = pyharp.ToonMcKay89Options() opt.wave_lower(wave_lo) opt.wave_upper(wave_hi) + opt.flags("planck") toon = pyharp.ToonMcKay89(opt) dop = pydisort.DisortOptions() @@ -273,11 +274,11 @@ def test_toon_lw_isothermal(): ) def test_toon_lw_isothermal_profile_matches_disort(): - """Isothermal longwave profiles should be flat and symmetric. + """The `planck` flag should gate thermal emission and match DISORT aloft. - With no scattering and matching top/surface temperatures, both upward and - downward longwave fluxes should be nearly level-independent, and Toon - should closely match the DISORT reference profile. + Passing `temf` without the `planck` flag should keep Toon on the + shortwave-only path. With `planck` enabled, the emitted upward longwave + profile should closely match the DISORT reference for an isothermal column. """ nwave = 1 ncol = 1 @@ -289,10 +290,16 @@ def test_toon_lw_isothermal_profile_matches_disort(): wave_lo = [500.0] wave_hi = [1000.0] - opt = pyharp.ToonMcKay89Options() - opt.wave_lower(wave_lo) - opt.wave_upper(wave_hi) - toon = pyharp.ToonMcKay89(opt) + sw_opt = pyharp.ToonMcKay89Options() + sw_opt.wave_lower(wave_lo) + sw_opt.wave_upper(wave_hi) + sw_toon = pyharp.ToonMcKay89(sw_opt) + + lw_opt = pyharp.ToonMcKay89Options() + lw_opt.wave_lower(wave_lo) + lw_opt.wave_upper(wave_hi) + lw_opt.flags("planck") + toon = pyharp.ToonMcKay89(lw_opt) dop = pydisort.DisortOptions() dop.upward(True) @@ -313,39 +320,31 @@ def test_toon_lw_isothermal_profile_matches_disort(): } temf = torch.ones(ncol, nlyr + 1) * temp_K + sw_result = sw_toon(prop, temf=temf, **bc) toon_result = toon(prop, temf=temf, **bc) disort_result = disort(prop, temf=temf, **bc) + assert torch.allclose(sw_result, torch.zeros_like(sw_result)) + toon_up = toon_result[0, 0, :, 0] - toon_dn = toon_result[0, 0, :, 1] disort_up = disort_result[0, 0, :, 0] - disort_dn = disort_result[0, 0, :, 1] toon_scale = float(torch.mean(toon_up)) disort_scale = float(torch.mean(disort_up)) - assert torch.max(torch.abs(toon_up - toon_dn)) / toon_scale < 1e-2, ( - f"Toon isothermal LW up/down mismatch too large: " - f"{float(torch.max(torch.abs(toon_up - toon_dn)) / toon_scale):.3e}" - ) assert torch.max(torch.abs(toon_up - torch.mean(toon_up))) / toon_scale < 1e-2, ( f"Toon isothermal LW upward profile is not flat enough: " f"{float(torch.max(torch.abs(toon_up - torch.mean(toon_up))) / toon_scale):.3e}" ) - assert torch.max(torch.abs(toon_dn - torch.mean(toon_dn))) / toon_scale < 1e-2, ( - f"Toon isothermal LW downward profile is not flat enough: " - f"{float(torch.max(torch.abs(toon_dn - torch.mean(toon_dn))) / toon_scale):.3e}" - ) - - assert torch.max(torch.abs(disort_up - disort_dn)) / disort_scale < 1e-2, ( - f"DISORT isothermal LW up/down mismatch too large: " - f"{float(torch.max(torch.abs(disort_up - disort_dn)) / disort_scale):.3e}" + assert torch.max(torch.abs(disort_up - torch.mean(disort_up))) / disort_scale < 1e-2, ( + f"DISORT isothermal LW upward profile is not flat enough: " + f"{float(torch.max(torch.abs(disort_up - torch.mean(disort_up))) / disort_scale):.3e}" ) - toon_disort_scale = max(float(torch.max(torch.abs(disort_result))), 1e-30) - assert torch.max(torch.abs(toon_result - disort_result)) / toon_disort_scale < 1.5e-1, ( - f"Toon/DISORT isothermal LW profile mismatch too large: " - f"{float(torch.max(torch.abs(toon_result - disort_result)) / toon_disort_scale):.3e}" + up_profile_scale = max(float(torch.max(torch.abs(disort_up))), 1e-30) + assert torch.max(torch.abs(toon_up - disort_up)) / up_profile_scale < 1.5e-6, ( + f"Toon/DISORT isothermal LW upward profile mismatch too large: " + f"{float(torch.max(torch.abs(toon_up - disort_up)) / up_profile_scale):.3e}" ) def test_toon_lw_isothermal_scattering(): @@ -367,6 +366,7 @@ def test_toon_lw_isothermal_scattering(): opt = pyharp.ToonMcKay89Options() opt.wave_lower(wave_lo) opt.wave_upper(wave_hi) + opt.flags("planck") toon = pyharp.ToonMcKay89(opt) dop = pydisort.DisortOptions() diff --git a/tests/toon_test.yaml b/tests/toon_test.yaml index ffe9673..b03cb56 100644 --- a/tests/toon_test.yaml +++ b/tests/toon_test.yaml @@ -30,8 +30,11 @@ bands: range: [200., 2000.] opacities: [scatterer] solver: toon + flags: planck,zenith_correction,hard_surface,delta_eddington_lw integration: wavenumber nwave: 3 + toon: + top_emission_flag: -1 - name: B_disort range: [200., 2000.]