Add API to normalise metrics

include/bout/coordinates.hxx

@@ -46,6 +46,7 @@
 
 class Mesh;
 class YBoundary;
+struct MetricNormaliser;
 
 /*!
  * Represents a coordinate system, and associated operators
@@ -154,8 +155,6 @@ public:
   /// get g_22 at the cell faces;
   const FieldMetric& g_22_ylow() const;
   const FieldMetric& g_22_yhigh() const;
-  FieldMetric& g_22_ylow();
-  FieldMetric& g_22_yhigh();
   // Cell Areas
   const FieldMetric& cell_area_xlow() const {
     if (_cell_area_xlow.has_value()) {
@@ -370,10 +369,16 @@ public:
   void setMetricTensor(const ContravariantMetricTensor& contravariant_metric_tensor,
                        const CovariantMetricTensor& covariant_metric_tensor);
 
+  void setMetricTensorJB(const ContravariantMetricTensor& contravariant_metric_tensor,
+                         const CovariantMetricTensor& covariant_metric_tensor,
+                         const FieldMetric& J, const FieldMetric& Bxy);
+
   void communicateMetricTensor();
 
   void communicateDz();
 
+  void normaliseMetric(const MetricNormaliser& norm);
+
   ///< Coordinate system Jacobian, so volume of cell is J*dx*dy*dz
   const FieldMetric& J() const;
 
@@ -543,4 +548,24 @@ namespace bout {
 std::string parallelSliceFieldName(std::string_view field, int offset);
 }
 
+/// Represents a way to normalise the coordinate system
+/// If a component returns nothing, no normalisation is performed.
+/// Coordinate values are divided by the respective component from
+/// MetricNormaliser, with the exception of the contravariant metric
+/// tensor, which is multiplied by the normalisation factor.
+struct MetricNormaliser {
+  std::optional<BoutReal> g = std::nullopt;
+  std::optional<BoutReal> g11 = std::nullopt;
+  std::optional<BoutReal> g22 = std::nullopt;
+  std::optional<BoutReal> g33 = std::nullopt;
+  std::optional<BoutReal> g12 = std::nullopt;
+  std::optional<BoutReal> g13 = std::nullopt;
+  std::optional<BoutReal> g23 = std::nullopt;
+  std::optional<BoutReal> dx = std::nullopt;
+  std::optional<BoutReal> dy = std::nullopt;
+  std::optional<BoutReal> dz = std::nullopt;
+  std::optional<BoutReal> J = std::nullopt;
+  std::optional<BoutReal> Bxy = std::nullopt;
+};
+
 #endif // BOUT_COORDINATES_H

include/bout/metric_tensor.hxx

@@ -17,6 +17,7 @@ using FieldMetric = Field2D;
 } // namespace bout
 
 class Coordinates;
+struct MetricNormaliser;
 
 class MetricTensor {
 public:
@@ -77,6 +78,9 @@ public:
 
   void communicate();
 
+  template <class F>
+  void normaliseMetric(const MetricNormaliser& norm, const F& op);
+
 private:
   FieldMetric g11_m, g22_m, g33_m, g12_m, g13_m, g23_m;
 };
@@ -90,6 +94,8 @@ public:
 
   auto inverse(const std::string& region = "RGN_ALL", bool communicate = true)
       -> ContravariantMetricTensor;
+
+  void normaliseMetric(const MetricNormaliser& norm);
 };
 
 class ContravariantMetricTensor : public MetricTensor {
@@ -98,6 +104,8 @@ public:
 
   auto inverse(const std::string& region = "RGN_ALL", bool communicate = true)
       -> CovariantMetricTensor;
+
+  void normaliseMetric(const MetricNormaliser& norm);
 };
 
 #endif //BOUT_METRIC_TENSOR_HXX

include/bout/tokamak_coordinates.hxx

@@ -45,6 +45,9 @@ TokamakCoordinates set_tokamak_coordinates(Mesh& mesh, BoutReal Lbar = 1.0,
                                            BoutReal Bbar = 1.0, bool no_shear = false,
                                            BoutReal shear_factor = 1.0);
 
+MetricNormaliser TokamakOrFCIMetricNormaliser(const Mesh* mesh, BoutReal Bnorm,
+                                              BoutReal rho_s0);
+
 } // namespace bout
 
 #endif //BOUT_TOKAMAK_COORDINATES_HXX

src/mesh/coordinates.cxx

@@ -823,6 +823,7 @@ Coordinates::FieldMetric Coordinates::recalculateJacobian() const {
 }
 
 Coordinates::FieldMetric Coordinates::recalculateBxy() const {
+  ASSERT2(not J().isFci());
   return sqrt(g_22()) / J();
 }
 
@@ -1268,6 +1269,16 @@ void Coordinates::setMetricTensor(
   setJ(recalculateJacobian());
   setBxy(recalculateBxy());
 }
+void Coordinates::setMetricTensorJB(
+    const ContravariantMetricTensor& contravariant_metric_tensor,
+    const CovariantMetricTensor& covariant_metric_tensor, const FieldMetric& J,
+    const FieldMetric& Bxy) {
+  contravariantMetricTensor = contravariant_metric_tensor;
+  covariantMetricTensor = covariant_metric_tensor;
+  setJ(J);
+  setBxy(Bxy);
+  invalidateMetricCaches();
+}
 
 void Coordinates::communicateMetricTensor() {
   contravariantMetricTensor.communicate();
@@ -1277,6 +1288,58 @@ void Coordinates::communicateMetricTensor() {
 void Coordinates::communicateDz() { localmesh->communicate(dz_); }
 
 void Coordinates::splitBxyParallelSlices() {
-  Bxy_.splitParallelSlices();
-  Bxy_.yup() = Bxy_.ydown() = Bxy_;
+  if (not Bxy_.hasParallelSlices()) {
+    auto copy = Bxy_;
+    Bxy_.splitParallelSlices();
+    Bxy_.yup() = Bxy_.ydown() = copy;
+  }
+}
+
+void Coordinates::normaliseMetric(const MetricNormaliser& norm) {
+  covariantMetricTensor.normaliseMetric(norm);
+  contravariantMetricTensor.normaliseMetric(norm);
+
+#if BOUT_USE_METRIC_3D
+  using FieldMetricParallel = Field3DParallel;
+#else
+  using FieldMetricParallel = Field2D;
+#endif
+
+  if (norm.J.has_value()) {
+    if (J().hasParallelSlices()) {
+      setJ(FieldMetricParallel{J() / *norm.J});
+    } else {
+      setJ(J() / *norm.J);
+    }
+  }
+  if (norm.Bxy.has_value()) {
+    if (Bxy().hasParallelSlices()) {
+      setBxy(FieldMetricParallel{Bxy() / *norm.Bxy});
+    } else {
+      setBxy(Bxy() / *norm.Bxy);
+    }
+  }
+  if (norm.dx.has_value()) {
+    setDx(dx() / *norm.dx);
+  }
+  if (norm.dy.has_value()) {
+    setDy(dy() / *norm.dy);
+  }
+  if (norm.dz.has_value()) {
+    setDz(dz() / *norm.dz);
+  }
+  invalidateMetricCaches();
+  if (norm.g.has_value() or norm.g22.has_value()) {
+    if (Bxy().isFci()) {
+      // No we compute g_22_* - they must not be cleared. If they get
+      // cleared, they will be recomputed, but not normalised!
+      auto g22 = norm.g.has_value() ? norm.g.value() : norm.g22.value();
+      g_22_ylow();
+      g_22_yhigh();
+      ASSERT2(_g_22_ylow.has_value());
+      (*_g_22_ylow) /= g22;
+      ASSERT2(_g_22_yhigh.has_value());
+      (*_g_22_yhigh) /= g22;
+    }
+  }
 }

src/mesh/difops.cxx

@@ -762,10 +762,13 @@ Field3D Laplace(const Field3D& f, CELL_LOC outloc,
  * Inverse of Laplacian operator in LaplaceXY solver
  *******************************************************************************/
 
-Field2D Laplace_perpXY(const Field2D& A, const Field2D& f) {
 #if BOUT_USE_METRIC_3D
+Field2D Laplace_perpXY([[maybe_unused]] const Field2D& A,
+                       [[maybe_unused]] const Field2D& f) {
   throw BoutException("Coordinates::Laplace_perpXY for 3D metric not implemented");
+}
 #else
+Field2D Laplace_perpXY(const Field2D& A, const Field2D& f) {
   const auto& coords = *f.getCoordinates();
 
   Field2D result;
@@ -821,8 +824,8 @@ Field2D Laplace_perpXY(const Field2D& A, const Field2D& f) {
   }
 
   return result;
-#endif
 }
+#endif
 
 /*******************************************************************************
 * b0xGrad_dot_Grad

src/mesh/metric_tensor.cxx

@@ -149,6 +149,48 @@ auto ContravariantMetricTensor::inverse(const std::string& region, bool communic
   return result;
 }
 
+template <class F>
+void MetricTensor::normaliseMetric(const MetricNormaliser& norm, const F& op) {
+  if (norm.g.has_value()) {
+    op(g11_m, norm.g);
+    op(g22_m, norm.g);
+    op(g33_m, norm.g);
+    op(g12_m, norm.g);
+    op(g13_m, norm.g);
+    op(g23_m, norm.g);
+  } else {
+    op(g11_m, norm.g11);
+    op(g22_m, norm.g22);
+    op(g33_m, norm.g33);
+    op(g12_m, norm.g12);
+    op(g13_m, norm.g13);
+    op(g23_m, norm.g23);
+  }
+}
+
+void ContravariantMetricTensor::normaliseMetric(const MetricNormaliser& norm) {
+  MetricTensor::normaliseMetric(norm, [](FieldMetric& f, auto fac) {
+    if (fac.has_value()) {
+      if (f.hasParallelSlices()) {
+        f.asField3DParallel() *= fac.value();
+      } else {
+        f *= fac.value();
+      }
+    }
+  });
+}
+void CovariantMetricTensor::normaliseMetric(const MetricNormaliser& norm) {
+  MetricTensor::normaliseMetric(norm, [](FieldMetric& f, auto fac) {
+    if (fac.has_value()) {
+      if (f.hasParallelSlices()) {
+        f.asField3DParallel() /= fac.value();
+      } else {
+        f /= fac.value();
+      }
+    }
+  });
+}
+
 void MetricTensor::communicate() {
   g11_m.getMesh()->communicate_no_slices(g11_m, g22_m, g33_m, g12_m, g13_m, g23_m);
 }

src/mesh/tokamak_coordinates.cxx

@@ -75,4 +75,21 @@ TokamakCoordinates set_tokamak_coordinates(Mesh& mesh, BoutReal Lbar, BoutReal B
 
   return {Rxy, Zxy, Bpxy, Btxy, Bxy, hthe, I, I_unnormalised};
 }
+
+MetricNormaliser TokamakOrFCIMetricNormaliser(const Mesh* mesh, BoutReal Bnorm,
+                                              BoutReal rho_s0) {
+  if (mesh->isFci()) {
+    return {.g{SQ(rho_s0)}, .J{rho_s0 * rho_s0 * rho_s0}, .Bxy{Bnorm}};
+  }
+  return {.g11{1 / SQ(Bnorm * rho_s0)},
+          .g22{SQ(rho_s0)},
+          .g33{SQ(rho_s0)},
+          .g12{1 / Bnorm},
+          .g13{1 / Bnorm},
+          .g23{SQ(rho_s0)},
+          .dx{rho_s0 * rho_s0 * Bnorm},
+          .J{rho_s0 / Bnorm},
+          .Bxy{Bnorm}};
+}
+
 } // namespace bout