tboudreaux/GridFire
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(KINSOL): Switch from Eigen to KINSOL

Browse Source 
Previously QSE solving was done using Eigen. While this worked we were
limited in the ability to use previous iterations to speed up later
steps. We have switched to KINSOL, from SUNDIALS, for linear solving.
This has drastically speed up the process of solving for QSE abundances,
primarily because the jacobian matrix does not need to be generated
every single time time a QSE abundance is requested.
This commit is contained in:
Emily Boudreaux
2025-11-19 12:06:21 -05:00
parent f7fbc6c1da
commit 442d4ed86c
12 changed files with 506 additions and 386 deletions
Download Patch File Download Diff File Expand all files Collapse all files

188
src/include/gridfire/engine/views/engine_multiscale.h
View File

@@ -3,6 +3,10 @@
#include "gridfire/engine/engine_abstract.h"
#include "gridfire/engine/views/engine_view_abstract.h"
#include "gridfire/engine/engine_graph.h"
#include "sundials/sundials_linearsolver.h"
#include "sundials/sundials_matrix.h"
#include "sundials/sundials_nvector.h"
#include "sundials/sundials_types.h"
#include "unsupported/Eigen/NonLinearOptimization"
@@ -78,6 +82,8 @@ namespace gridfire {
*/
explicit MultiscalePartitioningEngineView(DynamicEngine& baseEngine);
~MultiscalePartitioningEngineView() override;
/**
* @brief Gets the list of species in the network.
* @return A const reference to the vector of `Species` objects representing all species
@@ -611,120 +617,81 @@ namespace gridfire {
[[nodiscard]] [[maybe_unused]] std::string toString() const;
};
/**
* @brief Functor for solving QSE abundances using Eigen's nonlinear optimization.
*
* @par Purpose
* This struct provides the objective function (`operator()`) and its Jacobian
* (`df`) to Eigen's Levenberg-Marquardt solver. The goal is to find the abundances
* of algebraic species that make their time derivatives (`dY/dt`) equal to zero.
*
* @how
* - **`operator()`**: Takes a vector `v_qse` (scaled abundances of algebraic species) as input.
* It constructs a full trial abundance vector `y_trial`, calls the base engine's
* `calculateRHSAndEnergy`, and returns the `dY/dt` values for the algebraic species.
* The solver attempts to drive this return vector to zero.
* - **`df`**: Computes the Jacobian of the objective function. It calls the base engine's
* `generateJacobianMatrix` and extracts the sub-matrix corresponding to the algebraic
* species. It applies the chain rule to account for the `asinh` scaling used on the
* abundances.
*
* The abundances are scaled using `asinh` to handle the large dynamic range and ensure positivity.
*/
struct EigenFunctor {
using InputType = Eigen::Matrix<double, Eigen::Dynamic, 1>;
using OutputType = Eigen::Matrix<double, Eigen::Dynamic, 1>;
using JacobianType = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>;
enum {
InputsAtCompileTime = Eigen::Dynamic,
ValuesAtCompileTime = Eigen::Dynamic
class QSESolver {
public:
QSESolver(
const std::vector<fourdst::atomic::Species>& species,
const DynamicEngine& engine,
SUNContext sun_ctx
);
QSESolver(
const QSESolver& other
) = delete;
QSESolver& operator=(
const QSESolver& other
) = delete;
~QSESolver();
fourdst::composition::Composition solve(
const fourdst::composition::Composition& comp,
double T9,
double rho
) const;
size_t solves() const;
void log_diagnostics() const;
private:
static int sys_func(
N_Vector y,
N_Vector f,
void *user_data
);
static int sys_jac(
N_Vector y,
N_Vector fy,
SUNMatrix J,
void *user_data,
N_Vector tmp1,
N_Vector tmp2
);
struct UserData {
const DynamicEngine& engine;
double T9;
double rho;
fourdst::composition::Composition& comp;
const std::unordered_map<fourdst::atomic::Species, size_t>& qse_solve_species_index_map;
const std::vector<fourdst::atomic::Species>& qse_solve_species;
};
/**
* @brief Pointer to the MultiscalePartitioningEngineView instance.
*/
const MultiscalePartitioningEngineView& m_view;
/**
* @brief The set of species to solve for in the QSE group.
*/
const std::set<fourdst::atomic::Species>& m_qse_solve_species;
/**
* @brief Initial abundances of all species in the full network.
*/
const fourdst::composition::CompositionAbstract& m_initial_comp;
/**
* @brief Temperature in units of 10^9 K.
*/
const double m_T9;
/**
* @brief Density in g/cm^3.
*/
const double m_rho;
/**
* @brief Scaling factors for the species abundances, used to improve solver stability.
*/
const Eigen::VectorXd& m_Y_scale;
private:
mutable size_t m_solves = 0;
size_t m_N;
const DynamicEngine& m_engine;
std::vector<fourdst::atomic::Species> m_species;
std::unordered_map<fourdst::atomic::Species, size_t> m_speciesMap;
/**
* @brief Mapping from species to their indices in the QSE solve vector.
*/
const std::unordered_map<fourdst::atomic::Species, size_t> m_qse_solve_species_index_map;
// --- SUNDIALS / KINSOL Persistent resources
SUNContext m_sun_ctx = nullptr;
void* m_kinsol_mem = nullptr;
mutable std::optional<JacobianType> m_cached_jacobian = std::nullopt;
N_Vector m_Y = nullptr;
N_Vector m_scale = nullptr;
N_Vector m_f_scale = nullptr;
N_Vector m_constraints = nullptr;
N_Vector m_func_tmpl = nullptr;
/**
* @brief Constructs an EigenFunctor.
*
* @param view The MultiscalePartitioningEngineView instance.
* @param qse_solve_species Species to solve for in the QSE group.
* @param initial_comp Initial abundances of all species in the full network.
* @param T9 Temperature in units of 10^9 K.
* @param rho Density in g/cm^3.
* @param Y_scale Scaling factors for the species abundances.
* @param qse_solve_species_index_map Mapping from species to their indices in the QSE solve vector.
*/
EigenFunctor(
const MultiscalePartitioningEngineView& view,
const std::set<fourdst::atomic::Species>& qse_solve_species,
const fourdst::composition::CompositionAbstract& initial_comp,
const double T9,
const double rho,
const Eigen::VectorXd& Y_scale,
const std::unordered_map<fourdst::atomic::Species, size_t>& qse_solve_species_index_map
) :
m_view(view),
m_qse_solve_species(qse_solve_species),
m_initial_comp(initial_comp),
m_T9(T9),
m_rho(rho),
m_Y_scale(Y_scale),
m_qse_solve_species_index_map(qse_solve_species_index_map) {}
SUNMatrix m_J = nullptr;
SUNLinearSolver m_LS = nullptr;
/**
* @brief Gets the number of output values from the functor (size of the residual vector).
* @return The number of algebraic species being solved.
*/
[[nodiscard]] size_t values() const { return m_qse_solve_species.size(); }
/**
* @brief Gets the number of input values to the functor (size of the variable vector).
* @return The number of algebraic species being solved.
*/
[[nodiscard]] size_t inputs() const { return m_qse_solve_species.size(); }
static quill::Logger* getLogger() {
return LogManager::getInstance().getLogger("log");
}
/**
* @brief Evaluates the functor's residual vector `f_qse = dY_alg/dt`.
* @param v_qse The input vector of scaled algebraic abundances.
* @param f_qse The output residual vector.
* @return 0 on success.
*/
int operator()(const InputType& v_qse, OutputType& f_qse) const;
/**
* @brief Evaluates the Jacobian of the functor, `J_qse = d(f_qse)/d(v_qse)`.
* @param v_qse The input vector of scaled algebraic abundances.
* @param J_qse The output Jacobian matrix.
* @return 0 on success.
*/
int df(const InputType& v_qse, JacobianType& J_qse) const;
};
struct FluxValidationResult {
@@ -746,6 +713,11 @@ namespace gridfire {
* @brief The list of identified equilibrium groups.
*/
std::vector<QSEGroup> m_qse_groups;
/**
* @brief A set of solvers, one for each QSE group
*/
std::vector<std::unique_ptr<QSESolver>> m_qse_solvers;
/**
* @brief The simplified set of species presented to the solver (the "slow" species).
*/
@@ -771,6 +743,8 @@ namespace gridfire {
mutable std::unordered_map<uint64_t, fourdst::composition::Composition> m_composition_cache;
SUNContext m_sun_ctx = nullptr; // TODO: initialize and safely destroy this
private:
/**
* @brief Partitions the network by timescale.

66
src/include/gridfire/utils/sundials.h Normal file
View File

@@ -0,0 +1,66 @@
#pragma once
#include <unordered_map>
#include <nvector/nvector_serial.h>
#include "gridfire/exceptions/error_solver.h"
#include "sundials/sundials_nvector.h"
namespace gridfire::utils {
inline std::unordered_map<int, std::string> cvode_ret_code_map {
{0, "CV_SUCCESS: The solver succeeded."},
{1, "CV_TSTOP_RETURN: The solver reached the specified stopping time."},
{2, "CV_ROOT_RETURN: A root was found."},
{-99, "CV_WARNING: CVODE succeeded but in an unusual manner"},
{-1, "CV_TOO_MUCH_WORK: The solver took too many internal steps."},
{-2, "CV_TOO_MUCH_ACC: The solver could not satisfy the accuracy requested."},
{-3, "CV_ERR_FAILURE: The solver encountered a non-recoverable error."},
{-4, "CV_CONV_FAILURE: The solver failed to converge."},
{-5, "CV_LINIT_FAIL: The linear solver's initialization function failed."},
{-6, "CV_LSETUP_FAIL: The linear solver's setup function failed."},
{-7, "CV_LSOLVE_FAIL: The linear solver's solve function failed."},
{-8, "CV_RHSFUNC_FAIL: The right-hand side function failed in an unrecoverable manner."},
{-9, "CV_FIRST_RHSFUNC_ERR: The right-hand side function failed at the first call."},
{-10, "CV_REPTD_RHSFUNC_ERR: The right-hand side function repeatedly failed recoverable."},
{-11, "CV_UNREC_RHSFUNC_ERR: The right-hand side function failed unrecoverably."},
{-12, "CV_RTFUNC_FAIL: The rootfinding function failed in an unrecoverable manner."},
{-13, "CV_NLS_INIT_FAIL: The nonlinear solver's initialization function failed."},
{-14, "CV_NLS_SETUP_FAIL: The nonlinear solver's setup function failed."},
{-15, "CV_CONSTR_FAIL : The inequality constraint was violated and the solver was unable to recover."},
{-16, "CV_NLS_FAIL: The nonlinear solver's solve function failed."},
{-20, "CV_MEM_FAIL: Memory allocation failed."},
{-21, "CV_MEM_NULL: The CVODE memory structure is NULL."},
{-22, "CV_ILL_INPUT: An illegal input was detected."},
{-23, "CV_NO_MALLOC: The CVODE memory structure has not been allocated."},
{-24, "CV_BAD_K: The value of k is invalid."},
{-25, "CV_BAD_T: The value of t is invalid."},
{-26, "CV_BAD_DKY: The value of dky is invalid."},
{-27, "CV_TOO_CLOSE: The time points are too close together."},
{-28, "CV_VECTOROP_ERR: A vector operation failed."},
{-29, "CV_PROJ_MEM_NULL: The projection memory structure is NULL."},
{-30, "CV_PROJFUNC_FAIL: The projection function failed in an unrecoverable manner."},
{-31, "CV_REPTD_PROJFUNC_ERR: The projection function has repeated recoverable errors."}
};
inline void check_cvode_flag(const int flag, const std::string& func_name) {
if (flag < 0) {
if (!cvode_ret_code_map.contains(flag)) {
throw exceptions::CVODESolverFailureError("CVODE error in " + func_name + ": Unknown error code: " + std::to_string(flag));
}
throw exceptions::CVODESolverFailureError("CVODE error in " + func_name + ": " + cvode_ret_code_map.at(flag));
}
}
inline N_Vector init_sun_vector(uint64_t size, SUNContext sun_ctx) {
#ifdef SUNDIALS_HAVE_OPENMP
const N_Vector vec = N_VNew_OpenMP(size, 0, sun_ctx);
#elif SUNDIALS_HAVE_PTHREADS
const N_Vector vec = N_VNew_Pthreads(size, sun_ctx);
#else
const N_Vector vec = N_VNew_Serial(static_cast<long long>(size), sun_ctx);
#endif
check_cvode_flag(vec == nullptr ? -1 : 0, "N_VNew");
return vec;
}
}