use crate::foundation::{GeoError, GridGeometry, Result};
use indexmap::IndexMap;
use ndarray::Array2;
use std::path::Path;
#[derive(Clone)]
pub struct Surface {
pub geom: GridGeometry,
values: Array2<f64>,
attributes: IndexMap<String, Array2<f64>>,
}
impl Surface {
pub fn new(geom: GridGeometry, values: Array2<f64>) -> Result<Surface> {
check_shape(&geom, &values, "Surface::new")?;
Ok(Surface {
geom,
values,
attributes: IndexMap::new(),
})
}
pub(crate) fn from_values_unchecked(geom: GridGeometry, values: Array2<f64>) -> Surface {
Surface {
geom,
values,
attributes: IndexMap::new(),
}
}
pub fn constant(geom: GridGeometry, value: f64) -> Surface {
let values = Array2::from_elem((geom.ncol, geom.nrow), value);
Surface {
geom,
values,
attributes: IndexMap::new(),
}
}
pub fn load_irap_classic(path: impl AsRef<Path>) -> Result<Surface> {
let (geom, values) = crate::io::irap::load_irap_classic(path.as_ref())?;
Ok(Surface {
geom,
values,
attributes: IndexMap::new(),
})
}
pub fn save_irap_classic(&self, path: impl AsRef<Path>) -> Result<()> {
crate::io::irap::save_irap_classic(path.as_ref(), &self.geom, &self.values)
}
pub fn values(&self) -> &Array2<f64> {
&self.values
}
pub fn attr(&self, name: &str) -> Option<&Array2<f64>> {
self.attributes.get(name)
}
pub fn set_attr(&mut self, name: &str, values: Array2<f64>) -> Result<()> {
check_shape(&self.geom, &values, "Surface::set_attr")?;
self.attributes.insert(name.to_string(), values);
Ok(())
}
pub fn attr_names(&self) -> Vec<&str> {
self.attributes.keys().map(String::as_str).collect()
}
pub fn as_attr_surface(&self, name: &str) -> Option<Surface> {
self.attributes.get(name).map(|a| Surface {
geom: self.geom.clone(),
values: a.clone(),
attributes: IndexMap::new(),
})
}
pub fn sample(&self, x: f64, y: f64) -> Option<f64> {
let (ncol, nrow) = (self.geom.ncol, self.geom.nrow);
if ncol < 2 || nrow < 2 {
return None;
}
let (fi, fj) = self.geom.xy_to_ij(x, y)?;
if fi < 0.0 || fj < 0.0 {
return None;
}
let i = fi.floor() as usize;
let j = fj.floor() as usize;
if i >= ncol - 1 || j >= nrow - 1 {
return None;
}
let (t, u) = (fi - i as f64, fj - j as f64);
let c00 = self.values[[i, j]];
let c10 = self.values[[i + 1, j]];
let c01 = self.values[[i, j + 1]];
let c11 = self.values[[i + 1, j + 1]];
if c00.is_nan() || c10.is_nan() || c01.is_nan() || c11.is_nan() {
return None;
}
Some((1.0 - t) * (1.0 - u) * c00 + t * (1.0 - u) * c10 + (1.0 - t) * u * c01 + t * u * c11)
}
pub fn resample(&self, target: &GridGeometry) -> Surface {
let mut out = Array2::from_elem((target.ncol, target.nrow), f64::NAN);
for j in 0..target.nrow {
for i in 0..target.ncol {
let (x, y) = target.node_xy(i, j);
if let Some(v) = self.sample(x, y) {
out[[i, j]] = v;
}
}
}
Surface {
geom: target.clone(),
values: out,
attributes: IndexMap::new(),
}
}
}
fn check_shape(geom: &GridGeometry, values: &Array2<f64>, ctx: &str) -> Result<()> {
if values.dim() != (geom.ncol, geom.nrow) {
return Err(GeoError::GeometryMismatch(format!(
"{ctx}: values shape {:?} != grid (ncol={}, nrow={})",
values.dim(),
geom.ncol,
geom.nrow
)));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use approx::assert_relative_eq;
fn ramp() -> Surface {
let mut v = Array2::zeros((2, 2));
v[[0, 0]] = 0.0;
v[[1, 0]] = 10.0;
v[[0, 1]] = 20.0;
v[[1, 1]] = 30.0;
Surface::new(geom(), v).unwrap()
}
#[test]
fn bilinear_sample_hand_calc() {
let s = ramp();
assert_relative_eq!(s.sample(5.0, 5.0).unwrap(), 15.0); assert_relative_eq!(s.sample(2.0, 0.0).unwrap(), 2.0); assert_relative_eq!(s.sample(0.0, 0.0).unwrap(), 0.0); assert_eq!(s.sample(-1.0, 0.0), None); assert_eq!(s.sample(100.0, 100.0), None); }
#[test]
fn bilinear_nan_corner_is_none() {
let mut v = Array2::zeros((2, 2));
v[[0, 0]] = 0.0;
v[[1, 0]] = 10.0;
v[[0, 1]] = 20.0;
v[[1, 1]] = f64::NAN;
let s = Surface::new(geom(), v).unwrap();
assert_eq!(s.sample(5.0, 5.0), None);
}
#[test]
fn resample_interpolates_and_copies_geometry() {
let s = ramp();
let target = GridGeometry {
xori: 0.0,
yori: 0.0,
xinc: 5.0,
yinc: 5.0,
ncol: 2,
nrow: 2,
rotation_deg: 0.0,
yflip: false,
};
let r = s.resample(&target);
assert_eq!(r.geom, target);
assert_relative_eq!(r.values()[[0, 0]], 0.0);
assert_relative_eq!(r.values()[[1, 1]], 15.0); }
fn geom() -> GridGeometry {
GridGeometry {
xori: 0.0,
yori: 0.0,
xinc: 10.0,
yinc: 10.0,
ncol: 2,
nrow: 2,
rotation_deg: 0.0,
yflip: false,
}
}
#[test]
fn new_rejects_wrong_shape() {
let bad = Array2::from_elem((3, 3), 1.0);
assert!(Surface::new(geom(), bad).is_err());
}
#[test]
fn attributes_set_get_promote() {
let mut s = Surface::constant(geom(), 1.0);
s.set_attr("thickness", Array2::from_elem((2, 2), 5.0))
.unwrap();
assert_eq!(s.attr_names(), vec!["thickness"]);
assert_eq!(s.attr("thickness").unwrap()[[0, 0]], 5.0);
assert!(s.attr("missing").is_none());
let promoted = s.as_attr_surface("thickness").unwrap();
assert_eq!(promoted.values()[[1, 1]], 5.0);
assert!(s.set_attr("bad", Array2::from_elem((1, 1), 0.0)).is_err());
}
}