engine_graph.h

Go to the documentation of this file.

#pragma once
 
#include "fourdst/composition/atomicSpecies.h"
#include "fourdst/composition/composition.h"
#include "fourdst/logging/logging.h"
#include "fourdst/config/config.h"
 
#include "gridfire/network.h"
#include "gridfire/reaction/reaction.h"
#include "gridfire/engine/engine_abstract.h"
#include "gridfire/screening/screening_abstract.h"
#include "gridfire/screening/screening_types.h"
#include "gridfire/partition/partition_abstract.h"
#include "gridfire/engine/procedures/construction.h"
 
#include <string>
#include <unordered_map>
#include <vector>
#include <memory>
#include <ranges>
 
#include <boost/numeric/ublas/matrix_sparse.hpp>
 
#include "cppad/cppad.hpp"
#include "cppad/utility/sparse_rc.hpp"
#include "cppad/speed/sparse_jac_fun.hpp"
 
#include "procedures/priming.h"
 
 
#include "quill/LogMacros.h"
 
// PERF: The function getNetReactionStoichiometry returns a map of species to their stoichiometric coefficients for a given reaction.
//       this makes extra copies of the species, which is not ideal and could be optimized further.
//       Even more relevant is the member m_reactionIDMap which makes copies of a REACLIBReaction for each reaction ID.
//       REACLIBReactions are quite large data structures, so this could be a performance bottleneck.
// static bool isF17 = false;
namespace gridfire {
    typedef CppAD::AD<double> ADDouble;
 
    using fourdst::config::Config;
    using fourdst::logging::LogManager;
    using fourdst::constant::Constants;

    static constexpr double MIN_DENSITY_THRESHOLD = 1e-18;

    static constexpr double MIN_ABUNDANCE_THRESHOLD = 1e-18;

    static constexpr double MIN_JACOBIAN_THRESHOLD = 1e-24;
 

    class GraphEngine final : public DynamicEngine{
    public:
        explicit GraphEngine(
            const fourdst::composition::Composition &composition,
            const BuildDepthType = NetworkBuildDepth::Full
        );
 
        explicit GraphEngine(
            const fourdst::composition::Composition &composition,
            const partition::PartitionFunction& partitionFunction,
            const BuildDepthType buildDepth = NetworkBuildDepth::Full
        );

        explicit GraphEngine(const reaction::LogicalReactionSet &reactions);

        [[nodiscard]] std::expected<StepDerivatives<double>, expectations::StaleEngineError> calculateRHSAndEnergy(
            const std::vector<double>& Y,
            const double T9,
            const double rho
        ) const override;

        void generateJacobianMatrix(
            const std::vector<double>& Y_dynamic,
            const double T9,
            const double rho
        ) const override;
 
        void generateJacobianMatrix(
            const std::vector<double> &Y_dynamic,
            double T9,
            double rho,
            const SparsityPattern &sparsityPattern
        ) const override;

        void generateStoichiometryMatrix() override;

        [[nodiscard]] double calculateMolarReactionFlow(
            const reaction::Reaction& reaction,
            const std::vector<double>&Y,
            const double T9,
            const double rho
        ) const override;

        [[nodiscard]] const std::vector<fourdst::atomic::Species>& getNetworkSpecies() const override;

        [[nodiscard]] const reaction::LogicalReactionSet& getNetworkReactions() const override;
 
        void setNetworkReactions(const reaction::LogicalReactionSet& reactions) override;

        [[nodiscard]] double getJacobianMatrixEntry(
            const int i,
            const int j
        ) const override;

        [[nodiscard]] static std::unordered_map<fourdst::atomic::Species, int> getNetReactionStoichiometry(
            const reaction::Reaction& reaction
        );

        [[nodiscard]] int getStoichiometryMatrixEntry(
            const int speciesIndex,
            const int reactionIndex
        ) const override;

        [[nodiscard]] std::expected<std::unordered_map<fourdst::atomic::Species, double>, expectations::StaleEngineError> getSpeciesTimescales(
            const std::vector<double>& Y,
            double T9,
            double rho
        ) const override;
 
        [[nodiscard]] std::expected<std::unordered_map<fourdst::atomic::Species, double>, expectations::StaleEngineError> getSpeciesDestructionTimescales(
            const std::vector<double>& Y,
            double T9,
            double rho
        ) const override;
 
        fourdst::composition::Composition update(const NetIn &netIn) override;
 
        bool isStale(const NetIn &netIn) override;

        [[nodiscard]] bool involvesSpecies(
            const fourdst::atomic::Species& species
        ) const;

        void exportToDot(
            const std::string& filename
        ) const;

        void exportToCSV(
            const std::string& filename
        ) const;
 
        void setScreeningModel(screening::ScreeningType) override;
 
        [[nodiscard]] screening::ScreeningType getScreeningModel() const override;
 
        void setPrecomputation(bool precompute);
 
        [[nodiscard]] bool isPrecomputationEnabled() const;
 
        [[nodiscard]] const partition::PartitionFunction& getPartitionFunction() const;
 
        [[nodiscard]] double calculateReverseRate(
            const reaction::Reaction &reaction,
            double T9
        ) const;
 
        [[nodiscard]] double calculateReverseRateTwoBody(
            const reaction::Reaction &reaction,
            const double T9,
            const double forwardRate,
            const double expFactor
        ) const;
 
        [[nodiscard]] double calculateReverseRateTwoBodyDerivative(
            const reaction::Reaction &reaction,
            const double T9,
            const double reverseRate
        ) const;
 
        [[nodiscard]] bool isUsingReverseReactions() const;
 
        void setUseReverseReactions(bool useReverse);
 
        [[nodiscard]] int getSpeciesIndex(
            const fourdst::atomic::Species& species
        ) const override;
 
        [[nodiscard]] std::vector<double> mapNetInToMolarAbundanceVector(const NetIn &netIn) const override;
 
        [[nodiscard]] PrimingReport primeEngine(const NetIn &netIn) override;
 
        [[nodiscard]] BuildDepthType getDepth() const override;
 
        void rebuild(const fourdst::composition::Composition& comp, const BuildDepthType depth) override;
 
 
    private:
        struct PrecomputedReaction {
            // Forward cacheing
            size_t reaction_index;
            std::vector<size_t> unique_reactant_indices;
            std::vector<int> reactant_powers;
            double symmetry_factor;
            std::vector<size_t> affected_species_indices;
            std::vector<int> stoichiometric_coefficients;
 
            // Reverse cacheing
            std::vector<size_t> unique_product_indices; 
            std::vector<int> product_powers; 
            double reverse_symmetry_factor; 
        };
 
        struct constants {
            const double u = Constants::getInstance().get("u").value; 
            const double Na = Constants::getInstance().get("N_a").value; 
            const double c = Constants::getInstance().get("c").value; 
            const double kB = Constants::getInstance().get("kB").value; 
        };
    private:
        class AtomicReverseRate final : public CppAD::atomic_base<double> {
        public:
            AtomicReverseRate(
                const reaction::Reaction& reaction,
                const GraphEngine& engine
            ):
            atomic_base<double>("AtomicReverseRate"),
            m_reaction(reaction),
            m_engine(engine) {}
 
            bool forward(
                size_t p,
                size_t q,
                const CppAD::vector<bool>& vx,
                CppAD::vector<bool>& vy,
                const CppAD::vector<double>& tx,
                CppAD::vector<double>& ty
            ) override;
            bool reverse(
                size_t q,
                const CppAD::vector<double>& tx,
                const CppAD::vector<double>& ty,
                CppAD::vector<double>& px,
                const CppAD::vector<double>& py
            ) override;
            bool for_sparse_jac(
                size_t q,
                const CppAD::vector<std::set<size_t>>&r,
                CppAD::vector<std::set<size_t>>& s
            ) override;
            bool rev_sparse_jac(
                size_t q,
                const CppAD::vector<std::set<size_t>>&rt,
                CppAD::vector<std::set<size_t>>& st
            ) override;
 
        private:
            const reaction::Reaction& m_reaction;
            const GraphEngine& m_engine;
        };
    private:
        Config& m_config = Config::getInstance();
        quill::Logger* m_logger = LogManager::getInstance().getLogger("log");
 
        constants m_constants;
 
        reaction::LogicalReactionSet m_reactions; 
        std::unordered_map<std::string_view, reaction::Reaction*> m_reactionIDMap; 
 
        std::vector<fourdst::atomic::Species> m_networkSpecies; 
        std::unordered_map<std::string_view, fourdst::atomic::Species> m_networkSpeciesMap; 
        std::unordered_map<fourdst::atomic::Species, size_t> m_speciesToIndexMap; 
 
        boost::numeric::ublas::compressed_matrix<int> m_stoichiometryMatrix; 
 
        mutable boost::numeric::ublas::compressed_matrix<double> m_jacobianMatrix; 
        mutable CppAD::ADFun<double> m_rhsADFun; 
        mutable CppAD::sparse_jac_work m_jac_work; 
        CppAD::sparse_rc<std::vector<size_t>> m_full_jacobian_sparsity_pattern; 
 
        std::vector<std::unique_ptr<AtomicReverseRate>> m_atomicReverseRates;
 
        screening::ScreeningType m_screeningType = screening::ScreeningType::BARE; 
        std::unique_ptr<screening::ScreeningModel> m_screeningModel = screening::selectScreeningModel(m_screeningType);
 
        bool m_usePrecomputation = true; 
 
        bool m_useReverseReactions = true; 
 
        BuildDepthType m_depth;
 
        std::vector<PrecomputedReaction> m_precomputedReactions; 
        std::unique_ptr<partition::PartitionFunction> m_partitionFunction; 
 
    private:
        void syncInternalMaps();

        void collectNetworkSpecies();

        void populateReactionIDMap();

        void populateSpeciesToIndexMap();

        void reserveJacobianMatrix() const;

        void recordADTape();
 
        void collectAtomicReverseRateAtomicBases();
 
        void precomputeNetwork();
 

        [[nodiscard]] bool validateConservation() const;
 
 
        [[nodiscard]] StepDerivatives<double> calculateAllDerivativesUsingPrecomputation(
            const std::vector<double> &Y_in,
            const std::vector<double>& bare_rates,
            const std::vector<double> &bare_reverse_rates,
            double T9, double rho
        ) const;

        template <IsArithmeticOrAD T>
        T calculateMolarReactionFlow(
            const reaction::Reaction &reaction,
            const std::vector<T> &Y,
            const T T9,
            const T rho
        ) const;
 
        template<IsArithmeticOrAD T>
        T calculateReverseMolarReactionFlow(
            T T9,
            T rho,
            std::vector<T> screeningFactors,
            std::vector<T> Y,
            size_t reactionIndex,
            const reaction::LogicalReaction &reaction
        ) const;

        template<IsArithmeticOrAD T>
        [[nodiscard]] StepDerivatives<T> calculateAllDerivatives(
            const std::vector<T> &Y_in,
            T T9,
            T rho
        ) const;

        [[nodiscard]] StepDerivatives<double> calculateAllDerivatives(
            const std::vector<double>& Y_in,
            const double T9,
            const double rho
        ) const;

        [[nodiscard]] StepDerivatives<ADDouble> calculateAllDerivatives(
            const std::vector<ADDouble>& Y_in,
            const ADDouble &T9,
            const ADDouble &rho
        ) const;
    };
 
 
 
    template <IsArithmeticOrAD T>
    T GraphEngine::calculateReverseMolarReactionFlow(
        T T9,
        T rho,
        std::vector<T> screeningFactors,
        std::vector<T> Y,
        size_t reactionIndex,
        const reaction::LogicalReaction &reaction
    ) const {
        if (!m_useReverseReactions) {
            return static_cast<T>(0.0); // If reverse reactions are not used, return zero
        }
        T reverseMolarFlow = static_cast<T>(0.0);
 
        if (reaction.qValue() != 0.0) {
            T reverseRateConstant = static_cast<T>(0.0);
            if constexpr (std::is_same_v<T, ADDouble>) { // Check if T is an AD type at compile time
                const auto& atomic_func_ptr = m_atomicReverseRates[reactionIndex];
                if (atomic_func_ptr != nullptr) {
                    // A. Instantiate the atomic operator for the specific reaction
                    // B. Marshal the input vector
                    std::vector<T> ax = { T9 };
 
                    std::vector<T> ay(1);
                    (*atomic_func_ptr)(ax, ay);
                    reverseRateConstant = static_cast<T>(ay[0]);
                } else {
                    return reverseMolarFlow; // If no atomic function is available, return zero
                }
            } else {
                // A,B If not calling with an AD type, calculate the reverse rate directly
                reverseRateConstant = calculateReverseRate(reaction, T9);
            }
 
            // C. Get product multiplicities
            std::unordered_map<fourdst::atomic::Species, int> productCounts;
            for (const auto& product : reaction.products()) {
                productCounts[product]++;
            }
 
            // D. Calculate the symmetry factor
            T reverseSymmetryFactor = static_cast<T>(1.0);
            for (const auto &count: productCounts | std::views::values) {
                reverseSymmetryFactor /= static_cast<T>(std::tgamma(static_cast<double>(count + 1))); // Gamma function for factorial
            }
 
            // E. Calculate the abundance term
            T productAbundanceTerm = static_cast<T>(1.0);
            for (const auto& [species, count] : productCounts) {
                const unsigned long speciesIndex = m_speciesToIndexMap.at(species);
                productAbundanceTerm *= CppAD::pow(Y[speciesIndex], count);
            }
 
            // F. Determine the power for the density term
            const size_t num_products = reaction.products().size();
            const T rho_power = CppAD::pow(rho, static_cast<T>(num_products > 1 ? num_products - 1 : 0)); // Density raised to the power of (N-1) for N products
 
            // G. Assemble the reverse molar flow rate
            reverseMolarFlow = screeningFactors[reactionIndex] *
                               reverseRateConstant *
                               productAbundanceTerm *
                               reverseSymmetryFactor *
                               rho_power;
        }
        return reverseMolarFlow;
    }
 
    template<IsArithmeticOrAD T>
    StepDerivatives<T> GraphEngine::calculateAllDerivatives(
        const std::vector<T> &Y_in, T T9, T rho) const {
        std::vector<T> screeningFactors = m_screeningModel->calculateScreeningFactors(
            m_reactions,
            m_networkSpecies,
            Y_in,
            T9,
            rho
        );
 
        // --- Setup output derivatives structure ---
        StepDerivatives<T> result;
        result.dydt.resize(m_networkSpecies.size(), static_cast<T>(0.0));
 
        // --- AD Pre-setup (flags to control conditionals in an AD safe / branch aware manner) ---
        // ----- Constants for AD safe calculations ---
        const T zero = static_cast<T>(0.0);
        const T one = static_cast<T>(1.0);
 
        // ----- Initialize variables for molar concentration product and thresholds ---
        // Note: the logic here is that we use CppAD::CondExprLt to test thresholds and if they are less we set the flag
        //       to zero so that the final returned reaction flow is 0. This is as opposed to standard if statements
        //       which create branches that break the AD tape.
        const T rho_threshold = static_cast<T>(MIN_DENSITY_THRESHOLD);
 
        // --- Check if the density is below the threshold where we ignore reactions ---
        T threshold_flag = CppAD::CondExpLt(rho, rho_threshold, zero, one); // If rho < threshold, set flag to 0
 
        std::vector<T> Y = Y_in;
        for (size_t i = 0; i < m_networkSpecies.size(); ++i) {
            // We use CppAD::CondExpLt to handle AD taping and prevent branching
            // Note that while this is syntactically more complex this is equivalent to
            // if (Y[i] < 0) {Y[i] = 0;}
            // The issue is that this would introduce a branch which would require the auto diff tape to be re-recorded
            // each timestep, which is very inefficient.
            Y[i] = CppAD::CondExpLt(Y[i], zero, zero, Y[i]); // Ensure no negative abundances
        }
 
        const T u = static_cast<T>(m_constants.u); // Atomic mass unit in grams
        const T N_A = static_cast<T>(m_constants.Na); // Avogadro's number in mol^-1
        const T c = static_cast<T>(m_constants.c); // Speed of light in cm/s
 
        // --- SINGLE LOOP OVER ALL REACTIONS ---
        for (size_t reactionIndex = 0; reactionIndex < m_reactions.size(); ++reactionIndex) {
            const auto& reaction = m_reactions[reactionIndex];
 
            // 1. Calculate forward reaction rate
            const T forwardMolarReactionFlow = screeningFactors[reactionIndex] *
                calculateMolarReactionFlow<T>(reaction, Y, T9, rho);
 
            // 2. Calculate reverse reaction rate
            T reverseMolarFlow = calculateReverseMolarReactionFlow<T>(
                T9,
                rho,
                screeningFactors,
                Y,
                reactionIndex,
                reaction
            );
 
            const T molarReactionFlow = forwardMolarReactionFlow - reverseMolarFlow; // Net molar reaction flow
 
            // 3. Use the rate to update all relevant species derivatives (dY/dt)
            for (size_t speciesIndex = 0; speciesIndex < m_networkSpecies.size(); ++speciesIndex) {
                const T nu_ij = static_cast<T>(m_stoichiometryMatrix(speciesIndex, reactionIndex));
                result.dydt[speciesIndex] += threshold_flag * nu_ij * molarReactionFlow;
            }
        }
 
        T massProductionRate = static_cast<T>(0.0); // [mol][s^-1]
        for (const auto& [species, index] : m_speciesToIndexMap) {
            massProductionRate += result.dydt[index] * species.mass() * u;
        }
 
        result.nuclearEnergyGenerationRate = -massProductionRate * N_A * c * c; // [cm^2][s^-3] = [erg][s^-1][g^-1]
 
        return result;
    }
 
 
    template <IsArithmeticOrAD T>
    T GraphEngine::calculateMolarReactionFlow(
        const reaction::Reaction &reaction,
        const std::vector<T> &Y,
        const T T9,
        const T rho
    ) const {
 
        // --- Pre-setup (flags to control conditionals in an AD safe / branch aware manner) ---
        // ----- Constants for AD safe calculations ---
        const T zero = static_cast<T>(0.0);
 
        // --- Calculate the molar reaction rate (in units of [s^-1][cm^3(N-1)][mol^(1-N)] for N reactants) ---
        const T k_reaction = reaction.calculate_rate(T9);
 
        // --- Cound the number of each reactant species to account for species multiplicity ---
        std::unordered_map<std::string, int> reactant_counts;
        reactant_counts.reserve(reaction.reactants().size());
        for (const auto& reactant : reaction.reactants()) {
            reactant_counts[std::string(reactant.name())]++;
        }
        const int totalReactants = static_cast<int>(reaction.reactants().size());
 
        // --- Accumulator for the molar concentration ---
        auto molar_concentration_product = static_cast<T>(1.0);
 
        // --- Loop through each unique reactant species and calculate the molar concentration for that species then multiply that into the accumulator ---
        for (const auto& [species_name, count] : reactant_counts) {
            // --- Resolve species to molar abundance ---
            // PERF: Could probably optimize out this lookup
            const auto species_it = m_speciesToIndexMap.find(m_networkSpeciesMap.at(species_name));
            const size_t species_index = species_it->second;
            const T Yi = Y[species_index];
 
            // --- If count is > 1 , we need to raise the molar concentration to the power of count since there are really count bodies in that reaction ---
            molar_concentration_product *= CppAD::pow(Yi, static_cast<T>(count)); // ni^count
 
            // --- Apply factorial correction for identical reactions ---
            if (count > 1) {
                molar_concentration_product /= static_cast<T>(std::tgamma(static_cast<double>(count + 1))); // Gamma function for factorial
            }
        }
        // --- Final reaction flow calculation [mol][s^-1][g^-1] ---
        // Note: If the threshold flag ever gets set to zero this will return zero.
        //       This will result basically in multiple branches being written to the AD tape, which will make
        //       the tape more expensive to record, but it will also mean that we only need to record it once for
        //       the entire network.
        const T densityTerm = CppAD::pow(rho, totalReactants > 1 ? static_cast<T>(totalReactants - 1) : zero); // Density raised to the power of (N-1) for N reactants
        return molar_concentration_product * k_reaction * densityTerm;
    }
 
};