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