base_species.h

#pragma once
 
#include <cstddef>
#include <memory>
#include <optional>
#include <string>
#include <vector>
 
#include "shooting_target.h"
 
// Forward declarations to avoid circular includes
struct background;
struct SpeciesBuildContext;
struct precision;
struct perturb_vector;
struct perturb_workspace;
struct perturb_parameters_and_workspace;
 
#include "background_ic_context.h"
#include "common.h"  // class_store_columntitle, class_store_double, true, false
#include "perturb_source_context.h"
 
class BackgroundModule;      // forward declaration
class ThermodynamicsModule;  // forward declaration
 
// ── BackgroundColumnWriter ───────────────────────────────────────────────────
class BackgroundColumnWriter {
 public:
  explicit BackgroundColumnWriter(std::string& titles) : titles_(&titles) {}
  BackgroundColumnWriter(double* dataptr, int& storeidx)
      : dataptr_(dataptr), storeidx_(&storeidx) {}
 
  bool IsTitleMode() const {
    return titles_ != nullptr;
  }
 
  void Add(const char* title, double value, bool condition = true) {
    if (titles_) {
      class_store_columntitle(*titles_, title, condition ? true : false);
    }
    else if (dataptr_) {
      class_store_double(dataptr_, value, condition ? true : false, (*storeidx_));
    }
  }
  void Add(const std::string& title, double value, bool condition = true) {
    Add(title.c_str(), value, condition);
  }
 
 private:
  std::string* titles_ = nullptr;
  double* dataptr_     = nullptr;
  int* storeidx_       = nullptr;
};
struct perturbs;  // forward declaration
 
class BaseSpecies {
 public:
  enum class EnergyType { Radiation, Matter, DarkEnergy, Other };
  enum class TransferColumnSection { all, density, velocity };
 
  // ── Perturbation layout (per-pv per-species index storage) ───────────────
  struct PerturbLayout {
    virtual ~PerturbLayout() = default;
  };
 
  virtual std::unique_ptr<PerturbLayout> CreatePerturbLayout() const {
    return std::make_unique<PerturbLayout>();
  }
 
  virtual ~BaseSpecies()                     = default;
  BaseSpecies(const BaseSpecies&)            = delete;
  BaseSpecies& operator=(const BaseSpecies&) = delete;
 
  const std::string& name() const {
    return name_;
  }
  EnergyType energy_type() const {
    return energy_type_;
  }
 
  virtual void SetBackgroundModule(const BackgroundModule* /*bgm*/) {}
 
  virtual void SetThermodynamicsModule(const ThermodynamicsModule* /*thm*/) {}
 
  virtual void SetPerturbs(const perturbs* /*ppt*/) {}
 
  virtual std::optional<double> GetParam(const std::string& /*name*/) const {
    return std::nullopt;
  }
 
  // ── Background ────────────────────────────────────────────────────────────
 
  virtual void RegisterBackgroundIndices(int& index_bg) = 0;
 
  virtual void RegisterIntegrationIndices(int& index_bi) {}
 
  virtual void SetBackgroundInitialConditions(const BackgroundICContext& ctx) {}
 
  virtual void ComputeBackground(double a, const double* pvecback_B, double* pvecback) = 0;
 
  virtual void BackgroundDerivs(double tau, const double* y, double* dy, const double* pvecback) {}
 
  virtual double Rho(const double* pvecback) const = 0;
 
  virtual double P(const double* pvecback) const = 0;
 
  virtual double PPrime(double a,
                        double H,
                        const double* pvecback_B,
                        const double* pvecback) const {
    return 0.;
  }
 
  virtual void FinalizeBackground(double a, double H, const double* pvecback_B, double* pvecback) {}
 
  // ── Background output ─────────────────────────────────────────────────────
 
  virtual bool IsFreestreaming() const {
    return false;
  }
 
  virtual double FreestreamingRho(const double* pvecback) const {
    return IsFreestreaming() ? Rho(pvecback) : 0.;
  }
 
  bool IsPresent() const {
    return index_bg_rho_ >= 0;
  }
 
  virtual void WriteBackgroundColumnTitles(BackgroundColumnWriter& /*writer*/) const {}
 
  virtual void WriteBackgroundData(const double* /*pvecback*/,
                                   BackgroundColumnWriter& /*writer*/) const {}
 
  virtual void ProcessBackgroundTable(const double* /*background_table*/,
                                      int /*n_rows*/,
                                      int /*row_stride*/,
                                      const double* /*z_table*/) {}
 
  // ── Perturbations ─────────────────────────────────────────────────────────
 
  // ── Per-species layout signatures ───────────────────────────────────────
 
  virtual void RegisterTransferSourceIndices(int& /*index_tp*/,
                                             const SourceRequestContext& /*ctx*/) {}
 
  virtual void RegisterPerturbationIndices(PerturbLayout& /*layout*/,
                                           perturb_vector* /*pv*/,
                                           const precision* /*ppr*/,
                                           int& /*index_pt*/,
                                           const perturb_workspace* /*ppw*/,
                                           int /*gauge*/) {}
 
  virtual void RegisterVectorPerturbationIndices(PerturbLayout& /*layout*/,
                                                 perturb_vector* /*pv*/,
                                                 const precision* /*ppr*/,
                                                 int& /*index_pt*/,
                                                 const perturb_workspace* /*ppw*/,
                                                 int /*gauge*/) {}
 
  virtual void RegisterTensorPerturbationIndices(PerturbLayout& /*layout*/,
                                                 perturb_vector* /*pv*/,
                                                 const precision* /*ppr*/,
                                                 int& /*index_pt*/,
                                                 const perturb_workspace* /*ppw*/,
                                                 int /*gauge*/) {}
 
  virtual void PerturbDerivs(const PerturbLayout& layout,
                             double tau,
                             const double* y,
                             double* dy,
                             const perturb_parameters_and_workspace& ppaw) const = 0;
 
  virtual void PerturbVectorDerivs(const PerturbLayout& /*layout*/,
                                   double /*tau*/,
                                   const double* /*y*/,
                                   double* /*dy*/,
                                   const perturb_parameters_and_workspace& /*ppaw*/) const {}
 
  virtual void PerturbTensorDerivs(const PerturbLayout& /*layout*/,
                                   double /*tau*/,
                                   const double* /*y*/,
                                   double* /*dy*/,
                                   const perturb_parameters_and_workspace& /*ppaw*/) const {}
 
  virtual void WriteTensorOutputColumnTitles(std::string& /*tensor_titles*/) const {}
 
  virtual void ContributeTensorGwSource(const PerturbLayout& /*layout*/,
                                        double /*a*/,
                                        const double* /*y*/,
                                        perturb_workspace* /*ppw*/) const {}
 
  struct StressEnergyContribution {
    double rho              = 0.;  // background ρ
    double p                = 0.;  // background P
    double delta_rho        = 0.;  // δρ
    double rho_plus_p_theta = 0.;  // (ρ+P)θ
    double delta_p          = 0.;  // δp
    double rho_plus_p_shear = 0.;  // (ρ+P)σ
 
    StressEnergyContribution& operator+=(const StressEnergyContribution& o) {
      rho              += o.rho;
      p                += o.p;
      delta_rho        += o.delta_rho;
      rho_plus_p_theta += o.rho_plus_p_theta;
      delta_p          += o.delta_p;
      rho_plus_p_shear += o.rho_plus_p_shear;
      return *this;
    }
  };
 
  virtual StressEnergyContribution StressEnergy(const PerturbLayout& layout,
                                                const perturb_vector* pv,
                                                const double* y,
                                                const double* pvecback,
                                                const perturb_workspace* ppw) const = 0;
 
  // ── Stress-energy + matter tally (hot path) ───────────────────────────────
  // Non-virtual wrapper, inlined at the loop call site (e.species is BaseSpecies*).
  // Plain species: one virtual call (StressEnergy) + accumulate. Composites set
  // delegates_tally_ and forward the WHOLE call per child via DelegateTally, so each
  // child is evaluated exactly once. The three accumulators are owned by the module
  // loop; StressEnergyContribution::operator+= is reused as the bucket combiner.
  // ppw is only passed through to StressEnergy (never dereferenced), so this compiles
  // with perturb_workspace forward-declared.
  void TallyStressEnergy(const 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 {
    if (delegates_tally_) {
      DelegateTally(layout, pv, y, pvecback, ppw, total, total_cold, total_warm);
      return;
    }
    const StressEnergyContribution se  = StressEnergy(layout, pv, y, pvecback, ppw);
    total                             += se;
    // Actual ρ/δρ/(ρ+P)θ — no ρ−3P proxy. Radiation never reaches here (it does not
    // cluster as matter), so there is nothing to "zero out".
    if (clusters_as_matter_)
      (is_cold_ ? total_cold : total_warm) += se;
  }
 
  // Seam: only composites (delegates_tally_ == true) override this; plain species
  // never reach the base body.
  virtual void DelegateTally(const 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 {}
 
  // ── Stage 1: Output ──────────────────────────────────────────────────────
 
  virtual void WriteOutputColumns(
      PerturbColumnWriter& /*writer*/,
      const PerturbationsModule& /*mod*/,
      file_format /*fmt*/,
      TransferColumnSection /*section*/ = TransferColumnSection::all) const {}
 
  virtual void PrintVariables(PerturbColumnWriter& /*writer*/,
                              double /*tau*/,
                              const double* /*y*/,
                              const PerturbationsModule& /*mod*/,
                              const perturb_workspace* /*ppw*/) const {}
 
  // ── Stage 2: Source filling ───────────────────────────────────────────────
 
  virtual void FillSources(const PerturbLayout& /*layout*/,
                           const double* /*y*/,
                           const double* /*dy*/,
                           PerturbSourceContext& /*ctx*/) const {}
 
  // ── Stage 3: Initial conditions ───────────────────────────────────────────
 
  virtual void ApplyInitialConditions(const PerturbLayout& /*layout*/,
                                      double* /*y*/,
                                      const PerturbIcContext& /*ctx*/) {}
 
  virtual void CopyPerturbationsAcrossSwitch(const PerturbLayout& /*old_layout*/,
                                             const PerturbLayout& /*new_layout*/,
                                             const double* /*old_y*/,
                                             double* /*new_y*/,
                                             const PerturbSwitchContext& /*ctx*/) const {}
 
  // ── Synchronous → Newtonian gauge transformation ──────────────────────────
  virtual double RhoDotOverRho(const double* pvecback, double a_prime_over_a) const {
    return -3. * a_prime_over_a * (Rho(pvecback) + P(pvecback)) / Rho(pvecback);
  }
 
  virtual void PerturbSynchronousToNewtonian(const PerturbLayout& /*layout*/,
                                             double* /*y*/,
                                             const PerturbIcContext& /*ctx*/) {}
 
  // ── Shooter / root-finding ────────────────────────────────────────────────
  // Reported after construction; non-empty iff this species guessed its unknown
  // (target-form input set, direct unknown absent). Drives target detection, the
  // unknown vector, and the lazy-shooting trigger.
  virtual std::vector<ShootingTarget> GetShootingTargets() const {
    return {};
  }
  // Initial guess + Jacobian seed for each unknown, same order as GetShootingTargets.
  // Also called by CreateAll to obtain the value to build from when the unknown is absent.
  virtual void ComputeShootingGuess(const SpeciesBuildContext& /*ctx*/,
                                    std::vector<double>& /*guess*/,
                                    std::vector<double>& /*dxdy*/) const {}
  // computed-today minus the requested target. The authoritative target is supplied by
  // the shooter (collected once at discovery), so the species never re-derives it from
  // the file content — in iteration builds DoShooting has set the unknown, which for some
  // species (e.g. DCDM_DR's overlapping Omega_dcdmdr/Omega_ini_dcdm) makes the user's
  // target indistinguishable by key presence.
  virtual double ComputeShootingResidual(const ShootingResidualContext& /*ctx*/,
                                         const ShootingTarget& /*target*/) const {
    return 0.;
  }
 
  virtual void MarkUsedInSources(const PerturbLayout& /*layout*/,
                                 const perturb_workspace* /*ppw*/,
                                 int* /*used_in_sources*/) const {}
 
  virtual void MarkVectorUsedInSources(const PerturbLayout& /*layout*/,
                                       const perturb_workspace* /*ppw*/,
                                       int* /*used_in_sources*/) const {}
 
  virtual void MarkTensorUsedInSources(const PerturbLayout& /*layout*/,
                                       const perturb_workspace* /*ppw*/,
                                       int* /*used_in_sources*/) const {}
 
  virtual double GetOmega0() const = 0;
 
  virtual double GetRadiationOmega0() const {
    return 0.;
  }
 
  // ── NCDM/DR-family hooks (#308) ───────────────────────────────────────────
  // Neutral defaults let modules loop over all_species_ and call these directly
  // instead of downcasting; composites forward/sum over their children.
 
  virtual double DarkRadiationRhoToday(const double* /*pvecback_integration*/) const {
    return 0.;
  }
 
  virtual double NeutrinoOmega0() const {
    return 0.;
  }
 
  virtual double NeffContribution(double /*z*/) const {
    return 0.;
  }
 
  virtual void PrintNeffInfo() const {}
  virtual void PrintMassInfo() const {}
 
  virtual double TensorMasslessRelativisticRho(const double* /*pvecback*/) const {
    return 0.;
  }
 
  virtual void CheckUltraRelativisticAtIc(const double* /*pvecback*/, double /*tol*/) const {}
 
  virtual bool IsUltraRelativisticAtIc(const double* /*pvecback*/, double /*tol*/) const {
    return true;
  }
 
  virtual void WarnIfTooHeavyForHalofit(double /*m_ev_threshold*/) const {}
 
  virtual double BackgroundAIni(double a_proposed, double /*tol*/) const {
    return a_proposed;
  }
 
  bool ClustersAsMatterCached() const {
    return clusters_as_matter_;
  }
  bool IsColdCached() const {
    return is_cold_;
  }
 
  virtual void FinalizeMatterClassification() {
    clusters_as_matter_ = ClustersAsMatter();
    is_cold_            = IsColdMatterSpecies();
  }
 
  virtual bool ClustersAsMatter() const {
    return energy_type_ == EnergyType::Matter;
  }
 
  virtual bool IsColdMatterSpecies() const {
    return energy_type_ == EnergyType::Matter;
  }
 
  virtual bool HasWarmMatter() const {
    return ClustersAsMatter() && !IsColdMatterSpecies();
  }
 
 protected:
  BaseSpecies(std::string name, EnergyType energy_type)
      : name_(std::move(name)), energy_type_(energy_type) {}
 
  void ApplyFluidLikeNewtonianShift(double* y,
                                    int idx_delta,
                                    int idx_theta,
                                    const double* pvecback,
                                    const PerturbIcContext& ctx) const {
    if (idx_delta >= 0)
      y[idx_delta] += RhoDotOverRho(pvecback, ctx.a_prime_over_a) * ctx.alpha;
    if (idx_theta >= 0)
      y[idx_theta] += ctx.k * ctx.k * ctx.alpha;
  }
 
  std::string name_;
  EnergyType energy_type_;
 
  // Cached matter-tally classification, stamped once by FinalizeMatterClassification()
  // (driven by SpeciesCollection::freeze(); composites recurse into children). Read on
  // the stress-energy hot path, so it must be a plain member, not a per-step virtual.
  bool clusters_as_matter_ = false;
  bool is_cold_            = false;
  bool delegates_tally_    = false;  // true for composites (set in CompositeSpecies ctor)
 
  // Set by SpeciesCollection::freeze(); used by PrintVariables to look up layout.
  std::size_t collection_index_ = SIZE_MAX;
 
  // Set by RegisterBackgroundIndices(); -1 means "not registered / species absent"
  int index_bg_rho_ = -1;
  int index_bg_p_   = -1;  // only set by species that store p separately (e.g. NCDM)
 
  friend class SpeciesCollection;
};
 
// ── Background budget printout (background_verbose > 1) ───────────────────────
// Coarse presentation buckets for background_output_budget(). Distinct from
// EnergyType because NCDM-family species are EnergyType::Other yet warrant their
// own section, and composites (also EnergyType::Other) split across buckets by
// child.
enum class BudgetBucket { Radiation, NonRelativistic, Ncdm, Other };
 
// One printed budget line: a (sub-)species' today density fraction and its bucket.
struct BudgetLine {
  std::string label;
  double omega;
  BudgetBucket bucket;
};
 
// Classify a species into a budget bucket. NCDM-family -> Ncdm (the same family
// test GetNcdmSpecies uses); otherwise mapped from energy_type(). Defined in
// base_species.cpp.
BudgetBucket BudgetBucketOf(const BaseSpecies& s);