feat(python): added robust python bindings covering the entire codebase
This commit is contained in:
231
src/lib/solver/solver.cpp
Normal file
231
src/lib/solver/solver.cpp
Normal file
@@ -0,0 +1,231 @@
|
||||
#include "gridfire/solver/solver.h"
|
||||
#include "gridfire/engine/engine_graph.h"
|
||||
#include "gridfire/network.h"
|
||||
#include "gridfire/exceptions/error_engine.h"
|
||||
|
||||
#include "fourdst/composition/atomicSpecies.h"
|
||||
#include "fourdst/composition/composition.h"
|
||||
#include "fourdst/config/config.h"
|
||||
|
||||
#include "unsupported/Eigen/NonLinearOptimization"
|
||||
|
||||
#include <boost/numeric/odeint.hpp>
|
||||
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <stdexcept>
|
||||
#include <iomanip>
|
||||
|
||||
#include "quill/LogMacros.h"
|
||||
|
||||
namespace gridfire::solver {
|
||||
NetOut DirectNetworkSolver::evaluate(const NetIn &netIn) {
|
||||
namespace ublas = boost::numeric::ublas;
|
||||
namespace odeint = boost::numeric::odeint;
|
||||
using fourdst::composition::Composition;
|
||||
|
||||
|
||||
const double T9 = netIn.temperature / 1e9; // Convert temperature from Kelvin to T9 (T9 = T / 1e9)
|
||||
|
||||
const auto absTol = m_config.get<double>("gridfire:solver:DirectNetworkSolver:absTol", 1.0e-8);
|
||||
const auto relTol = m_config.get<double>("gridfire:solver:DirectNetworkSolver:relTol", 1.0e-8);
|
||||
|
||||
Composition equilibratedComposition = m_engine.update(netIn);
|
||||
size_t numSpecies = m_engine.getNetworkSpecies().size();
|
||||
ublas::vector<double> Y(numSpecies + 1);
|
||||
|
||||
RHSManager manager(m_engine, T9, netIn.density);
|
||||
JacobianFunctor jacobianFunctor(m_engine, T9, netIn.density);
|
||||
|
||||
auto populateY = [&](const Composition& comp) {
|
||||
const size_t numSpeciesInternal = m_engine.getNetworkSpecies().size();
|
||||
Y.resize(numSpeciesInternal + 1);
|
||||
for (size_t i = 0; i < numSpeciesInternal; i++) {
|
||||
const auto& species = m_engine.getNetworkSpecies()[i];
|
||||
if (!comp.contains(species)) {
|
||||
double lim = std::numeric_limits<double>::min();
|
||||
LOG_DEBUG(m_logger, "Species '{}' not found in composition. Setting abundance to {:0.3E}.", species.name(), lim);
|
||||
Y(i) = lim; // Species not in the composition, set to zero
|
||||
} else {
|
||||
Y(i) = comp.getMolarAbundance(species);
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: a good starting point to make the temperature, density, and energy self consistent would be to turn this into an accumulator
|
||||
Y(numSpeciesInternal) = 0.0; // Specific energy rate, initialized to zero
|
||||
};
|
||||
|
||||
// This is a quick debug that can be turned on. For solar code input parameters (T~1.5e7K, ρ~1.5e3 g/cm^3) this should be near 8e-17
|
||||
// std::cout << "D/H: " << equilibratedComposition.getMolarAbundance("H-2") / equilibratedComposition.getMolarAbundance("H-1") << std::endl;
|
||||
|
||||
populateY(equilibratedComposition);
|
||||
const auto stepper = odeint::make_controlled<odeint::rosenbrock4<double>>(absTol, relTol);
|
||||
|
||||
double current_time = 0.0;
|
||||
double current_initial_timestep = netIn.dt0;
|
||||
double accumulated_energy = 0.0;
|
||||
// size_t total_update_stages_triggered = 0;
|
||||
|
||||
while (current_time < netIn.tMax) {
|
||||
try {
|
||||
odeint::integrate_adaptive(
|
||||
stepper,
|
||||
std::make_pair(manager, jacobianFunctor),
|
||||
Y,
|
||||
current_time,
|
||||
netIn.tMax,
|
||||
current_initial_timestep,
|
||||
[&](const auto& state, double t) {
|
||||
current_time = t;
|
||||
manager.observe(state, t);
|
||||
}
|
||||
);
|
||||
current_time = netIn.tMax;
|
||||
} catch (const exceptions::StaleEngineTrigger &e) {
|
||||
LOG_INFO(m_logger, "Catching StaleEngineTrigger at t = {:0.3E} with T9 = {:0.3E}, rho = {:0.3E}. Triggering update stage (last stage took {} steps).", current_time, T9, netIn.density, e.totalSteps());
|
||||
exceptions::StaleEngineTrigger::state staleState = e.getState();
|
||||
accumulated_energy += e.energy(); // Add the specific energy rate to the accumulated energy
|
||||
// total_update_stages_triggered++;
|
||||
|
||||
Composition temp_comp;
|
||||
std::vector<double> mass_fractions;
|
||||
size_t num_species_at_stop = e.numSpecies();
|
||||
|
||||
if (num_species_at_stop != m_engine.getNetworkSpecies().size()) {
|
||||
throw std::runtime_error(
|
||||
"StaleEngineError state has a different number of species than the engine. This should not happen."
|
||||
);
|
||||
}
|
||||
mass_fractions.reserve(num_species_at_stop);
|
||||
|
||||
for (size_t i = 0; i < num_species_at_stop; ++i) {
|
||||
const auto& species = m_engine.getNetworkSpecies()[i];
|
||||
temp_comp.registerSpecies(species);
|
||||
mass_fractions.push_back(e.getMolarAbundance(i) * species.mass()); // Convert from molar abundance to mass fraction
|
||||
}
|
||||
temp_comp.setMassFraction(m_engine.getNetworkSpecies(), mass_fractions);
|
||||
temp_comp.finalize(true);
|
||||
|
||||
NetIn netInTemp = netIn;
|
||||
netInTemp.temperature = e.temperature();
|
||||
netInTemp.density = e.density();
|
||||
netInTemp.composition = std::move(temp_comp);
|
||||
|
||||
Composition currentComposition = m_engine.update(netInTemp);
|
||||
populateY(currentComposition);
|
||||
Y(Y.size() - 1) = e.energy(); // Set the specific energy rate from the stale state
|
||||
numSpecies = m_engine.getNetworkSpecies().size();
|
||||
|
||||
// current_initial_timestep = 0.001 * manager.m_last_step_time; // set the new timestep to the last successful timestep before repartitioning
|
||||
}
|
||||
}
|
||||
|
||||
accumulated_energy += Y(Y.size() - 1); // Add the specific energy rate to the accumulated energy
|
||||
|
||||
std::vector<double> finalMassFractions(numSpecies);
|
||||
for (size_t i = 0; i < numSpecies; ++i) {
|
||||
const double molarMass = m_engine.getNetworkSpecies()[i].mass();
|
||||
finalMassFractions[i] = Y(i) * molarMass; // Convert from molar abundance to mass fraction
|
||||
if (finalMassFractions[i] < MIN_ABUNDANCE_THRESHOLD) {
|
||||
finalMassFractions[i] = 0.0;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::string> speciesNames;
|
||||
speciesNames.reserve(numSpecies);
|
||||
for (const auto& species : m_engine.getNetworkSpecies()) {
|
||||
speciesNames.push_back(std::string(species.name()));
|
||||
}
|
||||
|
||||
Composition outputComposition(speciesNames);
|
||||
outputComposition.setMassFraction(speciesNames, finalMassFractions);
|
||||
outputComposition.finalize(true);
|
||||
|
||||
NetOut netOut;
|
||||
netOut.composition = std::move(outputComposition);
|
||||
netOut.energy = accumulated_energy; // Specific energy rate
|
||||
netOut.num_steps = manager.m_num_steps;
|
||||
|
||||
return netOut;
|
||||
}
|
||||
|
||||
void DirectNetworkSolver::RHSManager::operator()(
|
||||
const boost::numeric::ublas::vector<double> &Y,
|
||||
boost::numeric::ublas::vector<double> &dYdt,
|
||||
const double t
|
||||
) const {
|
||||
const size_t numSpecies = m_engine.getNetworkSpecies().size();
|
||||
if (t != m_cached_time || !m_cached_result.has_value() || m_cached_result.value().dydt.size() != numSpecies + 1) {
|
||||
compute_and_cache(Y, t);
|
||||
}
|
||||
const auto&[dydt, nuclearEnergyGenerationRate] = m_cached_result.value();
|
||||
dYdt.resize(numSpecies + 1);
|
||||
std::ranges::copy(dydt, dYdt.begin());
|
||||
dYdt(numSpecies) = nuclearEnergyGenerationRate; // Set the last element to the specific energy rate
|
||||
}
|
||||
|
||||
void DirectNetworkSolver::RHSManager::observe(
|
||||
const boost::numeric::ublas::vector<double> &state,
|
||||
const double t
|
||||
) const {
|
||||
double dt = t - m_last_observed_time;
|
||||
compute_and_cache(state, t);
|
||||
LOG_INFO(
|
||||
m_logger,
|
||||
"(Step {}) Observed state at t = {:0.3E} (dt = {:0.3E})",
|
||||
m_num_steps,
|
||||
t,
|
||||
dt
|
||||
);
|
||||
std::ostringstream oss;
|
||||
oss << std::scientific << std::setprecision(3);
|
||||
oss << "(Step: " << std::setw(10) << m_num_steps << ") t = " << t << " (dt = " << dt << ", eps_nuc: " << state(state.size() - 1) << " [erg])\n";
|
||||
std::cout << oss.str();
|
||||
m_last_observed_time = t;
|
||||
m_last_step_time = dt;
|
||||
|
||||
}
|
||||
|
||||
void DirectNetworkSolver::RHSManager::compute_and_cache(
|
||||
const boost::numeric::ublas::vector<double> &state,
|
||||
double t
|
||||
) const {
|
||||
std::vector<double> y_vec(state.begin(), state.end() - 1);
|
||||
std::ranges::replace_if(
|
||||
y_vec,
|
||||
[](const double yi){
|
||||
return yi < 0.0;
|
||||
},
|
||||
0.0 // Avoid negative abundances
|
||||
);
|
||||
|
||||
const auto result = m_engine.calculateRHSAndEnergy(y_vec, m_T9, m_rho);
|
||||
if (!result) {
|
||||
LOG_INFO(m_logger,
|
||||
"Triggering update stage due to stale engine detected at t = {:0.3E} with T9 = {:0.3E}, rho = {:0.3E}, y_vec (size: {})",
|
||||
t, m_T9, m_rho, y_vec.size());
|
||||
throw exceptions::StaleEngineTrigger({m_T9, m_rho, y_vec, t, m_num_steps, state(state.size() - 1)});
|
||||
}
|
||||
m_cached_result = result.value();
|
||||
m_cached_time = t;
|
||||
|
||||
m_num_steps++;
|
||||
}
|
||||
|
||||
void DirectNetworkSolver::JacobianFunctor::operator()(
|
||||
const boost::numeric::ublas::vector<double> &Y,
|
||||
boost::numeric::ublas::matrix<double> &J,
|
||||
double t,
|
||||
boost::numeric::ublas::vector<double> &dfdt
|
||||
) const {
|
||||
size_t numSpecies = m_engine.getNetworkSpecies().size();
|
||||
J.resize(numSpecies+1, numSpecies+1);
|
||||
J.clear();
|
||||
for (int i = 0; i < numSpecies; ++i) {
|
||||
for (int j = 0; j < numSpecies; ++j) {
|
||||
J(i, j) = m_engine.getJacobianMatrixEntry(i, j);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user