Actual source code: ex2.c

  1: static char help[] = "Deterministic LETKF example for the Lorenz-96 model. See "
  2:                      "Algorithm 6.4 of \n"
  3:                      "Asch, Bocquet, and Nodet (2016) \"Data Assimilation\" "
  4:                      "(SIAM, doi:10.1137/1.9781611974546).\n\n"
  5:                      "  Expected result: Similar to ETKF with full localization\n\n";

  7: /* Data assimilation framework header (provides PetscDA) */
  8: #include <petscda.h>
  9: /* PETSc DMDA header (provides DM, DMDA functionality) */
 10: #include <petscdmda.h>
 11: #include <petscts.h>
 12: #include <petscvec.h>

 14: /* Default parameter values */
 15: #define DEFAULT_N             40
 16: #define DEFAULT_STEPS         105000
 17: #define DEFAULT_BURN          5000
 18: #define DEFAULT_OBS_FREQ      5
 19: #define DEFAULT_RANDOM_SEED   12345
 20: #define DEFAULT_F             8.0
 21: #define DEFAULT_DT            0.05
 22: #define DEFAULT_OBS_ERROR_STD 1.0
 23: #define DEFAULT_ENSEMBLE_SIZE 30
 24: #define SPINUP_STEPS          1000 /* Spin up truth to Lorenz-96 attractor (~200 steps sufficient, 1000 for safety) */

 26: /* Minimum valid parameter values */
 27: #define MIN_N              1
 28: #define MIN_ENSEMBLE_SIZE  2
 29: #define MIN_OBS_FREQ       1
 30: #define PROGRESS_INTERVALS 10

 32: typedef struct {
 33:   DM        da; /* 1D periodic DM storing the Lorenz-96 state */
 34:   PetscInt  n;  /* State dimension (number of grid points) */
 35:   PetscReal F;  /* Constant forcing term in the Lorenz-96 equations */
 36:   PetscReal dt; /* Integration time step size */
 37:   TS        ts; /* Reusable time stepper for efficiency */
 38: } Lorenz96Ctx;

 40: /*
 41:   Lorenz96RHS - Compute the right-hand side of the Lorenz-96 equations
 42: */
 43: static PetscErrorCode Lorenz96RHS(TS ts, PetscReal t, Vec X, Vec F_vec, PetscCtx ctx)
 44: {
 45:   Lorenz96Ctx       *l95 = (Lorenz96Ctx *)ctx;
 46:   Vec                X_local;
 47:   const PetscScalar *x;
 48:   PetscScalar       *f;
 49:   PetscInt           xs, xm, i;

 51:   PetscFunctionBeginUser;
 52:   (void)ts;
 53:   (void)t;

 55:   PetscCall(DMDAGetCorners(l95->da, &xs, NULL, NULL, &xm, NULL, NULL));
 56:   PetscCall(DMGetLocalVector(l95->da, &X_local));
 57:   PetscCall(DMGlobalToLocalBegin(l95->da, X, INSERT_VALUES, X_local));
 58:   PetscCall(DMGlobalToLocalEnd(l95->da, X, INSERT_VALUES, X_local));
 59:   PetscCall(DMDAVecGetArrayRead(l95->da, X_local, &x));
 60:   PetscCall(DMDAVecGetArray(l95->da, F_vec, &f));

 62:   for (i = xs; i < xs + xm; i++) f[i] = (x[i + 1] - x[i - 2]) * x[i - 1] - x[i] + l95->F;

 64:   PetscCall(DMDAVecRestoreArrayRead(l95->da, X_local, &x));
 65:   PetscCall(DMDAVecRestoreArray(l95->da, F_vec, &f));
 66:   PetscCall(DMRestoreLocalVector(l95->da, &X_local));
 67:   PetscFunctionReturn(PETSC_SUCCESS);
 68: }

 70: /*
 71:   Lorenz96ContextCreate - Create and initialize a Lorenz96 context with reusable TS object
 72: */
 73: static PetscErrorCode Lorenz96ContextCreate(DM da, PetscInt n, PetscReal F, PetscReal dt, Lorenz96Ctx **ctx)
 74: {
 75:   Lorenz96Ctx *l95;

 77:   PetscFunctionBeginUser;
 78:   PetscCall(PetscNew(&l95));
 79:   l95->da = da;
 80:   l95->n  = n;
 81:   l95->F  = F;
 82:   l95->dt = dt;

 84:   PetscCall(TSCreate(PetscObjectComm((PetscObject)da), &l95->ts));
 85:   PetscCall(TSSetProblemType(l95->ts, TS_NONLINEAR));
 86:   PetscCall(TSSetRHSFunction(l95->ts, NULL, Lorenz96RHS, l95));
 87:   PetscCall(TSSetType(l95->ts, TSRK));
 88:   PetscCall(TSRKSetType(l95->ts, TSRK4));
 89:   PetscCall(TSSetTimeStep(l95->ts, dt));
 90:   PetscCall(TSSetMaxSteps(l95->ts, 1));
 91:   PetscCall(TSSetMaxTime(l95->ts, dt));
 92:   PetscCall(TSSetExactFinalTime(l95->ts, TS_EXACTFINALTIME_MATCHSTEP));
 93:   PetscCall(TSSetFromOptions(l95->ts));

 95:   *ctx = l95;
 96:   PetscFunctionReturn(PETSC_SUCCESS);
 97: }

 99: /*
100:   Lorenz96ContextDestroy - Destroy a Lorenz96 context and its TS object
101: */
102: static PetscErrorCode Lorenz96ContextDestroy(Lorenz96Ctx **ctx)
103: {
104:   PetscFunctionBeginUser;
105:   if (!ctx || !*ctx) PetscFunctionReturn(PETSC_SUCCESS);
106:   PetscCall(TSDestroy(&(*ctx)->ts));
107:   PetscCall(PetscFree(*ctx));
108:   PetscFunctionReturn(PETSC_SUCCESS);
109: }

111: /*
112:   Lorenz96Step - Advance state vector one time step using Lorenz-96 dynamics
113: */
114: static PetscErrorCode Lorenz96Step(Vec input, Vec output, PetscCtx ctx)
115: {
116:   Lorenz96Ctx *l95 = (Lorenz96Ctx *)ctx;

118:   PetscFunctionBeginUser;
119:   PetscCall(TSSetStepNumber(l95->ts, 0));
120:   PetscCall(TSSetTime(l95->ts, 0.0));
121:   if (input != output) PetscCall(VecCopy(input, output));
122:   PetscCall(TSSolve(l95->ts, output));
123:   PetscFunctionReturn(PETSC_SUCCESS);
124: }

126: /*
127:   CreateIdentityObservationMatrix - Create identity observation matrix H for Lorenz-96

129:   For the fully observed case, H is an nxn identity matrix representing y = H*x where
130:   each observation corresponds directly to a state variable.

132:   Input Parameter:
133: . n - State dimension (number of grid points)

135:   Output Parameter:
136: . H - Identity observation matrix (n x n), sparse AIJ format (H in P x N)
137: */
138: static PetscErrorCode CreateIdentityObservationMatrix(PetscInt n, Mat *H)
139: {
140:   PetscInt i;

142:   PetscFunctionBeginUser;
143:   /* Create identity observation matrix H (n x n) */
144:   PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, n, n, 1, NULL, 0, NULL, H));
145:   PetscCall(MatSetFromOptions(*H));

147:   /* Set diagonal entries to 1.0 for identity mapping */
148:   for (i = 0; i < n; i++) PetscCall(MatSetValue(*H, i, i, 1.0, INSERT_VALUES));

150:   PetscCall(MatAssemblyBegin(*H, MAT_FINAL_ASSEMBLY));
151:   PetscCall(MatAssemblyEnd(*H, MAT_FINAL_ASSEMBLY));
152:   PetscFunctionReturn(PETSC_SUCCESS);
153: }

155: /*
156:   ValidateParameters - Validate input parameters and apply constraints
157: */
158: static PetscErrorCode ValidateParameters(PetscInt *n, PetscInt *steps, PetscInt *burn, PetscInt *obs_freq, PetscInt *ensemble_size, PetscReal *dt, PetscReal *F, PetscReal *obs_error_std)
159: {
160:   PetscFunctionBeginUser;
161:   PetscCheck(*n > 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "State dimension n must be positive, got %" PetscInt_FMT, *n);
162:   PetscCheck(*steps >= 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number of steps must be non-negative, got %" PetscInt_FMT, *steps);
163:   PetscCheck(*ensemble_size >= MIN_ENSEMBLE_SIZE, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Ensemble size must be at least %" PetscInt_FMT " for meaningful statistics, got %" PetscInt_FMT, (PetscInt)MIN_ENSEMBLE_SIZE, *ensemble_size);

165:   if (*obs_freq < MIN_OBS_FREQ) {
166:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Observation frequency adjusted from %" PetscInt_FMT " to %" PetscInt_FMT "\n", *obs_freq, (PetscInt)MIN_OBS_FREQ));
167:     *obs_freq = MIN_OBS_FREQ;
168:   }
169:   if (*obs_freq > *steps && *steps > 0) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Observation frequency (%" PetscInt_FMT ") > total steps (%" PetscInt_FMT "), no observations will be assimilated.\n", *obs_freq, *steps));
170:   if (*burn > *steps) {
171:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Burn-in steps (%" PetscInt_FMT ") exceeds total steps (%" PetscInt_FMT "), setting burn = steps\n", *burn, *steps));
172:     *burn = *steps;
173:   }

175:   PetscCheck(*dt > 0.0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Time step dt must be positive, got %g", (double)*dt);
176:   PetscCheck(*obs_error_std > 0.0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Observation error std must be positive, got %g", (double)*obs_error_std);
177:   PetscCheck(PetscIsNormalReal(*F), PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Forcing parameter F must be a normal real number");
178:   PetscFunctionReturn(PETSC_SUCCESS);
179: }

181: /*
182:   ComputeRMSE - Compute root mean square error between two vectors
183: */
184: static PetscErrorCode ComputeRMSE(Vec v1, Vec v2, Vec work, PetscInt n, PetscReal *rmse)
185: {
186:   PetscReal norm;

188:   PetscFunctionBeginUser;
189:   PetscCall(VecWAXPY(work, -1.0, v2, v1));
190:   PetscCall(VecNorm(work, NORM_2, &norm));
191:   *rmse = norm / PetscSqrtReal((PetscReal)n);
192:   PetscFunctionReturn(PETSC_SUCCESS);
193: }

195: /*
196:   CreateLocalizationMatrix - Create and initialize full localization matrix Q

198:   For the fully observed case, Q is a dense nxn
199:   matrix with all entries = 1.0, meaning each vertex uses all observations.
200: */
201: static PetscErrorCode CreateLocalizationMatrix(PetscInt n, Mat *Q)
202: {
203:   PetscInt i, j;

205:   PetscFunctionBeginUser;
206:   /* Create Q matrix (n x n for identity observation operator)
207:      Each row will have exactly const non-zeros -- this can be relaxed */
208:   PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, n, n, n, NULL, 0, NULL, Q));
209:   PetscCall(MatSetFromOptions(*Q));

211:   /* Initialize with full localization (all weights = 1.0)
212:      Each vertex i uses all n observations */
213:   for (i = 0; i < n; i++) {
214:     for (j = 0; j < n; j++) PetscCall(MatSetValue(*Q, i, j, 1.0, INSERT_VALUES));
215:   }
216:   PetscCall(MatAssemblyBegin(*Q, MAT_FINAL_ASSEMBLY));
217:   PetscCall(MatAssemblyEnd(*Q, MAT_FINAL_ASSEMBLY));
218:   PetscFunctionReturn(PETSC_SUCCESS);
219: }

221: int main(int argc, char **argv)
222: {
223:   /* Configuration parameters */
224:   PetscInt  n             = DEFAULT_N;
225:   PetscInt  steps         = DEFAULT_STEPS;
226:   PetscInt  burn          = DEFAULT_BURN;
227:   PetscInt  obs_freq      = DEFAULT_OBS_FREQ;
228:   PetscInt  random_seed   = DEFAULT_RANDOM_SEED;
229:   PetscInt  ensemble_size = DEFAULT_ENSEMBLE_SIZE, n_obs_vertex = 7;
230:   PetscReal F                     = DEFAULT_F;
231:   PetscReal dt                    = DEFAULT_DT;
232:   PetscReal obs_error_std         = DEFAULT_OBS_ERROR_STD;
233:   PetscReal ensemble_init_std     = -1; /* Initial ensemble spread */
234:   PetscBool use_fake_localization = PETSC_FALSE, isletkf;
235:   PetscReal bd[3]                 = {DEFAULT_N, 0, 0};

237:   /* PETSc objects */
238:   Lorenz96Ctx *l95_ctx = NULL, *truth_ctx = NULL;
239:   DM           da_state;
240:   PetscDA      da;
241:   Vec          x0, x_mean, x_forecast;
242:   Vec          truth_state, rmse_work;
243:   Vec          observation, obs_noise, obs_error_var;
244:   PetscRandom  rng;
245:   Mat          Q = NULL; /* Localization matrix */
246:   Mat          H = NULL; /* Observation operator matrix */

248:   /* Statistics tracking */
249:   PetscReal rmse_forecast = 0.0, rmse_analysis = 0.0, spread = 0.0;
250:   PetscReal sum_rmse_forecast = 0.0, sum_rmse_analysis = 0.0;
251:   PetscInt  n_stat_steps = 0, n_obs_stat_steps = 0;
252:   PetscInt  obs_count = 0;
253:   PetscInt  step, progress_interval;

255:   PetscFunctionBeginUser;
256:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
257:   /* Kokkos initialization deferred to Phase 5 optimization */

259:   /* Parse command-line options */
260:   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Lorenz-96 Example", NULL);
261:   PetscCall(PetscOptionsInt("-n", "State dimension", "", n, &n, NULL));
262:   bd[0] = (PetscReal)n;
263:   PetscCall(PetscOptionsInt("-steps", "Number of time steps", "", steps, &steps, NULL));
264:   PetscCall(PetscOptionsInt("-burn", "Burn-in steps excluded from statistics", "", burn, &burn, NULL));
265:   PetscCall(PetscOptionsInt("-obs_freq", "Observation frequency", "", obs_freq, &obs_freq, NULL));
266:   PetscCall(PetscOptionsReal("-F", "Forcing parameter", "", F, &F, NULL));
267:   PetscCall(PetscOptionsReal("-dt", "Time step size", "", dt, &dt, NULL));
268:   PetscCall(PetscOptionsReal("-obs_error", "Observation error standard deviation", "", obs_error_std, &obs_error_std, NULL));
269:   PetscCall(PetscOptionsReal("-ensemble_init_std", "Initial ensemble spread standard deviation", "", ensemble_init_std, &ensemble_init_std, NULL));
270:   PetscCall(PetscOptionsInt("-random_seed", "Random seed for ensemble perturbations", "", random_seed, &random_seed, NULL));
271:   PetscCall(PetscOptionsBool("-use_fake_localization", "Use fake localization matrix", "", use_fake_localization, &use_fake_localization, NULL));
272:   if (!use_fake_localization) PetscCall(PetscOptionsInt("-n_obs_vertex", "Number of observations per vertex", "", n_obs_vertex, &n_obs_vertex, NULL));
273:   else n_obs_vertex = n; /* fully observed */
274:   PetscOptionsEnd();

276:   if (ensemble_init_std < 0) ensemble_init_std = obs_error_std;

278:   /* Validate and constrain parameters */
279:   PetscCall(ValidateParameters(&n, &steps, &burn, &obs_freq, &ensemble_size, &dt, &F, &obs_error_std));

281:   /* Calculate progress reporting interval */
282:   progress_interval = (steps >= PROGRESS_INTERVALS) ? (steps / PROGRESS_INTERVALS) : 1;

284:   /* Create 1D periodic DM for state space */
285:   PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, n, 1, 2, NULL, &da_state));
286:   PetscCall(DMSetFromOptions(da_state));
287:   PetscCall(DMSetUp(da_state));
288:   PetscCall(DMDASetUniformCoordinates(da_state, 0.0, (PetscReal)n, 0.0, 0.0, 0.0, 0.0));

290:   /* Create Lorenz96 context with reusable TS object */
291:   PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &l95_ctx));
292:   PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &truth_ctx));

294:   /* Initialize random number generator */
295:   PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rng));
296:   {
297:     PetscMPIInt rank;
298:     PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
299:     PetscCall(PetscRandomSetSeed(rng, (unsigned long)(random_seed + rank)));
300:   }
301:   PetscCall(PetscRandomSetFromOptions(rng));
302:   PetscCall(PetscRandomSeed(rng));

304:   /* Initialize state vectors */
305:   PetscCall(DMCreateGlobalVector(da_state, &x0));
306:   PetscCall(PetscRandomSetInterval(rng, -.1 * F, .1 * F));
307:   PetscCall(VecSetRandom(x0, rng));
308:   PetscCall(PetscRandomSetInterval(rng, 0, 1));

310:   /* Initialize truth trajectory */
311:   PetscCall(VecDuplicate(x0, &truth_state));
312:   PetscCall(VecCopy(x0, truth_state));
313:   PetscCall(VecDuplicate(x0, &rmse_work));

315:   /* Spin up truth to get onto attractor */
316:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Spinning up truth for %" PetscInt_FMT " steps...\n", (PetscInt)SPINUP_STEPS));
317:   for (int k = 0; k < SPINUP_STEPS; k++) PetscCall(Lorenz96Step(truth_state, truth_state, truth_ctx));

319:   /* Initialize observation vectors */
320:   PetscCall(VecDuplicate(x0, &observation));
321:   PetscCall(VecDuplicate(x0, &obs_noise));
322:   PetscCall(VecDuplicate(x0, &obs_error_var));
323:   PetscCall(VecSet(obs_error_var, obs_error_std * obs_error_std));

325:   /* Initialize ensemble statistics vectors */
326:   PetscCall(VecDuplicate(x0, &x_mean));
327:   PetscCall(VecDuplicate(x0, &x_forecast));

329:   /* Create identity observation matrix H */
330:   PetscCall(CreateIdentityObservationMatrix(n, &H));

332:   /* Create and configure PetscDA for ensemble data assimilation */
333:   PetscCall(PetscDACreate(PETSC_COMM_WORLD, &da));
334:   PetscCall(PetscDASetType(da, PETSCDALETKF)); /* Set LETKF type */
335:   /* Note: ndof defaults to 1 (scalar field) - perfect for Lorenz-96 */
336:   PetscCall(PetscDASetSizes(da, n, n));
337:   PetscCall(PetscDAEnsembleSetSize(da, ensemble_size));
338:   PetscCall(PetscDASetFromOptions(da));
339:   PetscCall(PetscDAEnsembleGetSize(da, &ensemble_size));
340:   PetscCall(PetscDASetUp(da));
341:   PetscCall(PetscDASetObsErrorVariance(da, obs_error_var));
342:   PetscCall(PetscObjectTypeCompare((PetscObject)da, PETSCDALETKF, &isletkf));

344:   /* Create and set localization matrix Q */
345:   if (!use_fake_localization && isletkf) {
346:     Vec      Vecxyz[3] = {NULL, NULL, NULL};
347:     Vec      coord;
348:     PetscInt d;

350:     PetscCall(DMGetCoordinates(da_state, &coord));
351:     for (d = 0; d < 1; d++) {
352:       PetscCall(DMCreateGlobalVector(da_state, &Vecxyz[d]));
353:       PetscCall(PetscObjectSetName((PetscObject)Vecxyz[d], "x_coordinate"));
354:       PetscCall(VecStrideGather(coord, d, Vecxyz[d], INSERT_VALUES));
355:     }
356:     PetscCall(PetscDALETKFGetLocalizationMatrix(n_obs_vertex, 1, Vecxyz, bd, H, &Q));
357:     PetscCall(PetscDALETKFSetObsPerVertex(da, n_obs_vertex));
358:     PetscCall(VecDestroy(&Vecxyz[0]));
359:   } else {
360:     PetscCall(CreateLocalizationMatrix(n, &Q));
361:     if (isletkf) {
362:       PetscCall(PetscDALETKFSetObsPerVertex(da, n_obs_vertex)); // fully observed
363:     }
364:   }
365:   PetscCall(PetscDALETKFSetLocalization(da, Q, H));
366:   if (isletkf) {
367:     PetscInt n_obs_vertex_actual;
368:     PetscCall(PetscDALETKFGetObsPerVertex(da, &n_obs_vertex_actual));
369:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Localization matrix Q created: %" PetscInt_FMT " x %" PetscInt_FMT "\n", n, n_obs_vertex_actual));
370:   }

372:   /* Initialize ensemble members from spun-up truth state */
373:   PetscCall(PetscDAEnsembleInitialize(da, truth_state, ensemble_init_std, rng));

375:   PetscCall(PetscDAViewFromOptions(da, NULL, "-petscda_view"));

377:   /* Print configuration summary */
378:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Lorenz-96 LETKF Example\n"));
379:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "======================\n"));
380:   PetscCall(PetscPrintf(PETSC_COMM_WORLD,
381:                         "  State dimension        : %" PetscInt_FMT "\n"
382:                         "  Ensemble size          : %" PetscInt_FMT "\n"
383:                         "  Forcing parameter (F)  : %.4f\n"
384:                         "  Time step (dt)         : %.4f\n"
385:                         "  Total steps            : %" PetscInt_FMT "\n"
386:                         "  Burn-in steps          : %" PetscInt_FMT "\n"
387:                         "  Spin-up steps          : %" PetscInt_FMT "\n"
388:                         "  Observation frequency  : %" PetscInt_FMT "\n"
389:                         "  Observation noise std  : %.3f\n"
390:                         "  Ensemble init std      : %.3f\n"
391:                         "  Random seed            : %" PetscInt_FMT "\n"
392:                         "  Localization (obs/vert): %" PetscInt_FMT " \n\n",
393:                         n, ensemble_size, (double)F, (double)dt, steps, burn, SPINUP_STEPS, obs_freq, (double)obs_error_std, (double)ensemble_init_std, random_seed, n_obs_vertex));

395:   /* Main assimilation cycle: forecast and analysis steps */
396:   for (step = 0; step <= steps; step++) {
397:     PetscReal time = step * dt;

399:     /* Forecast step: compute ensemble mean and forecast RMSE */
400:     PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
401:     PetscCall(VecCopy(x_mean, x_forecast));
402:     PetscCall(ComputeRMSE(x_forecast, truth_state, rmse_work, n, &rmse_forecast));
403:     rmse_analysis = rmse_forecast;

405:     /* Analysis step: assimilate observations when available */
406:     if (step % obs_freq == 0 && step > 0) {
407:       /* Generate synthetic noisy observations from truth */
408:       PetscCall(VecSetRandomGaussian(obs_noise, rng, 0.0, obs_error_std));
409:       PetscCall(VecWAXPY(observation, 1.0, obs_noise, truth_state));

411:       /* Perform LETKF analysis with observation matrix H */
412:       PetscCall(PetscDAEnsembleAnalysis(da, observation, H));

414:       /* Compute analysis RMSE */
415:       PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
416:       PetscCall(ComputeRMSE(x_mean, truth_state, rmse_work, n, &rmse_analysis));
417:       obs_count++;
418:     }

420:     /* Accumulate statistics after burn-in period.
421:        Forecast RMSE is accumulated every step; analysis RMSE only at observation times
422:        to avoid conflating forecast and analysis errors. */
423:     if (step >= burn) {
424:       sum_rmse_forecast += rmse_forecast;
425:       n_stat_steps++;
426:       if (step % obs_freq == 0 && step > 0) {
427:         sum_rmse_analysis += rmse_analysis;
428:         n_obs_stat_steps++;
429:       }
430:     }

432:     /* Progress reporting */
433:     if ((step % progress_interval == 0) || (step == steps) || (step == 0)) {
434:       Mat       X_anom;
435:       PetscReal norm_fro;

437:       PetscCall(PetscDAEnsembleComputeAnomalies(da, x_mean, &X_anom));
438:       PetscCall(MatNorm(X_anom, NORM_FROBENIUS, &norm_fro));
439:       spread = norm_fro / PetscSqrtReal((PetscReal)n);
440:       PetscCall(MatDestroy(&X_anom));
441:       PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Step %4" PetscInt_FMT ", time %6.3f  RMSE_forecast %.5f  RMSE_analysis %.5f  Spread %.5f%s\n", step, (double)time, (double)rmse_forecast, (double)rmse_analysis, (double)spread, (step < burn) ? " [burn-in]" : ""));
442:     }

444:     /* Propagate ensemble and truth trajectory */
445:     if (step < steps) {
446:       PetscCall(PetscDAEnsembleForecast(da, Lorenz96Step, l95_ctx));
447:       PetscCall(Lorenz96Step(truth_state, truth_state, truth_ctx));
448:     }
449:   }

451:   /* Report final statistics */
452:   if (n_stat_steps > 0) {
453:     PetscReal avg_rmse_forecast = sum_rmse_forecast / n_stat_steps;
454:     PetscReal avg_rmse_analysis = (n_obs_stat_steps > 0) ? sum_rmse_analysis / n_obs_stat_steps : 0.0;
455:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nStatistics (%" PetscInt_FMT " forecast steps, %" PetscInt_FMT " analysis steps post-burn-in):\n", n_stat_steps, n_obs_stat_steps));
456:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==================================================\n"));
457:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "  Mean RMSE (forecast) : %.5f\n", (double)avg_rmse_forecast));
458:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "  Mean RMSE (analysis) : %.5f\n", (double)avg_rmse_analysis));
459:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "  Observations used    : %" PetscInt_FMT "\n\n", obs_count));
460:   } else {
461:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nWarning: No post-burn-in statistics collected (burn >= steps)\n\n"));
462:   }

464:   /* Cleanup */
465:   PetscCall(MatDestroy(&H));
466:   PetscCall(MatDestroy(&Q));
467:   PetscCall(VecDestroy(&x_forecast));
468:   PetscCall(VecDestroy(&x_mean));
469:   PetscCall(VecDestroy(&obs_error_var));
470:   PetscCall(VecDestroy(&obs_noise));
471:   PetscCall(VecDestroy(&observation));
472:   PetscCall(VecDestroy(&rmse_work));
473:   PetscCall(VecDestroy(&truth_state));
474:   PetscCall(VecDestroy(&x0));
475:   PetscCall(PetscDADestroy(&da));
476:   PetscCall(DMDestroy(&da_state));
477:   PetscCall(Lorenz96ContextDestroy(&l95_ctx));
478:   PetscCall(Lorenz96ContextDestroy(&truth_ctx));
479:   PetscCall(PetscRandomDestroy(&rng));

481:   /* Kokkos finalization deferred to Phase 5 optimization */
482:   PetscCall(PetscFinalize());
483:   return 0;
484: }

486: /*TEST

488:   testset:
489:     requires: kokkos_kernels !complex
490:     args: -steps 112 -burn 10 -obs_freq 1 -obs_error 1 -petscda_view -petscda_ensemble_size 5

492:     test:
493:       suffix: chol
494:       args: -petscda_type letkf -petscda_ensemble_sqrt_type eigen

496:     test:
497:       nsize: 3
498:       suffix: letkf
499:       args: -petscda_type letkf -mat_type aijkokkos -dm_vec_type kokkos -info :vec -n_obs_vertex 5

501:     test:
502:       suffix: etkf
503:       args: -petscda_type etkf -petscda_ensemble_sqrt_type eigen

505:     test:
506:       suffix: etkf2
507:       args: -petscda_type etkf -petscda_ensemble_sqrt_type cholesky

509:   TEST*/