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          300 /* Spin up truth to Lorenz-96 attractor (~200 steps sufficient, 300 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:   PetscCall(DMDAGetCorners(l95->da, &xs, NULL, NULL, &xm, NULL, NULL));
 53:   PetscCall(DMGetLocalVector(l95->da, &X_local));
 54:   PetscCall(DMGlobalToLocalBegin(l95->da, X, INSERT_VALUES, X_local));
 55:   PetscCall(DMGlobalToLocalEnd(l95->da, X, INSERT_VALUES, X_local));
 56:   PetscCall(DMDAVecGetArrayRead(l95->da, X_local, &x));
 57:   PetscCall(DMDAVecGetArray(l95->da, F_vec, &f));

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

 61:   PetscCall(DMDAVecRestoreArrayRead(l95->da, X_local, &x));
 62:   PetscCall(DMDAVecRestoreArray(l95->da, F_vec, &f));
 63:   PetscCall(DMRestoreLocalVector(l95->da, &X_local));
 64:   PetscFunctionReturn(PETSC_SUCCESS);
 65: }

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

 74:   PetscFunctionBeginUser;
 75:   PetscCall(PetscNew(&l95));
 76:   l95->da = da;
 77:   l95->n  = n;
 78:   l95->F  = F;
 79:   l95->dt = dt;

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

 92:   *ctx = l95;
 93:   PetscFunctionReturn(PETSC_SUCCESS);
 94: }

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

108: /* Advance a single state vector one TS step. Used by the truth trajectory and as the per-column kernel of Lorenz96Step(). */
109: static PetscErrorCode Lorenz96StepVec(Lorenz96Ctx *l95, Vec x)
110: {
111:   PetscFunctionBeginUser;
112:   PetscCall(TSSetStepNumber(l95->ts, 0));
113:   PetscCall(TSSetTime(l95->ts, 0.0));
114:   PetscCall(TSSolve(l95->ts, x));
115:   PetscFunctionReturn(PETSC_SUCCESS);
116: }

118: /*
119:   Lorenz96Step - Advance every column of the ensemble matrix one time step using Lorenz-96 dynamics.
120:   TS only advances one state at a time, so loop over columns here.
121: */
122: static PetscErrorCode Lorenz96Step(Mat ensemble, PetscCtx ctx)
123: {
124:   Lorenz96Ctx *l95 = (Lorenz96Ctx *)ctx;
125:   PetscInt     n;

127:   PetscFunctionBeginUser;
128:   PetscCall(MatGetSize(ensemble, NULL, &n));
129:   for (PetscInt j = 0; j < n; j++) {
130:     Vec col;

132:     PetscCall(MatDenseGetColumnVec(ensemble, j, &col));
133:     PetscCall(Lorenz96StepVec(l95, col));
134:     PetscCall(MatDenseRestoreColumnVec(ensemble, j, &col));
135:   }
136:   PetscFunctionReturn(PETSC_SUCCESS);
137: }

139: /*
140:   CreateIdentityObservationMatrix - Create identity observation matrix H for Lorenz-96

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

145:   Input Parameter:
146: . n - State dimension (number of grid points)

148:   Output Parameter:
149: . H - Identity observation matrix (n x n), sparse AIJ format (H in P x N)
150: */
151: static PetscErrorCode CreateIdentityObservationMatrix(PetscInt n, Mat *H)
152: {
153:   PetscInt i;

155:   PetscFunctionBeginUser;
156:   /* Create identity observation matrix H (n x n) */
157:   PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, n, n, 1, NULL, 0, NULL, H));
158:   PetscCall(MatSetFromOptions(*H));

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

163:   PetscCall(MatAssemblyBegin(*H, MAT_FINAL_ASSEMBLY));
164:   PetscCall(MatAssemblyEnd(*H, MAT_FINAL_ASSEMBLY));
165:   PetscFunctionReturn(PETSC_SUCCESS);
166: }

168: /*
169:   ValidateParameters - Validate input parameters and apply constraints
170: */
171: static PetscErrorCode ValidateParameters(PetscInt *n, PetscInt *steps, PetscInt *burn, PetscInt *obs_freq, PetscInt *ensemble_size, PetscReal *dt, PetscReal *F, PetscReal *obs_error_std)
172: {
173:   PetscFunctionBeginUser;
174:   PetscCheck(*n > 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "State dimension n must be positive, got %" PetscInt_FMT, *n);
175:   PetscCheck(*steps >= 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number of steps must be non-negative, got %" PetscInt_FMT, *steps);
176:   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);

178:   if (*obs_freq < MIN_OBS_FREQ) {
179:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Observation frequency adjusted from %" PetscInt_FMT " to %" PetscInt_FMT "\n", *obs_freq, (PetscInt)MIN_OBS_FREQ));
180:     *obs_freq = MIN_OBS_FREQ;
181:   }
182:   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));
183:   if (*burn > *steps) {
184:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Burn-in steps (%" PetscInt_FMT ") exceeds total steps (%" PetscInt_FMT "), setting burn = steps\n", *burn, *steps));
185:     *burn = *steps;
186:   }

188:   PetscCheck(*dt > 0.0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Time step dt must be positive, got %g", (double)*dt);
189:   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);
190:   PetscCheck(PetscIsNormalReal(*F), PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Forcing parameter F must be a normal real number");
191:   PetscFunctionReturn(PETSC_SUCCESS);
192: }

194: /*
195:   ComputeRMSE - Compute root mean square error between two vectors
196: */
197: static PetscErrorCode ComputeRMSE(Vec v1, Vec v2, Vec work, PetscInt n, PetscReal *rmse)
198: {
199:   PetscReal norm;

201:   PetscFunctionBeginUser;
202:   PetscCall(VecWAXPY(work, -1.0, v2, v1));
203:   PetscCall(VecNorm(work, NORM_2, &norm));
204:   *rmse = norm / PetscSqrtReal((PetscReal)n);
205:   PetscFunctionReturn(PETSC_SUCCESS);
206: }

208: int main(int argc, char **argv)
209: {
210:   Lorenz96Ctx                 *l95_ctx = NULL, *truth_ctx = NULL;
211:   DM                           da_state;
212:   PetscDA                      da;
213:   Vec                          x0, x_mean, x_forecast;
214:   Vec                          truth_state, rmse_work;
215:   Vec                          observation, obs_noise, obs_error_var;
216:   Mat                          H = NULL; /* Observation operator matrix */
217:   PetscRandom                  rng;
218:   Vec                          xyz[3] = {NULL, NULL, NULL};
219:   Vec                          coord;
220:   PetscDALETKFLocalizationType loc_type;
221:   PetscBool                    radius_set;
222:   const char                  *da_prefix;
223:   PetscInt                     n = DEFAULT_N, steps = DEFAULT_STEPS, burn = DEFAULT_BURN, obs_freq = DEFAULT_OBS_FREQ;
224:   PetscInt                     random_seed = DEFAULT_RANDOM_SEED, ensemble_size = DEFAULT_ENSEMBLE_SIZE;
225:   PetscInt                     n_stat_steps = 0, n_obs_stat_steps = 0, obs_count = 0, step, progress_interval;
226:   PetscReal                    F = DEFAULT_F, dt = DEFAULT_DT, obs_error_std = DEFAULT_OBS_ERROR_STD;
227:   PetscReal                    ensemble_init_std = -1; /* Initial ensemble spread */
228:   PetscReal                    localization_radius;    /* Default 2*domain: effectively no localization with Gaspari-Cohn (max periodic distance is L/2) */
229:   PetscReal                    bd[3]         = {DEFAULT_N, 0, 0};
230:   PetscReal                    rmse_forecast = 0.0, rmse_analysis = 0.0, spread = 0.0;
231:   PetscReal                    sum_rmse_forecast = 0.0, sum_rmse_analysis = 0.0;

233:   PetscFunctionBeginUser;
234:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
235:   /* Kokkos initialization deferred to Phase 5 optimization */

237:   /* Parse command-line options */
238:   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Lorenz-96 Example", NULL);
239:   PetscCall(PetscOptionsInt("-n", "State dimension", "", n, &n, NULL));
240:   PetscCall(PetscOptionsInt("-steps", "Number of time steps", "", steps, &steps, NULL));
241:   PetscCall(PetscOptionsInt("-burn", "Burn-in steps excluded from statistics", "", burn, &burn, NULL));
242:   PetscCall(PetscOptionsInt("-obs_freq", "Observation frequency", "", obs_freq, &obs_freq, NULL));
243:   PetscCall(PetscOptionsReal("-F", "Forcing parameter", "", F, &F, NULL));
244:   PetscCall(PetscOptionsReal("-dt", "Time step size", "", dt, &dt, NULL));
245:   PetscCall(PetscOptionsReal("-obs_error", "Observation error standard deviation", "", obs_error_std, &obs_error_std, NULL));
246:   PetscCall(PetscOptionsReal("-ensemble_init_std", "Initial ensemble spread standard deviation", "", ensemble_init_std, &ensemble_init_std, NULL));
247:   PetscCall(PetscOptionsInt("-random_seed", "Random seed for ensemble perturbations", "", random_seed, &random_seed, NULL));
248:   PetscOptionsEnd();
249:   bd[0]               = (PetscReal)n;
250:   localization_radius = 2.0 * bd[0];

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

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

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

260:   /* Create 1D periodic DM for state space */
261:   PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, n, 1, 2, NULL, &da_state));
262:   PetscCall(DMSetFromOptions(da_state));
263:   PetscCall(DMSetUp(da_state));
264:   PetscCall(DMDASetUniformCoordinates(da_state, 0.0, (PetscReal)n, 0.0, 0.0, 0.0, 0.0));

266:   /* Create Lorenz96 context with reusable TS object */
267:   PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &l95_ctx));
268:   PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &truth_ctx));

270:   /* Initialize random number generator */
271:   PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rng));
272:   {
273:     PetscMPIInt rank;
274:     PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
275:     PetscCall(PetscRandomSetSeed(rng, (unsigned long)(random_seed + rank)));
276:   }
277:   PetscCall(PetscRandomSetFromOptions(rng));
278:   PetscCall(PetscRandomSeed(rng));

280:   /* Initialize state vectors */
281:   PetscCall(DMCreateGlobalVector(da_state, &x0));
282:   PetscCall(PetscRandomSetInterval(rng, -.1 * F, .1 * F));
283:   PetscCall(VecSetRandom(x0, rng));
284:   PetscCall(PetscRandomSetInterval(rng, 0, 1));

286:   /* Initialize truth trajectory */
287:   PetscCall(VecDuplicate(x0, &truth_state));
288:   PetscCall(VecCopy(x0, truth_state));
289:   PetscCall(VecDuplicate(x0, &rmse_work));

291:   /* Spin up truth to get onto attractor */
292:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Spinning up truth for %" PetscInt_FMT " steps...\n", (PetscInt)SPINUP_STEPS));
293:   for (PetscInt k = 0; k < SPINUP_STEPS; k++) PetscCall(Lorenz96StepVec(truth_ctx, truth_state));

295:   /* Initialize observation vectors */
296:   PetscCall(VecDuplicate(x0, &observation));
297:   PetscCall(VecDuplicate(x0, &obs_noise));
298:   PetscCall(VecDuplicate(x0, &obs_error_var));
299:   PetscCall(VecSet(obs_error_var, obs_error_std * obs_error_std));

301:   /* Initialize ensemble statistics vectors */
302:   PetscCall(VecDuplicate(x0, &x_mean));
303:   PetscCall(VecDuplicate(x0, &x_forecast));

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

308:   /* Create and configure PetscDA for ensemble data assimilation */
309:   PetscCall(PetscDACreate(PETSC_COMM_WORLD, &da));
310:   /* Note: ndof defaults to 1 (scalar field) - perfect for Lorenz-96 */
311:   PetscCall(PetscDASetSizes(da, n, n));
312:   PetscCall(PetscDAEnsembleSetSize(da, ensemble_size));
313:   PetscCall(PetscDASetFromOptions(da));
314:   PetscCall(PetscDAEnsembleGetSize(da, &ensemble_size));
315:   PetscCall(PetscDASetUp(da));
316:   PetscCall(PetscDASetObsErrorVariance(da, obs_error_var));

318:   /* Configure localization for LETKF (Q is built lazily on first analysis). PETSCDALETKF is
319:      now the only registered PetscDAType, so we can call the setters unconditionally. */
320:   PetscCall(DMGetCoordinates(da_state, &coord));
321:   PetscCall(DMCreateGlobalVector(da_state, &xyz[0]));
322:   PetscCall(PetscObjectSetName((PetscObject)xyz[0], "x_coordinate"));
323:   PetscCall(VecStrideGather(coord, 0, xyz[0], INSERT_VALUES));
324:   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)da, &da_prefix));
325:   PetscCall(PetscOptionsHasName(NULL, da_prefix, "-petscda_letkf_localization_radius", &radius_set));
326:   if (!radius_set) PetscCall(PetscDALETKFSetLocalizationRadius(da, localization_radius));
327:   PetscCall(PetscDALETKFGetLocalizationRadius(da, &localization_radius));
328:   PetscCall(PetscDALETKFSetLocalizationCoordinates(da, xyz, bd, H));
329:   PetscCall(VecDestroy(&xyz[0]));
330:   PetscCall(PetscDALETKFGetLocalizationType(da, &loc_type));
331:   if (loc_type != PETSCDA_LETKF_LOC_NONE) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Localization configured: %" PetscInt_FMT " vertices, radius=%g\n", n, (double)localization_radius));
332:   else PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Localization disabled (LETKF NONE; equivalent to global ETKF)\n"));

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

337:   /* Print configuration summary */
338:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Lorenz-96 LETKF Example\n"));
339:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "======================\n"));
340:   PetscCall(PetscPrintf(PETSC_COMM_WORLD,
341:                         "  State dimension        : %" PetscInt_FMT "\n"
342:                         "  Ensemble size          : %" PetscInt_FMT "\n"
343:                         "  Forcing parameter (F)  : %.4f\n"
344:                         "  Time step (dt)         : %.4f\n"
345:                         "  Total steps            : %" PetscInt_FMT "\n"
346:                         "  Burn-in steps          : %" PetscInt_FMT "\n"
347:                         "  Spin-up steps          : %" PetscInt_FMT "\n"
348:                         "  Observation frequency  : %" PetscInt_FMT "\n"
349:                         "  Observation noise std  : %.3f\n"
350:                         "  Ensemble init std      : %.3f\n"
351:                         "  Random seed            : %" PetscInt_FMT "\n",
352:                         n, ensemble_size, (double)F, (double)dt, steps, burn, (PetscInt)SPINUP_STEPS, obs_freq, (double)obs_error_std, (double)ensemble_init_std, random_seed));
353:   if (loc_type != PETSCDA_LETKF_LOC_NONE) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "  Localization radius    : %g\n", (double)localization_radius));
354:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n"));

356:   /* Main assimilation cycle: forecast and analysis steps */
357:   for (step = 0; step <= steps; step++) {
358:     PetscReal time = step * dt;

360:     /* Forecast step: compute ensemble mean and forecast RMSE */
361:     PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
362:     PetscCall(VecCopy(x_mean, x_forecast));
363:     PetscCall(ComputeRMSE(x_forecast, truth_state, rmse_work, n, &rmse_forecast));
364:     rmse_analysis = rmse_forecast;

366:     /* Analysis step: assimilate observations when available */
367:     if (step % obs_freq == 0 && step > 0) {
368:       /* Generate synthetic noisy observations from truth */
369:       PetscCall(VecSetRandomGaussian(obs_noise, rng, 0.0, obs_error_std));
370:       PetscCall(VecWAXPY(observation, 1.0, obs_noise, truth_state));

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

375:       /* Compute analysis RMSE */
376:       PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
377:       PetscCall(ComputeRMSE(x_mean, truth_state, rmse_work, n, &rmse_analysis));
378:       obs_count++;
379:     }

381:     /* Accumulate statistics after burn-in period.
382:        Forecast RMSE is accumulated every step; analysis RMSE only at observation times
383:        to avoid conflating forecast and analysis errors. */
384:     if (step >= burn) {
385:       sum_rmse_forecast += rmse_forecast;
386:       n_stat_steps++;
387:       if (step % obs_freq == 0 && step > 0) {
388:         sum_rmse_analysis += rmse_analysis;
389:         n_obs_stat_steps++;
390:       }
391:     }

393:     /* Progress reporting */
394:     if ((step % progress_interval == 0) || (step == steps) || (step == 0)) {
395:       Mat       X_anom;
396:       PetscReal norm_fro;

398:       PetscCall(PetscDAEnsembleComputeAnomalies(da, x_mean, &X_anom));
399:       PetscCall(MatNorm(X_anom, NORM_FROBENIUS, &norm_fro));
400:       spread = norm_fro / PetscSqrtReal((PetscReal)n);
401:       PetscCall(MatDestroy(&X_anom));
402:       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]" : ""));
403:     }

405:     /* Propagate ensemble and truth trajectory */
406:     if (step < steps) {
407:       PetscCall(PetscDAEnsembleForecast(da, Lorenz96Step, l95_ctx));
408:       PetscCall(Lorenz96StepVec(truth_ctx, truth_state));
409:     }
410:   }

412:   /* Report final statistics */
413:   if (n_stat_steps > 0) {
414:     PetscReal avg_rmse_forecast = sum_rmse_forecast / n_stat_steps;
415:     PetscReal avg_rmse_analysis = (n_obs_stat_steps > 0) ? sum_rmse_analysis / n_obs_stat_steps : 0.0;
416:     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));
417:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==================================================\n"));
418:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "  Mean RMSE (forecast) : %.5f\n", (double)avg_rmse_forecast));
419:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "  Mean RMSE (analysis) : %.5f\n", (double)avg_rmse_analysis));
420:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "  Observations used    : %" PetscInt_FMT "\n\n", obs_count));
421:   } else {
422:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nWarning: No post-burn-in statistics collected (burn >= steps)\n\n"));
423:   }

425:   PetscCall(PetscDAView(da, PETSC_VIEWER_STDOUT_WORLD));

427:   /* Cleanup */
428:   PetscCall(MatDestroy(&H));
429:   PetscCall(VecDestroy(&x_forecast));
430:   PetscCall(VecDestroy(&x_mean));
431:   PetscCall(VecDestroy(&obs_error_var));
432:   PetscCall(VecDestroy(&obs_noise));
433:   PetscCall(VecDestroy(&observation));
434:   PetscCall(VecDestroy(&rmse_work));
435:   PetscCall(VecDestroy(&truth_state));
436:   PetscCall(VecDestroy(&x0));
437:   PetscCall(PetscDADestroy(&da));
438:   PetscCall(DMDestroy(&da_state));
439:   PetscCall(Lorenz96ContextDestroy(&l95_ctx));
440:   PetscCall(Lorenz96ContextDestroy(&truth_ctx));
441:   PetscCall(PetscRandomDestroy(&rng));

443:   /* Kokkos finalization deferred to Phase 5 optimization */
444:   PetscCall(PetscFinalize());
445:   return 0;
446: }

448: /*TEST

450:   testset:
451:     requires: !complex
452:     args: -steps 20 -burn 5 -obs_freq 1 -obs_error 1 -petscda_ensemble_size 5 -petscda_type letkf

454:     test:
455:       suffix: letkf_serial

457:     test:
458:       suffix: letkf_loc_none
459:       args: -petscda_letkf_localization_type none

461:     test:
462:       nsize: 2
463:       suffix: letkf_cpu_2rank
464:       args: -petscda_letkf_localization_radius 5.0

466:     test:
467:       nsize: 2
468:       suffix: letkf_loc_none_2rank
469:       args: -petscda_letkf_localization_type none

471:   testset:
472:     requires: kokkos_kernels !complex
473:     args: -steps 20 -burn 5 -obs_freq 1 -obs_error 1 -petscda_ensemble_size 5 -petscda_type letkf -mat_type aijkokkos -dm_vec_type kokkos

475:     test:
476:       nsize: 3
477:       suffix: letkf
478:       args: -info :vec -petscda_letkf_localization_radius 5.0

480:     test:
481:       suffix: letkf_loc_none_kokkos
482:       args: -petscda_letkf_localization_type none

484:     test:
485:       nsize: 3
486:       suffix: letkf_loc_none_kokkos_3rank
487:       args: -petscda_letkf_localization_type none

489:   TEST*/