WIP: Petsc matrix operators for FCI

CMakeLists.txt

@@ -177,6 +177,7 @@ set(BOUT_SOURCES
     ./include/bout/parallel_boundary_region.hxx
     ./include/bout/paralleltransform.hxx
     ./include/bout/petsc_interface.hxx
+    ./include/bout/petsc_operators.hxx
     ./include/bout/petsclib.hxx
     ./include/bout/physicsmodel.hxx
     ./include/bout/rajalib.hxx
@@ -308,6 +309,7 @@ set(BOUT_SOURCES
     ./src/mesh/parallel/shiftedmetricinterp.hxx
     ./src/mesh/parallel_boundary_op.cxx
     ./src/mesh/parallel_boundary_region.cxx
+    ./src/mesh/petsc_operators.cxx
     ./src/mesh/surfaceiter.cxx
     ./src/physics/gyro_average.cxx
     ./src/physics/physicsmodel.cxx

examples/CMakeLists.txt

@@ -18,6 +18,7 @@ add_subdirectory(dalf3)
 add_subdirectory(eigen-box)
 add_subdirectory(elm-pb)
 add_subdirectory(elm-pb-outerloop)
+add_subdirectory(fci-operators)
 add_subdirectory(fci-wave)
 add_subdirectory(finite-volume/diffusion)
 add_subdirectory(finite-volume/fluid)

examples/fci-operators/CMakeLists.txt

@@ -0,0 +1,14 @@
+cmake_minimum_required(VERSION 3.13)
+
+project(fci-operators LANGUAGES CXX)
+
+if(NOT TARGET bout++::bout++)
+  find_package(bout++ REQUIRED)
+endif()
+
+bout_add_example(
+  fci-operators
+  SOURCES fci_operators_example.cxx
+  DATA_DIRS data
+  EXTRA_FILES rotating-ellipse.py makeplots.py
+)

examples/fci-operators/data/BOUT.inp

@@ -0,0 +1,5 @@
+[mesh]
+file = "rotating-ellipse-20x8x20.fci.nc"
+
+[mesh:paralleltransform]
+type = fci

examples/fci-operators/fci_operators_example.cxx

@@ -0,0 +1,48 @@
+#include "bout/bout.hxx"
+#include "bout/difops.hxx"
+#include "bout/field.hxx"
+#include "bout/field3d.hxx"
+#include "bout/field_factory.hxx"
+#include "bout/globals.hxx"
+#include "bout/options.hxx"
+#include "bout/petsc_operators.hxx"
+
+int main(int argc, char** argv) {
+  BoutInitialise(argc, argv);
+
+  const PetscOperators ops(bout::globals::mesh);
+
+  // Parallel operators
+  const auto parallel = ops.getParallel();
+
+  // Test field
+  Field3D f = FieldFactory::get()->create3D("sin(y) + cos(z)");
+  bout::globals::mesh->communicate(f);
+  f.applyParallelBoundary("parallel_dirichlet_o2");
+
+  Field3D f_neumann = FieldFactory::get()->create3D("sin(y) + cos(z)");
+  bout::globals::mesh->communicate(f_neumann);
+  f_neumann.applyParallelBoundary("parallel_neumann_o1");
+
+  Options dump;
+  dump["f"] = f;
+  dump["dV"] = parallel.dV;
+
+  dump["grad_par_op"] = parallel.Grad_par(f);
+  dump["grad_par_yud"] = Grad_par(f);
+
+  dump["div_par_op"] = parallel.Div_par(f);
+
+  auto* coords = bout::globals::mesh->getCoordinates();
+  bout::globals::mesh->communicate(coords->Bxy);
+  coords->Bxy.applyParallelBoundary("parallel_neumann_o1");
+  dump["div_par_yud"] = Div_par(f);
+
+  dump["div_par_grad_par_op"] = parallel.Div_par_Grad_par(f_neumann);
+  dump["div_par_grad_par_yud"] = Div_par_K_Grad_par(1.0, f_neumann);
+
+  bout::writeDefaultOutputFile(dump);
+
+  BoutFinalise();
+  return 0;
+}

examples/fci-operators/makeplots.py

@@ -0,0 +1,81 @@
+from boutdata import collect
+import matplotlib.pyplot as plt
+import numpy as np
+
+path = "data"
+
+grad_par_op = collect("grad_par_op", path=path)
+grad_par_yud = collect("grad_par_yud", path=path)
+
+div_par_op = collect("div_par_op", path=path)
+div_par_yud = collect("div_par_yud", path=path)
+
+dV = collect("dV", path=path)  # Cell volume
+
+# Integrate divergence
+op_div_int = np.sum((dV * div_par_op)[2:-2, :, :])
+yud_div_int = np.sum((dV * div_par_yud)[2:-2, :, :])
+
+print(f"Div_par volume integral Op: {op_div_int} Yup/down: {yud_div_int}")
+
+div_par_grad_par_op = collect("div_par_grad_par_op", path=path)
+div_par_grad_par_yud = collect("div_par_grad_par_yud", path=path)
+
+# Integrate divergence
+op_div2_int = np.sum((dV * div_par_grad_par_op)[2:-2, :, :])
+yud_div2_int = np.sum((dV * div_par_grad_par_yud)[2:-2, :, :])
+
+print(f"Div_par_Grad_par volume integral Op: {op_div2_int} Yup/down: {yud_div2_int}")
+
+
+def plot_comparison(a, alabel, b, blabel):
+    fig, axs = plt.subplots(1, 3, figsize=(10, 3))
+    vmax = max([np.amax(a), np.amax(b)])
+    vmin = min([np.amin(a), np.amin(b)])
+
+    cs = axs[0].contourf(a, vmin=vmin, vmax=vmax)
+    fig.colorbar(cs, ax=axs[0])
+    axs[0].set_title(alabel)
+
+    axs[1].contourf(b, vmin=vmin, vmax=vmax)
+    axs[1].set_title(blabel)
+
+    cs = axs[2].contourf(a - b)
+    axs[2].set_title("Difference")
+    fig.colorbar(cs, ax=axs[2])
+
+    return fig
+
+
+fig = plot_comparison(
+    grad_par_op[2:-2, 2, :], "Operator", grad_par_yud[2:-2, 2, :], "Yup/down"
+)
+fig.suptitle("Parallel Gradient")
+fig.tight_layout()
+fig.savefig("gradient.pdf")
+fig.savefig("gradient.png")
+
+fig = plot_comparison(
+    div_par_op[2:-2, 2, :],
+    f"Operator [integral {op_div_int:.2e}]",
+    div_par_yud[2:-2, 2, :],
+    f"Yup/down [integral {yud_div_int:.2e}]",
+)
+fig.suptitle("Parallel Divergence")
+fig.tight_layout()
+fig.savefig("divergence.pdf")
+fig.savefig("divergence.png")
+
+fig = plot_comparison(
+    div_par_grad_par_op[2:-2, 2, :],
+    f"Operator [integral {op_div2_int:.2e}]",
+    div_par_grad_par_yud[2:-2, 2, :],
+    f"Yup/down [integral {yud_div2_int:.2e}]",
+)
+fig.suptitle("Parallel diffusion")
+fig.tight_layout()
+fig.savefig("diffusion.pdf")
+fig.savefig("diffusion.png")
+
+
+plt.show()

examples/fci-operators/rotating-ellipse.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+
+# Create grids for a rotating ellipse stellarator, based on curvilinear grids
+
+import numpy as np
+
+import zoidberg
+
+#############################################################################
+# Define the magnetic field
+
+# Length in y after which the coils return to their starting (R,Z) locations
+yperiod = 2.0 * np.pi / 5
+
+xcentre = 2.0
+
+magnetic_field = zoidberg.field.RotatingEllipse(
+    xcentre=xcentre,  # Major radius of axis [m]
+    radius=0.8,  # Minor radius of the coils [m]
+    I_coil=0.4,  # Coil current
+    yperiod=yperiod,
+    Btor=1.0,  # Toroidal magnetic field
+)
+
+#############################################################################
+# Create the inner flux surface, starting at a point at phi=0
+# To do this we need to define the y locations of the poloidal points
+# where we will construct grids
+
+start_r = xcentre + 0.4
+start_z = 0.0
+
+nslices = 8  # Number of poloidal slices
+ycoords = np.linspace(0, yperiod, nslices)
+npoints = 20  # Points per poloidal slice
+
+rzcoord, ycoords = zoidberg.fieldtracer.trace_poincare(
+    magnetic_field, start_r, start_z, yperiod, y_slices=ycoords, revs=npoints, nover=1
+)
+
+inner_lines = []
+for i in range(nslices):
+    r = rzcoord[:, i, 0, 0]
+    z = rzcoord[:, i, 0, 1]
+    line = zoidberg.rzline.line_from_points(r, z)
+    # Re-map the points so they're approximately uniform in distance along the surface
+    # Note that this results in some motion of the line
+    line = line.equallySpaced()
+    inner_lines.append(line)
+
+# Now have a list of y coordinates (ycoords) and inner lines (inner_lines)
+
+#############################################################################
+# Generate a fixed circle for the outer boundary
+
+outer_line = zoidberg.rzline.circle(R0=xcentre, r=0.6)
+
+#############################################################################
+# Now have inner and outer boundaries for each poloidal grid
+# Generate a grid on each poloidal slice using the elliptic grid generator
+
+nx = 20
+nz = 20
+
+pol_grids = [
+    zoidberg.poloidal_grid.grid_elliptic(inner_line, outer_line, nx, nz)
+    for inner_line in inner_lines
+]
+
+#############################################################################
+# Create a grid, then calculate forward and backward maps
+
+grid = zoidberg.grid.Grid(pol_grids, ycoords, yperiod, yperiodic=True)
+
+samples_per_dim = 5  # Sub-sampling in each cell
+
+maps = zoidberg.make_maps(grid, magnetic_field, samples_per_dim=samples_per_dim)
+
+#############################################################################
+# Interpolation weights
+
+weight_vars = zoidberg.weights.calculate_parallel_map_operators(maps)
+maps.update(weight_vars)
+# Remove sub-sampled maps from output
+for var in [
+    "forward_xt_primes",
+    "forward_zt_primes",
+    "backward_xt_primes",
+    "backward_zt_primes",
+    "subcell_weights",
+]:
+    maps.pop(var, None)
+
+#############################################################################
+# Write grid file
+
+filename = f"rotating-ellipse-{nx}x{nslices}x{nz}-s{samples_per_dim}.fci.nc"
+
+print(f"Writing to grid file '{filename}'")
+zoidberg.write_maps(
+    grid, magnetic_field, maps, gridfile=filename, new_names=False, metric2d=False
+)

include/bout/array.hxx

@@ -27,6 +27,7 @@
 #define BOUT_ARRAY_H
 
 #include <algorithm>
+#include <initializer_list>
 #include <map>
 #include <memory>
 #include <tuple>
@@ -187,9 +188,22 @@ public:
   Array() noexcept : ptr(nullptr) {}
 
   /*!
-   * Create an array of given length
+   * Create an array of given length.
    */
-  Array(size_type len) { ptr = get(len); }
+  Array(size_type len) : ptr(get(len)) {}
+
+  /*!
+   * Create an array with initializer list.
+   * This is explicit to avoid confusion with (size_type) constructor.
+   */
+  static Array fromValues(std::initializer_list<T> init) {
+    if (init.size() == 0) {
+      return Array();
+    }
+    Array array(init.size());
+    std::copy(init.begin(), init.end(), array.begin());
+    return array;
+  }
 
   /*!
    * Destructor. Releases the underlying dataBlock
@@ -199,7 +213,7 @@ public:
   /*!
    * Copy constructor
    */
-  Array(const Array& other) noexcept { ptr = other.ptr; }
+  Array(const Array& other) noexcept : ptr(other.ptr) {}
 
   /*!
    * Assignment operator

include/bout/field3d.hxx

@@ -1,7 +1,7 @@
 /**************************************************************************
- * Copyright 2010 B.D.Dudson, S.Farley, M.V.Umansky, X.Q.Xu
+ * Copyright 2010 - 2026 BOUT++ contributors
  *
- * Contact: Ben Dudson, bd512@york.ac.uk
+ * Contact: Ben Dudson, dudson2@llnl.gov
  *
  * This file is part of BOUT++.
  *
@@ -349,6 +349,14 @@ public:
     return std::end(getRegion("RGN_ALL"));
   };
 
+  /// Convert (x, y, z) to an Ind3D
+  /// Requires mesh size.
+  Ind3D indexAt(int x, int y, int z) const {
+    const int ny = this->getNy();
+    const int nz = this->getNz();
+    return Ind3D{(((x * ny) + y) * nz) + z, ny, nz};
+  }
+
   BoutReal& operator[](const Ind3D& d) { return data[d.ind]; }
   const BoutReal& operator[](const Ind3D& d) const { return data[d.ind]; }