cosmology/correlation/
mod.rs

1use crate::power::{PowerFn, PowerSpectrum};
2use ouroboros::self_referencing;
3use quadrature::Output;
4use std::error::Error;
5use std::f64::consts::PI;
6use std::ops::Add;
7
8/// Default parameter for the lower k-bound of the correlation function integral
9pub const CORR_LOGK_MIN: f64 = -6.0;
10/// Default parameter for the upper k-bound of the correlation function integral
11pub const CORR_LOGK_MAX: f64 = 6.0;
12/// Default parameter for the abs error of the integrator
13pub const CORR_ABS_ERROR: f64 = 1e-10;
14// /// At r = 100, this gives 1/100th of a period
15// const MAX_K_STEP: f64 = 2.0 * PI / 100.0 / 100.0;
16
17#[self_referencing]
18pub struct CorrelationFunction {
19    z: f64,
20    params: CorrelationFunctionParameters,
21    lower_logk: f64,
22    upper_logk: f64,
23    target_error: f64,
24    #[borrows(params)]
25    #[covariant]
26    power_at_k: PowerFn<'this>,
27}
28
29impl CorrelationFunction {
30    /// Calculates the linear theory correlation function in real-space
31    pub fn correlation_function(&self, r: f64) -> f64 {
32        // Define prefactor
33        let prefactor = (2.0 * (PI).powi(2)).recip();
34
35        // Linear integrand
36        let integrand = |k: f64| k * self.borrow_power_at_k().power(k) * (k * r).sin() / r;
37
38        // Carry out integral
39        let intervals = r.clamp(1.0, 4.0) as usize;
40        let mut result = 0.0;
41        let lower_k = 1e-6 / r;
42        let upper_k = 1e4 / r;
43        let total_interval = upper_k - lower_k;
44        let subinterval_size = total_interval / intervals as f64;
45        for i in 0..intervals {
46            let interval_lower = lower_k + subinterval_size * i as f64;
47            let interval_upper = lower_k + subinterval_size * i.add(1) as f64;
48            let cf = quadrature::double_exponential::integrate(
49                integrand,
50                interval_lower,
51                interval_upper,
52                *self.borrow_target_error(),
53            );
54            check_integral(&cf);
55            result += cf.integral;
56        }
57        prefactor * result
58    }
59
60    // /// Returns the linear theory correlation function RSD monopole.
61    // /// Inspired by Kaiser 1987, Hamilton 1992 RSDs (in redshift-space).
62    // pub fn rsd_correlation_function_monopole(
63    //     &self,
64    //     r: f64,
65    //     b: f64,
66    // ) -> (f64, f64, f64) {
67
68    //     // Get linear theory real-space monopole
69    //     let corr_real = self.correlation_function(r);
70
71    //     // Calculate growth rate
72    //     let omega_matter = self.borrow_params().power.get_omega_matter();
73    //     let f = omega_matter.powf(5.0/9.0);
74
75    //     // Return result
76    //     (corr_real, prefactor * cf_2.integral, prefactor * cf_4.integral)
77    // }
78
79    pub fn get_correlation_function(
80        z: f64,
81        params: CorrelationFunctionParameters,
82    ) -> Result<CorrelationFunction, Box<dyn Error>> {
83        // Specify domain of integration, target error
84        let (lower_logk, upper_logk, target_error) = {
85            if let Some(ref acc_params) = params.accuracy_params {
86                // User specified bounds, error
87                (
88                    acc_params.lower_logk_bound,
89                    acc_params.upper_logk_bound,
90                    acc_params.target_error,
91                )
92            } else {
93                // Default parameters if None
94                (CORR_LOGK_MIN, CORR_LOGK_MAX, CORR_ABS_ERROR)
95            }
96        };
97
98        Ok(CorrelationFunctionBuilder {
99            z,
100            params,
101            lower_logk,
102            upper_logk,
103            target_error,
104            power_at_k_builder: |params: &CorrelationFunctionParameters| {
105                params.power.power_fn(z).unwrap()
106            },
107        }
108        .build())
109    }
110
111    pub fn get_params<'a>(&'a self) -> &'a CorrelationFunctionParameters {
112        &self.borrow_params()
113    }
114    pub fn get_redshift<'a>(&'a self) -> f64 {
115        *self.borrow_z()
116    }
117}
118
119#[allow(unused)]
120fn check_integral(cf: &Output) {
121    // TODO: check integral for convergence, perhaps just print warnings
122    // println!("evaluation: {}", cf.num_function_evaluations);
123}
124
125/// The parameters required for calculating the 2-point correlation function.
126#[derive(Clone)]
127pub struct CorrelationFunctionParameters {
128    /// The underlying power spectrum engine
129    pub power: PowerSpectrum,
130
131    /// Parameters controlling the accuracy of the correlation function
132    pub accuracy_params: Option<CorrFuncAccuracyParameters>,
133}
134
135// Parameters controlling the accuracy of the calculated correlation function.
136#[derive(Clone)]
137pub struct CorrFuncAccuracyParameters {
138    lower_logk_bound: f64,
139    upper_logk_bound: f64,
140    target_error: f64,
141}
142
143#[cfg(test)]
144#[cfg(feature = "colossus-python")]
145mod tests {
146
147    use super::*;
148    use crate::power::{PowerSpectrum, TransferFunction};
149
150    macro_rules! assert_eq_tol {
151        ($x:expr, $y:expr, $d:expr) => {
152            // Calculate fractional delta
153            let frac_delta = (($x - $y) / $y).abs();
154
155            // Compare frac_delta
156            let within = frac_delta < $d;
157
158            if !within {
159                // Construct err msg
160                let msg = format!(
161                    "Result {:.4e} not within {:.4e} of {:.4e}. frac_delta is {:.4e}",
162                    $x, $d, $y, frac_delta,
163                );
164
165                // Panic with err msg
166                panic!("{msg}");
167            }
168        };
169    }
170
171    macro_rules! eisenstein_corr(
172        ($z:ident, $h0:ident, $om0:ident, $ob0:ident, $t0:ident) => {
173
174          concat_idents::concat_idents!(test_name = test_eisen_corr, _, $z, _, $h0, _, $om0, _, $ob0, _, $t0, {
175            #[test]
176            fn test_name() {
177
178                let z: u32 = stringify!($z)[1..].parse::<u32>().unwrap();
179                let h: u32 = stringify!($h0)[1..].parse::<u32>().unwrap();
180                let om0: u32 = stringify!($om0)[1..].parse::<u32>().unwrap();
181                let ob0: u32 = stringify!($ob0)[1..].parse::<u32>().unwrap();
182                let t0: u32 = stringify!($t0)[1..].parse::<u32>().unwrap();
183
184                // Construct EisensteinHu model
185                let power = PowerSpectrum::new(TransferFunction::EisensteinHu {
186                    h: h as f64 / 100.0, // h
187                    omega_matter_0: om0 as f64 / 100.0, // omega_matter_0
188                    omega_baryon_0: ob0 as f64 / 100.0, // omega_baryon_0
189                    temp_cmb0: t0 as f64 / 100.0, // temp_cmb_0
190                    ns: 0.9665, // ns
191                    sigma_8: 0.8102, // sigma8
192                }).unwrap();
193
194                // Construct correlation function
195                let params = CorrelationFunctionParameters {
196                    power,
197                    accuracy_params: None, // default accuracy parameters
198                };
199                let corr = CorrelationFunction::get_correlation_function(
200                    z as f64,
201                    params
202                ).unwrap();
203
204              // Pick scales
205              let rs = [0.1, 1.0, 10.0, 100.0];
206
207              // Get result at redshift zero
208              let result = rs.map(|r| corr.correlation_function(r));
209
210              // Expected values, from COLOSSUS
211              let expected = {
212                use pyo3::prelude::*;
213                use pyo3::types::*;
214                Python::with_gil(|py| {
215
216                  // Get ks into python
217                  let list = PyList::new(py, &rs);
218                  let locals = PyDict::new(py);
219                  locals.set_item("rs", list).unwrap();
220
221                  py.run(format!(r#"from colossus.cosmology import cosmology
222import warnings
223warnings.filterwarnings("ignore")
224planck18 = cosmology.setCosmology("planck18")
225params = {{
226    "H0": {0},
227    "Om0": {1},
228    "Ob0": {2},
229    "Tcmb0": {3},
230    "ns": 0.9665,
231    "sigma8": 0.8102,
232}}
233cosmology.addCosmology("test", params=params)
234cosmo = cosmology.setCosmology("test")
235x = []
236for r in rs:
237    x.append(cosmo.correlationFunction(r, z={4}))
238                  "#, h as f64, om0 as f64 / 100.0, ob0 as f64 / 100.0,
239                  t0 as f64 / 100.0, z).as_str(), None, Some(locals)).unwrap();
240                  let x: Vec<_> = locals.get_item("x").unwrap().extract::<Vec<f64>>().unwrap();
241                  x
242                })
243              };
244
245              for i in 0..result.len() {
246                let abs_diff = (result[i]-expected[i]).abs();
247                let rel_diff = ((result[i]-expected[i])/expected[i]).abs();
248                assert!(abs_diff < 1e-2 || rel_diff < 5e-2, "result {:.3e} has abs_diff = {abs_diff:.3e}; rel_diff = {rel_diff:.3e} wrt expectation {:.3e}", result[i], expected[i]);
249              }
250            }
251          });
252        }
253      );
254    dry::macro_for!($H in [h50, h60, h70, h80, h90, h100] {
255        dry::macro_for!($M in [m10, m30, m50, m70, m90] {
256            dry::macro_for!($B in [b1, b2, b3] {
257                dry::macro_for!($T in [t270] {
258                    dry::macro_for!($Z in [z0, z1, z2, z10] {
259                        eisenstein_corr!($Z, $H, $M, $B, $T);
260                    });
261                });
262            });
263        });
264    });
265
266    macro_rules! eisenstein_corr_plot(
267        ($z:ident, $h0:ident, $om0:ident, $ob0:ident, $t0:ident) => {
268
269          concat_idents::concat_idents!(test_name = test_eisen_corr, _, $z, _, $h0, _, $om0, _, $ob0, _, $t0, {
270            #[test]
271            fn test_name() {
272
273                let z: u32 = stringify!($z)[1..].parse::<u32>().unwrap();
274                let h: u32 = stringify!($h0)[1..].parse::<u32>().unwrap();
275                let om0: u32 = stringify!($om0)[1..].parse::<u32>().unwrap();
276                let ob0: u32 = stringify!($ob0)[1..].parse::<u32>().unwrap();
277                let t0: u32 = stringify!($t0)[1..].parse::<u32>().unwrap();
278
279                // Construct EisensteinHu model
280                let power = PowerSpectrum::new(TransferFunction::EisensteinHu {
281                    h: h as f64 / 100.0, // h
282                    omega_matter_0: om0 as f64 / 100.0, // omega_matter_0
283                    omega_baryon_0: ob0 as f64 / 100.0, // omega_baryon_0
284                    temp_cmb0: t0 as f64 / 100.0, // temp_cmb_0
285                    ns: 0.9665, // ns
286                    sigma_8: 0.8102, // sigma8
287                }).unwrap();
288
289                // Construct correlation function
290                let params = CorrelationFunctionParameters {
291                    power,
292                    accuracy_params: None, // default accuracy parameters
293                };
294                let corr = CorrelationFunction::get_correlation_function(
295                    z as f64,
296                    params
297                ).unwrap();
298
299              // Pick scales
300              let rs = [0.1, 1.0, 10.0, 100.0];
301
302              // Get result at redshift zero
303              let result = rs.map(|r| corr.correlation_function(r));
304
305              // Expected values, from COLOSSUS
306              let expected = {
307                use pyo3::prelude::*;
308                use pyo3::types::*;
309                Python::with_gil(|py| {
310
311                  // Get ks into python
312                  let list = PyList::new(py, &rs);
313                  let locals = PyDict::new(py);
314                  locals.set_item("rs", list).unwrap();
315
316                  py.run(format!(r#"from colossus.cosmology import cosmology
317import warnings
318warnings.filterwarnings("ignore")
319planck18 = cosmology.setCosmology("planck18")
320params = {{
321    "H0": {0},
322    "Om0": {1},
323    "Ob0": {2},
324    "Tcmb0": {3},
325    "ns": 0.9665,
326    "sigma8": 0.8102,
327}}
328cosmology.addCosmology("test", params=params)
329cosmo = cosmology.setCosmology("test")
330x = []
331for r in rs:
332    x.append(cosmo.correlationFunction(r, z={4}))
333                  "#, h as f64, om0 as f64 / 100.0, ob0 as f64 / 100.0,
334                  t0 as f64 / 100.0, z).as_str(), None, Some(locals)).unwrap();
335                  let x: Vec<_> = locals.get_item("x").unwrap().extract::<Vec<f64>>().unwrap();
336                  x
337                })
338              };
339
340              for i in 0..result.len() {
341                let abs_diff = (result[i]-expected[i]).abs();
342                let rel_diff = ((result[i]-expected[i])/expected[i]).abs();
343                assert!(abs_diff < 1e-2 || rel_diff < 5e-2, "result {:.3e} has abs_diff = {abs_diff:.3e}; rel_diff = {rel_diff:.3e} wrt expectation {:.3e}", result[i], expected[i]);
344              }
345            }
346          });
347        }
348      );
349}
cosmology/correlation/mod.rs

cosmology/correlation/
mod.rs
