use crate::error::{LatticeError, Result};
use crate::lattice::PaddedTileLattice;
impl<T: Clone + Default> PaddedTileLattice<T> {
pub fn map<U, F>(&self, mut f: F) -> PaddedTileLattice<U>
where
U: Clone + Default,
F: FnMut(&T) -> U,
{
let mut out = PaddedTileLattice::<U>::zeroed(self.rows(), self.cols(), *self.geometry())
.expect("shape is preserved, so geometry stays valid");
for row in 0..self.rows() {
for col in 0..self.cols() {
let mapped = f(self.get(row, col).unwrap());
out.set(row, col, mapped).unwrap();
}
}
out
}
pub fn map_in_place<F>(&mut self, mut f: F)
where
F: FnMut(&T) -> T,
{
for row in 0..self.rows() {
for col in 0..self.cols() {
let mapped = f(self.get(row, col).unwrap());
self.set(row, col, mapped).unwrap();
}
}
}
pub fn zip_with<F>(
&self,
other: &PaddedTileLattice<T>,
mut f: F,
) -> Result<PaddedTileLattice<T>>
where
F: FnMut(&T, &T) -> T,
{
if self.geometry() != other.geometry() {
return Err(LatticeError::GeometryMismatch);
}
if self.rows() != other.rows() || self.cols() != other.cols() {
return Err(LatticeError::ContractionMismatch {
lhs_cols: self.cols(),
rhs_rows: other.rows(),
});
}
let mut out = PaddedTileLattice::<T>::zeroed(self.rows(), self.cols(), *self.geometry())?;
for row in 0..self.rows() {
for col in 0..self.cols() {
let combined = f(self.get(row, col).unwrap(), other.get(row, col).unwrap());
out.set(row, col, combined)?;
}
}
Ok(out)
}
}