From 9e2ce3ae245e772cc4c9b9e89fb4fed307d53193 Mon Sep 17 00:00:00 2001
From: Magnus Ulimoen <magnus@ulimoen.dev>
Date: Fri, 26 Mar 2021 16:27:45 +0100
Subject: [PATCH] Add Metrics iterator

---
 euler/src/lib.rs | 22 +++++++++-------------
 sbp/src/grid.rs  | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/euler/src/lib.rs b/euler/src/lib.rs
index d43339b..03d2c54 100644
--- a/euler/src/lib.rs
+++ b/euler/src/lib.rs
@@ -618,15 +618,11 @@ fn upwind_dissipation(
     let mut tmp0 = tmp.0 .0.view_mut().into_shape((4, n)).unwrap();
     let mut tmp1 = tmp.1 .0.view_mut().into_shape((4, n)).unwrap();
 
-    for ((((((y, mut tmp0), mut tmp1), detj_dxi_dx), detj_dxi_dy), detj_deta_dx), detj_deta_dy) in
-        yview
-            .axis_iter(ndarray::Axis(1))
-            .zip(tmp0.axis_iter_mut(ndarray::Axis(1)))
-            .zip(tmp1.axis_iter_mut(ndarray::Axis(1)))
-            .zip(metrics.detj_dxi_dx().iter())
-            .zip(metrics.detj_dxi_dy().iter())
-            .zip(metrics.detj_deta_dx().iter())
-            .zip(metrics.detj_deta_dy().iter())
+    for (((y, mut tmp0), mut tmp1), metric) in yview
+        .axis_iter(ndarray::Axis(1))
+        .zip(tmp0.axis_iter_mut(ndarray::Axis(1)))
+        .zip(tmp1.axis_iter_mut(ndarray::Axis(1)))
+        .zip(metrics.iter())
     {
         let rho = y[0];
         assert!(rho > 0.0);
@@ -637,8 +633,8 @@ fn upwind_dissipation(
         let u = rhou / rho;
         let v = rhov / rho;
 
-        let uhat = detj_dxi_dx * u + detj_dxi_dy * v;
-        let vhat = detj_deta_dx * u + detj_deta_dy * v;
+        let uhat = metric.detj_dxi_dx * u + metric.detj_dxi_dy * v;
+        let vhat = metric.detj_deta_dx * u + metric.detj_deta_dy * v;
 
         let p = pressure(GAMMA, rho, rhou, rhov, e);
         assert!(p > 0.0);
@@ -648,8 +644,8 @@ fn upwind_dissipation(
         // not that important in this case
         let hypot = |x: Float, y: Float| Float::sqrt(x * x + y * y);
 
-        let alpha_u = uhat.abs() + c * hypot(*detj_dxi_dx, *detj_dxi_dy);
-        let alpha_v = vhat.abs() + c * hypot(*detj_deta_dx, *detj_deta_dy);
+        let alpha_u = uhat.abs() + c * hypot(metric.detj_dxi_dx, metric.detj_dxi_dy);
+        let alpha_v = vhat.abs() + c * hypot(metric.detj_deta_dx, metric.detj_deta_dy);
 
         tmp0[0] = alpha_u * rho;
         tmp1[0] = alpha_v * rho;
diff --git a/sbp/src/grid.rs b/sbp/src/grid.rs
index fb56ced..f6ce1e4 100644
--- a/sbp/src/grid.rs
+++ b/sbp/src/grid.rs
@@ -162,4 +162,37 @@ impl Metrics {
     pub fn detj_deta_dy(&self) -> ArrayView2<Float> {
         self.detj_deta_dy.view()
     }
+    pub fn iter(
+        &self,
+    ) -> impl DoubleEndedIterator<Item = Metric> + ExactSizeIterator<Item = Metric> + '_ {
+        let n = self.nx() * self.ny();
+        let detj = &self.detj.as_slice().unwrap()[..n];
+        let detj_dxi_dx = &self.detj_dxi_dx.as_slice().unwrap()[..n];
+        let detj_dxi_dy = &self.detj_dxi_dy.as_slice().unwrap()[..n];
+        let detj_deta_dx = &self.detj_deta_dx.as_slice().unwrap()[..n];
+        let detj_deta_dy = &self.detj_deta_dy.as_slice().unwrap()[..n];
+
+        detj.iter()
+            .zip(detj_dxi_dx)
+            .zip(detj_dxi_dy)
+            .zip(detj_deta_dx)
+            .zip(detj_deta_dy)
+            .map(
+                |((((&detj, &detj_dxi_dx), &detj_dxi_dy), &detj_deta_dx), &detj_deta_dy)| Metric {
+                    detj,
+                    detj_dxi_dx,
+                    detj_dxi_dy,
+                    detj_deta_dx,
+                    detj_deta_dy,
+                },
+            )
+    }
+}
+
+pub struct Metric {
+    pub detj: Float,
+    pub detj_dxi_dx: Float,
+    pub detj_dxi_dy: Float,
+    pub detj_deta_dx: Float,
+    pub detj_deta_dy: Float,
 }