examples/fci-operators

Shows how to use the PetscOperators, and compares against the
yup/ydown operators. Includes a script to generate the mesh,
which currently requires a branch of Zoidberg to run.

The code to calculate Grad_par, Div_par and Div_par_Grad_par from
forward and backward operators is now in
`PetscOperators::getParallel()`. That calculates and returns a
collection of `Parallel` operators.

This example could be turned into a test:
- `Grad_par` should be close between the two methods
- Volume integrals of `Div_par` and `Div_par_Grad_par` results should be
  near zero (flux conserved).

Author: bendudson

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-20x8x20.fci.nc 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,91 @@
+#!/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)
+
+maps = zoidberg.make_maps(grid, magnetic_field)
+
+#############################################################################
+# Interpolation weights
+
+weight_vars = zoidberg.weights.calculate_parallel_map_operators(maps)
+maps.update(weight_vars)
+
+#############################################################################
+# Write grid file
+
+filename = f"rotating-ellipse-{nx}x{nslices}x{nz}.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/petsc_operators.hxx

@@ -419,6 +419,21 @@ public:
     return PetscOperator::diagonal(mapping, f);
   }
 
+  /// A standard collection of operators
+  /// parallel to the magnetic field
+  struct Parallel {
+    PetscOperator Grad_par;         ///< Gradient
+    PetscOperator Div_par;          ///< Divergence
+    PetscOperator Div_par_Grad_par; ///< Diffusion
+    Field3D dV;                     ///< Cell volume
+  };
+
+  /// Calculate a set of parallel operators
+  ///
+  /// This assumes that "forward" and "backward" operators
+  /// are stored in the mesh.
+  Parallel getParallel() const;
+
 private:
   /// Read a `Field3D` from the mesh.
   ///

src/mesh/petsc_operators.cxx

@@ -299,4 +299,110 @@ PetscOperator PetscOperators::get(const std::string& name) const {
                        this->meshGetArray<BoutReal>(this->mesh, name + "_weights"));
 }
 
+PetscOperators::Parallel PetscOperators::getParallel() const {
+  // Read maps from the mesh
+  auto forward = this->get("forward");
+  auto backward = this->get("backward");
+
+  // ---- Construct Grad_par ----
+  //
+  // Create a parallel gradient operator by combining the parallel
+  // length dl = dy * sqrt(g_22) with forward & backward operators
+  auto* coords = this->mesh->getCoordinates();
+  Field3D dl = coords->dy * sqrt(coords->g_22);
+  dl.splitParallelSlices();
+  dl.yup() = 0.0;
+  dl.ydown() = 0.0;
+  dl.applyParallelBoundary("parallel_neumann_o1");
+
+  auto inv_2dl = 0.5 / dl;
+  inv_2dl.splitParallelSlices();
+  inv_2dl.yup() = 0.0;
+  inv_2dl.ydown() = 0.0;
+  inv_2dl.applyParallelBoundary("parallel_neumann_o1");
+
+  inv_2dl.yup() *= 0.5;
+  inv_2dl.ydown() *= 0.5;
+
+  auto inv_2dl_op = this->diagonal(inv_2dl);
+
+  auto Grad_par = inv_2dl_op * (forward - backward);
+
+  // ---- Construct Div_par ----
+  //
+  // Use the Support Operator Method (SOM) to calculate
+  // Div_par from Grad_par and cell volumes.
+
+  // Cell volume
+  auto dV = coords->J * coords->dx * coords->dy * coords->dz;
+  dV.splitParallelSlices();
+  dV.yup() = 0.0;
+  dV.ydown() = 0.0;
+  dV.applyParallelBoundary("parallel_neumann_o1");
+  auto dV_op = this->diagonal(dV);
+
+  Field3D neg_inv_dV = -1. / dV;
+  neg_inv_dV.splitParallelSlices();
+  neg_inv_dV.yup() = 0.0;
+  neg_inv_dV.ydown() = 0.0;
+  neg_inv_dV.applyParallelBoundary("parallel_neumann_o1");
+  auto neg_inv_dV_op = this->diagonal(neg_inv_dV);
+
+  auto Div_par = neg_inv_dV_op * Grad_par.transpose() * dV_op;
+
+  // ---- Construct Div_par_Grad_par ----
+  //
+  // Requires gradients between planes, at +1/2 and -1/2, and interpolation
+  // operators to calculate quantities between cells.
+
+  // Identity operator
+  Field3D one{1.0};
+  one.splitParallelSlices();
+  one.yup() = 1.0;
+  one.ydown() = 1.0;
+  const auto identity = this->diagonal(one);
+
+  // Interpolate at + 1/2
+  const auto interp_plus_op = (identity + forward) * 0.5;
+
+  // dl averaged at +1/2
+  const Field3D dl_plus = interp_plus_op(dl);
+  Field3D inv_dl_plus = 1. / dl_plus;
+  inv_dl_plus.splitParallelSlices();
+  inv_dl_plus.yup() = 0.0;
+  inv_dl_plus.ydown() = 0.0;
+  inv_dl_plus.applyParallelBoundary("parallel_neumann_o1");
+  const auto inv_dl_plus_op = this->diagonal(inv_dl_plus);
+
+  // Gradient at + 1/2
+  const auto Grad_plus = inv_dl_plus_op * (forward - identity);
+
+  // Divergence at -1/2
+  const auto Div_minus = neg_inv_dV_op * Grad_plus.transpose() * dV_op;
+
+  // Interpolate at - 1/2
+  auto interp_minus_op = (identity + backward) * 0.5;
+
+  // dl averaged at -1/2
+  const Field3D dl_minus = interp_minus_op(dl);
+  Field3D inv_dl_minus = 1. / dl_minus;
+  inv_dl_minus.splitParallelSlices();
+  inv_dl_minus.yup() = 0.0;
+  inv_dl_minus.ydown() = 0.0;
+  inv_dl_minus.applyParallelBoundary("parallel_neumann_o1");
+  const auto inv_dl_minus_op = this->diagonal(inv_dl_minus);
+
+  // Gradient at - 1/2
+  auto Grad_minus = inv_dl_minus_op * (identity - backward);
+
+  // Divergence at +1/2
+  const auto Div_plus = neg_inv_dV_op * Grad_minus.transpose() * dV_op;
+
+  // Div(Grad_par()) operator
+  auto Div_par_Grad_par = ((Div_minus * Grad_plus) + (Div_plus * Grad_minus)) * 0.5;
+
+  return Parallel{std::move(Grad_par), std::move(Div_par), std::move(Div_par_Grad_par),
+                  std::move(dV)};
+}
+
 #endif // BOUT_HAS_PETSC