xref: /petsc/src/dm/tutorials/ex22.c (revision 9371c9d470a9602b6d10a8bf50c9b2280a79e45a)

static char help[] = "Extract a 2D slice in natural ordering from a 3D vector, Command line options : \n\
Mx/My/Mz - set the dimensions of the parent vector \n\
sliceaxis - string describing the axis along which the sice will be selected : X, Y, Z \n\
gp - global grid point number along the sliceaxis direction where the slice will be extracted from the parent vector \n";

/*
  This example shows to extract a 2D slice in natural ordering
  from a 3D DMDA vector (first by extracting the slice and then
  by converting it to natural ordering)
*/

#include <petscdmda.h>

const char *const sliceaxes[] = {"X", "Y", "Z", "sliceaxis", "DM_", NULL};

int main(int argc, char **argv) {
  DM                 da3D;                            /* 3D DMDA object */
  DM                 da2D;                            /* 2D DMDA object */
  Vec                vec_full;                        /* Parent vector */
  Vec                vec_extracted;                   /* Extracted slice vector (in DMDA ordering) */
  Vec                vec_slice;                       /* vector in natural ordering */
  Vec                vec_slice_g;                     /* aliased vector in natural ordering */
  IS                 patchis_3d;                      /* IS to select slice and extract subvector */
  IS                 patchis_2d;                      /* Patch IS for 2D vector, will be converted to application ordering */
  IS                 scatis_extracted_slice;          /* PETSc indexed IS for extracted slice */
  IS                 scatis_natural_slice;            /* natural/application ordered IS for slice*/
  IS                 scatis_natural_slice_g;          /* aliased natural/application ordered IS  for slice */
  VecScatter         vscat;                           /* scatter slice in DMDA ordering <-> slice in column major ordering */
  AO                 da2D_ao;                         /* AO associated with 2D DMDA */
  MPI_Comm           subset_mpi_comm = MPI_COMM_NULL; /* MPI communicator where the slice lives */
  PetscScalar     ***vecdata3d;                       /* Pointer to access 3d parent vector */
  const PetscScalar *array;                           /* pointer to create aliased Vec */
  PetscInt           Mx = 4, My = 4, Mz = 4;          /* Dimensions for 3D DMDA */
  const PetscInt    *l1, *l2;                         /* 3D DMDA layout */
  PetscInt           M1 = -1, M2 = -1;                /* Dimensions for 2D DMDA */
  PetscInt           m1 = -1, m2 = -1;                /* Layouts for 2D DMDA */
  PetscInt           gp        = 2;                   /* grid point along sliceaxis to pick the slice */
  DMDirection        sliceaxis = DM_X;                /* Select axis along which the slice will be extracted */
  PetscInt           i, j, k;                         /* Iteration indices */
  PetscInt           ixs, iys, izs;                   /* Corner indices for 3D vector */
  PetscInt           ixm, iym, izm;                   /* Widths of parent vector */
  PetscInt           low, high;                       /* ownership range indices */
  PetscInt           in;                              /* local size index for IS*/
  PetscInt           vn;                              /* local size index */
  const PetscInt    *is_array;                        /* pointer to create aliased IS */
  MatStencil         lower, upper;                    /* Stencils to select slice for Vec */
  PetscBool          patchis_offproc = PETSC_FALSE;   /* flag to DMDACreatePatchIS indicating that off-proc values are to be ignored */
  PetscMPIInt        rank, size;                      /* MPI rank and size */

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Initialize program and set problem parameters
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscFunctionBeginUser;
  PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
  PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
  PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));

  PetscOptionsBegin(PETSC_COMM_WORLD, "", "ex22 DMDA tutorial example options", "DMDA");
  PetscCall(PetscOptionsRangeInt("-Mx", "dimension along x-axis", "ex22.c", Mx, &Mx, NULL, 0, PETSC_MAX_INT));
  PetscCall(PetscOptionsRangeInt("-My", "dimension along y-axis", "ex22.c", My, &My, NULL, 0, PETSC_MAX_INT));
  PetscCall(PetscOptionsRangeInt("-Mz", "dimension along z-axis", "ex22.c", Mz, &Mz, NULL, 0, PETSC_MAX_INT));
  PetscCall(PetscOptionsEnum("-sliceaxis", "axis along which 2D slice is extracted from : X, Y, Z", "", sliceaxes, (PetscEnum)sliceaxis, (PetscEnum *)&sliceaxis, NULL));
  PetscCall(PetscOptionsRangeInt("-gp", "index along sliceaxis at which 2D slice is extracted", "ex22.c", gp, &gp, NULL, 0, PETSC_MAX_INT));
  PetscOptionsEnd();

  /* Ensure that the requested slice is not out of bounds for the selected axis */
  if (sliceaxis == DM_X) {
    PetscCheck(gp <= Mx, PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
  } else if (sliceaxis == DM_Y) {
    PetscCheck(gp <= My, PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
  } else if (sliceaxis == DM_Z) {
    PetscCheck(gp <= Mz, PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
  }

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Create 3D DMDA object.
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
  PetscCall(DMDACreate3d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, Mx, My, Mz, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, 1, 1, NULL, NULL, NULL, &da3D));
  PetscCall(DMSetFromOptions(da3D));
  PetscCall(DMSetUp(da3D));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Create the parent vector
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
  PetscCall(DMCreateGlobalVector(da3D, &vec_full));
  PetscCall(PetscObjectSetName((PetscObject)vec_full, "full_vector"));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Populate the 3D vector
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(DMDAGetCorners(da3D, &ixs, &iys, &izs, &ixm, &iym, &izm));
  PetscCall(DMDAVecGetArray(da3D, vec_full, &vecdata3d));
  for (k = izs; k < izs + izm; k++) {
    for (j = iys; j < iys + iym; j++) {
      for (i = ixs; i < ixs + ixm; i++) { vecdata3d[k][j][i] = ((i - Mx / 2.0) * (j + Mx / 2.0)) + k * 100; }
    }
  }
  PetscCall(DMDAVecRestoreArray(da3D, vec_full, &vecdata3d));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Get an IS corresponding to a 2D slice
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  if (sliceaxis == DM_X) {
    lower.i = gp;
    lower.j = 0;
    lower.k = 0;
    upper.i = gp;
    upper.j = My;
    upper.k = Mz;
  } else if (sliceaxis == DM_Y) {
    lower.i = 0;
    lower.j = gp;
    lower.k = 0;
    upper.i = Mx;
    upper.j = gp;
    upper.k = Mz;
  } else if (sliceaxis == DM_Z) {
    lower.i = 0;
    lower.j = 0;
    lower.k = gp;
    upper.i = Mx;
    upper.j = My;
    upper.k = gp;
  }
  PetscCall(DMDACreatePatchIS(da3D, &lower, &upper, &patchis_3d, patchis_offproc));
  PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n IS to select slice from 3D DMDA vector : \n"));
  PetscCall(ISView(patchis_3d, PETSC_VIEWER_STDOUT_WORLD));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Use the obtained IS to extract the slice as a subvector
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(VecGetSubVector(vec_full, patchis_3d, &vec_extracted));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     View the extracted subvector
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD, PETSC_VIEWER_ASCII_DENSE));
  PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in DMDA ordering : \n"));
  PetscCall(VecView(vec_extracted, PETSC_VIEWER_STDOUT_WORLD));
  PetscCall(PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Query 3D DMDA layout, get the subset MPI communicator
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
  if (sliceaxis == DM_X) {
    PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL));
    PetscCall(DMDAGetOwnershipRanges(da3D, NULL, &l1, &l2));
    M1 = My;
    M2 = Mz;
    PetscCall(DMDAGetProcessorSubset(da3D, DM_X, gp, &subset_mpi_comm));
  } else if (sliceaxis == DM_Y) {
    PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, NULL, &m2, NULL, NULL, NULL, NULL, NULL, NULL));
    PetscCall(DMDAGetOwnershipRanges(da3D, &l1, NULL, &l2));
    M1 = Mx;
    M2 = Mz;
    PetscCall(DMDAGetProcessorSubset(da3D, DM_Y, gp, &subset_mpi_comm));
  } else if (sliceaxis == DM_Z) {
    PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL, NULL));
    PetscCall(DMDAGetOwnershipRanges(da3D, &l1, &l2, NULL));
    M1 = Mx;
    M2 = My;
    PetscCall(DMDAGetProcessorSubset(da3D, DM_Z, gp, &subset_mpi_comm));
  }

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Create 2D DMDA object,
     vector (that will hold the slice as a column major flattened array) &
     index set (that will be used for scattering to the column major
     indexed slice vector)
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
  if (subset_mpi_comm != MPI_COMM_NULL) {
    PetscCallMPI(MPI_Comm_size(subset_mpi_comm, &size));
    PetscCall(PetscSynchronizedPrintf(subset_mpi_comm, "subset MPI subcomm size is : %d, includes global rank : %d \n", size, rank));
    PetscCall(PetscSynchronizedFlush(subset_mpi_comm, PETSC_STDOUT));
    PetscCall(DMDACreate2d(subset_mpi_comm, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, M1, M2, m1, m2, 1, 1, l1, l2, &da2D));
    PetscCall(DMSetFromOptions(da2D));
    PetscCall(DMSetUp(da2D));

    /* Create a 2D patch IS for the slice */
    lower.i = 0;
    lower.j = 0;
    upper.i = M1;
    upper.j = M2;
    PetscCall(DMDACreatePatchIS(da2D, &lower, &upper, &patchis_2d, patchis_offproc));

    /* Convert the 2D patch IS to natural indexing (column major flattened) */
    PetscCall(ISDuplicate(patchis_2d, &scatis_natural_slice));
    PetscCall(DMDAGetAO(da2D, &da2D_ao));
    PetscCall(AOPetscToApplicationIS(da2D_ao, scatis_natural_slice));
    PetscCall(ISGetIndices(scatis_natural_slice, &is_array));
    PetscCall(ISGetLocalSize(scatis_natural_slice, &in));

    /* Create an aliased IS on the 3D DMDA's communicator */
    PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, in, is_array, PETSC_USE_POINTER, &scatis_natural_slice_g));
    PetscCall(ISRestoreIndices(scatis_natural_slice, &is_array));

    /* Create a 2D DMDA global vector */
    PetscCall(DMCreateGlobalVector(da2D, &vec_slice));
    PetscCall(PetscObjectSetName((PetscObject)vec_slice, "slice_vector_natural"));
    PetscCall(VecGetLocalSize(vec_slice, &vn));
    PetscCall(VecGetArrayRead(vec_slice, &array));

    /* Create an aliased version of the above on the 3D DMDA's communicator */
    PetscCall(VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, vn, M1 * M2, array, &vec_slice_g));
    PetscCall(VecRestoreArrayRead(vec_slice, &array));
  } else {
    /* Ranks not part of the subset MPI communicator provide no entries, but the routines for creating
       the IS and Vec on the 3D DMDA's communicator still need to called, since they are collective routines */
    PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 0, NULL, PETSC_USE_POINTER, &scatis_natural_slice_g));
    PetscCall(VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, 0, M1 * M2, NULL, &vec_slice_g));
  }
  PetscCall(PetscBarrier(NULL));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Create IS that maps from the extracted slice vector
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(VecGetOwnershipRange(vec_extracted, &low, &high));
  PetscCall(ISCreateStride(PETSC_COMM_WORLD, high - low, low, 1, &scatis_extracted_slice));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Scatter extracted subvector -> natural 2D slice vector
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(VecScatterCreate(vec_extracted, scatis_extracted_slice, vec_slice_g, scatis_natural_slice_g, &vscat));
  PetscCall(VecScatterBegin(vscat, vec_extracted, vec_slice_g, INSERT_VALUES, SCATTER_FORWARD));
  PetscCall(VecScatterEnd(vscat, vec_extracted, vec_slice_g, INSERT_VALUES, SCATTER_FORWARD));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     View the natural 2D slice vector
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD, PETSC_VIEWER_ASCII_DENSE));
  PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in natural ordering : \n"));
  PetscCall(VecView(vec_slice_g, PETSC_VIEWER_STDOUT_WORLD));
  PetscCall(PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Restore subvector
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(VecRestoreSubVector(vec_full, patchis_3d, &vec_extracted));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Destroy data structures and exit.
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(VecDestroy(&vec_full));
  PetscCall(VecScatterDestroy(&vscat));
  PetscCall(ISDestroy(&scatis_extracted_slice));
  PetscCall(ISDestroy(&scatis_natural_slice_g));
  PetscCall(VecDestroy(&vec_slice_g));
  PetscCall(ISDestroy(&patchis_3d));
  PetscCall(DMDestroy(&da3D));

  if (subset_mpi_comm != MPI_COMM_NULL) {
    PetscCall(ISDestroy(&patchis_2d));
    PetscCall(ISDestroy(&scatis_natural_slice));
    PetscCall(VecDestroy(&vec_slice));
    PetscCall(DMDestroy(&da2D));
    PetscCallMPI(MPI_Comm_free(&subset_mpi_comm));
  }

  PetscCall(PetscFinalize());
  return 0;
}

/*TEST

    test:
      nsize: 1
      args: -sliceaxis X -gp 0

    test:
      suffix: 2
      nsize:  2
      args: -sliceaxis Y -gp 1
      filter: grep -v "subset MPI subcomm"

    test:
      suffix: 3
      nsize:  3
      args:  -sliceaxis Z -gp 2
      filter: grep -v "subset MPI subcomm"

    test:
      suffix: 4
      nsize:  4
      args: -sliceaxis X -gp 2
      filter: grep -v "subset MPI subcomm"

    test:
      suffix: 5
      nsize:  4
      args: -sliceaxis Z -gp 1
      filter: grep -v "subset MPI subcomm"

TEST*/