Computing a stencil on a t8code mesh.

Makefile

T8_DIR = $(HOME)/install/t8code/main
T8_INC = -I$(T8_DIR)/include
T8_LIB = -L$(T8_DIR)/lib -Wl,-rpath=$(T8_DIR)/lib

LIBS = -lz -lsc -lp4est -lt8

CXX = mpicxx

all: stencil

stencil: stencil.cxx
	$(CXX) -o $@ $< $(T8_INC) $(T8_LIB) $(LIBS)

stencil.cxx

#include <cmath>

#include <t8.h>                 /* General t8code header, always include this. */
#include <t8_vec.h>             /* Basic operations on 3D vectors. */
#include <t8_cmesh.h>           /* cmesh definition and basic interface. */
#include <t8_forest.h>          /* forest definition and basic interface. */
#include <t8_cmesh/t8_cmesh_examples.h> /* A collection of exemplary cmeshes */
#include <t8_schemes/t8_default/t8_default_cxx.hxx>     /* default refinement scheme. */

/* The data that we want to store for each element.
 * In this example we want to store the element's level and volume. */
struct data_per_element
{
  int                 level;
  double              midpoint[3];
  double              dx, dy;
  double              volume;
  double              height;
  double              schlieren;
  double              curvature;
};

static struct data_per_element *
create_element_data (t8_forest_t forest)
{
  t8_locidx_t         num_local_elements;
  t8_locidx_t         num_ghost_elements;
  struct data_per_element *element_data;

  t8_locidx_t         itree, num_local_trees;
  t8_locidx_t         current_index;
  t8_locidx_t         ielement, num_elements_in_tree;
  t8_eclass_t         tree_class;
  t8_eclass_scheme_c *eclass_scheme;
  const t8_element_t *element;

  /* Check that forest is a committed, that is valid and usable, forest. */
  T8_ASSERT (t8_forest_is_committed (forest));

  /* Get the number of local elements of forest. */
  num_local_elements = t8_forest_get_local_num_elements (forest);
  /* Get the number of ghost elements of forest. */
  num_ghost_elements = t8_forest_get_num_ghosts (forest);

  element_data = T8_ALLOC (struct data_per_element, num_local_elements + num_ghost_elements);

  /* Get the number of trees that have elements of this process. */
  num_local_trees = t8_forest_get_num_local_trees (forest);

  for (itree = 0, current_index = 0; itree < num_local_trees; ++itree) {
    tree_class = t8_forest_get_tree_class (forest, itree);
    eclass_scheme = t8_forest_get_eclass_scheme (forest, tree_class);

    num_elements_in_tree = t8_forest_get_tree_num_elements (forest, itree);

    for (ielement = 0; ielement < num_elements_in_tree; ++ielement, ++current_index) {
      element = t8_forest_get_element_in_tree (forest, itree, ielement);

      struct data_per_element *edat = &element_data[current_index];

      edat->level = eclass_scheme->t8_element_level (element);
      edat->volume = t8_forest_element_volume (forest, itree, element);
      t8_forest_element_centroid (forest, itree, element, edat->midpoint);

      edat->dx = sqrt(edat->volume);
      edat->dy = edat->dx;

      const double x = edat->midpoint[0] - 0.5;
      const double y = edat->midpoint[1] - 0.5;
      const double s = 0.25;

      const double r = sqrt(x*x + y*y)*20.0;

      // Some 'interesting' height function.
      edat->height = sin(2.0*r)/r;
    }
  }

  return element_data;
}

static void
compute_stencil (t8_forest_t forest, struct data_per_element *element_data)
{
  t8_locidx_t         num_local_elements;
  t8_locidx_t         num_ghost_elements;

  t8_locidx_t         itree, num_local_trees;
  t8_locidx_t         current_index;
  t8_locidx_t         ielement, num_elements_in_tree;
  t8_eclass_t         tree_class;
  t8_element_t       *element, **neighbors;
  int                 iface, ineigh;
  t8_eclass_scheme_c *eclass_scheme;
  t8_eclass_scheme_c *neigh_scheme;

  int                 num_faces; /**< The number of faces */
  int                 num_neighbors; /**< Number of neighbors for each face */
  int                *dual_faces; /**< The face indices of the neighbor elements */
  t8_locidx_t        *neighids; /**< Indices of the neighbor elements */
  int8_t              neigh_level; /**< The level of the face neighbors at this face. */

  /* Check that forest is a committed, that is valid and usable, forest. */
  T8_ASSERT (t8_forest_is_committed (forest));

  /* Get the number of local elements of forest. */
  num_local_elements = t8_forest_get_local_num_elements (forest);
  /* Get the number of ghost elements of forest. */
  num_ghost_elements = t8_forest_get_num_ghosts (forest);

  /* Get the number of trees that have elements of this process. */
  num_local_trees = t8_forest_get_num_local_trees (forest);

  double stencil[3][3] = {0};

  for (itree = 0, current_index = 0; itree < num_local_trees; ++itree) {
    tree_class = t8_forest_get_tree_class (forest, itree);
    eclass_scheme = t8_forest_get_eclass_scheme (forest, tree_class);

    num_elements_in_tree = t8_forest_get_tree_num_elements (forest, itree);

    for (ielement = 0; ielement < num_elements_in_tree; ++ielement, ++current_index) {
      element = t8_forest_get_element_in_tree (forest, itree, ielement);

      stencil[1][1] = element_data[current_index].height;

      num_faces = eclass_scheme->t8_element_num_faces (element);
      for (iface = 0; iface < num_faces; iface++) {

        t8_forest_leaf_face_neighbors (forest, itree, element,
                                      &neighbors, iface,
                                      &dual_faces,
                                      &num_neighbors,
                                      &neighids,
                                      &neigh_scheme, 1);

        double height = element_data[neighids[0]].height;

        switch (iface) {
          case 0: stencil[0][1] = height; break; // NORTH
          case 1: stencil[2][1] = height; break; // SOUTH
          case 2: stencil[1][0] = height; break; // WEST
          case 3: stencil[1][2] = height; break; // EAST
        }

        T8_FREE(neighbors);
        T8_FREE(dual_faces);
        T8_FREE(neighids);
      }

      const double dx = element_data[current_index].dx;
      const double dy = element_data[current_index].dy;

      const double xslope = 0.5/dx*(stencil[2][1] - stencil[0][1]);
      const double yslope = 0.5/dy*(stencil[1][2] - stencil[1][0]);

      const double xcurve = 1.0/(dx*dx)*(stencil[2][1] - 2.0*stencil[1][1] + stencil[0][1]);
      const double ycurve = 1.0/(dy*dy)*(stencil[1][2] - 2.0*stencil[1][1] + stencil[1][0]);

      element_data[current_index].schlieren = sqrt(xslope*xslope + yslope*yslope);
      element_data[current_index].curvature = sqrt(xcurve*xcurve + ycurve*ycurve);
    }
  }
}

/* Each process has computed the data entries for its local elements.
 * In order to get the values for the ghost elements, we use t8_forest_ghost_exchange_data.
 * Calling this function will fill all the ghost entries of our element data array with the
 * value on the process that owns the corresponding element. */
static void
exchange_ghost_data (t8_forest_t forest, struct data_per_element *data)
{
  sc_array           *sc_array_wrapper;
  t8_locidx_t         num_elements = t8_forest_get_local_num_elements (forest);
  t8_locidx_t         num_ghosts = t8_forest_get_num_ghosts (forest);

  /* t8_forest_ghost_exchange_data expects an sc_array (of length num_local_elements + num_ghosts).
   * We wrap our data array to an sc_array. */
  sc_array_wrapper = sc_array_new_data (data, sizeof (struct data_per_element), num_elements + num_ghosts);

  /* Carry out the data exchange. The entries with indices > num_local_elements will get overwritten. */
  t8_forest_ghost_exchange_data (forest, sc_array_wrapper);

  /* Destroy the wrapper array. This will not free the data memory since we used sc_array_new_data. */
  sc_array_destroy (sc_array_wrapper);
}

/* Write the forest as vtu and also write the element's volumes in the file.
 * 
 * t8code supports writing element based data to vtu as long as its stored
 * as doubles. Each of the data fields to write has to be provided in its own
 * array of length num_local_elements.
 * We support two types: T8_VTK_SCALAR - One double per element
 *                  and  T8_VTK_VECTOR - 3 doubles per element
 */
static void
output_data_to_vtu (t8_forest_t forest,
                             struct data_per_element *data,
                             const char *prefix)
{
  t8_locidx_t         num_elements = t8_forest_get_local_num_elements (forest);
  

  // We need to allocate a new array to store the data on their own.
  // These arrays have one entry per local element.
  double             *heights   = T8_ALLOC (double, num_elements);
  double             *schlieren = T8_ALLOC (double, num_elements);
  double             *curvature = T8_ALLOC (double, num_elements);

  // The number of user defined data fields to write.
  const int         num_data = 3;

  // For each user defined data field we need one t8_vtk_data_field_t variable.
  t8_vtk_data_field_t vtk_data[num_data];
  // Set the type of this variable. Since we have one value per element, we pick T8_VTK_SCALAR.
  vtk_data[0].type = T8_VTK_SCALAR;
  // The name of the field as should be written to the file.
  strcpy (vtk_data[0].description, "height");
  vtk_data[0].data = heights;
  // Copy the elment's volumes from our data array to the output array.
  for (t8_locidx_t ielem = 0; ielem < num_elements; ++ielem) {
    heights[ielem] = data[ielem].height;
  }

  vtk_data[1].type = T8_VTK_SCALAR;
  strcpy (vtk_data[1].description, "schlieren");
  vtk_data[1].data = schlieren;
  for (t8_locidx_t ielem = 0; ielem < num_elements; ++ielem) {
    schlieren[ielem] = data[ielem].schlieren;
  }

  vtk_data[2].type = T8_VTK_SCALAR;
  strcpy (vtk_data[2].description, "curvature");
  vtk_data[2].data = curvature;
  for (t8_locidx_t ielem = 0; ielem < num_elements; ++ielem) {
    curvature[ielem] = data[ielem].curvature;
  }

  {
    // To write user defined data.
    const int write_treeid = 1;
    const int write_mpirank = 1;
    const int write_level = 1;
    const int write_element_id = 1;
    const int write_ghosts = 0;
    t8_forest_write_vtk_ext (forest, prefix, write_treeid, write_mpirank,
                             write_level, write_element_id, write_ghosts,
                             0, 0, num_data, vtk_data);
  }

  T8_FREE (heights);
  T8_FREE (schlieren);
  T8_FREE (curvature);
}

int
main (int argc, char **argv)
{
  int                 mpiret;
  sc_MPI_Comm         comm;
  t8_forest_t         forest;

  /* The prefix for our output files. */
  const char         *prefix_forest_with_data = "example";

  /* The uniform refinement level of the forest. */
  const int           dim = 2;
  const int           level = 8;

  /* The array that will hold our per element data. */
  data_per_element *data;

  /*
   * Initialization.
   */

  // Initialize MPI. This has to happen before we initialize sc or t8code.
  mpiret = sc_MPI_Init (&argc, &argv);
  // Error check the MPI return value.
  SC_CHECK_MPI (mpiret);

  // Initialize the sc library, has to happen before we initialize t8code.
  sc_init (sc_MPI_COMM_WORLD, 1, 1, NULL, SC_LP_ESSENTIAL);
  // Initialize t8code with log level SC_LP_PRODUCTION. See sc.h for more info on the log levels.
  t8_init (SC_LP_PRODUCTION);

  // We will use MPI_COMM_WORLD as a communicator.
  comm = sc_MPI_COMM_WORLD;

  // Initialize a uniform forest with periodic boundaries.
  t8_cmesh_t          cmesh = t8_cmesh_new_periodic (comm, dim);
  t8_scheme_cxx_t    *scheme = t8_scheme_new_default_cxx ();

  const int do_face_ghost = 1;
  forest = t8_forest_new_uniform (cmesh, scheme, level, do_face_ghost, comm);

  /*
   * Data handling and computation.
   */

  // Build data array and gather data for the local elements.
  data = create_element_data (forest);

  // Exchange the neighboring data at MPI process boundaries.
  exchange_ghost_data (forest, data);

  // Compute stencil.
  compute_stencil (forest, data);

  // Output the data to vtu files.
  output_data_to_vtu (forest, data, prefix_forest_with_data);
  t8_global_productionf (" Wrote forest and data to %s*.\n", prefix_forest_with_data);

  /*
   * Clean-up
   */

  /* Free the data array. */
  T8_FREE (data);

  /* Destroy the forest. */
  t8_forest_unref (&forest);
  t8_global_productionf (" Destroyed forest.\n");

  sc_finalize ();

  mpiret = sc_MPI_Finalize ();
  SC_CHECK_MPI (mpiret);

  return 0;
}