Skip to content

Add extrap options #67

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chengcli merged 2 commits into from
Dec 16, 2025
Merged

Add extrap options #67

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b79bc51
done
chengcli Dec 16, 2025
1d5f3a1
wip
chengcli Dec 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 21 additions & 12 deletions python/csrc/pythermo.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,19 +72,28 @@ void bind_thermo(py::module &m) {
.def("effective_cp", &kintera::ThermoXImpl::effective_cp, py::arg("temp"),
py::arg("pres"), py::arg("xfrac"), py::arg("gain"),
py::arg("conc") = py::none())
.def("extrapolate_ad",
py::overload_cast<torch::Tensor, torch::Tensor, torch::Tensor,
double, double, bool>(
&kintera::ThermoXImpl::extrapolate_ad),
py::arg("temp"), py::arg("pres"), py::arg("xfrac"), py::arg("dlnp"),
py::arg("ds_dlnp") = 0., py::arg("verbose") = false)
.def(
"extrapolate_dlnp",
[](kintera::ThermoXImpl &self, torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac, double dlnp, double ds_dlnp, bool verbose) {
kintera::ExtrapOptions opts;
opts.dlnp(dlnp).ds_dlnp(ds_dlnp).verbose(verbose);
return self.extrapolate_dlnp(temp, pres, xfrac, opts);
},
py::arg("temp"), py::arg("pres"), py::arg("xfrac"), py::arg("dlnp"),
py::arg("ds_dlnp") = 0., py::arg("verbose") = false)

.def("extrapolate_ad",
py::overload_cast<torch::Tensor, torch::Tensor, torch::Tensor,
double, double, double, bool>(
&kintera::ThermoXImpl::extrapolate_ad),
py::arg("temp"), py::arg("pres"), py::arg("xfrac"), py::arg("grav"),
py::arg("dz"), py::arg("ds_dz") = 0., py::arg("verbose") = false);
.def(
"extrapolate_dz",
[](kintera::ThermoXImpl &self, torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac, double dz, double grav, double ds_dz,
bool verbose) {
kintera::ExtrapOptions opts;
opts.dz(dz).grav(grav).ds_dz(ds_dz).verbose(verbose);
return self.extrapolate_dz(temp, pres, xfrac, opts);
},
py::arg("temp"), py::arg("pres"), py::arg("xfrac"), py::arg("dz"),
py::arg("grav"), py::arg("ds_dz") = 0., py::arg("verbose") = false);
Comment on lines +89 to +96
Copy link

Copilot AI Dec 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The parameter order for extrapolate_dz has been changed from the original (temp, pres, xfrac, grav, dz, ds_dz, verbose) to (temp, pres, xfrac, dz, grav, ds_dz, verbose). This is a breaking API change - the positions of dz and grav have been swapped. Existing Python code using this function will silently get incorrect results (e.g., passing grav=9.8, dz=100 will be interpreted as dz=9.8, grav=100). Either restore the original parameter order or ensure this breaking change is intentional and well-documented.

Suggested change
torch::Tensor xfrac, double dz, double grav, double ds_dz,
bool verbose) {
kintera::ExtrapOptions opts;
opts.dz(dz).grav(grav).ds_dz(ds_dz).verbose(verbose);
return self.extrapolate_dz(temp, pres, xfrac, opts);
},
py::arg("temp"), py::arg("pres"), py::arg("xfrac"), py::arg("dz"),
py::arg("grav"), py::arg("ds_dz") = 0., py::arg("verbose") = false);
torch::Tensor xfrac, double grav, double dz, double ds_dz,
bool verbose) {
kintera::ExtrapOptions opts;
opts.grav(grav).dz(dz).ds_dz(ds_dz).verbose(verbose);
return self.extrapolate_dz(temp, pres, xfrac, opts);
},
py::arg("temp"), py::arg("pres"), py::arg("xfrac"), py::arg("grav"),
py::arg("dz"), py::arg("ds_dz") = 0., py::arg("verbose") = false);

Copilot uses AI. Check for mistakes.

m.def("relative_humidity", &kintera::relative_humidity, py::arg("temp"),
py::arg("conc"), py::arg("stoich"), py::arg("op"));
Expand Down
23 changes: 16 additions & 7 deletions src/thermo/extrapolate_ad.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,9 +45,13 @@ torch::Tensor ThermoXImpl::effective_cp(torch::Tensor temp, torch::Tensor pres,
return cp_normal + cp_latent;
}

void ThermoXImpl::extrapolate_ad(torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac, double dlnp,
double ds_dlnp, bool verbose) {
void ThermoXImpl::extrapolate_dlnp(torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac,
ExtrapOptions const& opts) {
double dlnp = opts.dlnp();
double ds_dlnp = opts.ds_dlnp();
bool verbose = opts.verbose();

if (verbose) {
std::cout << "Extrapolating adiabat with dlnp = " << dlnp << std::endl;
}
Expand Down Expand Up @@ -113,14 +117,19 @@ void ThermoXImpl::extrapolate_ad(torch::Tensor temp, torch::Tensor pres,
}

if (iter >= options->max_iter()) {
TORCH_WARN("extrapolate_ad does not converge after ", options->max_iter(),
TORCH_WARN("extrapolate_dlnp does not converge after ", options->max_iter(),
" iterations.");
}
}

void ThermoXImpl::extrapolate_ad(torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac, double grav, double dz,
double ds_dz, bool verbose) {
void ThermoXImpl::extrapolate_dz(torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac,
ExtrapOptions const& opts) {
double grav = opts.grav();
double dz = opts.dz();
double ds_dz = opts.ds_dz();
bool verbose = opts.verbose();

if (verbose) {
std::cout << "Extrapolating adiabat over dz = " << dz << " m"
<< " with gravity = " << grav << " m/s^2" << std::endl;
Expand Down
28 changes: 17 additions & 11 deletions src/thermo/thermo.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,19 @@ struct ThermoOptionsImpl final : public SpeciesThermoImpl {
};
using ThermoOptions = std::shared_ptr<ThermoOptionsImpl>;

struct ExtrapOptions {
// Logarithmic change in pressure (dlnp = ln(p_new / p_old))
ADD_ARG(double, dlnp) = 0.;
ADD_ARG(double, dz) = 0.;
// Gravitational acceleration (m/s^2), positive
ADD_ARG(double, grav) = 0;
Copy link

Copilot AI Dec 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The default value for grav should be 0. (with decimal point) to be consistent with other double-type fields in the struct (like dlnp, dz, ds_dlnp, and ds_dz which all use 0.). Currently it's 0 without the decimal point.

Suggested change
ADD_ARG(double, grav) = 0;
ADD_ARG(double, grav) = 0.;

Copilot uses AI. Check for mistakes.
ADD_ARG(double, ds_dlnp) = 0;
// Height change (m)
Comment on lines +86 to +90
Copy link

Copilot AI Dec 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comments for struct fields are misplaced. The comment "// Height change (m)" on line 90 should appear before the dz field on line 86, not before ds_dz. Similarly, the comment "// Gravitational acceleration (m/s^2), positive" on line 87-88 should appear before the grav field on line 88. Consider reorganizing the fields and comments so that each comment directly precedes its corresponding field definition.

Suggested change
ADD_ARG(double, dz) = 0.;
// Gravitational acceleration (m/s^2), positive
ADD_ARG(double, grav) = 0;
ADD_ARG(double, ds_dlnp) = 0;
// Height change (m)
// Height change (m)
ADD_ARG(double, dz) = 0.;
// Gravitational acceleration (m/s^2), positive
ADD_ARG(double, grav) = 0;
ADD_ARG(double, ds_dlnp) = 0;

Copilot uses AI. Check for mistakes.
ADD_ARG(double, ds_dz) = 0;
// If true, print debug information
ADD_ARG(bool, verbose) = false;
};

//! Mass Thermodynamics
class ThermoYImpl : public torch::nn::Cloneable<ThermoYImpl> {
public:
Expand Down Expand Up @@ -300,12 +313,9 @@ class ThermoXImpl : public torch::nn::Cloneable<ThermoXImpl> {
* \param[in,out] temp Temperature tensor (K)
* \param[in,out] pres Pressure tensor (Pa)
* \param[in,out] xfrac Mole fraction tensor
* \param[in] dlnp Logarithmic change in pressure (dlnp = ln(p_new / p_old))
* \param[in] verbose If true, print debug information
*/
void extrapolate_ad(torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac, double dlnp, double ds_dlnp = 0.,
bool verbose = false);
void extrapolate_dlnp(torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac, ExtrapOptions const& opts);

/*! \brief Extrapolate state TPX to a new height along an adiabat
* Extrapolates the state variables (temperature, pressure, and mole
Expand All @@ -314,13 +324,9 @@ class ThermoXImpl : public torch::nn::Cloneable<ThermoXImpl> {
* \param[in,out] temp Temperature tensor (K)
* \param[in,out] pres Pressure tensor (Pa)
* \param[in,out] xfrac Mole fraction tensor
* \param[in] grav Gravitational acceleration (m/s^2)
* \param[in] dz Height change (m)
* \param[in] verbose If true, print debug information
*/
void extrapolate_ad(torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac, double grav, double dz,
double ds_dz = 0., bool verbose = false);
void extrapolate_dz(torch::Tensor temp, torch::Tensor pres,
torch::Tensor xfrac, ExtrapOptions const& opts);

//! \brief Calculate the equilibrium state given temperature and pressure
/*!
Expand Down
2 changes: 1 addition & 1 deletion tests/test_earth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def test_case1():
xfrac[:, i, :] = xfrac[:, i - 1, :]

# adiabatic extrapolation
thermo.extrapolate_ad(temp[:, i], pres[:, i], xfrac[:, i, :], -dlnp);
thermo.extrapolate_dlnp(temp[:, i], pres[:, i], xfrac[:, i, :], -dlnp);

# compute molar concentration
conc = thermo.compute("TPX->V", [temp, pres, xfrac])
Expand Down
2 changes: 1 addition & 1 deletion tests/test_thermo.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -350,7 +350,7 @@ TEST_P(DeviceTest, extrapolate_ad) {
std::cout << "entropy before = " << entropy_mole0 << std::endl;
std::cout << "temp before = " << temp << std::endl;

thermo_x->extrapolate_ad(temp, pres, xfrac, -0.1);
thermo_x->extrapolate_dlnp(temp, pres, xfrac, ExtrapOptions().dlnp(-0.1));

conc = thermo_x->compute("TPX->V", {temp, pres, xfrac});
entropy_vol = thermo_x->compute("TPV->S", {temp, pres, conc});
Expand Down
Loading