Add PICASO CK table

python/opacity.pyi

@@ -59,7 +59,9 @@ class OpacityOptions:
         """
         Set the type of the opacity source format.
 
-        Valid options are: ``jit``, ``molecule-line``, ``molecule-cia``, ``fourcolumn``, ``wavetemp``, ``multiband-ck``, ``helios``.
+        Valid options are: ``jit``, ``molecule-line``, ``molecule-cia``,
+        ``fourcolumn``, ``wavetemp``, ``multiband-ck``, ``picaso-ck``,
+        ``helios``.
 
         Args:
             value (str): type of the opacity source

src/opacity/picaso_ck.cpp

@@ -0,0 +1,97 @@
+#include "picaso_ck.hpp"
+
+#include <harp/math/interpolation.hpp>
+#include <harp/utils/netcdf_opacity_utils.hpp>
+
+namespace harp {
+
+PicasoCKImpl::PicasoCKImpl(OpacityOptions const& options_) : options(options_) {
+  TORCH_CHECK(options->opacity_files().size() == 1,
+              "Only one opacity file is allowed");
+  TORCH_CHECK(options->type().empty() || options->type() == "picaso-ck",
+              "Mismatch opacity type: ", options->type(),
+              " expecting 'picaso-ck'");
+  reset();
+}
+
+void PicasoCKImpl::reset() {
+#ifdef NETCDFOUTPUT
+  int fileid = open_file(options->opacity_files()[0]);
+
+  auto qwave =
+      read_tensor_permuted(fileid, "quadrature_wavenumber", {"band", "gpoint"});
+  auto qweight =
+      read_tensor_permuted(fileid, "quadrature_weight", {"band", "gpoint"});
+  auto lower = read_1d_variable(fileid, "wavenumber_lower");
+  auto upper = read_1d_variable(fileid, "wavenumber_upper");
+  auto ng = qwave.size(1);
+
+  wavenumber = qwave.flatten();
+  weights = qweight.flatten();
+  wave_lower = lower.unsqueeze(1).expand({lower.size(0), ng}).flatten();
+  wave_upper = upper.unsqueeze(1).expand({upper.size(0), ng}).flatten();
+  ln_pressure = read_1d_variable(fileid, "pressure").log();
+  temperature_anomaly = read_1d_variable(fileid, "temperature_offset");
+  ln_temperature_base =
+      read_1d_variable(fileid, "nominal_temperature").log().unsqueeze(-1);
+  ln_sigma_cross =
+      read_tensor_permuted(fileid, "kappa",
+                           {"band", "gpoint", "pressure", "temperature_offset"})
+          .reshape({-1, ln_pressure.size(0), temperature_anomaly.size(0), 1})
+          .clamp_min(1.e-300)
+          .log();
+
+  check_nc(nc_close(fileid), "Failed to close NetCDF file");
+#else
+  TORCH_CHECK(false, "NetCDF support is not enabled");
+#endif
+
+  register_buffer("wavenumber", wavenumber);
+  register_buffer("weights", weights);
+  register_buffer("wave_lower", wave_lower);
+  register_buffer("wave_upper", wave_upper);
+  register_buffer("ln_pressure", ln_pressure);
+  register_buffer("temperature_anomaly", temperature_anomaly);
+  register_buffer("ln_temperature_base", ln_temperature_base);
+  register_buffer("ln_sigma_cross", ln_sigma_cross);
+}
+
+torch::Tensor PicasoCKImpl::forward(
+    torch::Tensor conc, std::map<std::string, torch::Tensor> const& kwargs) {
+  int ncol = conc.size(0);
+  int nlyr = conc.size(1);
+
+  TORCH_CHECK(kwargs.count("pres") > 0, "pres is required in kwargs");
+  TORCH_CHECK(kwargs.count("temp") > 0, "temp is required in kwargs");
+  auto const& pres = kwargs.at("pres");
+  auto const& temp = kwargs.at("temp");
+
+  TORCH_CHECK(pres.size(0) == ncol && pres.size(1) == nlyr,
+              "Invalid pres shape: ", pres.sizes(),
+              "; needs to be (ncol, nlyr)");
+  TORCH_CHECK(temp.size(0) == ncol && temp.size(1) == nlyr,
+              "Invalid temp shape: ", temp.sizes(),
+              "; needs to be (ncol, nlyr)");
+
+  auto wave_query =
+      kwargs.count("wavenumber") > 0 ? kwargs.at("wavenumber") : wavenumber;
+  TORCH_CHECK(wave_query.dim() == 1, "PicasoCK expects a 1D wavenumber grid");
+  auto lnp = pres.log();
+  auto temperature_base =
+      interpn({lnp}, {ln_pressure}, ln_temperature_base).squeeze(-1).exp();
+  auto tempa = temp - temperature_base;
+  int const nwave = wave_query.size(0);
+  auto wave = wave_query.unsqueeze(-1).unsqueeze(-1).expand(
+      {nwave, conc.size(0), conc.size(1)});
+  lnp = lnp.unsqueeze(0).expand_as(wave);
+  tempa = tempa.unsqueeze(0).expand_as(wave);
+  auto sigma =
+      interpn({wave, lnp, tempa},
+              {wavenumber, ln_pressure, temperature_anomaly}, ln_sigma_cross)
+          .exp();
+
+  // PICASO pre-mixed molecular opacity [cm^2/mol] times total concentration.
+  return 1.e-4 * sigma * conc.sum(-1).unsqueeze(0).unsqueeze(-1);
+}
+
+}  // namespace harp

src/opacity/picaso_ck.hpp

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <torch/nn/cloneable.h>
+#include <torch/nn/module.h>
+#include <torch/torch.h>
+
+#include "opacity_options.hpp"
+
+namespace harp {
+
+class PicasoCKImpl : public torch::nn::Cloneable<PicasoCKImpl> {
+ public:
+  torch::Tensor wavenumber, weights, wave_lower, wave_upper;
+  torch::Tensor ln_pressure, temperature_anomaly, ln_temperature_base;
+  torch::Tensor ln_sigma_cross;
+  OpacityOptions options;
+
+  PicasoCKImpl() : options(OpacityOptionsImpl::create()) {}
+  explicit PicasoCKImpl(OpacityOptions const& options_);
+  void reset() override;
+
+  torch::Tensor forward(torch::Tensor conc,
+                        std::map<std::string, torch::Tensor> const& kwargs);
+};
+TORCH_MODULE(PicasoCK);
+
+}  // namespace harp

src/radiation/radiation_band.cpp

@@ -12,6 +12,7 @@
 #include <harp/opacity/molecule_line.hpp>
 #include <harp/opacity/multiband.hpp>
 #include <harp/opacity/opacity_formatter.hpp>
+#include <harp/opacity/picaso_ck.hpp>
 #include <harp/opacity/wavetemp.hpp>
 #include <harp/utils/layer2level.hpp>
 #include <harp/utils/parse_yaml_input.hpp>
@@ -183,6 +184,10 @@ void RadiationBandImpl::reset() {
       auto a = MultiBand(op);
       nmax_prop = std::max(nmax_prop, 1);
       opacities[name] = torch::nn::AnyModule(a);
+    } else if (op->type() == "picaso-ck") {
+      auto a = PicasoCK(op);
+      nmax_prop = std::max(nmax_prop, 1);
+      opacities[name] = torch::nn::AnyModule(a);
     } else if (op->type() == "wavetemp") {
       auto a = WaveTemp(op);
       nmax_prop = std::max(nmax_prop, 1);
@@ -200,6 +205,7 @@ void RadiationBandImpl::reset() {
     }
     register_module(name, opacities[name].ptr());
   }
+  if (options->solver_name() == "toon") nmax_prop = std::max(nmax_prop, 3);
 
   // create rtsolver
   auto [uphi, umu] = get_direction_grids<double>(ray_out);
@@ -324,8 +330,10 @@ torch::Tensor RadiationBandImpl::forward(
   }
 
   // run rt solver
-  if (kwargs->find("tempf") != kwargs->end()) {
-    int nlyr = prop.size(-1);
+  bool use_planck =
+      options->solver_name() != "toon" || options->toon()->planck();
+  if (use_planck && kwargs->find("tempf") != kwargs->end()) {
+    int nlyr = prop.size(-2);
     int nlev = kwargs->at("tempf").size(-1);
     TORCH_CHECK(nlev == nlyr + 1, "'tempf' size must be nlyr + 1 = ", nlyr + 1,
                 ", got ", nlev);
@@ -340,7 +348,7 @@ torch::Tensor RadiationBandImpl::forward(
       std::cout << "  Done running rt solver with level temperatures."
                 << std::endl;
     }
-  } else if (kwargs->find("temp") != kwargs->end()) {
+  } else if (use_planck && kwargs->find("temp") != kwargs->end()) {
     Layer2LevelOptions l2l;
     l2l.order(options->l2l_order());
     l2l.lower(kExtrapolate).upper(kExtrapolate).check_positivity(true);

src/radiation/radiation_band.hpp

@@ -35,6 +35,7 @@ using OpacityDict = std::map<std::string, OpacityOptions>;
  *  - "molecule-cia": CIA opacity from NetCDF binary absorption coefficient
  * dumps
  *  - "multiband-ck": multi-band correlated-k opacity table
+ *  - "picaso-ck": PICASO correlated-k NetCDF table
  *  - "wavetemp": opacity table defined on wavenumber and temperature grid (CIA)
  *  - "fourcolumn": Four-column opacity table (aerosol)
  *  - "helios": Helios opacity table

src/rtsolver/toon_mckay89.hpp

@@ -113,6 +113,11 @@ class ToonMcKay89Impl : public torch::nn::Cloneable<ToonMcKay89Impl> {
                         std::map<std::string, torch::Tensor>* bc,
                         std::string bname = "",
                         torch::optional<torch::Tensor> temf = torch::nullopt);
+
+ protected:
+  FORWARD_HAS_DEFAULT_ARGS(
+      {2, torch::nn::AnyValue(std::string(""))},
+      {3, torch::nn::AnyValue(torch::optional<torch::Tensor>(torch::nullopt))})
 };
 TORCH_MODULE(ToonMcKay89);