Actual source code: dalocalizationletkf.kokkos.cxx

      1: #include <petsc.h>
      2: #include <petscmat.h>
      3: #include <petsc_kokkos.hpp>

      5: typedef struct {
      6:   PetscReal distance;
      7:   PetscInt  obs_index;
      8: } DistObsPair;

     10: KOKKOS_INLINE_FUNCTION
     11: static PetscReal GaspariCohn(PetscReal distance, PetscReal radius)
     12: {
     13:   if (radius <= 0.0) return 0.0;
     14:   const PetscReal r = distance / radius;

     16:   if (r >= 2.0) return 0.0;

     18:   const PetscReal r2 = r * r;
     19:   const PetscReal r3 = r2 * r;
     20:   const PetscReal r4 = r3 * r;
     21:   const PetscReal r5 = r4 * r;

     23:   if (r <= 1.0) {
     24:     // Region [0, 1]
     25:     return -0.25 * r5 + 0.5 * r4 + 0.625 * r3 - (5.0 / 3.0) * r2 + 1.0;
     26:   } else {
     27:     // Region [1, 2]
     28:     return (1.0 / 12.0) * r5 - 0.5 * r4 + 0.625 * r3 + (5.0 / 3.0) * r2 - 5.0 * r + 4.0 - (2.0 / 3.0) / r;
     29:   }
     30: }

     32: #define RADIUS_FACTOR 1.1

     34: template <class ViewType>
     35: struct RadiusStatsFunctor {
     36:   ViewType best_dists;
     37:   PetscInt n_obs_vertex;

     39:   struct value_type {
     40:     double sum, sq_sum;
     41:   };

     43:   KOKKOS_INLINE_FUNCTION void operator()(const PetscInt i, value_type &update) const
     44:   {
     45:     PetscReal r2 = best_dists(i, n_obs_vertex - 1);
     46:     PetscReal r  = std::sqrt(r2);
     47:     r *= RADIUS_FACTOR;
     48:     if (r == 0.0) r = 1.0;
     49:     update.sum += r;
     50:     update.sq_sum += r * r;
     51:   }

     53:   KOKKOS_INLINE_FUNCTION void init(value_type &update) const
     54:   {
     55:     update.sum    = 0.0;
     56:     update.sq_sum = 0.0;
     57:   }

     59:   KOKKOS_INLINE_FUNCTION void join(value_type &dest, const value_type &src) const
     60:   {
     61:     dest.sum += src.sum;
     62:     dest.sq_sum += src.sq_sum;
     63:   }
     64: };

     66: /*@
     67:   PetscDALETKFGetLocalizationMatrix - Compute localization weight matrix for LETKF [move to ml/da/interface]

     69:   Collective

     71:   Input Parameters:
     72: + n_obs_vertex - Number of observations to localize to per vertex
     73: . n_dof        - Number of degrees of freedom
     74: . Vecxyz       - Array of vectors containing the vertex coordinates
     75: . bd           - Array of boundary extents per dimension (used for periodicity)
     76: - H            - Observation operator matrix

     78:   Output Parameter:
     79: . Q - Localization weight matrix (sparse, AIJ format)

     81:   Level: intermediate

     83:   Notes:
     84:   The output matrix Q has dimensions (n_vert_global x n_obs_global) where
     85:   n_vert_global is the number of vertices in the DMPlex. Each row contains
     86:   exactly n_obs_vertex non-zero entries corresponding to the nearest
     87:   observations, weighted by the Gaspari-Cohn fifth-order piecewise
     88:   rational function.

     90:   The observation locations are computed as H * V where V is the vector
     91:   of vertex coordinates. The localization weights ensure smooth tapering
     92:   of observation influence with distance.

     94:   Kokkos is required for this routine.

     96: .seealso: [](ch_da), `PetscDALETKFSetLocalization()`
     97: @*/
     98: PetscErrorCode PetscDALETKFGetLocalizationMatrix(const PetscInt n_obs_vertex, const PetscInt n_dof, Vec Vecxyz[3], PetscReal bd[3], Mat H, Mat *Q)
     99: {
    100:   PetscInt dim = 0, n_vert_local, d, n_obs_global, n_obs_local;
    101:   Vec     *obs_vecs;
    102:   MPI_Comm comm;

    104:   PetscFunctionBegin;
    106:   PetscAssertPointer(Q, 6);

    108:   PetscCall(PetscKokkosInitializeCheck());
    109:   PetscCall(PetscObjectGetComm((PetscObject)H, &comm));
    110:   PetscCall(MatGetLocalSize(H, &n_obs_local, NULL));
    111:   PetscCall(MatGetSize(H, &n_obs_global, NULL));
    112:   /* Infer dim from the number of vectors in Vecxyz */
    113:   for (d = 0; d < 3; ++d) {
    114:     if (Vecxyz[d]) dim++;
    115:     else break;
    116:   }
    117:   PetscCall(VecGetLocalSize(Vecxyz[0], &n_vert_local));

    119:   /* Check H dimensions */
    120:   // If n_obs_global < n_obs_vertex, we will pad with -1 indices and 0.0 weights. ???
    121:   // This is not an error condition, but rather a case where we have fewer observations than requested neighbors.
    122:   PetscCheck(dim > 0, comm, PETSC_ERR_ARG_WRONG, "Dim must be > 0");
    123:   PetscCheck(n_obs_vertex > 0 && n_obs_vertex <= n_obs_global, comm, PETSC_ERR_ARG_WRONG, "n_obs_vertex must be > 0 and <= n_obs_global");

    125:   /* Allocate storage for observation locations */
    126:   PetscCall(PetscMalloc1(dim, &obs_vecs));

    128:   /* Compute observation locations per dimension */
    129:   for (d = 0; d < dim; ++d) {
    130:     PetscCall(MatCreateVecs(H, NULL, &obs_vecs[d]));
    131:     PetscCall(MatMult(H, Vecxyz[d], obs_vecs[d]));
    132:   }

    134:   /* Create output matrix Q in N/n_dof x P */
    135:   PetscCall(MatCreate(comm, Q));
    136:   PetscCall(MatSetSizes(*Q, n_vert_local, n_obs_local, PETSC_DETERMINE, n_obs_global));
    137:   PetscCall(MatSetType(*Q, MATAIJ));
    138:   PetscCall(MatSeqAIJSetPreallocation(*Q, n_obs_vertex, NULL));
    139:   PetscCall(MatMPIAIJSetPreallocation(*Q, n_obs_vertex, NULL, n_obs_vertex, NULL));
    140:   PetscCall(MatSetFromOptions(*Q));
    141:   PetscCall(MatSetUp(*Q));

    143:   PetscCall(PetscInfo((PetscObject)*Q, "Computing LETKF localization matrix: %" PetscInt_FMT " vertices, %" PetscInt_FMT " observations, %" PetscInt_FMT " observations per vertex\n", n_vert_local, n_obs_global, n_obs_vertex));

    145:   /* Prepare Kokkos Views */
    146:   using ExecSpace = Kokkos::DefaultExecutionSpace;
    147:   using MemSpace  = ExecSpace::memory_space;

    149:   /* Vertex Coordinates */
    150:   // Use LayoutLeft for coalesced access on GPU (i is contiguous)
    151:   Kokkos::View<PetscScalar **, Kokkos::LayoutLeft, MemSpace> vertex_coords_dev("vertex_coords", n_vert_local, dim);
    152:   {
    153:     // Host view must match the data layout from VecGetArray (d-major, i-minor implies LayoutLeft for (i,d) view)
    154:     Kokkos::View<PetscScalar **, Kokkos::LayoutLeft, Kokkos::HostSpace> vertex_coords_host("vertex_coords_host", n_vert_local, dim);
    155:     for (d = 0; d < dim; ++d) {
    156:       const PetscScalar *local_coords_array;
    157:       PetscCall(VecGetArrayRead(Vecxyz[d], &local_coords_array));
    158:       // Copy data. Since vertex_coords_host is LayoutLeft, &vertex_coords_host(0, d) is the start of column d.
    159:       for (PetscInt i = 0; i < n_vert_local; ++i) vertex_coords_host(i, d) = local_coords_array[i];
    160:       PetscCall(VecRestoreArrayRead(Vecxyz[d], &local_coords_array));
    161:     }
    162:     Kokkos::deep_copy(vertex_coords_dev, vertex_coords_host);
    163:   }

    165:   /* Observation Coordinates */
    166:   Kokkos::View<PetscReal **, Kokkos::LayoutRight, MemSpace> obs_coords_dev("obs_coords", n_obs_global, dim);
    167:   {
    168:     Kokkos::View<PetscReal **, Kokkos::LayoutRight, Kokkos::HostSpace> obs_coords_host("obs_coords_host", n_obs_global, dim);
    169:     for (d = 0; d < dim; ++d) {
    170:       VecScatter         ctx;
    171:       Vec                seq_vec;
    172:       const PetscScalar *array;

    174:       PetscCall(VecScatterCreateToAll(obs_vecs[d], &ctx, &seq_vec));
    175:       PetscCall(VecScatterBegin(ctx, obs_vecs[d], seq_vec, INSERT_VALUES, SCATTER_FORWARD));
    176:       PetscCall(VecScatterEnd(ctx, obs_vecs[d], seq_vec, INSERT_VALUES, SCATTER_FORWARD));

    178:       PetscCall(VecGetArrayRead(seq_vec, &array));
    179:       for (PetscInt j = 0; j < n_obs_global; ++j) obs_coords_host(j, d) = PetscRealPart(array[j]);
    180:       PetscCall(VecRestoreArrayRead(seq_vec, &array));
    181:       PetscCall(VecScatterDestroy(&ctx));
    182:       PetscCall(VecDestroy(&seq_vec));
    183:     }
    184:     Kokkos::deep_copy(obs_coords_dev, obs_coords_host);
    185:   }

    187:   PetscInt rstart;
    188:   PetscCall(VecGetOwnershipRange(Vecxyz[0], &rstart, NULL));

    190:   /* Output Views */
    191:   // LayoutLeft for coalesced access on GPU
    192:   Kokkos::View<PetscInt **, Kokkos::LayoutLeft, MemSpace>    indices_dev("indices", n_vert_local, n_obs_vertex);
    193:   Kokkos::View<PetscScalar **, Kokkos::LayoutLeft, MemSpace> values_dev("values", n_vert_local, n_obs_vertex);

    195:   /* Temporary storage for top-k per vertex */
    196:   // LayoutLeft for coalesced access on GPU.
    197:   // Note: For the insertion sort within a thread, LayoutRight would offer better cache locality for the thread's private list.
    198:   // However, LayoutLeft is preferred for coalesced access across threads during the final weight computation and initialization.
    199:   // Given the random access nature of the sort (divergence), we stick to the default GPU layout (Left).
    200:   Kokkos::View<PetscReal **, Kokkos::LayoutLeft, MemSpace> best_dists_dev("best_dists", n_vert_local, n_obs_vertex);
    201:   Kokkos::View<PetscInt **, Kokkos::LayoutLeft, MemSpace>  best_idxs_dev("best_idxs", n_vert_local, n_obs_vertex);

    203:   /* Copy boundary data to device */
    204:   Kokkos::View<PetscReal *, MemSpace> bd_dev("bd_dev", dim);
    205:   {
    206:     Kokkos::View<PetscReal *, Kokkos::HostSpace> bd_host("bd_host", dim);
    207:     for (PetscInt d = 0; d < dim; ++d) bd_host(d) = bd[d];
    208:     Kokkos::deep_copy(bd_dev, bd_host);
    209:   }

    211:   /* Main Kernel */
    212:   Kokkos::parallel_for(
    213:     "ComputeLocalization", Kokkos::RangePolicy<ExecSpace>(0, n_vert_local), KOKKOS_LAMBDA(const PetscInt i) {
    214:       PetscReal current_max_dist = PETSC_MAX_REAL;

    216:       // Cache vertex coordinates in registers to avoid repeated global memory access
    217:       // dim is small (<= 3), so this fits easily in registers
    218:       PetscReal v_coords[3] = {0.0, 0.0, 0.0};
    219:       for (PetscInt d = 0; d < dim; ++d) v_coords[d] = PetscRealPart(vertex_coords_dev(i, d));

    221:       // Initialize with infinity
    222:       for (PetscInt k = 0; k < n_obs_vertex; ++k) {
    223:         best_dists_dev(i, k) = PETSC_MAX_REAL;
    224:         best_idxs_dev(i, k)  = -1;
    225:       }

    227:       // Iterate over all observations
    228:       for (PetscInt j = 0; j < n_obs_global; ++j) {
    229:         PetscReal dist2 = 0.0;
    230:         for (PetscInt d = 0; d < dim; ++d) {
    231:           PetscReal diff = v_coords[d] - obs_coords_dev(j, d);
    232:           if (bd_dev(d) != 0) { // Periodic boundary
    233:             PetscReal domain_size = bd_dev(d);
    234:             if (diff > 0.5 * domain_size) diff -= domain_size;
    235:             else if (diff < -0.5 * domain_size) diff += domain_size;
    236:           }
    237:           dist2 += diff * diff;
    238:         }

    240:         // Check if this observation is closer than the furthest stored observation
    241:         if (dist2 < current_max_dist) {
    242:           // Insert sorted
    243:           PetscInt pos = n_obs_vertex - 1;
    244:           while (pos > 0 && best_dists_dev(i, pos - 1) > dist2) {
    245:             best_dists_dev(i, pos) = best_dists_dev(i, pos - 1);
    246:             best_idxs_dev(i, pos)  = best_idxs_dev(i, pos - 1);
    247:             pos--;
    248:           }
    249:           best_dists_dev(i, pos) = dist2;
    250:           best_idxs_dev(i, pos)  = j;

    252:           // Update current max distance
    253:           current_max_dist = best_dists_dev(i, n_obs_vertex - 1);
    254:         }
    255:       }
    256:       // Compute weights
    257:       PetscReal radius2 = best_dists_dev(i, n_obs_vertex - 1);
    258:       PetscReal radius  = std::sqrt(radius2);
    259:       radius *= RADIUS_FACTOR;
    260:       if (radius == 0.0) radius = 1.0;

    262:       for (PetscInt k = 0; k < n_obs_vertex; ++k) {
    263:         if (best_idxs_dev(i, k) != -1) {
    264:           PetscReal dist    = std::sqrt(best_dists_dev(i, k));
    265:           indices_dev(i, k) = best_idxs_dev(i, k);
    266:           values_dev(i, k)  = GaspariCohn(dist, radius);
    267:         } else {
    268:           indices_dev(i, k) = -1; // Ignore this entry
    269:           values_dev(i, k)  = 0.0;
    270:         }
    271:       }
    272:     });

    274:   /* Copy back to host and fill matrix */
    275:   // Host views must be LayoutRight for MatSetValues (row-major)
    276:   Kokkos::View<PetscInt **, Kokkos::LayoutRight, Kokkos::HostSpace>    indices_host("indices_host", n_vert_local, n_obs_vertex);
    277:   Kokkos::View<PetscScalar **, Kokkos::LayoutRight, Kokkos::HostSpace> values_host("values_host", n_vert_local, n_obs_vertex);

    279:   // Deep copy will handle layout conversion (transpose) if device views are LayoutLeft
    280:   // Note: Kokkos::deep_copy cannot copy between different layouts if the memory spaces are different (e.g. GPU to Host).
    281:   // We need an intermediate mirror view on the host with the same layout as the device view.
    282:   Kokkos::View<PetscInt **, Kokkos::LayoutLeft, Kokkos::HostSpace>    indices_host_left = Kokkos::create_mirror_view(indices_dev);
    283:   Kokkos::View<PetscScalar **, Kokkos::LayoutLeft, Kokkos::HostSpace> values_host_left  = Kokkos::create_mirror_view(values_dev);

    285:   Kokkos::deep_copy(indices_host_left, indices_dev);
    286:   Kokkos::deep_copy(values_host_left, values_dev);

    288:   // Now copy from LayoutLeft host view to LayoutRight host view
    289:   Kokkos::deep_copy(indices_host, indices_host_left);
    290:   Kokkos::deep_copy(values_host, values_host_left);

    292:   for (PetscInt i = 0; i < n_vert_local; ++i) {
    293:     PetscInt globalRow = rstart + i;
    294:     PetscCall(MatSetValues(*Q, 1, &globalRow, n_obs_vertex, &indices_host(i, 0), &values_host(i, 0), INSERT_VALUES));
    295:   }

    297:   /* Compute mean and std dev of localization radius */
    298:   {
    299:     using FunctorType = RadiusStatsFunctor<decltype(best_dists_dev)>;
    300:     typename FunctorType::value_type result;
    301:     Kokkos::parallel_reduce("ComputeRadiusStats", Kokkos::RangePolicy<ExecSpace>(0, n_vert_local), FunctorType{best_dists_dev, n_obs_vertex}, result);

    303:     if (n_vert_local > 0) {
    304:       double mean   = result.sum / n_vert_local;
    305:       double var    = (result.sq_sum / n_vert_local) - (mean * mean);
    306:       double stddev = (var > 1e-1 * PETSC_SQRT_MACHINE_EPSILON) ? std::sqrt(var) : 0.0;
    307:       PetscCall(PetscInfo((PetscObject)obs_vecs[0], "LETKF localization radius: mean %g, std dev %g\n", mean, stddev));
    308:     }
    309:   }

    311:   /* Cleanup storage */
    312:   for (d = 0; d < dim; ++d) PetscCall(VecDestroy(&obs_vecs[d]));
    313:   PetscCall(PetscFree(obs_vecs));

    315:   /* Assemble matrix */
    316:   PetscCall(MatAssemblyBegin(*Q, MAT_FINAL_ASSEMBLY));
    317:   PetscCall(MatAssemblyEnd(*Q, MAT_FINAL_ASSEMBLY));
    318:   PetscFunctionReturn(PETSC_SUCCESS);
    319: }