Actual source code: letkf.h

      1: #pragma once
      2: #include <petsc/private/daimpl.h>
      3: #include <petsc/private/daensembleimpl.h>

      5: typedef struct {
      6:   PetscDA_Ensemble en;
      7:   Vec              mean;
      8:   Vec              y_mean;
      9:   Vec              delta_scaled;
     10:   Vec              w;
     11:   Vec              r_inv_sqrt;
     12:   Mat              Z;
     13:   Mat              S;
     14:   Mat              T_sqrt;
     15:   Mat              w_ones;
     16:   Mat              Q;            /* Localization matrix (n_grid x n_observations_total) Each row has exactly n_obs_vertex non-zeros */
     17:   PetscInt         n_obs_vertex; /* number of local observations per grid point (const now) */
     18:   PetscInt         n_grid;       /* Number of grid points (n_grid = state_size / da->ndof) */
     19:   PetscInt         batch_size;   /* Batch size for GPU processing */

     21:   /* Localization support for MPI */
     22:   IS         obs_is_local;    // Indices of observations needed by this process
     23:   VecScatter obs_scat;        // Scatter context for observations
     24:   Vec        obs_work;        // Local work vector for observations
     25:   Vec        y_mean_work;     // Local work vector for y_mean
     26:   Vec        r_inv_sqrt_work; // Local work vector for r_inv_sqrt
     27:   Mat        Z_work;          // Local work matrix for Z (SeqDense)
     28:   PetscHMapI obs_g2l;         // Map global observation index to local index in obs_work

     30:   /* Device pointers for Q (Kokkos views cast to void*) */
     31:   void *Q_device_i;
     32:   void *Q_device_j; // Holds LOCAL indices into obs_work, not global indices
     33:   void *Q_device_a;

     35:   /* Persistent solver handles and workspace */
     36:   void *solver_handle; // cusolverDnHandle_t / rocblas_handle / sycl::queue*
     37:   void *eigen_work;    // EigenWorkspace*
     38: } PetscDA_LETKF;

     40: #if defined(PETSC_HAVE_KOKKOS_KERNELS)
     41: PETSC_EXTERN PetscErrorCode PetscDALETKFLocalAnalysis(PetscDA, PetscDA_LETKF *, PetscInt, PetscInt, Mat, Vec, Mat, Vec, Vec);
     42: PETSC_EXTERN PetscErrorCode PetscDALETKFLocalAnalysis_GPU(PetscDA, PetscDA_LETKF *, PetscInt, PetscInt, Mat, Vec, Mat, Vec, Vec);
     43: PETSC_EXTERN PetscErrorCode PetscDALETKFSetupLocalization_Kokkos(PetscDA_LETKF *, Mat);
     44: PETSC_EXTERN PetscErrorCode PetscDALETKFDestroyLocalization_Kokkos(PetscDA_LETKF *);
     45: #endif