Add ability to use variables from grid file in input file expressions

include/bout/sys/generator_context.hxx

@@ -24,12 +24,8 @@ public:
       : Context(i.x(), i.y(), 0, (loc == CELL_ZLOW) ? CELL_CENTRE : loc, msh, t) {}
 
   /// Specify a cell index, together with the cell location, mesh and time
-  ///
   Context(int ix, int iy, int iz, CELL_LOC loc, Mesh* msh, BoutReal t);
 
-  /// Specify the values directly
-  Context(BoutReal x, BoutReal y, BoutReal z, Mesh* msh, BoutReal t);
-
   /// If constructed without parameters, contains no values (null).
   /// Requesting x,y,z or t throws an exception
   Context() = default;
@@ -44,6 +40,11 @@ public:
   BoutReal z() const { return get("z"); }
   BoutReal t() const { return get("t"); }
 
+  /// Cell indices
+  int ix() const { return ix_; }
+  int jy() const { return jy_; }
+  int kz() const { return kz_; }
+
   /// Set the value of a parameter with given name
   Context& set(const std::string& name, BoutReal value) {
     parameters[name] = value;
@@ -76,6 +77,10 @@ public:
   }
 
 private:
+  int ix_{0};
+  int jy_{0};
+  int kz_{0};
+
   Mesh* localmesh{nullptr}; ///< The mesh on which the position is defined
 
   /// Contains user-set values which can be set and retrieved

manual/sphinx/user_docs/input_grids.rst

@@ -155,6 +155,21 @@ The only quantities which are required are the sizes of the grid. If
 these are the only quantities specified, then the coordinates revert to
 Cartesian.
 
+You can read additional quantities from the grid and make them available in
+expressions in the input file by listing them in the ``input:grid_variables``
+section, with the key being the name in the grid file (``mesh:file``) and the
+value being the type (one of ``field3d``, ``field2d``, ``boutreal``):
+
+.. code-block:: cfg
+
+   [input:grid_variables]
+   rho = field2d
+   theta = field2d
+   scale = boutreal
+
+   [mesh]
+   B = (scale / rho) * cos(theta)
+
 This section describes how to generate inputs for tokamak equilibria. If
 you’re not interested in tokamaks then you can skip to the next section.
 

src/field/field_factory.cxx

@@ -19,20 +19,28 @@
  * along with BOUT++.  If not, see <http://www.gnu.org/licenses/>.
  *
  **************************************************************************/
-#include <bout/globals.hxx>
 
 #include <bout/field_factory.hxx>
 
-#include <cmath>
-
+#include <bout/bout_enum_class.hxx>
+#include <bout/bout_types.hxx>
+#include <bout/boutexception.hxx>
 #include <bout/constants.hxx>
+#include <bout/field2d.hxx>
+#include <bout/field3d.hxx>
+#include <bout/fieldperp.hxx>
+#include <bout/globals.hxx>
 #include <bout/output.hxx>
+#include <bout/sys/expressionparser.hxx>
+#include <bout/traits.hxx>
 #include <bout/utils.hxx>
 
-#include "bout/constants.hxx"
-
 #include "fieldgenerators.hxx"
 
+#include <cmath>
+#include <memory>
+#include <string>
+
 using bout::generator::Context;
 
 /// Helper function to create a FieldValue generator from a BoutReal
@@ -45,6 +53,8 @@ FieldGeneratorPtr generator(BoutReal* ptr) {
   return std::make_shared<FieldValuePtr>(ptr);
 }
 
+BOUT_ENUM_CLASS(GridVariableFunction, field3d, field2d, boutreal);
+
 namespace {
 /// Provides a placeholder whose target can be changed after creation.
 /// This enables recursive FieldGenerator expressions to be generated
@@ -80,6 +90,48 @@ class FieldIndirect : public FieldGenerator {
 
   FieldGeneratorPtr target;
 };
+
+// Read variables from the grid file and make them available in expressions
+template <class T>
+auto add_grid_variable(FieldFactory& factory, Mesh& mesh, const std::string& name) {
+  T var;
+  mesh.get(var, name);
+  factory.addGenerator(name, std::make_shared<GridVariable<T>>(var, name));
+}
+
+auto read_grid_variables(FieldFactory& factory, Mesh& mesh, Options& options) {
+  auto& field_variables = options["input"]["grid_variables"].doc(
+      "Variables to read from the grid file and make available in expressions");
+
+  for (const auto& [name, value] : field_variables) {
+    if (not mesh.isDataSourceGridFile()) {
+      throw BoutException(
+          "A grid file ('mesh:file') is required for `input:grid_variables`");
+    }
+
+    if (not mesh.sourceHasVar(name)) {
+      const auto filename = Options::root()["mesh"]["file"].as<std::string>();
+      throw BoutException(
+          "Grid file '{}' missing `{}` specified in `input:grid_variables`", filename,
+          name);
+    }
+
+    const auto func = value.as<GridVariableFunction>();
+    switch (func) {
+    case GridVariableFunction::field3d:
+      add_grid_variable<Field3D>(factory, mesh, name);
+      break;
+    case GridVariableFunction::field2d:
+      add_grid_variable<Field2D>(factory, mesh, name);
+      break;
+    case GridVariableFunction::boutreal:
+      BoutReal var{};
+      mesh.get(var, name);
+      factory.addGenerator(name, std::make_shared<FieldValue>(var));
+      break;
+    }
+  }
+}
 } // namespace
 
 //////////////////////////////////////////////////////////
@@ -176,6 +228,9 @@ FieldFactory::FieldFactory(Mesh* localmesh, Options* opt)
 
   // Where switch function
   addGenerator("where", std::make_shared<FieldWhere>(nullptr, nullptr, nullptr));
+
+  // Variables from the grid file
+  read_grid_variables(*this, *localmesh, nonconst_options);
 }
 
 Field2D FieldFactory::create2D(const std::string& value, const Options* opt,

src/field/fieldgenerators.hxx

@@ -9,9 +9,13 @@
 
 #include <bout/boutexception.hxx>
 #include <bout/field_factory.hxx>
+#include <bout/sys/expressionparser.hxx>
+#include <bout/traits.hxx>
 #include <bout/unused.hxx>
 
 #include <cmath>
+#include <memory>
+#include <string>
 
 //////////////////////////////////////////////////////////
 // Generators from values
@@ -352,4 +356,29 @@ private:
   FieldGeneratorPtr test, gt0, lt0;
 };
 
+/// A `Field3D` that can be used in expressions
+template <class T, typename = bout::utils::EnableIfField<T>>
+class GridVariable : public FieldGenerator {
+public:
+  GridVariable(T var, std::string name)
+      : variable(std::move(var)), name(std::move(name)) {}
+
+  double generate(const bout::generator::Context& ctx) override {
+    return variable(ctx.ix(), ctx.jy(), ctx.kz());
+  }
+
+  FieldGeneratorPtr clone(const std::list<FieldGeneratorPtr> args) override {
+    if (args.size() != 0) {
+      throw ParseException("Variable '{}' takes no arguments but got {:d}", args.size());
+    }
+    return std::make_shared<GridVariable<T>>(variable, name);
+  }
+
+  std::string str() const override { return name; }
+
+private:
+  T variable;
+  std::string name;
+};
+
 #endif // BOUT_FIELDGENERATORS_H

src/sys/generator_context.cxx

@@ -7,7 +7,7 @@ namespace bout {
 namespace generator {
 
 Context::Context(int ix, int iy, int iz, CELL_LOC loc, Mesh* msh, BoutReal t)
-    : localmesh(msh) {
+    : ix_(ix), jy_(iy), kz_(iz), localmesh(msh) {
 
   parameters["x"] = (loc == CELL_XLOW) ? 0.5 * (msh->GlobalX(ix) + msh->GlobalX(ix - 1))
                                        : msh->GlobalX(ix);
@@ -23,20 +23,17 @@ 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)
-    : localmesh(msh) {
-
-  // Add one to X index if boundary is in -x direction, so that XLOW is on the boundary
-  int ix = (bndry->bx < 0) ? bndry->x + 1 : bndry->x;
+    : // 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) {
 
   parameters["x"] = ((loc == CELL_XLOW) || (bndry->bx != 0))
-                        ? 0.5 * (msh->GlobalX(ix) + msh->GlobalX(ix - 1))
-                        : msh->GlobalX(ix);
-
-  int iy = (bndry->by < 0) ? bndry->y + 1 : bndry->y;
+                        ? 0.5 * (msh->GlobalX(ix_) + msh->GlobalX(ix_ - 1))
+                        : msh->GlobalX(ix_);
 
-  parameters["y"] = ((loc == CELL_YLOW) || bndry->by)
-                        ? PI * (msh->GlobalY(iy) + msh->GlobalY(iy - 1))
-                        : TWOPI * msh->GlobalY(iy);
+  parameters["y"] = ((loc == CELL_YLOW) || (bndry->by != 0))
+                        ? PI * (msh->GlobalY(jy_) + msh->GlobalY(jy_ - 1))
+                        : TWOPI * msh->GlobalY(jy_);
 
   parameters["z"] = (loc == CELL_ZLOW)
                         ? TWOPI * (iz - 0.5) / static_cast<BoutReal>(msh->LocalNz)
@@ -45,8 +42,5 @@ Context::Context(const BoundaryRegion* bndry, int iz, CELL_LOC loc, BoutReal t,
   parameters["t"] = t;
 }
 
-Context::Context(BoutReal x, BoutReal y, BoutReal z, Mesh* msh, BoutReal t)
-    : localmesh(msh), parameters{{"x", x}, {"y", y}, {"z", z}, {"t", t}} {}
-
 } // namespace generator
 } // namespace bout

tests/unit/fake_mesh.hxx

@@ -133,7 +133,7 @@ public:
   MPI_Comm getYcomm(int UNUSED(jx)) const override { return BoutComm::get(); }
   bool periodicY(int UNUSED(jx)) const override { return true; }
   bool periodicY(int UNUSED(jx), BoutReal& UNUSED(ts)) const override { return true; }
-  int numberOfYBoundaries() const override { return 1; }
+  int numberOfYBoundaries() const override { return 0; }
   std::pair<bool, BoutReal> hasBranchCutLower(int UNUSED(jx)) const override {
     return std::make_pair(false, 0.);
   }