From 97a7fd05d27d92eaceb4b7680bb19c88a2760348 Mon Sep 17 00:00:00 2001 From: Emily Boudreaux Date: Wed, 10 Dec 2025 12:50:35 -0500 Subject: [PATCH] feat(SpectralSolver): Began work on multizone spectral solver The single zone solver we have is too slow for a true high resolution multi-zone enviroment. Began work on a spectral element method multi-zone solver --- src/include/gridfire/config/config.h | 19 + src/include/gridfire/engine/engine_abstract.h | 4 + src/include/gridfire/engine/engine_graph.h | 5 + .../solver/strategies/CVODE_solver_strategy.h | 3 +- .../strategies/SpectralSolverStrategy.h | 196 +++++ .../gridfire/solver/strategies/strategies.h | 3 +- .../solver/strategies/strategy_abstract.h | 32 +- src/lib/engine/engine_graph.cpp | 30 +- src/lib/engine/views/engine_adaptive.cpp | 44 +- src/lib/engine/views/engine_defined.cpp | 18 +- src/lib/engine/views/engine_multiscale.cpp | 44 +- .../strategies/CVODE_solver_strategy.cpp | 2 +- .../strategies/SpectralSolverStrategy.cpp | 770 ++++++++++++++++++ src/meson.build | 2 + tests/graphnet_sandbox/meson.build | 6 + tests/graphnet_sandbox/spectral_main.cpp | 13 +- 16 files changed, 1100 insertions(+), 91 deletions(-) diff --git a/src/include/gridfire/config/config.h b/src/include/gridfire/config/config.h index 9c807503..ee729f01 100644 --- a/src/include/gridfire/config/config.h +++ b/src/include/gridfire/config/config.h @@ -8,8 +8,27 @@ namespace gridfire::config { double relTol = 1.0e-5; }; + + struct SpectralSolverConfig { + struct MonitorFunctionConfig { + double structure_weight = 1.0; + double abundance_weight = 10.0; + double alpha = 0.2; + double beta = 0.8; + }; + struct BasisConfig { + size_t num_elements = 50; + }; + double absTol = 1.0e-8; + double relTol = 1.0e-5; + size_t degree = 3; + MonitorFunctionConfig monitorFunction; + BasisConfig basis; + }; + struct SolverConfig { CVODESolverConfig cvode; + SpectralSolverConfig spectral; }; struct AdaptiveEngineViewConfig { diff --git a/src/include/gridfire/engine/engine_abstract.h b/src/include/gridfire/engine/engine_abstract.h index 4ee96930..732dd6f4 100644 --- a/src/include/gridfire/engine/engine_abstract.h +++ b/src/include/gridfire/engine/engine_abstract.h @@ -514,5 +514,9 @@ namespace gridfire::engine { */ [[nodiscard]] virtual SpeciesStatus getSpeciesStatus(const fourdst::atomic::Species& species) const = 0; + [[nodiscard]] virtual std::optional> getMostRecentRHSCalculation() const { + return std::nullopt; + } + }; } \ No newline at end of file diff --git a/src/include/gridfire/engine/engine_graph.h b/src/include/gridfire/engine/engine_graph.h index 0f57ca5c..77e786bf 100644 --- a/src/include/gridfire/engine/engine_graph.h +++ b/src/include/gridfire/engine/engine_graph.h @@ -14,6 +14,8 @@ #include "gridfire/engine/procedures/construction.h" #include "gridfire/config/config.h" +#include "ankerl/unordered_dense.h" + #include #include #include @@ -764,6 +766,8 @@ namespace gridfire::engine { m_store_intermediate_reaction_contributions = value; } + [[nodiscard]] std::optional> getMostRecentRHSCalculation() const override; + private: struct PrecomputedReaction { @@ -887,6 +891,7 @@ namespace gridfire::engine { mutable std::unordered_map> m_stepDerivativesCache; mutable std::unordered_map, std::vector>> m_jacobianSubsetCache; mutable std::unordered_map m_jacWorkCache; + mutable std::optional> m_most_recent_rhs_calculation; bool m_has_been_primed = false; ///< Flag indicating if the engine has been primed. diff --git a/src/include/gridfire/solver/strategies/CVODE_solver_strategy.h b/src/include/gridfire/solver/strategies/CVODE_solver_strategy.h index babcc5b7..8db5000b 100644 --- a/src/include/gridfire/solver/strategies/CVODE_solver_strategy.h +++ b/src/include/gridfire/solver/strategies/CVODE_solver_strategy.h @@ -37,7 +37,6 @@ #ifdef SUNDIALS_HAVE_PTHREADS #include #endif -// Default to serial if no parallelism is enabled #ifndef SUNDIALS_HAVE_OPENMP #ifndef SUNDIALS_HAVE_PTHREADS #include @@ -79,7 +78,7 @@ namespace gridfire::solver { * std::cout << "Final energy: " << out.energy << " erg/g\n"; * @endcode */ - class CVODESolverStrategy final : public DynamicNetworkSolverStrategy { + class CVODESolverStrategy final : public SingleZoneDynamicNetworkSolverStrategy { public: /** * @brief Construct the CVODE strategy and create a SUNDIALS context. diff --git a/src/include/gridfire/solver/strategies/SpectralSolverStrategy.h b/src/include/gridfire/solver/strategies/SpectralSolverStrategy.h index e69de29b..260c0deb 100644 --- a/src/include/gridfire/solver/strategies/SpectralSolverStrategy.h +++ b/src/include/gridfire/solver/strategies/SpectralSolverStrategy.h @@ -0,0 +1,196 @@ +#pragma once + +#include "gridfire/solver/strategies/strategy_abstract.h" +#include "gridfire/engine/engine_abstract.h" +#include "gridfire/types/types.h" +#include "gridfire/config/config.h" + +#include "fourdst/logging/logging.h" +#include "fourdst/constants/const.h" + +#include +#include +#include + +#ifdef SUNDIALS_HAVE_OPENMP + #include +#endif +#ifdef SUNDIALS_HAVE_PTHREADS + #include +#endif +#ifndef SUNDIALS_HAVE_OPENMP + #ifndef SUNDIALS_HAVE_PTHREADS + #include + #endif +#endif + +namespace gridfire::solver { + class SpectralSolverStrategy final : public MultiZoneDynamicNetworkSolverStrategy { + public: + explicit SpectralSolverStrategy(engine::DynamicEngine& engine); + ~SpectralSolverStrategy() override; + + std::vector evaluate( + const std::vector &netIns, + const std::vector& mass_coords + ) override; + + void set_callback(const std::any &callback) override; + [[nodiscard]] std::vector> describe_callback_context() const override; + + [[nodiscard]] bool get_stdout_logging_enabled() const; + void set_stdout_logging_enabled(bool logging_enabled); + + public: + struct TimestepContext final : public SolverContextBase { + TimestepContext( + const double t, + const N_Vector &state, + const double dt, + const double last_time_step, + const engine::DynamicEngine &engine + ) : + t(t), + state(state), + dt(dt), + last_time_step(last_time_step), + engine(engine) {} + + [[nodiscard]] std::vector> describe() const override; + + const double t; + const N_Vector& state; + const double dt; + const double last_time_step; + const engine::DynamicEngine& engine; + }; + + struct BasisEval { + size_t start_idx; + std::vector phi; + }; + + struct SplineBasis { + std::vector knots; + std::vector quadrature_nodes; + std::vector quadrature_weights; + int degree = 3; + + + std::vector quad_evals; + }; + public: + using TimestepCallback = std::function; + private: + + struct SpectralCoefficients { + size_t num_sets; + size_t num_coefficients; + std::vector coefficients; + + double operator()(size_t i, size_t j) const; + }; + + struct GridPoint { + double T9; + double rho; + fourdst::composition::Composition composition; + }; + + struct Constants { + const double c = fourdst::constant::Constants::getInstance().get("c").value; + const double N_a = fourdst::constant::Constants::getInstance().get("N_a").value; + const double c2 = c * c; + }; + + struct DenseLinearSolver { + SUNMatrix A; + SUNLinearSolver LS; + N_Vector temp_vector; + SUNContext ctx; + + DenseLinearSolver(size_t size, SUNContext sun_ctx); + ~DenseLinearSolver(); + + DenseLinearSolver(const DenseLinearSolver&) = delete; + DenseLinearSolver& operator=(const DenseLinearSolver&) = delete; + + void setup() const; + void zero() const; + + void init_from_cache(size_t num_basis_funcs, const std::vector& shell_cache) const; + void init_from_basis(size_t num_basis_funcs, const SplineBasis& basis) const; + + void solve_inplace(N_Vector x, size_t num_vars, size_t basis_size) const; + }; + + struct CVODEUserData { + SpectralSolverStrategy* solver_instance{}; + engine::DynamicEngine* engine; + + DenseLinearSolver* mass_matrix_solver_instance; + const SplineBasis* basis; + }; + + private: + fourdst::config::Config m_config; + quill::Logger* m_logger = fourdst::logging::LogManager::getInstance().getLogger("log"); + + SUNContext m_sun_ctx = nullptr; ///< SUNDIALS context (lifetime of the solver). + void* m_cvode_mem = nullptr; ///< CVODE memory block. + N_Vector m_Y = nullptr; ///< CVODE state vector (species + energy accumulator). + SUNMatrix m_J = nullptr; ///< Dense Jacobian matrix. + SUNLinearSolver m_LS = nullptr; ///< Dense linear solver. + + + std::optional m_callback; ///< Optional per-step callback. + int m_num_steps = 0; ///< CVODE step counter (used for diagnostics and triggers). + + bool m_stdout_logging_enabled = true; ///< If true, print per-step logs and use CV_ONE_STEP. + + N_Vector m_constraints = nullptr; ///< CVODE constraints vector (>= 0 for species entries). + + std::optional m_absTol; ///< User-specified absolute tolerance. + std::optional m_relTol; ///< User-specified relative tolerance. + + bool m_detailed_step_logging = false; ///< If true, log detailed step diagnostics (error ratios, Jacobian, species balance). + + mutable size_t m_last_size = 0; + mutable size_t m_last_composition_hash = 0ULL; + mutable sunrealtype m_last_good_time_step = 0ULL; + + SplineBasis m_current_basis; + + Constants m_constants; + + N_Vector m_T_coeffs = nullptr; + N_Vector m_rho_coeffs = nullptr; + + + private: + std::vector evaluate_monitor_function(const std::vector& current_shells) const; + SplineBasis generate_basis_from_monitor(const std::vector& monitor_values, const std::vector& mass_coordinates) const; + + GridPoint reconstruct_at_quadrature(const N_Vector y_coeffs, size_t quad_index, const SplineBasis &basis) const; + + std::vector reconstruct_solution(const std::vector& original_inputs, const std::vector& mass_coordinates, const N_Vector final_coeffs, const SplineBasis& basis, double dt) const; + + static int cvode_rhs_wrapper(sunrealtype t, N_Vector y, N_Vector, void* user_data); + static int cvode_jac_wrapper(sunrealtype t, N_Vector y, N_Vector ydot, SUNMatrix J, void* user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3); + + int calculate_rhs(sunrealtype t, N_Vector y_coeffs, N_Vector ydot_coeffs, CVODEUserData* data) const; + + static void project_specific_variable( + const std::vector& current_shells, + const std::vector& mass_coordinates, + const std::vector& shell_cache, + const DenseLinearSolver& linear_solver, + N_Vector output_vec, + size_t output_offset, + const std::function &getter, + bool use_log + ); + + }; + +} diff --git a/src/include/gridfire/solver/strategies/strategies.h b/src/include/gridfire/solver/strategies/strategies.h index 3e59d2ef..f296dd3e 100644 --- a/src/include/gridfire/solver/strategies/strategies.h +++ b/src/include/gridfire/solver/strategies/strategies.h @@ -2,4 +2,5 @@ #include "gridfire/solver/strategies/triggers/triggers.h" #include "gridfire/solver/strategies/strategy_abstract.h" -#include "gridfire/solver/strategies/CVODE_solver_strategy.h" \ No newline at end of file +#include "gridfire/solver/strategies/CVODE_solver_strategy.h" +#include "gridfire/solver/strategies/SpectralSolverStrategy.h" \ No newline at end of file diff --git a/src/include/gridfire/solver/strategies/strategy_abstract.h b/src/include/gridfire/solver/strategies/strategy_abstract.h index 3562012e..376a9508 100644 --- a/src/include/gridfire/solver/strategies/strategy_abstract.h +++ b/src/include/gridfire/solver/strategies/strategy_abstract.h @@ -10,6 +10,9 @@ #include namespace gridfire::solver { + template + concept IsEngine = std::is_base_of_v; + /** * @struct SolverContextBase * @brief Base class for solver callback contexts. @@ -34,7 +37,7 @@ namespace gridfire::solver { [[nodiscard]] virtual std::vector> describe() const = 0; }; /** - * @class NetworkSolverStrategy + * @class SingleZoneNetworkSolverStrategy * @brief Abstract base class for network solver strategies. * * This class defines the interface for network solver strategies, which are responsible @@ -43,19 +46,19 @@ namespace gridfire::solver { * * @tparam EngineT The type of engine to use with this solver strategy. Must inherit from Engine. */ - template - class NetworkSolverStrategy { + template + class SingleZoneNetworkSolverStrategy { public: /** * @brief Constructor for the NetworkSolverStrategy. * @param engine The engine to use for evaluating the network. */ - explicit NetworkSolverStrategy(EngineT& engine) : m_engine(engine) {}; + explicit SingleZoneNetworkSolverStrategy(EngineT& engine) : m_engine(engine) {}; /** * @brief Virtual destructor. */ - virtual ~NetworkSolverStrategy() = default; + virtual ~SingleZoneNetworkSolverStrategy() = default; /** * @brief Evaluates the network for a given timestep. @@ -92,8 +95,25 @@ namespace gridfire::solver { EngineT& m_engine; ///< The engine used by this solver strategy. }; + template + class MultiZoneNetworkSolverStrategy { + public: + explicit MultiZoneNetworkSolverStrategy(EngineT& engine) : m_engine(engine) {}; + virtual ~MultiZoneNetworkSolverStrategy() = default; + + virtual std::vector evaluate( + const std::vector& netIns, + const std::vector& mass_coords + ) = 0; + virtual void set_callback(const std::any& callback) = 0; + [[nodiscard]] virtual std::vector> describe_callback_context() const = 0; + protected: + EngineT& m_engine; ///< The engine used by this solver strategy. + }; + /** * @brief Type alias for a network solver strategy that uses a DynamicEngine. */ - using DynamicNetworkSolverStrategy = NetworkSolverStrategy; + using SingleZoneDynamicNetworkSolverStrategy = SingleZoneNetworkSolverStrategy; + using MultiZoneDynamicNetworkSolverStrategy = MultiZoneNetworkSolverStrategy; } diff --git a/src/lib/engine/engine_graph.cpp b/src/lib/engine/engine_graph.cpp index d29474bf..d3641340 100644 --- a/src/lib/engine/engine_graph.cpp +++ b/src/lib/engine/engine_graph.cpp @@ -147,6 +147,7 @@ namespace gridfire::engine { ) const { LOG_TRACE_L3(m_logger, "Calculating RHS and Energy in GraphEngine at T9 = {}, rho = {}.", T9, rho); const double Ye = comp.getElectronAbundance(); + const std::vector molarAbundances = comp.getMolarAbundanceVector(); if (m_usePrecomputation) { const std::size_t state_hash = utils::hash_state(comp, T9, rho, activeReactions); if (m_stepDerivativesCache.contains(state_hash)) { @@ -158,9 +159,10 @@ namespace gridfire::engine { bare_rates.reserve(activeReactions.size()); bare_reverse_rates.reserve(activeReactions.size()); + for (const auto& reaction: activeReactions) { assert(m_reactions.contains(*reaction)); // A bug which results in this failing indicates a serious internal inconsistency and should only be present during development. - bare_rates.push_back(reaction->calculate_rate(T9, rho, Ye, 0.0, comp.getMolarAbundanceVector(), m_indexToSpeciesMap)); + bare_rates.push_back(reaction->calculate_rate(T9, rho, Ye, 0.0, molarAbundances, m_indexToSpeciesMap)); if (reaction->type() != reaction::ReactionType::WEAK) { bare_reverse_rates.push_back(calculateReverseRate(*reaction, T9, rho, comp)); } @@ -171,11 +173,12 @@ namespace gridfire::engine { // --- The public facing interface can always use the precomputed version since taping is done internally --- StepDerivatives result = calculateAllDerivativesUsingPrecomputation(comp, bare_rates, bare_reverse_rates, T9, rho, activeReactions); m_stepDerivativesCache.insert(std::make_pair(state_hash, result)); + m_most_recent_rhs_calculation = result; return result; } else { LOG_TRACE_L2(m_logger, "Not using precomputation for reaction rates in GraphEngine calculateRHSAndEnergy."); StepDerivatives result = calculateAllDerivatives( - comp.getMolarAbundanceVector(), + molarAbundances, T9, rho, Ye, @@ -191,6 +194,7 @@ namespace gridfire::engine { return false; } ); + m_most_recent_rhs_calculation = result; return result; } } @@ -435,7 +439,8 @@ namespace gridfire::engine { double factor; if (power == 1) { factor = abundance; } else if (power == 2) { factor = abundance * abundance; } - else { factor = std::pow(abundance, power); } + else if (power == 3) { factor = abundance * abundance * abundance; } + else { factor = std::pow(abundance, static_cast(power)); } if (!std::isfinite(factor)) { LOG_CRITICAL(m_logger, "Non-finite factor encountered in forward abundance product for reaction '{}'. Check input abundances for validity.", reaction.id()); @@ -527,10 +532,13 @@ namespace gridfire::engine { molarReactionFlows.reserve(m_precomputedReactions.size()); size_t reactionCounter = 0; + std::vector reactionIndices; + reactionIndices.reserve(m_precomputedReactions.size()); for (const auto& reaction : activeReactions) { - uint64_t reactionHash = utils::hash_reaction(*reaction); + uint64_t reactionHash = reaction->hash(0); const size_t reactionIndex = m_precomputedReactionIndexMap.at(reactionHash); + reactionIndices.push_back(reactionIndex); const PrecomputedReaction& precomputedReaction = m_precomputedReactions[reactionIndex]; double netFlow = compute_reaction_flow( @@ -556,8 +564,7 @@ namespace gridfire::engine { LOG_TRACE_L3(m_logger, "Computed {} molar reaction flows for active reactions. Assembling these into RHS", molarReactionFlows.size()); reactionCounter = 0; - for (const auto& reaction: activeReactions) { - const size_t j = m_precomputedReactionIndexMap.at(utils::hash_reaction(*reaction)); + for (const auto& [reaction, j]: std::views::zip(activeReactions, reactionIndices)) { const auto& precomp = m_precomputedReactions[j]; const double R_j = molarReactionFlows[reactionCounter]; @@ -828,6 +835,13 @@ namespace gridfire::engine { } + std::optional> GraphEngine::getMostRecentRHSCalculation() const { + if (!m_most_recent_rhs_calculation.has_value()) { + return std::nullopt; + } + return m_most_recent_rhs_calculation.value(); + } + StepDerivatives GraphEngine::calculateAllDerivativesUsingPrecomputation( const fourdst::composition::CompositionAbstract &comp, @@ -1439,7 +1453,7 @@ namespace gridfire::engine { PrecomputedReaction precomp; precomp.reaction_index = i; precomp.reaction_type = reaction.type(); - uint64_t reactionHash = utils::hash_reaction(reaction); + uint64_t reactionHash = reaction.hash(0); precomp.reaction_hash = reactionHash; m_precomputedReactionIndexMap[reactionHash] = i; @@ -1617,7 +1631,7 @@ namespace gridfire::engine { for (size_t k = 0; k < activeReactions.size(); ++k) { int t_id = omp_get_thread_num(); const auto& reaction = activeReactions[k]; - const size_t reactionIndex = m_precomputedReactionIndexMap.at(utils::hash_reaction(reaction)); + const size_t reactionIndex = m_precomputedReactionIndexMap.at(reaction.hash(0)); const PrecomputedReaction& precomputedReaction = m_precomputedReactions[reactionIndex]; double netFlow = compute_reaction_flow( diff --git a/src/lib/engine/views/engine_adaptive.cpp b/src/lib/engine/views/engine_adaptive.cpp index d8eb27c5..e90d846b 100644 --- a/src/lib/engine/views/engine_adaptive.cpp +++ b/src/lib/engine/views/engine_adaptive.cpp @@ -85,47 +85,9 @@ namespace gridfire::engine { ) const { LOG_TRACE_L2(m_logger, "Calculating RHS and Energy in AdaptiveEngineView at T9 = {}, rho = {}.", T9, rho); validateState(); - LOG_TRACE_L2( - m_logger, - "Adaptive engine view state validated prior to composition collection. Input Composition: {}", - [&comp]() -> std::string { - std::stringstream ss; - size_t i = 0; - for (const auto& [species, abundance] : comp) { - ss << species.name() << ": " << abundance; - if (i < comp.size() - 1) { - ss << ", "; - } - i++; - } - return ss.str(); - }()); - fourdst::composition::Composition collectedComp; - std::size_t state_hash = utils::hash_state(comp, T9, rho, m_activeReactions); - if (m_collected_composition_cache.contains(state_hash)) { - collectedComp = m_collected_composition_cache.at(state_hash); - } else { - collectedComp = collectComposition(comp, T9, rho); - m_collected_composition_cache[state_hash] = collectedComp; - } - LOG_TRACE_L2( - m_logger, - "Composition Collected prior to passing to base engine. Collected Composition: {}", - [&comp, &collectedComp]() -> std::string { - std::stringstream ss; - size_t i = 0; - for (const auto& [species, abundance] : collectedComp) { - ss << species.name() << ": " << abundance; - if (comp.contains(species)) { - ss << " (input: " << comp.getMolarAbundance(species) << ")"; - } - if (i < collectedComp.size() - 1) { - ss << ", "; - } - i++; - } - return ss.str(); - }()); + + const fourdst::composition::Composition collectedComp = collectComposition(comp, T9, rho); + auto result = m_baseEngine.calculateRHSAndEnergy(collectedComp, T9, rho, true); LOG_TRACE_L2(m_logger, "Base engine calculation of RHS and Energy complete."); diff --git a/src/lib/engine/views/engine_defined.cpp b/src/lib/engine/views/engine_defined.cpp index 69b1624c..c8b50601 100644 --- a/src/lib/engine/views/engine_defined.cpp +++ b/src/lib/engine/views/engine_defined.cpp @@ -15,8 +15,6 @@ #include #include -#include "fourdst/composition/exceptions/exceptions_composition.h" - namespace gridfire::engine { using fourdst::atomic::Species; @@ -47,7 +45,7 @@ namespace gridfire::engine { ) const { validateNetworkState(); - const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies); + const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies | std::ranges::to()); const auto result = m_baseEngine.calculateRHSAndEnergy(masked, T9, rho, m_activeReactions); if (!result) { @@ -64,7 +62,7 @@ namespace gridfire::engine { ) const { validateNetworkState(); - const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies); + const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies | std::ranges::to()); return m_baseEngine.calculateEpsDerivatives(masked, T9, rho, m_activeReactions); } @@ -78,7 +76,7 @@ namespace gridfire::engine { if (!m_activeSpeciesVectorCache.has_value()) { m_activeSpeciesVectorCache = std::vector(m_activeSpecies.begin(), m_activeSpecies.end()); } - const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies); + const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies | std::ranges::to()); return m_baseEngine.generateJacobianMatrix(masked, T9, rho, m_activeSpeciesVectorCache.value()); } @@ -95,7 +93,7 @@ namespace gridfire::engine { activeSpecies.end() ); - const fourdst::composition::MaskedComposition masked(comp, activeSpeciesSet); + const fourdst::composition::MaskedComposition masked(comp, activeSpeciesSet | std::ranges::to()); return m_baseEngine.generateJacobianMatrix(masked, T9, rho, activeSpecies); } @@ -106,7 +104,7 @@ namespace gridfire::engine { const SparsityPattern &sparsityPattern ) const { validateNetworkState(); - const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies); + const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies | std::ranges::to()); return m_baseEngine.generateJacobianMatrix(masked, T9, rho, sparsityPattern); } @@ -151,7 +149,7 @@ namespace gridfire::engine { throw std::runtime_error("Reaction not found in active reactions: " + std::string(reaction.id())); } - const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies); + const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies | std::ranges::to()); return m_baseEngine.calculateMolarReactionFlow(reaction, masked, T9, rho); } @@ -176,7 +174,7 @@ namespace gridfire::engine { const double rho ) const { validateNetworkState(); - const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies); + const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies | std::ranges::to()); const auto result = m_baseEngine.getSpeciesTimescales(masked, T9, rho, m_activeReactions); if (!result) { @@ -199,7 +197,7 @@ namespace gridfire::engine { const double rho ) const { validateNetworkState(); - const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies); + const fourdst::composition::MaskedComposition masked(comp, m_activeSpecies | std::ranges::to()); const auto result = m_baseEngine.getSpeciesDestructionTimescales(masked, T9, rho, m_activeReactions); diff --git a/src/lib/engine/views/engine_multiscale.cpp b/src/lib/engine/views/engine_multiscale.cpp index e3a55f63..9e372769 100644 --- a/src/lib/engine/views/engine_multiscale.cpp +++ b/src/lib/engine/views/engine_multiscale.cpp @@ -232,8 +232,7 @@ namespace gridfire::engine { } return ss.str(); }()); - // TODO: Figure out why setting trust -> trust causes issues. The only place I think I am setting that to true is in AdaptiveEngineView which has just called getNormalizedEquilibratedComposition... - const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(comp, T9, rho, false); + const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(comp, T9, rho, trust); LOG_TRACE_L2(m_logger, "Equilibrated composition prior to calling base engine is {}", [&qseComposition, &comp]() -> std::string { std::stringstream ss; size_t i = 0; @@ -1095,13 +1094,8 @@ namespace gridfire::engine { LOG_TRACE_L3(m_logger, "Cache Miss in Multiscale Partitioning Engine View for composition at T9 = {}, rho = {}. Solving QSE abundances...", T9, rho); // Only solve if the composition and thermodynamic conditions have not been cached yet - fourdst::composition::Composition qseComposition = solveQSEAbundances(comp, T9, rho); + fourdst::composition::Composition qseComposition(solveQSEAbundances(comp, T9, rho)); - for (const auto &[sp, y]: qseComposition) { - if (y < 0.0 && std::abs(y) < 1e-20) { - qseComposition.setMolarAbundance(sp, 0.0); // normalize small negative abundances to zero - } - } m_composition_cache[composite_hash] = qseComposition; return qseComposition; @@ -1530,8 +1524,13 @@ namespace gridfire::engine { fourdst::composition::Composition outputComposition(comp); + std::vector species; + std::vector abundances; + species.reserve(m_algebraic_species.size()); + abundances.reserve(m_algebraic_species.size()); + for (const auto& [group, solver]: std::views::zip(m_qse_groups, m_qse_solvers)) { - const fourdst::composition::Composition groupResult = solver->solve(outputComposition, T9, rho); + const fourdst::composition::Composition& groupResult = solver->solve(outputComposition, T9, rho); for (const auto& [sp, y] : groupResult) { if (!std::isfinite(y)) { LOG_CRITICAL(m_logger, "Non-finite abundance {} computed for species {} in QSE group solve at T9 = {}, rho = {}.", @@ -1539,10 +1538,16 @@ namespace gridfire::engine { m_logger->flush_log(); throw exceptions::EngineError("Non-finite abundance computed for species " + std::string(sp.name()) + " in QSE group solve."); } - outputComposition.setMolarAbundance(sp, y); + if (y < 0.0 && std::abs(y) < 1e-20) { + abundances.push_back(0.0); + } else { + abundances.push_back(y); + } + species.emplace_back(sp); } - solver->log_diagnostics(group, outputComposition); + // solver->log_diagnostics(group, outputComposition); } + outputComposition.setMolarAbundance(species, abundances); LOG_TRACE_L2(m_logger, "Done solving for QSE abundances!"); return outputComposition; } @@ -1821,7 +1826,7 @@ namespace gridfire::engine { utils::check_sundials_flag(KINSetMaxSetupCalls(m_kinsol_mem, 20), "KINSetMaxSetupCalls", utils::SUNDIALS_RET_CODE_TYPES::KINSOL); - utils::check_sundials_flag(KINSetFuncNormTol(m_kinsol_mem, 1e-6), "KINSetFuncNormTol", utils::SUNDIALS_RET_CODE_TYPES::KINSOL); + utils::check_sundials_flag(KINSetFuncNormTol(m_kinsol_mem, 1e-8), "KINSetFuncNormTol", utils::SUNDIALS_RET_CODE_TYPES::KINSOL); utils::check_sundials_flag(KINSetNumMaxIters(m_kinsol_mem, 200), "KINSetNumMaxIters", utils::SUNDIALS_RET_CODE_TYPES::KINSOL); utils::check_sundials_flag(KINSetScaledStepTol(m_kinsol_mem, 1e-10), "KINSetScaledStepTol", utils::SUNDIALS_RET_CODE_TYPES::KINSOL); @@ -1899,15 +1904,22 @@ namespace gridfire::engine { scale_data[i] = 1.0 / Y; } - auto initial_rhs = m_engine.calculateRHSAndEnergy(result, T9, rho, false); - if (!initial_rhs) { - throw std::runtime_error("In QSE solver failed to calculate initial RHS"); + StepDerivatives rhsGuess; + auto cached_rhs = m_engine.getMostRecentRHSCalculation(); + if (!cached_rhs) { + const auto initial_rhs = m_engine.calculateRHSAndEnergy(result, T9, rho, false); + if (!initial_rhs) { + throw std::runtime_error("In QSE solver failed to calculate initial RHS for caching"); + } + rhsGuess = initial_rhs.value(); + } else { + rhsGuess = cached_rhs.value(); } sunrealtype* f_scale_data = N_VGetArrayPointer(m_f_scale); for (size_t i = 0; i < m_N; ++i) { const auto& species = m_species[i]; - double dydt = std::abs(initial_rhs.value().dydt.at(species)); + double dydt = std::abs(rhsGuess.dydt.at(species)); f_scale_data[i] = 1.0 / (dydt + 1e-15); } diff --git a/src/lib/solver/strategies/CVODE_solver_strategy.cpp b/src/lib/solver/strategies/CVODE_solver_strategy.cpp index b9e3c412..f7fc5ed0 100644 --- a/src/lib/solver/strategies/CVODE_solver_strategy.cpp +++ b/src/lib/solver/strategies/CVODE_solver_strategy.cpp @@ -74,7 +74,7 @@ namespace gridfire::solver { } - CVODESolverStrategy::CVODESolverStrategy(DynamicEngine &engine): NetworkSolverStrategy(engine) { + CVODESolverStrategy::CVODESolverStrategy(DynamicEngine &engine): SingleZoneNetworkSolverStrategy(engine) { // TODO: In order to support MPI this function must be changed const int flag = SUNContext_Create(SUN_COMM_NULL, &m_sun_ctx); if (flag < 0) { diff --git a/src/lib/solver/strategies/SpectralSolverStrategy.cpp b/src/lib/solver/strategies/SpectralSolverStrategy.cpp index e69de29b..b2b48afd 100644 --- a/src/lib/solver/strategies/SpectralSolverStrategy.cpp +++ b/src/lib/solver/strategies/SpectralSolverStrategy.cpp @@ -0,0 +1,770 @@ +#include "gridfire/solver/strategies/SpectralSolverStrategy.h" + +#include + +#include "gridfire/utils/sundials.h" + +#include "quill/LogMacros.h" +#include "sunmatrix/sunmatrix_dense.h" + +namespace { + std::pair> evaluate_bspline( + double x, + const gridfire::solver::SpectralSolverStrategy::SplineBasis& basis + ) { + const int p = basis.degree; + const std::vector& t = basis.knots; + + auto it = std::ranges::upper_bound(t, x); + size_t i = std::distance(t.begin(), it) - 1; + + if (i < static_cast(p)) i = p; + if (i >= t.size() - 1 - p) i = t.size() - 2 - p; + + if (x >= t.back()) { + i = t.size() - p - 2; + } + + // Cox-de Boor algorithm + std::vector N(p + 1); + std::vector left(p + 1); + std::vector right(p + 1); + + N[0] = 1.0; + + for (int j = 1; j <= p; ++j) { + left[j] = x - t[i + 1 - j]; + right[j] = t[i + j] - x; + double saved = 0.0; + + for (int r = 0; r < j; ++r) { + double temp = N[r] / (right[r + 1] + left[j - r]); + N[r] = saved + right[r + 1] * temp; + + saved = left[j - r] * temp; + } + N[j] = saved; + } + return {i - p, N}; + } +} + +namespace gridfire::solver { + + SpectralSolverStrategy::SpectralSolverStrategy(engine::DynamicEngine& engine) : MultiZoneNetworkSolverStrategy (engine) { + LOG_INFO(m_logger, "Initializing SpectralSolverStrategy"); + + utils::check_sundials_flag(SUNContext_Create(SUN_COMM_NULL, &m_sun_ctx), "SUNContext_Create", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + + m_absTol = m_config->solver.spectral.absTol; + m_relTol = m_config->solver.spectral.relTol; + + LOG_INFO(m_logger, "SpectralSolverStrategy initialized successfully"); + } + + SpectralSolverStrategy::~SpectralSolverStrategy() { + LOG_INFO(m_logger, "Destroying SpectralSolverStrategy"); + + + if (m_cvode_mem) { + CVodeFree(&m_cvode_mem); + m_cvode_mem = nullptr; + } + + if (m_LS) SUNLinSolFree(m_LS); + if (m_J) SUNMatDestroy(m_J); + if (m_Y) N_VDestroy(m_Y); + if (m_constraints) N_VDestroy(m_constraints); + + if (m_sun_ctx) { + SUNContext_Free(&m_sun_ctx); + m_sun_ctx = nullptr; + } + + if (m_T_coeffs) N_VDestroy(m_T_coeffs); + if (m_rho_coeffs) N_VDestroy(m_rho_coeffs); + + LOG_INFO(m_logger, "SpectralSolverStrategy destroyed successfully"); + } + + //////////////////////////////////////////////////////////////////////////////// + /// Main Evaluation Loop + ///////////////////////////////////////////////////////////////////////////////// + + std::vector SpectralSolverStrategy::evaluate(const std::vector& netIns, const std::vector& mass_coords) { + LOG_INFO(m_logger, "Starting spectral solver evaluation for {} zones", netIns.size()); + + assert(std::ranges::all_of(netIns, [&netIns](const NetIn& in) { return in.tMax == netIns[0].tMax; }) && "All NetIn entries must have the same tMax for spectral solver evaluation."); + + std::vector updatedNetIns = netIns; + for (auto& netIn : updatedNetIns) { + netIn.composition = m_engine.update(netIn); + } + + + ///////////////////////////////////// + /// Evaluate the monitor function /// + ///////////////////////////////////// + + const std::vector monitor_function = evaluate_monitor_function(updatedNetIns); + m_current_basis = generate_basis_from_monitor(monitor_function, mass_coords); + + size_t num_basis_funcs = m_current_basis.knots.size() - m_current_basis.degree - 1; + + std::vector shell_cache(updatedNetIns.size()); + for (size_t shellID = 0; shellID < shell_cache.size(); ++shellID) { + auto [start, phi] = evaluate_bspline(mass_coords[shellID], m_current_basis); + shell_cache[shellID] = {.start_idx=start, .phi=phi}; + } + + DenseLinearSolver proj_solver(num_basis_funcs, m_sun_ctx); + proj_solver.init_from_cache(num_basis_funcs, shell_cache); + + if (m_T_coeffs) N_VDestroy(m_T_coeffs); + m_T_coeffs = N_VNew_Serial(static_cast(num_basis_funcs), m_sun_ctx); + project_specific_variable(updatedNetIns, mass_coords, shell_cache, proj_solver, m_T_coeffs, 0, [](const NetIn& s) { return s.temperature; }, true); + + if (m_rho_coeffs) N_VDestroy(m_rho_coeffs); + m_rho_coeffs = N_VNew_Serial(static_cast(num_basis_funcs), m_sun_ctx); + project_specific_variable(updatedNetIns, mass_coords, shell_cache, proj_solver, m_rho_coeffs, 0, [](const NetIn& s) { return s.density; }, true); + + size_t num_species = m_engine.getNetworkSpecies().size(); + size_t current_offset = 0; + + size_t total_coefficients = num_basis_funcs * (num_species + 1); + + if (m_Y) N_VDestroy(m_Y); + if (m_constraints) N_VDestroy(m_constraints); + + m_Y = N_VNew_Serial(static_cast(total_coefficients), m_sun_ctx); + m_constraints = N_VClone(m_Y); + N_VConst(0.0, m_constraints); // For now no constraints on coefficients + + for (const auto& sp : m_engine.getNetworkSpecies()) { + project_specific_variable( + updatedNetIns, + mass_coords, + shell_cache, + proj_solver, + m_Y, + current_offset, + [&sp](const NetIn& s) { return s.composition.getMolarAbundance(sp); }, + false + ); + current_offset += num_basis_funcs; + } + + sunrealtype* y_data = N_VGetArrayPointer(m_Y); + const size_t energy_offset = num_species * num_basis_funcs; + + assert(energy_offset == current_offset && "Energy offset calculation mismatch in spectral solver initialization."); + + for (size_t i = 0; i < num_basis_funcs; ++i) { + y_data[energy_offset + i] = 0.0; + } + + DenseLinearSolver mass_solver(num_basis_funcs, m_sun_ctx); + mass_solver.init_from_basis(num_basis_funcs, m_current_basis); + + ///////////////////////////////////// + /// CVODE Initialization /// + ///////////////////////////////////// + + CVODEUserData data; + data.solver_instance = this; + data.engine = &m_engine; + data.mass_matrix_solver_instance = &mass_solver; + data.basis = &m_current_basis; + + const double absTol = m_absTol.value_or(1e-10); + const double relTol = m_relTol.value_or(1e-6); + + const bool size_changed = m_last_size != total_coefficients; + m_last_size = total_coefficients; + + if (m_cvode_mem == nullptr || size_changed) { + if (m_cvode_mem) { + CVodeFree(&m_cvode_mem); + m_cvode_mem = nullptr; + } + if (m_LS) { + SUNLinSolFree(m_LS); + m_LS = nullptr; + } + if (m_J) { + SUNMatDestroy(m_J); + m_J = nullptr; + } + + m_cvode_mem = CVodeCreate(CV_BDF, m_sun_ctx); + utils::check_sundials_flag(m_cvode_mem == nullptr ? -1 : 0, "CVodeCreate", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + + utils::check_sundials_flag(CVodeInit(m_cvode_mem, cvode_rhs_wrapper, 0.0, m_Y), "CVodeInit", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + m_J = SUNDenseMatrix(static_cast(total_coefficients), static_cast(total_coefficients), m_sun_ctx); + m_LS = SUNLinSol_Dense(m_Y, m_J, m_sun_ctx); + utils::check_sundials_flag(CVodeSetLinearSolver(m_cvode_mem, m_LS, m_J), "CVodeSetLinearSolver", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + + // For now, we will not attach a Jacobian function, using finite differences + } else { + utils::check_sundials_flag(CVodeReInit(m_cvode_mem, 0.0, m_Y), "CVodeReInit", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + } + + utils::check_sundials_flag(CVodeSStolerances(m_cvode_mem, relTol, absTol), "CVodeSStolerances", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + utils::check_sundials_flag(CVodeSetUserData(m_cvode_mem, &data), "CVodeSetUserData", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + + ///////////////////////////////////// + /// Time Integration Loop /// + ///////////////////////////////////// + + const double target_time = updatedNetIns[0].tMax; + double current_time = 0.0; + + while (current_time < target_time) { + int flag = CVode(m_cvode_mem, target_time, m_Y, ¤t_time, CV_ONE_STEP); + utils::check_sundials_flag(flag, "CVode", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + + std::println("Advanced to time: {:10.4e} / {:10.4e}", current_time, target_time); + } + + std::vector results = reconstruct_solution(updatedNetIns, mass_coords, m_Y, m_current_basis, target_time); + return results; + } + + void SpectralSolverStrategy::set_callback(const std::any &callback) { + m_callback = std::any_cast(callback); + } + + std::vector> SpectralSolverStrategy::describe_callback_context() const { + throw std::runtime_error("SpectralSolverStrategy does not yet implement describe_callback_context."); + } + + bool SpectralSolverStrategy::get_stdout_logging_enabled() const { + return m_stdout_logging_enabled; + } + + void SpectralSolverStrategy::set_stdout_logging_enabled(bool logging_enabled) { + m_stdout_logging_enabled = logging_enabled; + } + + //////////////////////////////////////////////////////////////////////////////// + /// Static Wrappers for SUNDIALS Callbacks + //////////////////////////////////////////////////////////////////////////////// + + int SpectralSolverStrategy::cvode_rhs_wrapper( + const sunrealtype t, + const N_Vector y_coeffs, + const N_Vector ydot_coeffs, + void *user_data + ) { + auto *data = static_cast(user_data); + const auto *instance = data->solver_instance; + + try { + return instance -> calculate_rhs(t, y_coeffs, ydot_coeffs, data); + } catch (const std::exception& e) { + LOG_CRITICAL(instance->m_logger, "Uncaught exception in Spectral Solver RHS wrapper at time {}: {}", t, e.what()); + return -1; + } catch (...) { + LOG_CRITICAL(instance->m_logger, "Unknown uncaught exception in Spectral Solver RHS wrapper at time {}", t); + return -1; + } + } + + int SpectralSolverStrategy::cvode_jac_wrapper( + const sunrealtype t, + const N_Vector y, + const N_Vector ydot, + const SUNMatrix J, + void *user_data, + const N_Vector tmp1, + const N_Vector tmp2, + const N_Vector tmp3 + ) { + const auto *data = static_cast(user_data); + const auto *instance = data->solver_instance; + + try { + LOG_WARNING_LIMIT_EVERY_N(1000, instance->m_logger, "Analytic Jacobian Generation not yet implemented, using finite difference approximation"); + return 0; + } catch (const std::exception& e) { + LOG_CRITICAL(instance->m_logger, "Uncaught exception in Spectral Solver Jacobian wrapper at time {}: {}", t, e.what()); + return -1; + } catch (...) { + LOG_CRITICAL(instance->m_logger, "Unknown uncaught exception in Spectral Solver Jacobian wrapper at time {}", t); + return -1; + } + } + + //////////////////////////////////////////////////////////////////////////////// + /// RHS implementation + //////////////////////////////////////////////////////////////////////////////// + + int SpectralSolverStrategy::calculate_rhs( + sunrealtype t, + N_Vector y_coeffs, + N_Vector ydot_coeffs, + CVODEUserData* data + ) const { + const auto& basis = m_current_basis; + DenseLinearSolver* mass_solver = data->mass_matrix_solver_instance; + const auto& species_list = m_engine.getNetworkSpecies(); + + const size_t num_basis_funcs = basis.knots.size() - basis.degree - 1; + const size_t num_species = species_list.size(); + + sunrealtype* rhs_data = N_VGetArrayPointer(ydot_coeffs); + N_VConst(0.0, ydot_coeffs); + + // PERF: In future we can use openMP to parallelize over these basis functions once we make the engines thread safe + for (size_t q = 0; q < basis.quadrature_nodes.size(); ++q) { + double w_q = basis.quadrature_weights[q]; + + const auto& [start_idx, phi] = basis.quad_evals[q]; + + GridPoint gp = reconstruct_at_quadrature(y_coeffs, q, basis); + std::expected, engine::EngineStatus> results = m_engine.calculateRHSAndEnergy(gp.composition, gp.T9, gp.rho, false); + + // PERF: When switching to parallel execution, we will need to protect this section with a mutex or use atomic operations since we cannot throw safely from multiple threads + if (!results) { + LOG_CRITICAL(m_logger, "Engine failed to calculate RHS at time {}: {}", t, EngineStatus_to_string(results.error())); + return -1; + } + + const auto& [dydt, eps_nuc, contributions, nu_loss, nu_flux] = results.value(); + + for (size_t s = 0; s < num_species; ++s) { + double rate = dydt.at(species_list[s]); + size_t species_offset = s * num_basis_funcs; + + for (size_t k = 0; k < phi.size(); ++k) { + size_t global_idx = species_offset + start_idx + k; + rhs_data[global_idx] += w_q * phi[k] * rate; + } + } + + size_t energy_offset = num_species * num_basis_funcs; + + for (size_t k = 0; k < phi.size(); ++k) { + size_t global_idx = energy_offset + start_idx + k; + rhs_data[global_idx] += eps_nuc * w_q * phi[k]; + } + } + + size_t total_vars = num_species + 1; + mass_solver->solve_inplace(ydot_coeffs, total_vars, num_basis_funcs); + + return 0; + } + + + //////////////////////////////////////////////////////////////////////////////// + /// Spectral Utilities + /// These include basis generation, monitor function evaluation + /// projection and reconstruction routines. + //////////////////////////////////////////////////////////////////////////////// + + std::vector SpectralSolverStrategy::evaluate_monitor_function(const std::vector& current_shells) const { + const size_t n_shells = current_shells.size(); + if (n_shells < 3) { + return std::vector(n_shells, 1.0); // NOLINT(*-return-braced-init-list) + } + + std::vector M(n_shells, 1.0); + + auto accumulate_variable = [&](auto getter, double weight, bool use_log) { + std::vector data(n_shells); + double min_val = std::numeric_limits::max(); + double max_val = std::numeric_limits::lowest(); + + for (size_t i = 0 ; i < n_shells; ++i) { + double val = getter(current_shells[i]); + if (use_log) { + val = std::log10(std::max(val, 1e-100)); + } + + data[i] = val; + + if (val < min_val) min_val = val; + if (val > max_val) max_val = val; + } + + const double scale = max_val - min_val; + if (scale < 1e-10) return; + + for (size_t i = 1; i < n_shells - 1; ++i) { + const double v_prev = data[i-1]; + const double v_curr = data[i]; + const double v_next = data[i+1]; + + // Finite difference estimates for first and second derivatives + double d1 = std::abs(v_next - v_prev) / 2.0; + double d2 = std::abs(v_next - 2.0 * v_curr + v_prev); + + d1 /= scale; + d2 /= scale; + + const double alpha = m_config->solver.spectral.monitorFunction.alpha; + const double beta = m_config->solver.spectral.monitorFunction.beta; + + M[i] += weight * (alpha * d1 + beta * d2); + } + }; + + const double structure_weight = m_config->solver.spectral.monitorFunction.structure_weight; + double abundance_weight = m_config->solver.spectral.monitorFunction.abundance_weight; + accumulate_variable([](const NetIn& s) { return s.temperature; }, structure_weight, true); + accumulate_variable([](const NetIn& s) { return s.density; }, structure_weight, true); + + for (const auto& sp : m_engine.getNetworkSpecies()) { + accumulate_variable([&sp](const NetIn& s) { return s.composition.getMolarAbundance(sp); }, abundance_weight, false); + } + + ////////////////////////////// + /// Smoothing the Monitor /// + ////////////////////////////// + + std::vector M_smooth = M; + for (size_t i = 1; i < n_shells - 1; ++i) { + M_smooth[i] = (M[i-1] + 2.0 * M[i] + M[i+1]) / 4.0; + } + + M_smooth[0] = M_smooth[1]; + M_smooth[n_shells-1] = M_smooth[n_shells-2]; + + return M_smooth; + + } + + SpectralSolverStrategy::SplineBasis SpectralSolverStrategy::generate_basis_from_monitor( + const std::vector& monitor_values, + const std::vector& mass_coordinates + ) const { + SplineBasis basis; + basis.degree = 3; // Cubic Spline + + const size_t n_shells = monitor_values.size(); + + std::vector I(n_shells, 0.0); + double current_integral = 0.0; + + for (size_t i = 1; i < n_shells; ++i) { + const double dx = mass_coordinates[i] - mass_coordinates[i-1]; + double dI = 0.5 * (monitor_values[i] + monitor_values[i-1]) * dx; + + dI = std::max(dI, 1e-30); + current_integral += dI; + I[i] = current_integral; + } + + const double total_integral = I.back(); + for (size_t i = 0; i < n_shells; ++i) { + I[i] /= total_integral; + } + + const size_t num_elements = m_config->solver.spectral.basis.num_elements; + basis.knots.reserve(num_elements + 1 + 2 * basis.degree); + + // Note that these imply that mass_coordinates must be sorted in increasing order + double min_mass = mass_coordinates.front(); + double max_mass = mass_coordinates.back(); + + for (int i = 0; i < basis.degree; ++i) { + basis.knots.push_back(min_mass); + } + + for (size_t k = 1; k < num_elements; ++k) { + double target_I = static_cast(k) / static_cast(num_elements); + + auto it = std::ranges::lower_bound(I, target_I); + size_t idx = std::distance(I.begin(), it); + + if (idx == 0) idx = 1; + if (idx >= n_shells) idx = n_shells - 1; + + double I0 = I[idx-1]; + double I1 = I[idx]; + double m0 = mass_coordinates[idx-1]; + double m1 = mass_coordinates[idx]; + + double fraction = (target_I - I0) / (I1 - I0); + double knot_location = m0 + fraction * (m1 - m0); + + basis.knots.push_back(knot_location); + } + + for (int i = 0; i < basis.degree; ++i) { + basis.knots.push_back(max_mass); + } + + constexpr double sqrt_3_over_5 = 0.77459666924; + constexpr double five_over_nine = 5.0 / 9.0; + constexpr double eight_over_nine = 8.0 / 9.0; + static constexpr std::array gl_nodes = {-sqrt_3_over_5, 0.0, sqrt_3_over_5}; + static constexpr std::array gl_weights = {five_over_nine, eight_over_nine, five_over_nine}; + + basis.quadrature_nodes.clear(); + basis.quadrature_weights.clear(); + + for (size_t i = basis.degree; i < basis.knots.size() - basis.degree - 1; ++i) { + double a = basis.knots[i]; + double b = basis.knots[i+1]; + + if ( b - a < 1e-14) continue; + + double mid = 0.5 * (a + b); + double half_width = 0.5 * (b - a); + + for (size_t j = 0; j < gl_nodes.size(); ++j) { + double phys_node = mid + gl_nodes[j] * half_width; + double phys_weight = gl_weights[j] * half_width; + + basis.quadrature_nodes.push_back(phys_node); + basis.quadrature_weights.push_back(phys_weight); + + auto [start, phi] = evaluate_bspline(phys_node, basis); + basis.quad_evals.push_back({start, phi}); + } + } + + return basis; + } + + + SpectralSolverStrategy::GridPoint SpectralSolverStrategy::reconstruct_at_quadrature( + const N_Vector y_coeffs, + const size_t quad_index, + const SplineBasis &basis + ) const { + auto [start_idx, vals] = basis.quad_evals[quad_index]; + + const sunrealtype* T_ptr = N_VGetArrayPointer(m_T_coeffs); + const sunrealtype* rho_ptr = N_VGetArrayPointer(m_rho_coeffs); + const sunrealtype* y_data = N_VGetArrayPointer(y_coeffs); + + const size_t num_basis_funcs = basis.knots.size() - basis.degree - 1; + const std::vector& species_list = m_engine.getNetworkSpecies(); + const size_t num_species = species_list.size(); + + double logT = 0.0; + double logRho = 0.0; + + for (size_t k = 0; k < vals.size(); ++k) { + size_t idx = start_idx + k; + logT += T_ptr[idx] * vals[k]; + logRho += rho_ptr[idx] * vals[k]; + } + + GridPoint result; + result.T9 = std::pow(10.0, logT) / 1e9; + result.rho = std::pow(10.0, logRho); + + for (size_t s = 0; s < num_species; ++s) { + const fourdst::atomic::Species& species = species_list[s]; + double abundance = 0.0; + const size_t offset = s * num_basis_funcs; + + for (size_t k = 0; k < vals.size(); ++k) { + abundance += y_data[offset + start_idx + k] * vals[k]; + } + + // Note: It is possible this will lead to a loss of mass conservation. In future we may want to implement a better way to handle this. + if (abundance < 0.0) abundance = 0.0; + + result.composition.registerSpecies(species); + result.composition.setMolarAbundance(species, abundance); + } + + return result; + } + + std::vector SpectralSolverStrategy::reconstruct_solution( + const std::vector& original_inputs, + const std::vector& mass_coordinates, + const N_Vector final_coeffs, + const SplineBasis& basis, + const double dt + ) const { + const size_t n_shells = original_inputs.size(); + const size_t num_basis_funcs = basis.knots.size() - basis.degree - 1; + + std::vector outputs; + outputs.reserve(n_shells); + + const sunrealtype* c_data = N_VGetArrayPointer(final_coeffs); + + const auto& species_list = m_engine.getNetworkSpecies(); + for (size_t shellID = 0; shellID < n_shells; ++shellID) { + const double x = mass_coordinates[shellID]; + + auto [start_idx, vals] = evaluate_bspline(x, basis); + + auto reconstruct_var = [&](const size_t coeff_offset) -> double { + double result = 0.0; + for (size_t i = 0; i < vals.size(); ++i) { + result += c_data[coeff_offset + start_idx + i] * vals[i]; + } + return result; + }; + + fourdst::composition::Composition comp_new; + + for (size_t s_idx = 0; s_idx < species_list.size(); ++s_idx) { + const fourdst::atomic::Species& sp = species_list[s_idx]; + comp_new.registerSpecies(sp); + + const size_t current_offset = s_idx * num_basis_funcs; + double Y_val = reconstruct_var(current_offset); + + if (Y_val < 0.0 && Y_val > -1.0e-16) { + Y_val = 0.0; + } + + if (Y_val < 0.0 && Y_val > -1e-16) Y_val = 0.0; + if (Y_val >= 0.0) { + comp_new.setMolarAbundance(sp, Y_val); + } + } + + const double energy = reconstruct_var(species_list.size() * num_basis_funcs); + + NetOut netOut; + netOut.composition = comp_new; + netOut.energy = energy; + netOut.num_steps = -1; // Not tracked in spectral solver + + outputs.push_back(std::move(netOut)); + } + + return outputs; + } + + void SpectralSolverStrategy::project_specific_variable( + const std::vector ¤t_shells, + const std::vector &mass_coordinates, + const std::vector &shell_cache, + const DenseLinearSolver &linear_solver, + N_Vector output_vec, + size_t output_offset, + const std::function &getter, + bool use_log + ) { + const size_t n_shells = current_shells.size(); + + sunrealtype* out_ptr = N_VGetArrayPointer(output_vec); + size_t basis_size = N_VGetLength(linear_solver.temp_vector); + + for (size_t i = 0; i < basis_size; ++i ) { + out_ptr[output_offset + i] = 0.0; + } + + for (size_t shellID = 0; shellID < n_shells; ++shellID) { + double val = getter(current_shells[shellID]); + if (use_log) val = std::log10(std::max(val, 1e-100)); + + const auto& eval = shell_cache[shellID]; + + for (size_t i = 0; i < eval.phi.size(); ++i) { + out_ptr[output_offset + eval.start_idx + i] += val * eval.phi[i]; + } + } + + sunrealtype* tmp_data = N_VGetArrayPointer(linear_solver.temp_vector); + for (size_t i = 0; i < basis_size; ++i) tmp_data[i] = out_ptr[output_offset + i]; + + SUNLinSolSolve(linear_solver.LS, linear_solver.A, linear_solver.temp_vector, linear_solver.temp_vector, 0.0); + + for (size_t i = 0; i < basis_size; ++i) out_ptr[output_offset + i] = tmp_data[i]; + } + + /////////////////////////////////////////////////////////////////////////////// + /// SpectralSolverStrategy::MassMatrixSolver Implementation + /////////////////////////////////////////////////////////////////////////////// + + SpectralSolverStrategy::DenseLinearSolver::DenseLinearSolver( + size_t size, + SUNContext sun_ctx + ) : ctx(sun_ctx) { + A = SUNDenseMatrix(size, size, sun_ctx); + temp_vector = N_VNew_Serial(size, sun_ctx); + + LS = SUNLinSol_Dense(temp_vector, A, sun_ctx); + + if (!A || !temp_vector || !LS) { + throw std::runtime_error("Failed to create MassMatrixSolver components."); + } + + zero(); + } + + SpectralSolverStrategy::DenseLinearSolver::~DenseLinearSolver() { + if (LS) SUNLinSolFree(LS); + if (A) SUNMatDestroy(A); + if (temp_vector) N_VDestroy(temp_vector); + } + + void SpectralSolverStrategy::DenseLinearSolver::zero() const { + SUNMatZero(A); + } + + void SpectralSolverStrategy::DenseLinearSolver::init_from_cache( + const size_t num_basis_funcs, + const std::vector &shell_cache + ) const { + sunrealtype* a_data = SUNDenseMatrix_Data(A); + for (const auto&[start_idx, phi] : shell_cache) { + for (size_t i = 0; i < phi.size(); ++i) { + const size_t row = start_idx + i; + + for (size_t j = 0; j < phi.size(); ++j) { + const size_t col = start_idx + j; + + a_data[col * num_basis_funcs + row] += phi[i] * phi[j]; + } + } + } + setup(); + } + + void SpectralSolverStrategy::DenseLinearSolver::init_from_basis( + const size_t num_basis_funcs, + const SplineBasis &basis + ) const { + sunrealtype* m_data = SUNDenseMatrix_Data(A); + for (size_t q = 0; q < basis.quadrature_nodes.size(); ++q) { + double w_q = basis.quadrature_weights[q]; + const auto& eval = basis.quad_evals[q]; + + for (size_t i = 0; i < eval.phi.size(); ++i) { + size_t row = eval.start_idx + i; + for (size_t j = 0; j < eval.phi.size(); ++j) { + size_t col = eval.start_idx + j; + m_data[col * num_basis_funcs + row] += w_q * eval.phi[j] * eval.phi[i]; + } + } + } + setup(); + } + + void SpectralSolverStrategy::DenseLinearSolver::setup() const { + utils::check_sundials_flag(SUNLinSolSetup(LS, A), "SUNLinSolSetup - Mass Matrix Solver", utils::SUNDIALS_RET_CODE_TYPES::CVODE); + } + + // ReSharper disable once CppMemberFunctionMayBeConst + void SpectralSolverStrategy::DenseLinearSolver::solve_inplace(const N_Vector x, const size_t num_vars, const size_t basis_size) const { + sunrealtype* x_data = N_VGetArrayPointer(x); + sunrealtype* tmp_data = N_VGetArrayPointer(temp_vector); + + for (size_t v = 0; v < num_vars; ++v) { + const size_t offset = v * basis_size; + + for (size_t i = 0; i < basis_size; ++i) { + tmp_data[i] = x_data[offset + i]; + } + SUNLinSolSolve(LS, A, temp_vector, temp_vector, 0.0); + + for (size_t i = 0; i < basis_size; ++i) { + x_data[offset + i] = tmp_data[i]; + } + } + } +} diff --git a/src/meson.build b/src/meson.build index 17ea2b21..f7af84f3 100644 --- a/src/meson.build +++ b/src/meson.build @@ -16,6 +16,7 @@ gridfire_sources = files( 'lib/io/network_file.cpp', 'lib/io/generative/python.cpp', 'lib/solver/strategies/CVODE_solver_strategy.cpp', + 'lib/solver/strategies/SpectralSolverStrategy.cpp', 'lib/solver/strategies/triggers/engine_partitioning_trigger.cpp', 'lib/screening/screening_types.cpp', 'lib/screening/screening_weak.cpp', @@ -40,6 +41,7 @@ gridfire_build_dependencies = [ eigen_dep, sundials_dep, json_dep, + uod_dep, ] if get_option('use_mimalloc') diff --git a/tests/graphnet_sandbox/meson.build b/tests/graphnet_sandbox/meson.build index ef538396..65054041 100644 --- a/tests/graphnet_sandbox/meson.build +++ b/tests/graphnet_sandbox/meson.build @@ -3,3 +3,9 @@ executable( 'main.cpp', dependencies: [gridfire_dep, cli11_dep], ) + +executable( + 'spectral_sandbox', + 'spectral_main.cpp', + dependencies: [gridfire_dep, cli11_dep] +) diff --git a/tests/graphnet_sandbox/spectral_main.cpp b/tests/graphnet_sandbox/spectral_main.cpp index 7cb80f45..9f0212e0 100644 --- a/tests/graphnet_sandbox/spectral_main.cpp +++ b/tests/graphnet_sandbox/spectral_main.cpp @@ -36,11 +36,12 @@ std::vector linspace(const double start, const double end, const size_t return result; } -std::vector init(const double tMin, const double tMax, const double rhoMin, const double rhoMax, const double nShells, const double tMax) { +std::vector init(const double tMin, const double tMax, const double rhoMin, const double rhoMax, const double nShells, const double evolveTime) { std::setlocale(LC_ALL, ""); g_previousHandler = std::set_terminate(quill_terminate_handler); quill::Logger* logger = fourdst::logging::LogManager::getInstance().getLogger("log"); logger->set_log_level(quill::LogLevel::TraceL2); + LOG_INFO(logger, "Initializing GridFire Spectral Solver Sandbox..."); using namespace gridfire; const std::vector X = {0.7081145999999999, 2.94e-5, 0.276, 0.003, 0.0011, 9.62e-3, 1.62e-3, 5.16e-4}; @@ -57,7 +58,7 @@ std::vector init(const double tMin, const double tMax, const do netIn.density = ρ; netIn.energy = 0; - netIn.tMax = tMax; + netIn.tMax = evolveTime; netIn.dt0 = 1e-12; netIns.push_back(netIn); @@ -83,19 +84,19 @@ int main(int argc, char** argv) { double tMax = 2.5e7; double rhoMin = 1.0e2; double rhoMax = 1.0e4; - double nShells = 500; - double tMax = 3.1536e+16; + double nShells = 5; + double evolveTime = 3.1536e+16; app.add_option("--tMin", tMin, "Minimum time in seconds"); app.add_option("--tMax", tMax, "Maximum time in seconds"); app.add_option("--rhoMin", rhoMin, "Minimum density in g/cm^3"); app.add_option("--rhoMax", rhoMax, "Maximum density in g/cm^3"); app.add_option("--nShells", nShells, "Number of shells"); - app.add_option("--tMax", tMax, "Maximum time in seconds"); + app.add_option("--evolveTime", evolveTime, "Maximum time in seconds"); CLI11_PARSE(app, argc, argv); - std::vector netIns = init(tMin, tMax, rhoMin, rhoMax, nShells, tMax); + const std::vector netIns = init(tMin, tMax, rhoMin, rhoMax, nShells, evolveTime); policy::MainSequencePolicy stellarPolicy(netIns[0].composition); stellarPolicy.construct();