New shifted metric 2

include/bout/mesh.hxx

@@ -322,7 +322,15 @@ class Mesh {
   /// \param[in] jx   The local (on this processor) index in X
   /// \param[out] ts  The Twist-Shift angle if periodic
   virtual bool periodicY(int jx, BoutReal &ts) const = 0;
+
+  /// Is there a branch cut at the lower y edge of this processor's grid points
+  /// at x-index jx?
+  virtual bool hasBranchCutDown(int jx) const = 0;
 
+  /// Is there a branch cut at the upper y edge of this processor's grid points
+  /// at x-index jx?
+  virtual bool hasBranchCutUp(int jx) const = 0;
+
   virtual int ySize(int jx) const; ///< The number of points in Y at fixed X index \p jx
 
   // Y communications

include/bout/paralleltransform.hxx

@@ -173,17 +173,17 @@ private:
    * @param[in] zangle  The angle (z coordinate) to shift by
    * @param[out] out  A 1D array of length \p len, already allocated
    */
-  void shiftZ(const BoutReal *in, int len, BoutReal zangle,  BoutReal *out);
+  void shiftZ(const BoutReal* in, int len, BoutReal zangle, BoutReal* out);
 
   /*!
    * Shift a given 1D array, assumed to be in Z, by the given \p zangle
    *
    * @param[in] in  A 1D array of length mesh.LocalNz
-   * @param[in] phs Phase shift, assumed to have length (mesh.LocalNz/2 + 1) i.e. the number of modes
+   * @param[in] phs Phase shift, assumed to have length (mesh.LocalNz/2 + 1)
+   * i.e. the number of modes
    * @param[out] out  A 1D array of length mesh.LocalNz, already allocated
    */
   void shiftZ(const BoutReal* in, const dcomplex* phs, BoutReal* out);
 };
 
-
 #endif // __PARALLELTRANSFORM_H__

src/mesh/impls/bout/boutmesh.hxx

@@ -105,6 +105,20 @@ class BoutMesh : public Mesh {
   /// \param[in] jx   The local (on this processor) index in X
   bool periodicY(int jx) const;
 
+  /// Is there a branch cut at the lower y edge of this processor's grid points
+  /// at x-index jx?
+  bool hasBranchCutDown(int jx) const {
+    return (TS_down_in && (DDATA_INDEST != -1) && (jx < DDATA_XSPLIT))
+      || (TS_down_out && (DDATA_OUTDEST != -1) && (jx >= DDATA_XSPLIT));
+  }
+
+  /// Is there a branch cut at the upper y edge of this processor's grid points
+  /// at x-index jx?
+  bool hasBranchCutUp(int jx) const {
+    return (TS_up_in && (UDATA_INDEST != -1) && (jx < UDATA_XSPLIT))
+      || (TS_up_out && (UDATA_OUTDEST != -1) && (jx >= UDATA_XSPLIT));
+  }
+
   int ySize(int jx) const; ///< The number of points in Y at fixed X index \p jx
 
   /////////////////////////////////////////////

src/mesh/mesh.cxx

@@ -13,6 +13,7 @@
 #include <output.hxx>
 
 #include "parallel/fci.hxx"
+#include "parallel/shifttofieldaligned.hxx"
 
 Mesh* Mesh::create(GridDataSource *s, Options *opt) {
   return MeshFactory::getInstance()->createMesh(s, opt);
@@ -299,8 +300,12 @@ void Mesh::setParallelTransform() {
 
   }else if(ptstr == "shifted") {
     // Shifted metric method
-  transform = bout::utils::make_unique<ShiftedMetric>(*this);
+    transform = bout::utils::make_unique<ShiftedMetric>(*this);
 
+  }else if(ptstr == "shifttofieldaligned") {
+    // Shifted metric method, shifting to globally field-aligned coordinates
+    transform = bout::utils::make_unique<ShiftToFieldAligned>(*this);
+
   }else if(ptstr == "fci") {
 
     Options *fci_options = Options::getRoot()->getSection("fci");

src/mesh/parallel/makefile

@@ -2,7 +2,7 @@
 BOUT_TOP = ../../..
 
 DIRS            = 
-SOURCEC		= shiftedmetric.cxx fci.cxx
+SOURCEC		= shiftedmetric.cxx fci.cxx shifttofieldaligned.cxx
 TARGET		= lib
 
 include $(BOUT_TOP)/make.config

src/mesh/parallel/shifttofieldaligned.cxx

@@ -0,0 +1,172 @@
+/*
+ * Implements the shifted metric method for parallel derivatives by
+ * transforming to globally field-aligned coordinates (in contrast to
+ * ShiftedMetric where the coordinates are only locally field-aligned).
+ *
+ * By default fields are stored so that X-Z are orthogonal,
+ * and so not aligned in Y.
+ *
+ */
+
+#include "shifttofieldaligned.hxx"
+
+#include <bout/constants.hxx>
+#include <bout/mesh.hxx>
+#include <fft.hxx>
+#include <interpolation.hxx>
+
+namespace {
+void fixzShiftBoundaries(Field2D& zShift) {
+  Mesh* localmesh = zShift.getMesh();
+
+  localmesh->communicate(zShift);
+
+  for (int x = 0; x < localmesh->LocalNx; x++) {
+    if (localmesh->hasBranchCutDown(x)) {
+      // ShiftAngle is the total toroidal shift of a field line after one poloidal turn
+      BoutReal ShiftAngle = 0.;
+      localmesh->periodicY(x, ShiftAngle);
+      for (int y = 0; y < localmesh->ystart; y++) {
+        zShift(x, y) -= ShiftAngle;
+      }
+    }
+    if (localmesh->hasBranchCutUp(x)) {
+      // ShiftAngle is the total toroidal shift of a field line after one poloidal turn
+      BoutReal ShiftAngle = 0.;
+      localmesh->periodicY(x, ShiftAngle);
+      for (int y = localmesh->yend + 1; y < localmesh->LocalNy; y++) {
+        zShift(x, y) += ShiftAngle;
+      }
+    }
+  }
+
+  zShift.applyBoundary("free_o3");
+}
+} // namespace
+
+ShiftToFieldAligned::ShiftToFieldAligned(Mesh& mesh_in) {
+  // Must *not* twist-shift when communicating
+  bool twistshift = Options::root()["TwistShift"].withDefault(false);
+  if (twistshift) {
+    throw BoutException("ShiftToFieldAligned requires TwistShift=false");
+  }
+
+  Field2D zShift(&mesh_in);
+  // Read the zShift angle from the mesh
+  if (mesh_in.get(zShift, "zShift")) {
+    // No zShift variable. Try qinty in BOUT grid files
+    mesh_in.get(zShift, "qinty");
+  }
+
+  init(mesh_in, zShift);
+}
+
+void ShiftToFieldAligned::init(Mesh& mesh_in, const Field2D& zShift) {
+  // always include CELL_CENTRE implementation
+  implementations.emplace(CELL_CENTRE, Implementation(mesh_in, CELL_CENTRE, zShift));
+
+  if (mesh_in.StaggerGrids) {
+    // Don't know which locations we will need, so populate all those that can
+    // be interpolated (for now at least)
+    implementations.emplace(CELL_XLOW, Implementation(mesh_in, CELL_XLOW, zShift));
+    implementations.emplace(CELL_XLOW, Implementation(mesh_in, CELL_YLOW, zShift));
+    implementations.emplace(CELL_XLOW, Implementation(mesh_in, CELL_ZLOW, zShift));
+  }
+}
+
+ShiftToFieldAligned::Implementation::Implementation(Mesh& mesh_in,
+                                                    const CELL_LOC location_in,
+                                                    const Field2D& zShift_in)
+    : mesh(mesh_in), location(location_in), zShift(&mesh_in) {
+
+  zShift = interp_to(zShift_in, location);
+  fixzShiftBoundaries(zShift);
+  cachePhases();
+}
+
+void ShiftToFieldAligned::Implementation::cachePhases() {
+  // If we wanted to be efficient we could move the following cached phase setup
+  // into the relevant shifting routines (with a bool to record if the phases
+  // have already been cached) so that we only calculate the phase if we
+  // actually call a relevant shift routine -- however that would make calls to
+  // toFieldAligned and fromFieldAligned non-const, so it's better to keep this
+  // in the constructor
+
+  // As we're attached to a mesh we can expect the z direction to
+  // not change once we've been created so precalculate the complex
+  // phases used in transformations
+  nmodes = mesh.LocalNz / 2 + 1;
+  BoutReal zlength = mesh.getCoordinates()->zlength();
+
+  // Allocate storage for our 3d vector structures.
+  fromAlignedPhs = Tensor<dcomplex>(mesh.LocalNx, mesh.LocalNy, mesh.LocalNz);
+  toAlignedPhs = Tensor<dcomplex>(mesh.LocalNx, mesh.LocalNy, mesh.LocalNz);
+
+  // To/From field aligned phases
+  BOUT_FOR(i, mesh.getRegion2D("RGN_ALL")) {
+    for (int jz = 0; jz < nmodes; jz++) {
+      BoutReal kwave = jz * 2.0 * PI / zlength; // wave number is 1/[rad]
+      fromAlignedPhs(i.x(), i.y(), jz) =
+        dcomplex(cos(kwave * zShift[i]), -sin(kwave * zShift[i]));
+      toAlignedPhs(i.x(), i.y(), jz) =
+        dcomplex(cos(kwave * zShift[i]), sin(kwave * zShift[i]));
+    }
+  }
+}
+
+/*!
+ * Get the shifted field so that X-Z is not orthogonal,
+ * and Y is then field aligned.
+ */
+const Field3D ShiftToFieldAligned::Implementation::toFieldAligned(const Field3D& f,
+                                                                  const REGION region) {
+  return shiftZ(f, toAlignedPhs, region);
+}
+
+/*!
+ * Shift back, so that X-Z is orthogonal,
+ * but Y is not field aligned.
+ */
+const Field3D ShiftToFieldAligned::Implementation::fromFieldAligned(const Field3D& f,
+                                                                    const REGION region) {
+  return shiftZ(f, fromAlignedPhs, region);
+}
+
+const Field3D ShiftToFieldAligned::Implementation::shiftZ(
+    const Field3D& f, const Tensor<dcomplex>& phs, const REGION region) {
+  ASSERT1(&mesh == f.getMesh());
+  ASSERT1(f.getLocation() == location);
+
+  if (mesh.LocalNz == 1)
+    return f; // Shifting makes no difference
+
+  Field3D result(&mesh);
+  result.allocate();
+  result.setLocation(location);
+
+  BOUT_FOR(i, mesh.getRegion2D(REGION_STRING(region))) {
+    shiftZ(&f(i, 0), &phs(i.x(), i.y(), 0), &result(i, 0));
+  }
+
+  return result;
+}
+
+void ShiftToFieldAligned::Implementation::shiftZ(const BoutReal* in,
+                                                 const dcomplex* phs,
+                                                 BoutReal* out) {
+  Array<dcomplex> cmplx(nmodes);
+
+  // Take forward FFT
+  rfft(in, mesh.LocalNz, &cmplx[0]);
+
+  // Following is an algorithm approach to write a = a*b where a and b are
+  // vectors of dcomplex.
+  //  std::transform(cmplxOneOff.begin(),cmplxOneOff.end(), ptr.begin(),
+  //                 cmplxOneOff.begin(), std::multiplies<dcomplex>());
+
+  for (int jz = 1; jz < nmodes; jz++) {
+    cmplx[jz] *= phs[jz];
+  }
+
+  irfft(&cmplx[0], mesh.LocalNz, out); // Reverse FFT
+}

src/mesh/parallel/shifttofieldaligned.hxx

@@ -0,0 +1,133 @@
+/**************************************************************************
+ * Parallel derivatives use shifted-metric scheme implemented by shifting
+ * to globally field-aligned coordinates
+ *
+ **************************************************************************
+ * Copyright 2018 B.D.Dudson, P. Hill, J. Omotani
+ *
+ * Contact: Ben Dudson, bd512@york.ac.uk
+ *
+ * This file is part of BOUT++.
+ *
+ * BOUT++ is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * BOUT++ is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with BOUT++.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ **************************************************************************/
+
+#ifndef __SHIFTTOFIELDALIGNED_H__
+#define __SHIFTTOFIELDALIGNED_H__
+
+#include <bout/paralleltransform.hxx>
+
+class ShiftToFieldAligned : public ParallelTransform {
+public:
+  ShiftToFieldAligned(Mesh& mesh_in);
+  /// Use an existing zShift
+  ShiftToFieldAligned(Mesh& mesh_in, const Field2D& zShift) { init(mesh_in, zShift); }
+
+  /*!
+   * Do nothing here since we don't cache field-aligned fields
+   */
+  void calcYUpDown(Field3D& UNUSED(f)) override {}
+
+  /*!
+   * Uses FFTs and a phase shift to align the grid points
+   * with the y coordinate (along magnetic field usually).
+   *
+   * Note that the returned field will no longer be orthogonal
+   * in X-Z, and the metric tensor will need to be changed
+   * if X derivatives are used.
+   */
+  const Field3D toFieldAligned(const Field3D& f, const REGION region = RGN_NOX) override {
+    return implementations.at(f.getLocation()).toFieldAligned(f, region);
+  }
+
+  /*!
+   * Converts a field back to X-Z orthogonal coordinates
+   * from field aligned coordinates.
+   */
+  const Field3D fromFieldAligned(const Field3D& f,
+                                 const REGION region = RGN_NOX) override {
+    return implementations.at(f.getLocation()).fromFieldAligned(f, region);
+  }
+
+  bool canToFromFieldAligned() override { return true; }
+
+private:
+  class Implementation {
+  public:
+    Implementation(Mesh& mesh_in, const CELL_LOC location_in, const Field2D& zShift_in);
+
+    /*!
+     * Uses FFTs and a phase shift to align the grid points
+     * with the y coordinate (along magnetic field usually).
+     *
+     * Note that the returned field will no longer be orthogonal
+     * in X-Z, and the metric tensor will need to be changed
+     * if X derivatives are used.
+     */
+    const Field3D toFieldAligned(const Field3D& f, const REGION region = RGN_NOX);
+
+    /*!
+     * Converts a field back to X-Z orthogonal coordinates
+     * from field aligned coordinates.
+     */
+    const Field3D fromFieldAligned(const Field3D& f, const REGION region = RGN_NOX);
+
+  private:
+    Mesh& mesh;
+    const CELL_LOC location;
+    Field2D zShift;
+
+    /// Calculate and store the phases for to/from field aligned using zShift
+    void cachePhases();
+
+    int nmodes;
+
+    Tensor<dcomplex> toAlignedPhs; ///< Cache of phase shifts for transforming from X-Z
+                                   ///  orthogonal coordinates to field-aligned
+                                   ///  coordinates.
+    Tensor<dcomplex> fromAlignedPhs; ///< Cache of phase shifts for transforming from
+                                     ///  field-aligned coordinates to X-Z orthogonal
+                                     ///  coordinates.
+
+    /*!
+     * Shift a 3D field \p f by the given phase \p phs in Z
+     *
+     * Calculates FFT in Z, multiplies by the complex phase
+     * and inverse FFTS.
+     *
+     * @param[in] f  The field to shift
+     * @param[in] phs  The phase to shift by
+     */
+    const Field3D shiftZ(const Field3D& f, const Tensor<dcomplex>& phs,
+                         const REGION region = RGN_NOX);
+
+    /*!
+     * Shift a given 1D array, assumed to be in Z, by the given \p zangle
+     *
+     * @param[in] in  A 1D array of length mesh.LocalNz
+     * @param[in] phs Phase shift, assumed to have length (mesh.LocalNz/2 + 1)
+     * i.e. the number of modes
+     * @param[out] out  A 1D array of length mesh.LocalNz, already allocated
+     */
+    void shiftZ(const BoutReal* in, const dcomplex* phs, BoutReal* out);
+  };
+
+  std::map<CELL_LOC, Implementation> implementations;
+
+  // populate the implementations map
+  void init(Mesh& mesh_in, const Field2D& zShift);
+};
+
+#endif // __SHIFTTOFIELDALIGNED_H__