use crate::ionex::cell;
use crate::portable_distr::{Distribution, StandardNormal};
use crate::portable_math::PortableFloat;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;
pub const EMAG2_VOID: f64 = 99_999.0;
#[derive(Clone, Debug, PartialEq)]
pub struct GeoGrid {
pub lat0_deg: f64,
pub lon0_deg: f64,
pub dlat_deg: f64,
pub dlon_deg: f64,
pub n_lat: usize,
pub n_lon: usize,
pub values: Vec<f64>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Emag2Column {
SeaLevel,
UpCont,
}
impl GeoGrid {
pub fn from_fn<F>(
south: f64,
north: f64,
west: f64,
east: f64,
step_deg: f64,
f: F,
) -> Result<Self, String>
where
F: Fn(f64, f64) -> f64,
{
let (n_lat, n_lon) = box_shape(south, north, west, east, step_deg)?;
let mut values = Vec::with_capacity(n_lat * n_lon);
for i in 0..n_lat {
let lat = south + i as f64 * step_deg;
for j in 0..n_lon {
values.push(f(lat, west + j as f64 * step_deg));
}
}
Ok(Self {
lat0_deg: south,
lon0_deg: west,
dlat_deg: step_deg,
dlon_deg: step_deg,
n_lat,
n_lon,
values,
})
}
pub fn node(&self, i: usize, j: usize) -> f64 {
self.values[i * self.n_lon + j]
}
pub fn north_deg(&self) -> f64 {
self.lat0_deg + (self.n_lat - 1) as f64 * self.dlat_deg
}
pub fn east_deg(&self) -> f64 {
self.lon0_deg + (self.n_lon - 1) as f64 * self.dlon_deg
}
pub fn contains(&self, lat_deg: f64, lon_deg: f64) -> bool {
lat_deg >= self.lat0_deg
&& lat_deg <= self.north_deg()
&& lon_deg >= self.lon0_deg
&& lon_deg <= self.east_deg()
}
pub fn value_at(&self, lat_deg: f64, lon_deg: f64) -> f64 {
let (i0, fi) = cell(lat_deg, self.lat0_deg, self.dlat_deg, self.n_lat);
let (j0, fj) = cell(lon_deg, self.lon0_deg, self.dlon_deg, self.n_lon);
let v00 = self.node(i0, j0);
let v01 = self.node(i0, j0 + 1);
let v10 = self.node(i0 + 1, j0);
let v11 = self.node(i0 + 1, j0 + 1);
[
((1.0 - fi) * (1.0 - fj), v00),
((1.0 - fi) * fj, v01),
(fi * (1.0 - fj), v10),
(fi * fj, v11),
]
.iter()
.filter(|(w, _)| *w != 0.0)
.map(|(w, v)| w * v)
.sum()
}
pub fn sampler_deg(&self) -> impl Fn(f64, f64) -> f64 + '_ {
move |lat_deg: f64, lon_deg: f64| self.value_at(lat_deg, lon_deg)
}
pub fn void_count(&self) -> usize {
self.values.iter().filter(|v| v.is_nan()).count()
}
pub fn mean_std(&self) -> (f64, f64) {
let vals: Vec<f64> = self
.values
.iter()
.copied()
.filter(|v| !v.is_nan())
.collect();
if vals.is_empty() {
return (0.0, 0.0);
}
let n = vals.len() as f64;
let mean = vals.iter().sum::<f64>() / n;
let var = vals.iter().map(|v| (v - mean) * (v - mean)).sum::<f64>() / n;
(mean, var.sqrt())
}
pub fn plus(&self, other: &GeoGrid) -> GeoGrid {
let mut out = self.clone();
for i in 0..self.n_lat {
let lat = self.lat0_deg + i as f64 * self.dlat_deg;
for j in 0..self.n_lon {
let lon = self.lon0_deg + j as f64 * self.dlon_deg;
out.values[i * self.n_lon + j] += other.value_at(lat, lon);
}
}
out
}
#[allow(clippy::too_many_arguments)]
pub fn gravity_anomaly_from_sh(
field: &crate::gravity_sh::SphericalHarmonicField,
ell: &crate::gravity_sh::LevelEllipsoid,
zero_degree: bool,
south: f64,
north: f64,
west: f64,
east: f64,
step_deg: f64,
) -> Result<Self, String> {
let (n_lat, n_lon) = box_shape(south, north, west, east, step_deg)?;
let lons: Vec<f64> = (0..n_lon)
.map(|j| (west + j as f64 * step_deg).to_radians())
.collect();
let mut values = Vec::with_capacity(n_lat * n_lon);
for i in 0..n_lat {
let lat = (south + i as f64 * step_deg).to_radians();
values.extend(field.gravity_anomaly_sa_parallel_mgal(
ell,
zero_degree,
lat,
0.0,
&lons,
));
}
Ok(Self {
lat0_deg: south,
lon0_deg: west,
dlat_deg: step_deg,
dlon_deg: step_deg,
n_lat,
n_lon,
values,
})
}
pub fn from_emag2_csv(text: &str, column: Emag2Column) -> Result<Self, String> {
let col = match column {
Emag2Column::SeaLevel => 4,
Emag2Column::UpCont => 5,
};
let mut rows: Vec<(f64, f64, f64)> = Vec::new();
for (ln, line) in text.lines().enumerate() {
let t = line.trim();
if t.is_empty() || t.starts_with('#') {
continue;
}
let p: Vec<&str> = t.split(',').map(str::trim).collect();
if p.len() < 8 {
return Err(format!(
"EMAG2 line {}: expected 8 comma-separated columns, found {}",
ln + 1,
p.len()
));
}
let num = |k: usize, what: &str| {
p[k].parse::<f64>()
.map_err(|_| format!("EMAG2 line {}: bad {what} `{}`", ln + 1, p[k]))
};
let lon = num(2, "longitude")?;
let lat = num(3, "latitude")?;
let v = num(col, "anomaly")?;
if !(lon.is_finite() && lat.is_finite() && v.is_finite()) {
return Err(format!("EMAG2 line {}: non-finite value", ln + 1));
}
rows.push((lat, lon, if v == EMAG2_VOID { f64::NAN } else { v }));
}
if rows.is_empty() {
return Err("EMAG2: no data rows".into());
}
let lats = distinct_sorted(rows.iter().map(|r| r.0));
let lons = distinct_sorted(rows.iter().map(|r| r.1));
let dlat = regular_step(&lats, "latitude")?;
let dlon = regular_step(&lons, "longitude")?;
let (n_lat, n_lon) = (lats.len(), lons.len());
if rows.len() != n_lat * n_lon {
return Err(format!(
"EMAG2: {} rows do not fill the {n_lat} x {n_lon} lattice they span",
rows.len()
));
}
let (lat0, lon0) = (lats[0], lons[0]);
let mut values = vec![f64::NAN; n_lat * n_lon];
let mut seen = vec![false; n_lat * n_lon];
for (lat, lon, v) in rows {
let i = ((lat - lat0) / dlat).round() as usize;
let j = ((lon - lon0) / dlon).round() as usize;
let k = i * n_lon + j;
if seen[k] {
return Err(format!("EMAG2: duplicate node at {lat}, {lon}"));
}
seen[k] = true;
values[k] = v;
}
Ok(Self {
lat0_deg: lat0,
lon0_deg: lon0,
dlat_deg: dlat,
dlon_deg: dlon,
n_lat,
n_lon,
values,
})
}
}
fn box_shape(
south: f64,
north: f64,
west: f64,
east: f64,
step: f64,
) -> Result<(usize, usize), String> {
if !(step.is_finite() && step > 0.0) {
return Err(format!("grid step must be finite and positive, got {step}"));
}
if !(south.is_finite() && north.is_finite() && west.is_finite() && east.is_finite())
|| north <= south
|| east <= west
{
return Err(format!(
"grid box must have north > south and east > west, got [{south}, {north}] x [{west}, {east}]"
));
}
let n = |span: f64| ((span / step - 1e-9).ceil() as usize + 1).max(2);
let (n_lat, n_lon) = (n(north - south), n(east - west));
if n_lat.saturating_mul(n_lon) > 50_000_000 {
return Err(format!("grid of {n_lat} x {n_lon} nodes is too large"));
}
Ok((n_lat, n_lon))
}
fn distinct_sorted(it: impl Iterator<Item = f64>) -> Vec<f64> {
let mut v: Vec<f64> = it.collect();
v.sort_by(|a, b| a.total_cmp(b));
v.dedup_by(|a, b| (*a - *b).abs() < 1e-7);
v
}
fn regular_step(axis: &[f64], what: &str) -> Result<f64, String> {
if axis.len() < 2 {
return Err(format!(
"EMAG2: the {what} axis needs at least 2 distinct values"
));
}
let step = (axis[axis.len() - 1] - axis[0]) / (axis.len() - 1) as f64;
for w in axis.windows(2) {
if ((w[1] - w[0]) - step).abs() > 1e-6 {
return Err(format!(
"EMAG2: the {what} axis is not regular ({} to {} against a mean step of {step})",
w[0], w[1]
));
}
}
Ok(step)
}
#[allow(clippy::too_many_arguments)]
pub fn correlated_error_field(
south: f64,
north: f64,
west: f64,
east: f64,
step_deg: f64,
sigma: f64,
corr_len_deg: f64,
seed: u64,
) -> Result<GeoGrid, String> {
if !(sigma.is_finite() && sigma >= 0.0) {
return Err(format!(
"error sigma must be finite and non-negative, got {sigma}"
));
}
if !(corr_len_deg.is_finite() && corr_len_deg > 0.0) {
return Err(format!(
"correlation length must be finite and positive, got {corr_len_deg}"
));
}
let (n_lat, n_lon) = box_shape(south, north, west, east, step_deg)?;
let cos_mid = ((south + north) / 2.0).to_radians().pcos().max(1e-3);
let w_lat = corr_len_deg / std::f64::consts::SQRT_2 / step_deg;
let w_lon = w_lat / cos_mid;
let kernel = |w: f64| -> Vec<f64> {
let half = (4.0 * w).ceil().max(1.0) as isize;
(-half..=half)
.map(|k| (-(k as f64) * (k as f64) / (2.0 * w * w)).pexp())
.collect()
};
let (k_lat, k_lon) = (kernel(w_lat), kernel(w_lon));
let (h_lat, h_lon) = (k_lat.len() / 2, k_lon.len() / 2);
let ss = |k: &[f64]| k.iter().map(|x| x * x).sum::<f64>();
let scale = sigma / (ss(&k_lat) * ss(&k_lon)).sqrt();
let (p_lat, p_lon) = (n_lat + 2 * h_lat, n_lon + 2 * h_lon);
let mut rng = ChaCha8Rng::seed_from_u64(seed);
let white: Vec<f64> = (0..p_lat * p_lon)
.map(|_| StandardNormal.sample(&mut rng))
.collect();
let mut rows_smoothed = vec![0.0; p_lat * n_lon];
for i in 0..p_lat {
for j in 0..n_lon {
let mut acc = 0.0;
for (k, w) in k_lon.iter().enumerate() {
acc += w * white[i * p_lon + j + k];
}
rows_smoothed[i * n_lon + j] = acc;
}
}
let mut values = vec![0.0; n_lat * n_lon];
for i in 0..n_lat {
for j in 0..n_lon {
let mut acc = 0.0;
for (k, w) in k_lat.iter().enumerate() {
acc += w * rows_smoothed[(i + k) * n_lon + j];
}
values[i * n_lon + j] = scale * acc;
}
}
Ok(GeoGrid {
lat0_deg: south,
lon0_deg: west,
dlat_deg: step_deg,
dlon_deg: step_deg,
n_lat,
n_lon,
values,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn tiny() -> GeoGrid {
GeoGrid {
lat0_deg: 10.0,
lon0_deg: 20.0,
dlat_deg: 1.0,
dlon_deg: 1.0,
n_lat: 2,
n_lon: 3,
values: vec![0.0, 10.0, 20.0, 100.0, 110.0, f64::NAN],
}
}
#[test]
fn bilinear_is_exact_at_nodes_and_midpoints() {
let g = tiny();
assert_eq!(g.value_at(10.0, 20.0), 0.0);
assert_eq!(g.value_at(11.0, 21.0), 110.0);
assert!((g.value_at(10.5, 20.5) - 55.0).abs() < 1e-12);
assert!(g.value_at(10.5, 21.5).is_nan());
assert_eq!(g.void_count(), 1);
}
#[test]
fn from_fn_reproduces_a_plane_exactly() {
let g = GeoGrid::from_fn(1.0, 2.0, 3.0, 4.0, 0.1, |la, lo| 3.0 * la - 2.0 * lo).unwrap();
assert_eq!((g.n_lat, g.n_lon), (11, 11));
let (la, lo) = (1.234, 3.567);
assert!((g.value_at(la, lo) - (3.0 * la - 2.0 * lo)).abs() < 1e-12);
assert!(GeoGrid::from_fn(1.0, 1.0, 3.0, 4.0, 0.1, |_, _| 0.0).is_err());
assert!(GeoGrid::from_fn(1.0, 2.0, 3.0, 4.0, 0.0, |_, _| 0.0).is_err());
}
const EMAG2_SAMPLE: &str = "# synthetic test rows in the EMAG2v3 column layout\n\
0,0,20.0,10.0,1.0,2.0,019,5.0\n\
1,0,20.5,10.0,3.0,4.0,019,5.0\n\
0,1,20.0,10.5,5.0,6.0,019,5.0\n\
1,1,20.5,10.5,99999,99999,999,-999\n";
#[test]
fn emag2_reader_lays_out_the_lattice_and_masks_no_data() {
let g = GeoGrid::from_emag2_csv(EMAG2_SAMPLE, Emag2Column::SeaLevel).unwrap();
assert_eq!((g.n_lat, g.n_lon), (2, 2));
assert_eq!((g.lat0_deg, g.lon0_deg), (10.0, 20.0));
assert!((g.dlat_deg - 0.5).abs() < 1e-12 && (g.dlon_deg - 0.5).abs() < 1e-12);
assert_eq!(g.node(0, 0), 1.0);
assert_eq!(g.node(0, 1), 3.0);
assert_eq!(g.node(1, 0), 5.0);
assert!(g.node(1, 1).is_nan());
let up = GeoGrid::from_emag2_csv(EMAG2_SAMPLE, Emag2Column::UpCont).unwrap();
assert_eq!(up.node(0, 1), 4.0);
}
#[test]
fn emag2_reader_refuses_broken_lattices() {
let short = "0,0,20.0,10.0,1,2,19,5\n1,0,20.5,10.0,1,2,19,5\n0,1,20.0,10.5,1,2,19,5\n";
assert!(GeoGrid::from_emag2_csv(short, Emag2Column::SeaLevel).is_err());
assert!(GeoGrid::from_emag2_csv("0,0,20.0,10.0,1\n", Emag2Column::SeaLevel).is_err());
let irregular = "0,0,20.0,10.0,1,2,19,5\n1,0,20.5,10.0,1,2,19,5\n2,0,21.5,10.0,1,2,19,5\n\
0,1,20.0,10.5,1,2,19,5\n1,1,20.5,10.5,1,2,19,5\n2,1,21.5,10.5,1,2,19,5\n";
assert!(GeoGrid::from_emag2_csv(irregular, Emag2Column::SeaLevel).is_err());
assert!(GeoGrid::from_emag2_csv("# only a comment\n", Emag2Column::SeaLevel).is_err());
}
#[test]
fn error_field_has_the_configured_sigma_and_correlation() {
let (s, n, w, e, step) = (0.0, 6.0, 0.0, 6.0, 0.02);
let sigma = 3.0;
let l = 0.1;
let g = correlated_error_field(s, n, w, e, step, sigma, l, 7).unwrap();
let (mean, sd) = g.mean_std();
let mean_se = sigma * (2.0 * std::f64::consts::PI * l * l / ((n - s) * (e - w))).sqrt();
assert!(mean.abs() < 3.0 * mean_se, "mean {mean} vs 3 x {mean_se}");
assert!((sd - sigma).abs() / sigma < 0.05, "sd {sd} vs {sigma}");
let lag = (l / step).round() as usize;
let (mut num, mut den) = (0.0, 0.0);
for i in 0..g.n_lat - lag {
for j in 0..g.n_lon {
let a = g.node(i, j) - mean;
num += a * (g.node(i + lag, j) - mean);
den += a * a;
}
}
let rho = num / den;
assert!((rho - (-0.5_f64).pexp()).abs() < 0.05, "rho(L) = {rho}");
let again = correlated_error_field(s, n, w, e, step, sigma, l, 7).unwrap();
assert_eq!(g, again);
let other = correlated_error_field(s, n, w, e, step, sigma, l, 8).unwrap();
assert_ne!(g.values, other.values);
let zero = correlated_error_field(s, 1.0, w, 1.0, 0.1, 0.0, l, 7).unwrap();
assert!(zero.values.iter().all(|v| *v == 0.0));
}
}