Merge pull request E3SM-Project#62 from mwarusz/mw/omega/operators

Add mesh operators

components/omega/doc/devGuide/HorzOperators.md

@@ -0,0 +1,43 @@
+(omega-dev-horz-operators)=
+
+# Horizontal Operators
+
+The TRiSK scheme horizontal operators are implemented in C++ as functors, which
+are classes that can be called with input arguments in the same way as functions
+can. However, in contrast to plain functions, functors can contain internal
+state. In the case of Omega operators they contain shallow copies of various
+`HorzMesh` arrays that are needed to compute the operator, but are not strictly
+inputs of the operator.
+
+Operators are constructed from an instance of the `HorzMesh` class, for example
+```c++
+    auto mesh = OMEGA::HorzMesh::getDefault();
+    DivergenceOnCell DivOnCell(mesh);
+```
+which sets up the previously mentioned internal state.
+
+Each Omega operator provides a C++ call operator, which computes its value on a
+mesh element given the element index and operator-specific input arrays.
+Typically, operators are created outside of a parallel region and are used
+inside a parallel loop over mesh elements, for example
+```c++
+    auto mesh = OMEGA::HorzMesh::getDefault();
+    DivergenceOnCell DivOnCell(mesh);
+    parallelFor({mesh->NCellsOwned}, KOKKOS_LAMBDA(Int ICell) {
+        Real Div = DivOnCell(ICell, Vec); // computes divergence of Vec over cell with index ICell
+    });
+```
+
+Currently, the following operators are implemented:
+- `DivergenceOnCell`
+- `GradientOnEdge`
+- `CurlOnVertex`
+- `TangentialReconOnEdge`
+
+Some tendency terms in the Omega PDE solver could in principle be constructed
+using these operators as building blocks. However, very often tendency terms
+require evaluation of slightly modified operators. Moreover, there is a
+potential performance cost of nesting operators within classes. Hence, the
+primary motivation for introducing these classes is to provide reference
+implementation of the basic TRiSK operators and for diagnostic and debugging
+purposes.

components/omega/doc/devGuide/QuickStart.md

@@ -130,18 +130,29 @@ script that you can run to build Omega:
 ./omega_build.sh
 ```
 
+(omega-dev-quick-start-getting-meshes)=
+### Getting test meshes
+
+Some tests require a valid Omega mesh file. Different tests require different meshes. At the moment, mesh files need
+to be copied or linked to specifically named files under the `test` directory. Appropriate mesh files can be downloaded from:
+- [Ocean Mesh](https://web.lcrc.anl.gov/public/e3sm/inputdata/ocn/mpas-o/oQU240/ocean.QU.240km.151209.nc)
+- [Global Mesh](https://web.lcrc.anl.gov/public/e3sm/polaris/ocean/polaris_cache/global_convergence/icos/cosine_bell/Icos480/init/initial_state.230220.nc)
+- [Planar Mesh](https://gist.github.com/mwarusz/f8caf260398dbe140d2102ec46a41268/raw/e3c29afbadc835797604369114321d93fd69886d/PlanarPeriodic48x48.nc)
+```sh
+wget -O ocean_test_mesh.nc https://web.lcrc.anl.gov/public/e3sm/inputdata/ocn/mpas-o/oQU240/ocean.QU.240km.151209.nc
+wget -O global_test_mesh.nc https://web.lcrc.anl.gov/public/e3sm/polaris/ocean/polaris_cache/global_convergence/icos/cosine_bell/Icos480/init/initial_state.230220.nc
+wget -O planar_test_mesh.nc https://gist.github.com/mwarusz/f8caf260398dbe140d2102ec46a41268/raw/e3c29afbadc835797604369114321d93fd69886d/PlanarPeriodic48x48.nc
+ln -sf  ocean_test_mesh.nc test/OmegaMesh.nc
+ln -sf  global_test_mesh.nc test/OmegaSphereMesh.nc
+ln -sf  planar_test_mesh.nc test/OmegaPlanarMesh.nc
+```
+
 ### Running CTests
 
 Omega includes several unit tests that run through CTest.  These need to be
-run on a compute node.  Some tests also require a valid Omega mesh file
-called `test/OmegaMesh.nc`.  An appropriate mesh file can be downloaded from
-[mesh.230220.nc](https://web.lcrc.anl.gov/public/e3sm/polaris/ocean/omega_ctest/ocean.QU.240km.151209.nc).
-```sh
-wget https://web.lcrc.anl.gov/public/e3sm/polaris/ocean/omega_ctest/ocean.QU.240km.151209.nc
-mv ocean.QU.240km.151209.nc test/OmegaMesh.nc
-```
+run on a compute node.
 
-Then, run the tests:
+To run the tests:
 ```sh
 ./omega_ctest.sh
 ```

components/omega/doc/index.md

@@ -33,6 +33,7 @@ userGuide/IO
 userGuide/IOField
 userGuide/Halo
 userGuide/HorzMesh
+userGuide/HorzOperators
 userGuide/TimeMgr
 ```
 
@@ -57,6 +58,7 @@ devGuide/IO
 devGuide/IOField
 devGuide/Halo
 devGuide/HorzMesh
+devGuide/HorzOperators
 devGuide/TimeMgr
 ```
 

components/omega/doc/userGuide/HorzOperators.md

@@ -0,0 +1,19 @@
+(omega-user-horz-operators)=
+
+# Horizontal Operators
+
+The horizontal discretization in Omega is a staggered numerical scheme known as
+TRiSK. It defines discrete versions of basic differential operators (divergence,
+gradient, and curl) as well as other operators that are needed by the scheme,
+such as reconstruction of tangential velocity from its normal components. Omega
+provides reference implementations of these operators, each in a separate C++
+class. The class name describes the operator and the mesh element associated
+with its result.
+
+The following operators are currently implemented:
+- `DivergenceOnCell`
+- `GradientOnEdge`
+- `CurlOnVertex`
+- `TangentialReconOnEdge`
+
+There are no user-configurable options.

components/omega/doc/userGuide/OmegaBuild.md

@@ -132,7 +132,8 @@ parallel job launcher such as SLURM srun should be available. You may first
 get allocation of an interactive computing node or use batch system.
 In addition, you must either copy (or soft link) a mesh file into the same
 directory as the unit test executables or specify in a configuration file
-(not implemented yet) the path to that mesh input file.
+(not implemented yet) the path to that mesh input file. For detailed instructions
+on how to obtain appropiate test meshes see {ref}`omega-dev-quick-start-getting-meshes`.
 
 For example, if you are on an interactive computing node, you can run
 Omega unit test by running `omega_ctest.sh` as shown below.

components/omega/src/base/Decomp.cpp

@@ -321,14 +321,13 @@ int readMesh(const int MeshFileID, // file ID for open mesh file
 // Initialize the decomposition and create the default decomposition with
 // (currently) one partition per MPI task using a ParMetis KWay method.
 
-int Decomp::init() {
+int Decomp::init(const std::string &MeshFileName) {
 
    int Err = 0; // default successful return code
 
    // TODO: retrieve from Config when available - currently hardwired
    // Initialize decomposition options
    I4 InHaloWidth           = 3;
-   std::string MeshFileName = "OmegaMesh.nc";
    std::string DecompMethod = "MetisKWay";
    PartMethod Method        = getPartMethodFromStr(DecompMethod);
 

components/omega/src/base/Decomp.h

@@ -219,7 +219,7 @@ class Decomp {
    /// Initializes Omega decomposition info and creates the default
    /// decomposition based on the default MachEnv and configuration
    /// options.
-   static int init();
+   static int init(const std::string &MeshFileName = "OmegaMesh.nc");
 
    /// Construct a new decomposition across an input MachEnv with
    /// NPart partitions of a mesh that is read from a mesh file.

components/omega/src/infra/OmegaKokkos.h

@@ -10,6 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "DataTypes.h"
+#include <utility>
 
 namespace OMEGA {
 
@@ -86,24 +87,26 @@ inline void parallelFor(const int (&upper_bounds)[N], const F &f,
 // parallelReduce: with label
 template <int N, class F, class R, class... Args>
 inline void parallelReduce(const std::string &label,
-                           const int (&upper_bounds)[N], const F &f, R &reducer,
+                           const int (&upper_bounds)[N], const F &f,
+                           R &&reducer,
                            const int (&tile)[N] = DefaultTile<N>::value) {
    if constexpr (N == 1) {
       const auto policy = Kokkos::RangePolicy<Args...>(0, upper_bounds[0]);
-      Kokkos::parallel_reduce(label, policy, f, reducer);
+      Kokkos::parallel_reduce(label, policy, f, std::forward<R>(reducer));
 
    } else {
       const int lower_bounds[N] = {0};
       const auto policy = Bounds<N, Args...>(lower_bounds, upper_bounds, tile);
-      Kokkos::parallel_reduce(label, policy, f, reducer);
+      Kokkos::parallel_reduce(label, policy, f, std::forward<R>(reducer));
    }
 }
 
 // parallelReduce: without label
 template <int N, class F, class R, class... Args>
-inline void parallelReduce(const int (&upper_bounds)[N], const F &f, R &reducer,
+inline void parallelReduce(const int (&upper_bounds)[N], const F &f,
+                           R &&reducer,
                            const int (&tile)[N] = DefaultTile<N>::value) {
-   parallelReduce("", upper_bounds, f, tile, reducer);
+   parallelReduce("", upper_bounds, f, std::forward<R>(reducer), tile);
 }
 
 } // end namespace OMEGA

components/omega/src/ocn/HorzMesh.h

@@ -13,6 +13,7 @@
 #include "DataTypes.h"
 #include "Decomp.h"
 #include "MachEnv.h"
+#include "OmegaKokkos.h"
 
 #include <string>
 

components/omega/src/ocn/HorzOperators.cpp

@@ -0,0 +1,24 @@
+#include "HorzOperators.h"
+#include "DataTypes.h"
+#include "HorzMesh.h"
+
+namespace OMEGA {
+
+DivergenceOnCell::DivergenceOnCell(HorzMesh const *Mesh)
+    : NEdgesOnCell(Mesh->NEdgesOnCell), EdgesOnCell(Mesh->EdgesOnCell),
+      DvEdge(Mesh->DvEdge), AreaCell(Mesh->AreaCell),
+      EdgeSignOnCell(Mesh->EdgeSignOnCell) {}
+
+GradientOnEdge::GradientOnEdge(HorzMesh const *Mesh)
+    : CellsOnEdge(Mesh->CellsOnEdge), DcEdge(Mesh->DcEdge) {}
+
+CurlOnVertex::CurlOnVertex(HorzMesh const *Mesh)
+    : VertexDegree(Mesh->VertexDegree), EdgesOnVertex(Mesh->EdgesOnVertex),
+      DcEdge(Mesh->DcEdge), AreaTriangle(Mesh->AreaTriangle),
+      EdgeSignOnVertex(Mesh->EdgeSignOnVertex) {}
+
+TangentialReconOnEdge::TangentialReconOnEdge(HorzMesh const *Mesh)
+    : NEdgesOnEdge(Mesh->NEdgesOnEdge), EdgesOnEdge(Mesh->EdgesOnEdge),
+      WeightsOnEdge(Mesh->WeightsOnEdge) {}
+
+} // namespace OMEGA

components/omega/src/ocn/HorzOperators.h

@@ -0,0 +1,98 @@
+#ifndef OMEGA_HORZOPERATORS_H
+#define OMEGA_HORZOPERATORS_H
+
+#include "DataTypes.h"
+#include "HorzMesh.h"
+
+namespace OMEGA {
+
+class DivergenceOnCell {
+ public:
+   DivergenceOnCell(HorzMesh const *Mesh);
+
+   KOKKOS_FUNCTION Real operator()(int ICell,
+                                   const Array1DReal &VecEdge) const {
+      Real DivCell = 0;
+      for (int J = 0; J < NEdgesOnCell(ICell); ++J) {
+         const int JEdge = EdgesOnCell(ICell, J);
+         DivCell -= DvEdge(JEdge) * EdgeSignOnCell(ICell, J) * VecEdge(JEdge);
+      }
+      const Real InvAreaCell = 1. / AreaCell(ICell);
+      DivCell *= InvAreaCell;
+      return DivCell;
+   }
+
+ private:
+   Array1DI4 NEdgesOnCell;
+   Array2DI4 EdgesOnCell;
+   Array1DR8 DvEdge;
+   Array1DR8 AreaCell;
+   Array2DR8 EdgeSignOnCell;
+};
+
+class GradientOnEdge {
+ public:
+   GradientOnEdge(HorzMesh const *Mesh);
+
+   KOKKOS_FUNCTION Real operator()(int IEdge,
+                                   const Array1DReal &ScalarCell) const {
+      const auto JCell0    = CellsOnEdge(IEdge, 0);
+      const auto JCell1    = CellsOnEdge(IEdge, 1);
+      const Real InvDcEdge = 1. / DcEdge(IEdge);
+      const Real GradEdge =
+          InvDcEdge * (ScalarCell(JCell1) - ScalarCell(JCell0));
+      return GradEdge;
+   }
+
+ private:
+   Array2DI4 CellsOnEdge;
+   Array1DR8 DcEdge;
+};
+
+class CurlOnVertex {
+ public:
+   CurlOnVertex(HorzMesh const *Mesh);
+
+   KOKKOS_FUNCTION Real operator()(int IVertex,
+                                   const Array1DReal &VecEdge) const {
+      Real CurlVertex = 0;
+      for (int J = 0; J < VertexDegree; ++J) {
+         const int JEdge = EdgesOnVertex(IVertex, J);
+         CurlVertex +=
+             DcEdge(JEdge) * EdgeSignOnVertex(IVertex, J) * VecEdge(JEdge);
+      }
+      const Real InvAreaTriangle = 1. / AreaTriangle(IVertex);
+      CurlVertex *= InvAreaTriangle;
+      return CurlVertex;
+   }
+
+ private:
+   I4 VertexDegree;
+   Array2DI4 EdgesOnVertex;
+   Array1DR8 DcEdge;
+   Array1DR8 AreaTriangle;
+   Array2DR8 EdgeSignOnVertex;
+};
+
+class TangentialReconOnEdge {
+ public:
+   TangentialReconOnEdge(HorzMesh const *Mesh);
+
+   KOKKOS_FUNCTION Real operator()(int IEdge,
+                                   const Array1DReal &VecEdge) const {
+      Real ReconEdge = 0;
+      for (int J = 0; J < NEdgesOnEdge(IEdge); ++J) {
+         const int JEdge = EdgesOnEdge(IEdge, J);
+         ReconEdge += WeightsOnEdge(IEdge, J) * VecEdge(JEdge);
+      }
+      return ReconEdge;
+   }
+
+ private:
+   Array1DI4 NEdgesOnEdge;
+   Array2DI4 EdgesOnEdge;
+   Array2DR8 WeightsOnEdge;
+};
+
+} // namespace OMEGA
+#endif

components/omega/test/CMakeLists.txt

@@ -111,6 +111,34 @@ add_omega_test(
     "-n 8;--cpu-bind=cores"
 )
 
+#######################
+# HorzOperators tests
+#######################
+
+add_omega_test(
+    HORZOPERATORS_PLANE_TEST
+    testHorzOperatorsPlane.exe
+    ocn/HorzOperatorsTest.cpp
+    "-n 8;--cpu-bind=cores"
+)
+target_compile_definitions(
+  testHorzOperatorsPlane.exe
+  PRIVATE
+  HORZOPERATORS_TEST_PLANE
+)
+
+add_omega_test(
+  HORZOPERATORS_SPHERE_TEST
+    testHorzOperatorsSphere.exe
+    ocn/HorzOperatorsTest.cpp
+    "-n 8;--cpu-bind=cores"
+)
+target_compile_definitions(
+  testHorzOperatorsSphere.exe
+  PRIVATE
+  HORZOPERATORS_TEST_SPHERE_1
+)
+
 #############
 # IO test
 #############
@@ -190,6 +218,8 @@ set_tests_properties(
   DECOMP_TEST
   HALO_TEST
   HORZMESH_TEST
+  HORZOPERATORS_PLANE_TEST
+  HORZOPERATORS_SPHERE_TEST
   IO_TEST
   METADATA_TEST
   TIMEMGR_TEST