From 27fd2fce5fce520ee181caefd754a2de5ce60e3a Mon Sep 17 00:00:00 2001 From: stormy/cli Date: Tue, 9 Jun 2026 17:56:35 -0400 Subject: [PATCH 1/5] Add PICASO CK opacity --- python/opacity.pyi | 4 +- src/opacity/picaso_ck.cpp | 86 ++++++++++++++++++++++++++++++++ src/opacity/picaso_ck.hpp | 27 ++++++++++ src/radiation/radiation_band.cpp | 11 +++- src/radiation/radiation_band.hpp | 1 + src/rtsolver/toon_mckay89.hpp | 5 ++ 6 files changed, 131 insertions(+), 3 deletions(-) create mode 100644 src/opacity/picaso_ck.cpp create mode 100644 src/opacity/picaso_ck.hpp diff --git a/python/opacity.pyi b/python/opacity.pyi index de5b8da..3eb1af2 100644 --- a/python/opacity.pyi +++ b/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 diff --git a/src/opacity/picaso_ck.cpp b/src/opacity/picaso_ck.cpp new file mode 100644 index 0000000..589050b --- /dev/null +++ b/src/opacity/picaso_ck.cpp @@ -0,0 +1,86 @@ +#include "picaso_ck.hpp" + +#include +#include + +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 const& kwargs) { + 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"); + auto wave_query = + kwargs.count("wavenumber") > 0 ? kwargs.at("wavenumber") : wavenumber; + + 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 diff --git a/src/opacity/picaso_ck.hpp b/src/opacity/picaso_ck.hpp new file mode 100644 index 0000000..bdcf9f2 --- /dev/null +++ b/src/opacity/picaso_ck.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include +#include +#include + +#include "opacity_options.hpp" + +namespace harp { + +class PicasoCKImpl : public torch::nn::Cloneable { + 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 const& kwargs); +}; +TORCH_MODULE(PicasoCK); + +} // namespace harp diff --git a/src/radiation/radiation_band.cpp b/src/radiation/radiation_band.cpp index 48fe80a..689cf7a 100644 --- a/src/radiation/radiation_band.cpp +++ b/src/radiation/radiation_band.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -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(ray_out); @@ -324,7 +330,8 @@ torch::Tensor RadiationBandImpl::forward( } // run rt solver - if (kwargs->find("tempf") != kwargs->end()) { + bool use_planck = options->solver_name() != "toon" || options->toon()->planck(); + if (use_planck && kwargs->find("tempf") != kwargs->end()) { int nlyr = prop.size(-1); int nlev = kwargs->at("tempf").size(-1); TORCH_CHECK(nlev == nlyr + 1, "'tempf' size must be nlyr + 1 = ", nlyr + 1, @@ -340,7 +347,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); diff --git a/src/radiation/radiation_band.hpp b/src/radiation/radiation_band.hpp index 11403b6..6de429f 100644 --- a/src/radiation/radiation_band.hpp +++ b/src/radiation/radiation_band.hpp @@ -35,6 +35,7 @@ using OpacityDict = std::map; * - "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 diff --git a/src/rtsolver/toon_mckay89.hpp b/src/rtsolver/toon_mckay89.hpp index fedd88f..635caa1 100644 --- a/src/rtsolver/toon_mckay89.hpp +++ b/src/rtsolver/toon_mckay89.hpp @@ -113,6 +113,11 @@ class ToonMcKay89Impl : public torch::nn::Cloneable { std::map* bc, std::string bname = "", torch::optional temf = torch::nullopt); + + protected: + FORWARD_HAS_DEFAULT_ARGS( + {2, torch::nn::AnyValue(std::string(""))}, + {3, torch::nn::AnyValue(torch::optional(torch::nullopt))}) }; TORCH_MODULE(ToonMcKay89); From 1284111a032593685c4d510bcfde365735a7dc76 Mon Sep 17 00:00:00 2001 From: stormy/cli Date: Tue, 9 Jun 2026 17:58:08 -0400 Subject: [PATCH 2/5] wip --- src/opacity/picaso_ck.cpp | 12 ++++++------ src/radiation/radiation_band.cpp | 3 ++- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/opacity/picaso_ck.cpp b/src/opacity/picaso_ck.cpp index 589050b..e3d4239 100644 --- a/src/opacity/picaso_ck.cpp +++ b/src/opacity/picaso_ck.cpp @@ -18,10 +18,10 @@ 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 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); @@ -70,8 +70,8 @@ torch::Tensor PicasoCKImpl::forward( 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)}); + 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 = diff --git a/src/radiation/radiation_band.cpp b/src/radiation/radiation_band.cpp index 689cf7a..84a76a8 100644 --- a/src/radiation/radiation_band.cpp +++ b/src/radiation/radiation_band.cpp @@ -330,7 +330,8 @@ torch::Tensor RadiationBandImpl::forward( } // run rt solver - bool use_planck = options->solver_name() != "toon" || options->toon()->planck(); + bool use_planck = + options->solver_name() != "toon" || options->toon()->planck(); if (use_planck && kwargs->find("tempf") != kwargs->end()) { int nlyr = prop.size(-1); int nlev = kwargs->at("tempf").size(-1); From a126c8ca6779987f016d2342c9614f3e7bd2c69f Mon Sep 17 00:00:00 2001 From: Cheng Li <69489965+chengcli@users.noreply.github.com> Date: Tue, 9 Jun 2026 19:50:22 -0400 Subject: [PATCH 3/5] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/radiation/radiation_band.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/radiation/radiation_band.cpp b/src/radiation/radiation_band.cpp index 84a76a8..927949c 100644 --- a/src/radiation/radiation_band.cpp +++ b/src/radiation/radiation_band.cpp @@ -333,7 +333,7 @@ torch::Tensor RadiationBandImpl::forward( bool use_planck = options->solver_name() != "toon" || options->toon()->planck(); if (use_planck && kwargs->find("tempf") != kwargs->end()) { - int nlyr = prop.size(-1); + 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); From ee234aa7db13b3ebcce2942aa35e726f7c8c1aa3 Mon Sep 17 00:00:00 2001 From: Cheng Li <69489965+chengcli@users.noreply.github.com> Date: Tue, 9 Jun 2026 21:39:37 -0400 Subject: [PATCH 4/5] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/opacity/picaso_ck.cpp | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/opacity/picaso_ck.cpp b/src/opacity/picaso_ck.cpp index e3d4239..b2b5520 100644 --- a/src/opacity/picaso_ck.cpp +++ b/src/opacity/picaso_ck.cpp @@ -58,13 +58,23 @@ void PicasoCKImpl::reset() { torch::Tensor PicasoCKImpl::forward( torch::Tensor conc, std::map 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(); From 564461e731b5b12f735b511ab261648afc788f28 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 10 Jun 2026 01:56:38 +0000 Subject: [PATCH 5/5] Fix clang-format violations in picaso_ck.cpp --- src/opacity/picaso_ck.cpp | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/opacity/picaso_ck.cpp b/src/opacity/picaso_ck.cpp index b2b5520..7c10ab2 100644 --- a/src/opacity/picaso_ck.cpp +++ b/src/opacity/picaso_ck.cpp @@ -67,14 +67,15 @@ torch::Tensor PicasoCKImpl::forward( 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)"); + "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)"); + "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"); + 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();