From 6888f04672e34592b01113c8df6caacd8dd0ae29 Mon Sep 17 00:00:00 2001
From: tomc271 <tomc271@gmail.com>
Date: Wed, 22 Nov 2023 21:39:30 +0000
Subject: [PATCH] Add MetricTensor::map() method and use to call
 interpolateAndExtrapolate() on each component.

---
 include/bout/metricTensor.hxx |  3 +++
 src/mesh/coordinates.cxx      | 32 +++++++++-----------------------
 src/mesh/metricTensor.cxx     |  9 +++++++++
 3 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/include/bout/metricTensor.hxx b/include/bout/metricTensor.hxx
index 739d76afd3..6a71ed9a5c 100644
--- a/include/bout/metricTensor.hxx
+++ b/include/bout/metricTensor.hxx
@@ -40,6 +40,9 @@ public:
   MetricTensor oppositeRepresentation(const CELL_LOC location, Mesh* mesh,
                                       const std::string& region = "RGN_ALL");
 
+  // Transforms the MetricTensor by applying the given function to every component
+  void map(const std::function<const Field2D(const FieldMetric)> function);
+
   MetricTensor applyToComponents(
       const std::function<const FieldMetric(const FieldMetric)> function) const;
 
diff --git a/src/mesh/coordinates.cxx b/src/mesh/coordinates.cxx
index 4bc5d895ae..dda1285c15 100644
--- a/src/mesh/coordinates.cxx
+++ b/src/mesh/coordinates.cxx
@@ -634,30 +634,16 @@ Coordinates::Coordinates(Mesh* mesh, Options* options, const CELL_LOC loc,
       }
     }
 
-    FieldMetric g_11, g_22, g_33, g_12, g_13, g_23;
-
     // More robust to extrapolate derived quantities directly, rather than
     // deriving from extrapolated covariant metric components
-    g_11 =
-        interpolateAndExtrapolate(covariantMetricTensor.Getg11(), location, extrapolate_x,
-                                  extrapolate_y, false, transform.get());
-    g_22 =
-        interpolateAndExtrapolate(covariantMetricTensor.Getg22(), location, extrapolate_x,
-                                  extrapolate_y, false, transform.get());
-    g_33 =
-        interpolateAndExtrapolate(covariantMetricTensor.Getg33(), location, extrapolate_x,
-                                  extrapolate_y, false, transform.get());
-    g_12 =
-        interpolateAndExtrapolate(covariantMetricTensor.Getg12(), location, extrapolate_x,
-                                  extrapolate_y, false, transform.get());
-    g_13 =
-        interpolateAndExtrapolate(covariantMetricTensor.Getg13(), location, extrapolate_x,
-                                  extrapolate_y, false, transform.get());
-    g_23 =
-        interpolateAndExtrapolate(covariantMetricTensor.Getg23(), location, extrapolate_x,
-                                  extrapolate_y, false, transform.get());
-    covariantMetricTensor.setMetricTensor(
-        MetricTensor(g_11, g_22, g_33, g_12, g_13, g_23));
+
+    std::function<const FieldMetric(const FieldMetric)>
+        interpolateAndExtrapolate_function =
+            [this, extrapolate_y, extrapolate_x](const FieldMetric component) {
+              return interpolateAndExtrapolate(component, location, extrapolate_x,
+                                               extrapolate_y, false, transform.get());
+            };
+    covariantMetricTensor.map(interpolateAndExtrapolate_function);
 
     // Check covariant metrics
     checkCovariant();
@@ -684,7 +670,7 @@ Coordinates::Coordinates(Mesh* mesh, Options* options, const CELL_LOC loc,
       communicate(this_J);
 
       // Re-evaluate Bxy using new J
-      this_Bxy = sqrt(g_22) / this_J;
+      this_Bxy = sqrt(g_22()) / this_J;
     }
 
     // Check jacobian
diff --git a/src/mesh/metricTensor.cxx b/src/mesh/metricTensor.cxx
index 8f18de9e24..7f639fb813 100644
--- a/src/mesh/metricTensor.cxx
+++ b/src/mesh/metricTensor.cxx
@@ -175,6 +175,15 @@ std::vector<MetricTensor::FieldMetric> MetricTensor::getComponents() const {
   return std::vector<FieldMetric>{g11, g22, g33, g12, g13, g23};
 }
 
+void MetricTensor::map(
+    const std::function<const FieldMetric(const FieldMetric)> function) {
+
+  const auto components_out = applyToComponents(function);
+
+  auto [g_11, g_22, g_33, g_12, g_13, g_23] = components_out;
+  setMetricTensor(MetricTensor(g_11, g_22, g_33, g_12, g_13, g_23));
+}
+
 MetricTensor MetricTensor::applyToComponents(
     const std::function<const FieldMetric(const FieldMetric)> function) const {