Expression unit_integral to normalize sources

include/bout/sys/generator_context.hxx

@@ -39,6 +39,7 @@ public:
   BoutReal y() const { return get("y"); }
   BoutReal z() const { return get("z"); }
   BoutReal t() const { return get("t"); }
+  CELL_LOC location() const { return location_; }
 
   /// Cell indices
   int ix() const { return ix_; }
@@ -80,6 +81,7 @@ private:
   int ix_{0};
   int jy_{0};
   int kz_{0};
+  CELL_LOC location_{CELL_CENTRE};
 
   Mesh* localmesh{nullptr}; ///< The mesh on which the position is defined
 

manual/sphinx/user_docs/field_expressions.rst

@@ -25,7 +25,6 @@ sense:
 - Unary algebraic operators such as ``sqrt``, ``abs``, ``exp``, ``log``,
   ``sin``, ``cos``, ``tan``, ``sinh``, ``cosh``, ``tanh``, and ``SQ``
 - Simple conditionals with ``if_else`` and ``if_else_zero``
-- Reductions such as ``min``, ``max``, and ``mean``
 
 For example::
 

manual/sphinx/user_docs/variable_init.rst

@@ -227,6 +227,10 @@ expressions.
    +------------------------------------------+------------------------------------------------------+
    | ``tan(x)``                               | Tangent                                              |
    +------------------------------------------+------------------------------------------------------+
+   | ``unit_integral(x)``                     | Scale ``x`` so that its mesh-weighted integral       |
+   |                                          | :math:`\sum x J\,dx\,dy\,dz` over ``RGN_NOBNDRY`` is |
+   |                                          | :math:`1`                                            |
+   +------------------------------------------+------------------------------------------------------+
    | ``erf(x)``                               | The error function                                   |
    +------------------------------------------+------------------------------------------------------+
    | ``TanhHat(x, width, centre, steepness)`` | The hat function                                     |

src/field/field_factory.cxx

@@ -220,6 +220,7 @@ FieldFactory::FieldFactory(Mesh* localmesh, Options* opt)
   addGenerator("power", std::make_shared<FieldGenTwoArg<pow>>(nullptr, nullptr));
 
   addGenerator("round", std::make_shared<FieldRound>(nullptr));
+  addGenerator("unit_integral", std::make_shared<FieldUnitIntegral>());
 
   // Ballooning transform
   addGenerator("ballooning", std::make_shared<FieldBallooning>(fieldmesh));

src/field/fieldgenerators.cxx

@@ -1,6 +1,7 @@
 
 #include "fieldgenerators.hxx"
 
+#include <bout/boutcomm.hxx>
 #include <memory>
 
 #include <bout/constants.hxx>
@@ -47,6 +48,69 @@ BoutReal FieldHeaviside::generate(const Context& ctx) {
   return (gen->generate(ctx) > 0.0) ? 1.0 : 0.0;
 }
 
+FieldGeneratorPtr FieldUnitIntegral::clone(const std::list<FieldGeneratorPtr> args) {
+  if (args.size() != 1) {
+    throw ParseException(
+        "Incorrect number of arguments to unit_integral function. Expecting 1, got {:d}",
+        args.size());
+  }
+
+  return std::make_shared<FieldUnitIntegral>(args.front());
+}
+
+bool FieldUnitIntegral::cacheMatches(const Context& ctx) const {
+  return cache_valid && (cached_mesh == ctx.getMesh()) && (cached_time == ctx.t())
+         && (cached_location == ctx.location());
+}
+
+void FieldUnitIntegral::populateCache(const Context& ctx) {
+  Mesh* localmesh = ctx.getMesh();
+  ASSERT0(localmesh != nullptr);
+  Coordinates* coords = localmesh->getCoordinates(ctx.location());
+  if (coords == nullptr) {
+    throw BoutException("unit_integral function needs coordinates at {}",
+                        toString(ctx.location()));
+  }
+
+  cached_values = Field3D(localmesh).setLocation(ctx.location()).allocate();
+
+  BOUT_FOR(i, cached_values.getRegion("RGN_ALL")) {
+    cached_values[i] = gen->generate(Context(i, ctx.location(), localmesh, ctx.t()));
+  }
+
+  BoutReal local_integral = 0.0;
+  BOUT_FOR(i, cached_values.getRegion("RGN_NOBNDRY")) {
+    local_integral += cached_values[i] * coords->J()[i] * coords->dx()[i]
+                      * coords->dy()[i] * coords->dz()[i];
+  }
+
+  BoutReal integral = 0.0;
+  MPI_Allreduce(&local_integral, &integral, 1, MPI_DOUBLE, MPI_SUM, BoutComm::get());
+
+  if (integral == 0.0) {
+    throw BoutException("unit_integral function integral is zero");
+  }
+
+  BOUT_FOR(i, cached_values.getRegion("RGN_ALL")) { cached_values[i] /= integral; }
+
+  cached_mesh = localmesh;
+  cached_time = ctx.t();
+  cached_location = ctx.location();
+  cache_valid = true;
+}
+
+BoutReal FieldUnitIntegral::generate(const Context& ctx) {
+  if (!cacheMatches(ctx)) {
+    std::lock_guard<std::mutex> guard(cache_mutex);
+    if (!cacheMatches(ctx)) {
+      populateCache(ctx);
+    }
+  }
+
+  ASSERT1(ctx.location() == cached_location);
+  return cached_values(ctx.ix(), ctx.jy(), ctx.kz());
+}
+
 //////////////////////////////////////////////////////////
 // Ballooning transform
 // Use a truncated Ballooning transform to enforce periodicity in y and z

src/field/fieldgenerators.hxx

@@ -265,6 +265,30 @@ private:
   FieldGeneratorPtr gen;
 };
 
+/// Scale an expression so its volume integral over RGN_NOBNDRY is 1
+class FieldUnitIntegral : public FieldGenerator {
+public:
+  explicit FieldUnitIntegral(FieldGeneratorPtr g = nullptr) : gen(std::move(g)) {}
+
+  FieldGeneratorPtr clone(const std::list<FieldGeneratorPtr> args) override;
+  BoutReal generate(const bout::generator::Context& pos) override;
+  std::string str() const override {
+    return std::string("unit_integral(") + gen->str() + std::string(")");
+  }
+
+private:
+  void populateCache(const bout::generator::Context& ctx);
+  bool cacheMatches(const bout::generator::Context& ctx) const;
+
+  FieldGeneratorPtr gen;
+  Field3D cached_values{};
+  Mesh* cached_mesh{nullptr};
+  BoutReal cached_time{0.0};
+  CELL_LOC cached_location{CELL_CENTRE};
+  bool cache_valid{false};
+  std::mutex cache_mutex;
+};
+
 //////////////////////////////////////////////////////////
 // Ballooning transform
 // Use a truncated Ballooning transform to enforce periodicity

src/sys/generator_context.cxx

@@ -9,7 +9,7 @@ namespace bout {
 namespace generator {
 
 Context::Context(int ix, int iy, int iz, CELL_LOC loc, Mesh* msh, BoutReal t)
-    : ix_(ix), jy_(iy), kz_(iz), localmesh(msh) {
+    : ix_(ix), jy_(iy), kz_(iz), location_(loc), localmesh(msh) {
 
   parameters["x"] = (loc == CELL_XLOW) ? 0.5 * (msh->GlobalX(ix) + msh->GlobalX(ix - 1))
                                        : msh->GlobalX(ix);
@@ -26,7 +26,8 @@ Context::Context(int ix, int iy, int iz, CELL_LOC loc, Mesh* msh, BoutReal t)
 Context::Context(const BoundaryRegion* bndry, int iz, CELL_LOC loc, BoutReal t, Mesh* msh)
     : // Add one to X index if boundary is in -x direction, so that XLOW is on the boundary
       ix_((bndry->bx < 0) ? bndry->x + 1 : bndry->x),
-      jy_((bndry->by < 0) ? bndry->y + 1 : bndry->y), kz_(iz), localmesh(msh) {
+      jy_((bndry->by < 0) ? bndry->y + 1 : bndry->y), kz_(iz), location_(loc),
+      localmesh(msh) {
 
   parameters["x"] = ((loc == CELL_XLOW) || (bndry->bx != 0))
                         ? 0.5 * (msh->GlobalX(ix_) + msh->GlobalX(ix_ - 1))

tests/unit/field/test_field_factory.cxx

@@ -3,6 +3,7 @@
 #include "fake_mesh.hxx"
 #include "test_extras.hxx"
 #include "bout/bout_types.hxx"
+#include "bout/boutcomm.hxx"
 #include "bout/boutexception.hxx"
 #include "bout/constants.hxx"
 #include "bout/coordinates.hxx"
@@ -29,6 +30,43 @@
 using namespace bout::globals;
 using bout::generator::Context;
 
+namespace {
+BoutReal volumeIntegral(const Field3D& field) {
+  const auto* coords = field.getCoordinates();
+  if (coords == nullptr) {
+    throw BoutException("Field has no coordinates");
+  }
+
+  BoutReal local = 0.0;
+  BOUT_FOR(i, field.getRegion("RGN_NOBNDRY")) {
+    local +=
+        field[i] * coords->J()[i] * coords->dx()[i] * coords->dy()[i] * coords->dz()[i];
+  }
+
+  BoutReal global = 0.0;
+  MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_SUM, BoutComm::get());
+  return global;
+}
+
+BoutReal totalVolume(Mesh* localmesh, CELL_LOC location = CELL_CENTRE) {
+  auto* coords = localmesh->getCoordinates(location);
+  if (coords == nullptr) {
+    throw BoutException("Mesh has no coordinates");
+  }
+
+  BoutReal local = 0.0;
+  Field3D ones{localmesh};
+  ones.allocate();
+  BOUT_FOR(i, ones.getRegion("RGN_NOBNDRY")) {
+    local += coords->J()[i] * coords->dx()[i] * coords->dy()[i] * coords->dz()[i];
+  }
+
+  BoutReal global = 0.0;
+  MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_SUM, BoutComm::get());
+  return global;
+}
+} // namespace
+
 // Reuse the "standard" fixture for FakeMesh
 template <typename T>
 class FieldFactoryCreationTest : public FakeMeshFixture {
@@ -60,6 +98,8 @@ using Fields = ::testing::Types<Field2D, Field3D>;
 
 TYPED_TEST_SUITE(FieldFactoryCreationTest, Fields);
 
+using FieldFactory3DCreationTest = FieldFactoryCreationTest<Field3D>;
+
 TYPED_TEST(FieldFactoryCreationTest, CreateFromValueGenerator) {
   auto value = BoutReal{4.};
   auto output = this->create(generator(value));
@@ -179,6 +219,40 @@ TYPED_TEST(FieldFactoryCreationTest, CreateZ) {
   EXPECT_TRUE(IsFieldEqual(output, expected));
 }
 
+TEST_F(FieldFactory3DCreationTest, CreateUnitIntegralConstant) {
+  const auto expected = 1.0 / totalVolume(mesh);
+
+  auto output = this->factory.create3D("unit_integral(1)");
+
+  EXPECT_TRUE(IsFieldEqual(output, expected));
+  EXPECT_NEAR(volumeIntegral(output), 1.0, 1e-12);
+}
+
+TEST_F(FieldFactory3DCreationTest, CreateUnitIntegralUsesCellVolume) {
+  Field2D dy{mesh};
+  dy.allocate();
+  BOUT_FOR(i, dy.getRegion("RGN_ALL")) { dy[i] = static_cast<BoutReal>(i.y() + 1); }
+  mesh->getCoordinates()->setDy(dy);
+
+  const auto expected = 1.0 / totalVolume(mesh);
+  auto output = this->factory.create3D("unit_integral(1)");
+
+  EXPECT_TRUE(IsFieldEqual(output, expected));
+  EXPECT_NEAR(volumeIntegral(output), 1.0, 1e-12);
+}
+
+TEST_F(FieldFactory3DCreationTest, CreateUnitIntegralRefreshesCachedParseAtNewTime) {
+  auto output_t0 =
+      this->factory.create3D("unit_integral(y + t)", nullptr, mesh, CELL_CENTRE, 0.0);
+  auto output_t1 =
+      this->factory.create3D("unit_integral(y + t)", nullptr, mesh, CELL_CENTRE, 1.0);
+
+  EXPECT_NEAR(volumeIntegral(output_t0), 1.0, 1e-12);
+  EXPECT_NEAR(volumeIntegral(output_t1), 1.0, 1e-12);
+  EXPECT_NE(output_t0(mesh->xstart, mesh->ystart, 0),
+            output_t1(mesh->xstart, mesh->ystart, 0));
+}
+
 TYPED_TEST(FieldFactoryCreationTest, CreateXStaggered) {
   // Need this->mesh_staggered to access member of base FakeMeshFixture because
   // derived FieldFactoryCreationTest is a template clas