composite_species.h

#pragma once
#include <memory>
#include <string>
#include <vector>
 
#include "base_species.h"
 
class CompositeSpecies : public BaseSpecies {
 public:
  CompositeSpecies(std::string name, EnergyType energy_type)
      : BaseSpecies(std::move(name), energy_type) {
    delegates_tally_ = true;
  }
 
  struct PerturbLayout : BaseSpecies::PerturbLayout {
    std::vector<std::unique_ptr<BaseSpecies::PerturbLayout>> child_layouts;
  };
 
  std::unique_ptr<BaseSpecies::PerturbLayout> CreatePerturbLayout() const override;
 
  void RegisterPerturbationIndices(BaseSpecies::PerturbLayout& layout,
                                   perturb_vector* pv,
                                   const precision* ppr,
                                   int& index_pt,
                                   const perturb_workspace* ppw,
                                   int gauge) override;
  void ApplyInitialConditions(const BaseSpecies::PerturbLayout& layout,
                              double* y,
                              const PerturbIcContext& ctx) override;
  StressEnergyContribution StressEnergy(const BaseSpecies::PerturbLayout& layout,
                                        const perturb_vector* pv,
                                        const double* y,
                                        const double* pvecback,
                                        const perturb_workspace* ppw) const override;
  void DelegateTally(const BaseSpecies::PerturbLayout& layout,
                     const perturb_vector* pv,
                     const double* y,
                     const double* pvecback,
                     const perturb_workspace* ppw,
                     StressEnergyContribution& total,
                     StressEnergyContribution& total_cold,
                     StressEnergyContribution& total_warm) const override;
  void PerturbDerivs(const BaseSpecies::PerturbLayout& layout,
                     double tau,
                     const double* y,
                     double* dy,
                     const perturb_parameters_and_workspace& ppaw) const override;
  void FillSources(const BaseSpecies::PerturbLayout& layout,
                   const double* y,
                   const double* dy,
                   PerturbSourceContext& ctx) const override;
  void PerturbSynchronousToNewtonian(const BaseSpecies::PerturbLayout& layout,
                                     double* y,
                                     const PerturbIcContext& ctx) override;
  void CopyPerturbationsAcrossSwitch(const BaseSpecies::PerturbLayout& old_layout,
                                     const BaseSpecies::PerturbLayout& new_layout,
                                     const double* old_y,
                                     double* new_y,
                                     const PerturbSwitchContext& ctx) const override;
 
  // ── Registration ────────────────────────────────────────────────────────
  void RegisterBackgroundIndices(int& index_bg) override;
  void RegisterIntegrationIndices(int& index_bi) override;
  void RegisterTransferSourceIndices(int& index_tp, const SourceRequestContext& ctx) override;
  // ── Omega0 (closure) ────────────────────────────────────────────────────
  double GetOmega0() const override {
    double sum = 0.0;
    for (const auto& c : children_)
      sum += c->GetOmega0();
    return sum;
  }
 
  void AppendBudgetLines(const double* pvecback_today,
                         double rho_crit,
                         std::vector<BudgetLine>& out) const {
    for (const auto& c : children_)
      out.push_back({c->name(), c->Rho(pvecback_today) / rho_crit, BudgetBucketOf(*c)});
  }
 
  double GetRadiationOmega0() const override {
    double sum = 0.0;
    for (const auto& c : children_)
      sum += c->GetRadiationOmega0();
    return sum;
  }
 
  double DarkRadiationRhoToday(const double* pvecback_integration) const override {
    double sum = 0.0;
    for (const auto& c : children_)
      sum += c->DarkRadiationRhoToday(pvecback_integration);
    return sum;
  }
 
  double NeutrinoOmega0() const override {
    double sum = 0.0;
    for (const auto& c : children_)
      sum += c->NeutrinoOmega0();
    return sum;
  }
 
  double NeffContribution(double z) const override {
    double sum = 0.0;
    for (const auto& c : children_)
      sum += c->NeffContribution(z);
    return sum;
  }
 
  void PrintNeffInfo() const override {
    for (const auto& c : children_)
      c->PrintNeffInfo();
  }
  void PrintMassInfo() const override {
    for (const auto& c : children_)
      c->PrintMassInfo();
  }
 
  double TensorMasslessRelativisticRho(const double* pvecback) const override {
    double sum = 0.0;
    for (const auto& c : children_)
      sum += c->TensorMasslessRelativisticRho(pvecback);
    return sum;
  }
  void CheckUltraRelativisticAtIc(const double* pvecback, double tol) const override {
    for (const auto& c : children_)
      c->CheckUltraRelativisticAtIc(pvecback, tol);
  }
  bool IsUltraRelativisticAtIc(const double* pvecback, double tol) const override {
    for (const auto& c : children_)
      if (!c->IsUltraRelativisticAtIc(pvecback, tol))
        return false;
    return true;
  }
  void WarnIfTooHeavyForHalofit(double m_ev_threshold) const override {
    for (const auto& c : children_)
      c->WarnIfTooHeavyForHalofit(m_ev_threshold);
  }
 
  double BackgroundAIni(double a_proposed, double tol) const override {
    for (const auto& c : children_)
      a_proposed = c->BackgroundAIni(a_proposed, tol);
    return a_proposed;
  }
 
  // ── Background ──────────────────────────────────────────────────────────
  void SetBackgroundModule(const BackgroundModule* bgm) override;
  void SetThermodynamicsModule(const ThermodynamicsModule* thm) override;
  void SetPerturbs(const perturbs* ppt) override;
  void SetBackgroundInitialConditions(const BackgroundICContext& ctx) override;
  void ComputeBackground(double a, const double* pvecback_B, double* pvecback) override;
  void BackgroundDerivs(double tau, const double* y, double* dy, const double* pvecback) override;
  void FinalizeBackground(double a, double H, const double* pvecback_B, double* pvecback) override;
  double Rho(const double* pvecback) const override;
  double P(const double* pvecback) const override;
  double PPrime(double a,
                double H,
                const double* pvecback_B,
                const double* pvecback) const override;
  double FreestreamingRho(const double* pvecback) const override;
 
  // ── Matter tally ────────────────────────────────────────────────────────
  bool ClustersAsMatter() const override;
  bool IsColdMatterSpecies() const override;
  bool HasWarmMatter() const override;
  void FinalizeMatterClassification() override;
 
 protected:
  std::vector<std::unique_ptr<BaseSpecies>> children_;
 
  virtual void AddCouplingDerivs(double tau,
                                 const double* y,
                                 double* dy,
                                 const perturb_parameters_and_workspace& ppaw) const;
};