caustic 0.0.12

A General-Purpose 6D Collisionless Gravitational Dynamics Solver
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

//! Cosmological comoving-coordinate Strang splitting.
//!
//! Adapts second-order Strang splitting for an expanding universe by
//! incorporating scale-factor-dependent coefficients and Hubble drag.
//! The scale factor a(t) is evolved via Euler integration each step.
//!
//! Coordinate modifications relative to standard Strang splitting:
//! - Drift: displacement scaled by 1/a (comoving coordinates)
//! - Kick: acceleration scaled by a (momentum in comoving frame)
//! - Poisson: effective G_eff = G * a^2 for the comoving Poisson equation
//! - CFL: time step also limited by Hubble time 0.1/H(t)

use std::sync::Arc;

use super::super::{
    advecator::Advector,
    integrator::{StepProducts, StepTimings, TimeIntegrator},
    phasespace::PhaseSpaceRepr,
    progress::{StepPhase, StepProgress},
    solver::PoissonSolver,
    types::*,
};
use super::helpers;
use crate::CausticError;

/// Cosmological Strang splitting with scale-factor-dependent coefficients.
pub struct CosmologicalStrangSplitting {
    pub g: f64,
    /// Current scale factor a(t).
    pub scale_factor: f64,
    /// Hubble parameter H(t).
    pub hubble: f64,
    /// Matter density parameter Omega_m.
    pub omega_m: f64,
    last_timings: StepTimings,
    progress: Option<Arc<StepProgress>>,
}

impl CosmologicalStrangSplitting {
    pub fn new(g: f64, scale_factor: f64, hubble: f64, omega_m: f64) -> Self {
        Self {
            g,
            scale_factor,
            hubble,
            omega_m,
            last_timings: StepTimings::default(),
            progress: None,
        }
    }

    /// Update the scale factor and Hubble parameter for the current time.
    pub fn set_cosmology(&mut self, scale_factor: f64, hubble: f64) {
        self.scale_factor = scale_factor;
        self.hubble = hubble;
    }
}

impl TimeIntegrator for CosmologicalStrangSplitting {
    fn advance(
        &mut self,
        repr: &mut dyn PhaseSpaceRepr,
        solver: &dyn PoissonSolver,
        advector: &dyn Advector,
        dt: f64,
    ) -> Result<StepProducts, CausticError> {
        let _span = tracing::info_span!("cosmo_strang_advance").entered();
        let mut timings = StepTimings::default();
        let a = self.scale_factor;

        if let Some(ref p) = self.progress {
            p.start_step();
        }
        helpers::report_phase!(self.progress, StepPhase::DriftHalf1, 0, 5);

        // Comoving drift: displacement scaled by 1/a
        helpers::time_ms!(timings, drift_ms, advector.drift(repr, dt / (2.0 * a)));

        helpers::report_phase!(self.progress, StepPhase::PoissonSolve, 1, 5);

        // Poisson solve with cosmological factor (G_eff = G * a^2)
        let accel = helpers::time_ms!(timings, poisson_ms, {
            let g_eff = self.g * a * a;
            let density = repr.compute_density();
            let potential = solver.solve(&density, g_eff);
            solver.compute_acceleration(&potential)
        });

        helpers::report_phase!(self.progress, StepPhase::Kick, 2, 5);

        // Kick with cosmological factor
        helpers::time_ms!(timings, kick_ms, advector.kick(repr, &accel, dt * a));

        helpers::report_phase!(self.progress, StepPhase::DriftHalf2, 3, 5);

        helpers::time_ms!(timings, drift_ms, advector.drift(repr, dt / (2.0 * a)));

        // Update scale factor via simple Euler step: a' = a * H * dt
        self.scale_factor += self.scale_factor * self.hubble * dt;

        helpers::report_phase!(self.progress, StepPhase::StepComplete, 4, 5);

        // Compute end-of-step products for caller reuse (use updated scale factor)
        let (density, potential, acceleration) = helpers::time_ms!(timings, density_ms, {
            let g_eff = self.g * self.scale_factor * self.scale_factor;
            let density = repr.compute_density();
            let potential = solver.solve(&density, g_eff);
            let acceleration = solver.compute_acceleration(&potential);
            (density, potential, acceleration)
        });

        self.last_timings = timings;

        Ok(StepProducts {
            density,
            potential,
            acceleration,
        })
    }

    fn max_dt(&self, repr: &dyn PhaseSpaceRepr, cfl_factor: f64) -> f64 {
        let density = repr.compute_density();
        let rho_max = density.data.iter().cloned().fold(0.0_f64, f64::max);
        if rho_max <= 0.0 || self.g <= 0.0 {
            return 1e10;
        }
        let t_dyn = 1.0 / (self.g * rho_max).sqrt();
        // Also limit by Hubble time
        let t_hubble = if self.hubble.abs() > 1e-30 {
            0.1 / self.hubble.abs()
        } else {
            1e10
        };
        cfl_factor * t_dyn.min(t_hubble)
    }

    fn last_step_timings(&self) -> Option<&StepTimings> {
        Some(&self.last_timings)
    }

    fn set_progress(&mut self, progress: Arc<StepProgress>) {
        self.progress = Some(progress);
    }
}