Integrate 1D toon

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.

python/csrc/pyrtsolver.cpp

@@ -26,10 +26,16 @@ void bind_rtsolver(py::module& m) {
            })
       .ADD_OPTION(std::vector<double>, harp::ToonMcKay89OptionsImpl, wave_lower)
       .ADD_OPTION(std::vector<double>, 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<harp::ToonMcKay89Impl>(m, "ToonMcKay89")
       .def(py::init<>())

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<std::string>()));
+    }
+    if (band["toon"]) {
+      auto const toon = band["toon"];
+      op->toon()->top_emission_flag(
+          toon["top_emission_flag"].as<int>(op->toon()->top_emission_flag()));
+      if (toon["flags"]) {
+        op->toon()->flags(trim_copy(toon["flags"].as<std::string>()));
+      }
+      std::vector<std::string> legacy_flags;
+      if (toon["zenith_correction"].as<bool>(false)) {
+        legacy_flags.push_back("zenith_correction");
+      }
+      if (toon["hard_surface"].as<bool>(false)) {
+        legacy_flags.push_back("hard_surface");
+      }
+      if (toon["delta_eddington_lw"].as<bool>(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<double>(op->nwave(), wmin));
     op->set_wave_upper(std::vector<double>(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());
   }

src/rtsolver/toon_mckay89.cpp

@@ -22,20 +22,20 @@ torch::Tensor ToonMcKay89Impl::forward(torch::Tensor prop,
                                        torch::optional<torch::Tensor> 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) {

src/rtsolver/toon_mckay89.hpp

@@ -7,6 +7,11 @@
 #include <torch/nn/modules/common.h>
 #include <torch/nn/modules/container/any.h>
 
+// C/C++
+#include <cctype>
+#include <sstream>
+#include <string>
+
 // harp
 #include <harp/add_arg.h>
 
@@ -21,7 +26,36 @@ struct ToonMcKay89OptionsImpl {
   std::shared_ptr<ToonMcKay89OptionsImpl> clone() const {
     return std::make_shared<ToonMcKay89OptionsImpl>(*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<char>(std::tolower(static_cast<unsigned char>(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,24 +76,15 @@ struct ToonMcKay89OptionsImpl {
   //! set upper wavenumber(length) at each bin
   ADD_ARG(std::vector<double>, 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)
   //!   1 = full Planck at TOA (infinite isothermal slab above)
   //!  -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<ToonMcKay89OptionsImpl>;

tests/test_toon.cpp

@@ -1,14 +1,89 @@
 // external
 #include <gtest/gtest.h>
 
+// C/C++
+#include <filesystem>
+
 // harp
+#include <harp/radiation/radiation.hpp>
+#include <harp/radiation/radiation_band.hpp>
 #include <harp/rtsolver/toon_mckay89.hpp>
 
 // tests
 #include "device_testing.hpp"
 
 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<double>{200.0, 200.0, 200.0}));
+  EXPECT_EQ(op->toon()->wave_upper(),
+            (std::vector<double>{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<double>{200.0, 500.0};
+  auto wave_upper = std::vector<double>{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<std::string, torch::Tensor> sw_bc;
+  std::map<std::string, torch::Tensor> 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<double>(), 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)}) {