Actual source code: ex1.c
1: /* Data assimilation framework header (provides PetscDA) */
2: #include <petscda.h>
3: /* PETSc DMDA header (provides DM, DMDA functionality) */
4: #include <petscdmda.h>
5: #include <petscts.h>
6: #include <petscvec.h>
8: static char help[] = "Deterministic ETKF example for the Lorenz-96 model. See "
9: "Algorithm 6.4 of \n"
10: "Asch, Bocquet, and Nodet (2016) \"Data Assimilation\" "
11: "(SIAM, doi:10.1137/1.9781611974546).\n\n"
12: " Note: Asch et al. run Lorenz-96 for 100,000 steps, but our ensemble collapses.\n\n";
14: /* \begin{algorithm}
15: \caption{Ensemble Transform Kalman Filter (ETKF) - Deterministic}
16: \begin{algorithmic}[1]
17: \State \textbf{Initialize:} Ensemble $\mathbf{E}_0 = [\mathbf{x}_0^{(1)}, \ldots, \mathbf{x}_0^{(m)}]$ \Comment{$\mathbf{E}_0 \in \mathbb{R}^{n \times m}$}
18: \For{$k = 1, 2, \ldots, K$}
19: \State \textbf{// Forecast Step}
20: \For{$i = 1$ to $m$}
21: \State $\mathbf{x}_k^{f,(i)} = \mathcal{M}_{k-1}(\mathbf{x}_{k-1}^{a,(i)})$ \Comment{$\mathbf{x}_k^{f,(i)} \in \mathbb{R}^{n}$}
22: \EndFor
23: \State Assemble forecast matrix: $\mathbf{E}_k^f = [\mathbf{x}_k^{f,(1)}, \ldots, \mathbf{x}_k^{f,(m)}]$ \Comment{$\mathbf{E}_k^f \in \mathbb{R}^{n \times m}$}
24: \State
25: \State \textbf{// Analysis Step (if observation available)}
26: \If{observation $\mathbf{y}_k$ available}
27: \State Compute ensemble mean: $\overline{\mathbf{x}}_k^f = \frac{1}{m}\sum_{i=1}^m \mathbf{x}_k^{f,(i)}$ \Comment{$\overline{\mathbf{x}}_k^f \in \mathbb{R}^{n}$}
28: \State Compute anomalies: $\mathbf{X}_k = \frac{1}{\sqrt{m-1}}(\mathbf{E}_k^f - \overline{\mathbf{x}}_k^f \mathbf{1}^T)$ \Comment{$\mathbf{X}_k \in \mathbb{R}^{n \times m}$}
29: \State Apply observation operator: $\mathbf{Z}_k = \mathcal{H}(\mathbf{E}_k^f)$ \Comment{$\mathbf{Z}_k \in \mathbb{R}^{b \times m}$}
30: \State Compute obs ensemble mean: $\overline{\mathbf{y}}_k = \frac{1}{m}\sum_{i=1}^m \mathcal{H}(\mathbf{x}_k^{f,(i)})$ \Comment{$\overline{\mathbf{y}}_k \in \mathbb{R}^{b}$}
31: \State Compute obs anomalies: $\mathbf{S}_k = \frac{1}{\sqrt{m-1}}\mathbf{R}^{-1/2}(\mathbf{Z}_k - \overline{\mathbf{y}}_k\mathbf{1}^T)$ \Comment{$\mathbf{S}_k \in \mathbb{R}^{b \times m}$}
32: \State Compute raw innovation: $\boldsymbol{\delta}_k = \mathbf{y}_k - \overline{\mathbf{y}}_k$ \Comment{$\boldsymbol{\delta}_k \in \mathbb{R}^{b}$}
33: \State Whiten innovation: $\tilde{\boldsymbol{\delta}}_k = \mathbf{R}^{-1/2}\boldsymbol{\delta}_k$ \Comment{$\tilde{\boldsymbol{\delta}}_k \in \mathbb{R}^{b}$}
34: \State Compute transform matrix: $\mathbf{T}_k = (\mathbf{I}_m + \mathbf{S}_k^T\mathbf{S}_k)^{-1}$ \Comment{$\mathbf{T}_k \in \mathbb{R}^{m \times m}$}
35: \State Compute weight vector: $\mathbf{w}_k = \mathbf{T}_k \mathbf{S}_k^T \tilde{\boldsymbol{\delta}}_k$ \Comment{$\mathbf{w}_k \in \mathbb{R}^{m}$}
36: \State Form deterministic map: $\mathbf{G}_k = \mathbf{w}_k\mathbf{1}^T + \sqrt{m-1}\,\mathbf{T}_k^{1/2}$ \Comment{$\mathbf{G}_k \in \mathbb{R}^{m \times m}$, $\mathbf{U} = \mathbf{I}_m$}
37: \State Update ensemble: $\mathbf{E}_k^a = \overline{\mathbf{x}}_k^f\mathbf{1}^T + \mathbf{X}_k \mathbf{G}_k$ \Comment{$\mathbf{E}_k^a \in \mathbb{R}^{n \times m}$}
38: \Else
39: \State $\mathbf{E}_k^a = \mathbf{E}_k^f$ \Comment{$\mathbf{E}_k^a \in \mathbb{R}^{n \times m}$}
40: \EndIf
41: \EndFor
42: \end{algorithmic}
43: \end{algorithm}
44: */
46: /* Default parameter values */
47: #define DEFAULT_N 40
48: #define DEFAULT_STEPS 105000
49: #define DEFAULT_BURN 5000
50: #define DEFAULT_OBS_FREQ 1
51: #define DEFAULT_RANDOM_SEED 12345
52: #define DEFAULT_F 8.0
53: #define DEFAULT_DT 0.05
54: #define DEFAULT_OBS_ERROR_STD 1.0
55: #define DEFAULT_ENSEMBLE_SIZE 30
56: #define SPINUP_STEPS 0
58: /* Minimum valid parameter values */
59: #define MIN_N 1
60: #define MIN_ENSEMBLE_SIZE 2
61: #define MIN_OBS_FREQ 1
62: #define PROGRESS_INTERVALS 10
64: typedef struct {
65: DM da; /* 1D periodic DM storing the Lorenz-96 state */
66: PetscInt n; /* State dimension (number of grid points) */
67: PetscReal F; /* Constant forcing term in the Lorenz-96 equations */
68: PetscReal dt; /* Integration time step size */
69: TS ts; /* Reusable time stepper for efficiency */
70: } Lorenz96Ctx;
72: /*
73: Lorenz96RHS - Compute the right-hand side of the Lorenz-96 equations
75: Input Parameters:
76: + ts - The time-stepping context (unused but required by interface)
77: . t - Current time (unused but required by interface)
78: . X - State vector
79: - ctx - User context (Lorenz96Ctx)
81: Output Parameter:
82: . F_vec - RHS vector (tendency)
83: */
84: static PetscErrorCode Lorenz96RHS(TS ts, PetscReal t, Vec X, Vec F_vec, PetscCtx ctx)
85: {
86: Lorenz96Ctx *l95 = (Lorenz96Ctx *)ctx;
87: Vec X_local;
88: const PetscScalar *x;
89: PetscScalar *f;
90: PetscInt xs, xm, i;
92: PetscFunctionBeginUser;
93: (void)ts; /* Mark as intentionally unused to avoid compiler warnings */
94: (void)t;
96: /* Work with a local (ghosted) vector so the Lorenz-96 stencil has the
97: * required neighbors for periodic boundary conditions. */
98: PetscCall(DMDAGetCorners(l95->da, &xs, NULL, NULL, &xm, NULL, NULL));
99: PetscCall(DMGetLocalVector(l95->da, &X_local));
100: PetscCall(DMGlobalToLocalBegin(l95->da, X, INSERT_VALUES, X_local));
101: PetscCall(DMGlobalToLocalEnd(l95->da, X, INSERT_VALUES, X_local));
102: PetscCall(DMDAVecGetArrayRead(l95->da, X_local, &x));
103: PetscCall(DMDAVecGetArray(l95->da, F_vec, &f));
105: /* Standard Lorenz-96 tendency: (x_{i+1} - x_{i-2}) * x_{i-1} - x_i + F. */
106: for (i = xs; i < xs + xm; i++) f[i] = (x[i + 1] - x[i - 2]) * x[i - 1] - x[i] + l95->F;
108: PetscCall(DMDAVecRestoreArrayRead(l95->da, X_local, &x));
109: PetscCall(DMDAVecRestoreArray(l95->da, F_vec, &f));
110: PetscCall(DMRestoreLocalVector(l95->da, &X_local));
111: PetscFunctionReturn(PETSC_SUCCESS);
112: }
114: /*
115: Lorenz96ContextCreate - Create and initialize a Lorenz96 context with reusable TS object
117: Input Parameters:
118: + da - DM for state space
119: . n - State dimension
120: . F - Forcing parameter
121: . dt - Time step size
123: Output Parameter:
124: . ctx - Initialized Lorenz96 context
125: */
126: static PetscErrorCode Lorenz96ContextCreate(DM da, PetscInt n, PetscReal F, PetscReal dt, Lorenz96Ctx **ctx)
127: {
128: Lorenz96Ctx *l95;
130: PetscFunctionBeginUser;
131: PetscCall(PetscNew(&l95));
132: l95->da = da;
133: l95->n = n;
134: l95->F = F;
135: l95->dt = dt;
137: /* Create and configure a reusable time stepper to avoid repeated allocation/deallocation */
138: PetscCall(TSCreate(PetscObjectComm((PetscObject)da), &l95->ts));
139: PetscCall(TSSetProblemType(l95->ts, TS_NONLINEAR));
140: PetscCall(TSSetRHSFunction(l95->ts, NULL, Lorenz96RHS, l95));
141: PetscCall(TSSetType(l95->ts, TSRK));
142: PetscCall(TSRKSetType(l95->ts, TSRK4));
143: PetscCall(TSSetTimeStep(l95->ts, dt));
144: PetscCall(TSSetMaxTime(l95->ts, dt));
145: PetscCall(TSSetExactFinalTime(l95->ts, TS_EXACTFINALTIME_MATCHSTEP));
146: PetscCall(TSSetFromOptions(l95->ts));
148: *ctx = l95;
149: PetscFunctionReturn(PETSC_SUCCESS);
150: }
152: /*
153: Lorenz96ContextDestroy - Destroy a Lorenz96 context and its TS object
154: */
155: static PetscErrorCode Lorenz96ContextDestroy(Lorenz96Ctx **ctx)
156: {
157: PetscFunctionBeginUser;
158: if (!ctx || !*ctx) PetscFunctionReturn(PETSC_SUCCESS);
159: PetscCall(TSDestroy(&(*ctx)->ts));
160: PetscCall(PetscFree(*ctx));
161: PetscFunctionReturn(PETSC_SUCCESS);
162: }
164: /*
165: Lorenz96Step - Advance state vector one time step using Lorenz-96 dynamics
167: Input Parameters:
168: + input - state vector to be advanced
169: - ctx - Lorenz96 context (contains reusable TS)
171: Output Parameter:
172: . output - state vector after one time step
174: Notes:
175: Uses a single explicit RK4 step with the pre-configured TS object for efficiency.
176: */
177: static PetscErrorCode Lorenz96Step(Vec input, Vec output, PetscCtx ctx)
178: {
179: Lorenz96Ctx *l95 = (Lorenz96Ctx *)ctx;
181: PetscFunctionBeginUser;
182: /* Reset the TS time for each integration */
183: PetscCall(TSSetTime(l95->ts, 0.0));
184: if (input != output) PetscCall(VecCopy(input, output));
185: PetscCall(TSSolve(l95->ts, output));
186: PetscFunctionReturn(PETSC_SUCCESS);
187: }
189: /*
190: CreateIdentityObservationMatrix - Create identity observation matrix H for Lorenz-96
192: For the fully observed case, H is an nxn identity matrix representing y = H*x where
193: each observation corresponds directly to a state variable.
195: Input Parameter:
196: . n - State dimension (number of grid points)
198: Output Parameter:
199: . H - Identity observation matrix (n x n), sparse AIJ format
200: */
201: static PetscErrorCode CreateIdentityObservationMatrix(PetscInt n, Mat *H)
202: {
203: PetscInt i;
205: PetscFunctionBeginUser;
206: /* Create identity observation matrix H (n x n) */
207: PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, n, n, 1, NULL, 0, NULL, H));
208: PetscCall(MatSetFromOptions(*H));
210: /* Set diagonal entries to 1.0 for identity mapping */
211: for (i = 0; i < n; i++) PetscCall(MatSetValue(*H, i, i, 1.0, INSERT_VALUES));
213: PetscCall(MatAssemblyBegin(*H, MAT_FINAL_ASSEMBLY));
214: PetscCall(MatAssemblyEnd(*H, MAT_FINAL_ASSEMBLY));
215: PetscFunctionReturn(PETSC_SUCCESS);
216: }
218: /*
219: ValidateParameters - Validate input parameters and apply constraints
221: Input/Output Parameters:
222: + n - State dimension
223: . steps - Number of time steps
224: . burn - Burn-in steps
225: . obs_freq - Observation frequency
226: . ensemble_size - Ensemble size
227: . dt - Time step
228: . F - Forcing parameter
229: - obs_error_std - Observation error standard deviation
231: This routine should be reduced/eliminated. It is the job of PETSc functions to validate input, not the tutorials code
232: */
233: static PetscErrorCode ValidateParameters(PetscInt *n, PetscInt *steps, PetscInt *burn, PetscInt *obs_freq, PetscInt *ensemble_size, PetscReal *dt, PetscReal *F, PetscReal *obs_error_std)
234: {
235: PetscFunctionBeginUser;
236: /* Validate and constrain integer parameters */
237: PetscCheck(*n > 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "State dimension n must be positive, got %" PetscInt_FMT, *n);
238: PetscCheck(*steps >= 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number of steps must be non-negative, got %" PetscInt_FMT, *steps);
239: 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);
241: /* Apply constraints */
242: if (*obs_freq < MIN_OBS_FREQ) {
243: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Observation frequency adjusted from %" PetscInt_FMT " to %" PetscInt_FMT "\n", *obs_freq, (PetscInt)MIN_OBS_FREQ));
244: *obs_freq = MIN_OBS_FREQ;
245: }
246: 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));
247: if (*burn > *steps) {
248: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Burn-in steps (%" PetscInt_FMT ") exceeds total steps (%" PetscInt_FMT "), setting burn = steps\n", *burn, *steps));
249: *burn = *steps;
250: }
252: /* Validate real-valued parameters */
253: PetscCheck(*dt > 0.0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Time step dt must be positive, got %g", (double)*dt);
254: 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);
255: PetscCheck(PetscIsNormalReal(*F), PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Forcing parameter F must be a normal real number");
256: PetscFunctionReturn(PETSC_SUCCESS);
257: }
259: /*
260: ComputeRMSE - Compute root mean square error between two vectors
262: Input Parameters:
263: + v1 - First vector
264: . v2 - Second vector
265: - n - Vector dimension (for normalization)
267: Output Parameter:
268: . rmse - Root mean square error
269: */
270: static PetscErrorCode ComputeRMSE(Vec v1, Vec v2, Vec work, PetscInt n, PetscReal *rmse)
271: {
272: PetscReal norm;
274: PetscFunctionBeginUser;
275: PetscCall(VecWAXPY(work, -1.0, v2, v1));
276: PetscCall(VecNorm(work, NORM_2, &norm));
277: *rmse = norm / PetscSqrtReal((PetscReal)n);
278: PetscFunctionReturn(PETSC_SUCCESS);
279: }
281: int main(int argc, char **argv)
282: {
283: /* Configuration parameters */
284: PetscInt n = DEFAULT_N;
285: PetscInt steps = DEFAULT_STEPS;
286: PetscInt burn = DEFAULT_BURN;
287: PetscInt obs_freq = DEFAULT_OBS_FREQ;
288: PetscInt random_seed = DEFAULT_RANDOM_SEED;
289: PetscInt ensemble_size = DEFAULT_ENSEMBLE_SIZE;
290: PetscReal F = DEFAULT_F;
291: PetscReal dt = DEFAULT_DT;
292: PetscReal obs_error_std = DEFAULT_OBS_ERROR_STD;
293: PetscReal ensemble_init_std = 1; /* Initial ensemble spread */
295: /* PETSc objects */
296: Lorenz96Ctx *l95_ctx = NULL, *truth_ctx = NULL;
297: DM da_state;
298: PetscDA da;
299: Vec x0, x_mean, x_forecast;
300: Vec truth_state, rmse_work;
301: Vec observation, obs_noise, obs_error_var;
302: PetscRandom rng;
303: Mat H = NULL; /* Observation operator matrix */
305: /* Statistics tracking */
306: PetscReal rmse_forecast = 0.0, rmse_analysis = 0.0;
307: PetscReal sum_rmse_forecast = 0.0, sum_rmse_analysis = 0.0;
308: PetscInt n_stat_steps = 0;
309: PetscInt obs_count = 0;
310: PetscInt step, progress_interval;
312: PetscFunctionBeginUser;
313: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
315: /* Parse command-line options */
316: PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Lorenz-96 ETKF Quick Example", NULL);
317: PetscCall(PetscOptionsInt("-n", "State dimension", "", n, &n, NULL));
318: PetscCall(PetscOptionsInt("-steps", "Number of time steps", "", steps, &steps, NULL));
319: PetscCall(PetscOptionsInt("-burn", "Burn-in steps excluded from statistics", "", burn, &burn, NULL));
320: PetscCall(PetscOptionsInt("-obs_freq", "Observation frequency", "", obs_freq, &obs_freq, NULL));
321: PetscCall(PetscOptionsReal("-F", "Forcing parameter", "", F, &F, NULL));
322: PetscCall(PetscOptionsReal("-dt", "Time step size", "", dt, &dt, NULL));
323: PetscCall(PetscOptionsReal("-obs_error", "Observation error standard deviation", "", obs_error_std, &obs_error_std, NULL));
324: PetscCall(PetscOptionsReal("-ensemble_init_std", "Initial ensemble spread standard deviation", "", ensemble_init_std, &ensemble_init_std, NULL));
325: PetscCall(PetscOptionsInt("-random_seed", "Random seed for ensemble perturbations", "", random_seed, &random_seed, NULL));
326: PetscOptionsEnd();
328: /* Validate and constrain parameters */
329: PetscCall(ValidateParameters(&n, &steps, &burn, &obs_freq, &ensemble_size, &dt, &F, &obs_error_std));
331: /* Calculate progress reporting interval (avoid division by zero) */
332: progress_interval = (steps >= PROGRESS_INTERVALS) ? (steps / PROGRESS_INTERVALS) : 1;
334: /* Create 1D periodic DM for state space */
335: PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, n, 1, 2, NULL, &da_state));
336: PetscCall(DMSetFromOptions(da_state));
337: PetscCall(DMSetUp(da_state));
338: PetscCall(DMDASetUniformCoordinates(da_state, 0.0, (PetscReal)n, 0.0, 0.0, 0.0, 0.0));
340: /* Create Lorenz96 context with reusable TS object */
341: PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &l95_ctx));
342: PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &truth_ctx));
344: /* Initialize random number generator */
345: PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rng));
346: PetscCall(PetscRandomSetSeed(rng, (unsigned long)random_seed));
347: PetscCall(PetscRandomSetFromOptions(rng));
348: PetscCall(PetscRandomSeed(rng));
350: /* Initialize state vectors */
351: PetscCall(DMCreateGlobalVector(da_state, &x0));
352: PetscCall(PetscRandomSetInterval(rng, -.1 * F, .1 * F));
353: PetscCall(VecSetRandom(x0, rng));
354: PetscCall(PetscRandomSetInterval(rng, 0, 1));
356: /* Initialize truth trajectory */
357: PetscCall(VecDuplicate(x0, &truth_state));
358: PetscCall(VecCopy(x0, truth_state));
359: PetscCall(VecDuplicate(x0, &rmse_work));
361: /* Spin up truth to get onto attractor */
362: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Spinning up truth for %" PetscInt_FMT " steps...\n", (PetscInt)SPINUP_STEPS));
363: for (PetscInt k = 0; k < SPINUP_STEPS; k++) PetscCall(Lorenz96Step(truth_state, truth_state, truth_ctx));
365: /* Initialize observation vectors */
366: PetscCall(VecDuplicate(x0, &observation));
367: PetscCall(VecDuplicate(x0, &obs_noise));
368: PetscCall(VecDuplicate(x0, &obs_error_var));
369: PetscCall(VecSet(obs_error_var, obs_error_std * obs_error_std));
371: /* Initialize ensemble statistics vectors */
372: PetscCall(VecDuplicate(x0, &x_mean));
373: PetscCall(VecDuplicate(x0, &x_forecast));
375: /* Create identity observation matrix H */
376: PetscCall(CreateIdentityObservationMatrix(n, &H));
378: /* Create and configure PetscDA for ensemble data assimilation */
379: PetscCall(PetscDACreate(PETSC_COMM_WORLD, &da));
380: PetscCall(PetscDASetType(da, PETSCDAETKF)); /* Set ETKF type */
381: PetscCall(PetscDASetSizes(da, n, n));
382: PetscCall(PetscDAEnsembleSetSize(da, ensemble_size));
383: PetscCall(PetscDASetFromOptions(da));
384: PetscCall(PetscDAEnsembleGetSize(da, &ensemble_size));
385: PetscCall(PetscDASetUp(da));
386: PetscCall(PetscDAViewFromOptions(da, NULL, "-petscda_view"));
387: PetscCall(PetscDASetObsErrorVariance(da, obs_error_var));
389: /* Initialize ensemble members from spun-up truth state with appropriate spread */
390: PetscCall(PetscDAEnsembleInitialize(da, truth_state, ensemble_init_std, rng));
392: /* Print configuration summary */
393: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Lorenz-96 ETKF Example\n"));
394: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "======================\n"));
395: PetscCall(PetscPrintf(PETSC_COMM_WORLD,
396: " State dimension : %" PetscInt_FMT "\n"
397: " Ensemble size : %" PetscInt_FMT "\n"
398: " Forcing parameter (F) : %.4f\n"
399: " Time step (dt) : %.4f\n"
400: " Total steps : %" PetscInt_FMT "\n"
401: " Burn-in steps : %" PetscInt_FMT "\n"
402: " Observation frequency : %" PetscInt_FMT "\n"
403: " Observation noise std : %.3f\n"
404: " Ensemble init std : %.3f\n"
405: " Random seed : %" PetscInt_FMT "\n\n",
406: n, ensemble_size, (double)F, (double)dt, steps, burn, obs_freq, (double)obs_error_std, (double)ensemble_init_std, random_seed));
408: /* Main assimilation cycle: forecast and analysis steps */
409: for (step = 0; step <= steps; step++) {
410: PetscReal time = step * dt;
412: /* Forecast step: compute ensemble mean and forecast RMSE */
413: PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
414: PetscCall(VecCopy(x_mean, x_forecast));
415: PetscCall(ComputeRMSE(x_forecast, truth_state, rmse_work, n, &rmse_forecast));
416: rmse_analysis = rmse_forecast; /* Default to forecast RMSE if no analysis */
418: /* Analysis step: assimilate observations when available */
419: if (step % obs_freq == 0 && step > 0) {
420: /* Generate synthetic noisy observations from truth */
421: PetscCall(VecSetRandomGaussian(obs_noise, rng, 0.0, obs_error_std));
422: PetscCall(VecWAXPY(observation, 1.0, obs_noise, truth_state));
424: /* Perform ETKF analysis with observation matrix H */
425: PetscCall(PetscDAEnsembleAnalysis(da, observation, H));
427: /* Compute analysis RMSE */
428: PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
429: PetscCall(ComputeRMSE(x_mean, truth_state, rmse_work, n, &rmse_analysis));
430: obs_count++;
431: }
433: /* Accumulate statistics after burn-in period */
434: if (step >= burn) {
435: sum_rmse_forecast += rmse_forecast;
436: sum_rmse_analysis += rmse_analysis;
437: n_stat_steps++;
438: }
440: /* Progress reporting */
441: if ((step % progress_interval == 0) || (step == steps) || (step == 0)) {
442: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Step %4" PetscInt_FMT ", time %6.3f RMSE_forecast %.5f RMSE_analysis %.5f%s\n", step, (double)time, (double)rmse_forecast, (double)rmse_analysis, (step < burn) ? " [burn-in]" : ""));
443: }
445: /* Propagate ensemble and truth trajectory */
446: if (step < steps) {
447: PetscCall(PetscDAEnsembleForecast(da, Lorenz96Step, l95_ctx));
448: PetscCall(Lorenz96Step(truth_state, truth_state, truth_ctx));
449: }
450: }
452: /* Report final statistics */
453: if (n_stat_steps > 0) {
454: PetscReal avg_rmse_forecast = sum_rmse_forecast / n_stat_steps;
455: PetscReal avg_rmse_analysis = sum_rmse_analysis / n_stat_steps;
456: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nStatistics (%" PetscInt_FMT " post-burn-in steps):\n", n_stat_steps));
457: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==================================================\n"));
458: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean RMSE (forecast) : %.5f\n", (double)avg_rmse_forecast));
459: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean RMSE (analysis) : %.5f\n", (double)avg_rmse_analysis));
460: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Observations used : %" PetscInt_FMT "\n\n", obs_count));
461: } else {
462: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nWarning: No post-burn-in statistics collected (burn >= steps)\n\n"));
463: }
465: /* Test VecSetRandomGaussian to verify Gaussian distribution */
466: {
467: Vec test_vec;
468: PetscInt test_size = 10000; /* Large sample for statistical testing */
469: PetscScalar *array;
470: PetscReal mean_target = 2.0, std_target = 1.5;
471: PetscReal sample_mean = 0.0, sample_variance = 0.0, sample_std;
472: PetscReal skewness = 0.0, kurtosis = 0.0;
473: PetscInt i;
474: PetscBool test_gaussian = PETSC_FALSE;
476: PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_gaussian", &test_gaussian, NULL));
478: if (test_gaussian) {
479: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n==============================================\n"));
480: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Testing VecSetRandomGaussian\n"));
481: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==============================================\n"));
483: /* Create test vector */
484: PetscCall(VecCreate(PETSC_COMM_WORLD, &test_vec));
485: PetscCall(VecSetSizes(test_vec, PETSC_DECIDE, test_size));
486: PetscCall(VecSetFromOptions(test_vec));
488: /* Generate Gaussian random numbers */
489: PetscCall(VecSetRandomGaussian(test_vec, rng, mean_target, std_target));
491: /* Get array for statistical analysis */
492: PetscCall(VecGetArray(test_vec, &array));
494: /* Compute sample mean */
495: for (i = 0; i < test_size; i++) sample_mean += PetscRealPart(array[i]);
496: sample_mean /= test_size;
498: /* Compute sample variance and higher moments */
499: for (i = 0; i < test_size; i++) {
500: PetscReal diff = PetscRealPart(array[i]) - sample_mean;
501: PetscReal diff2 = diff * diff;
502: sample_variance += diff2;
503: skewness += diff * diff2;
504: kurtosis += diff2 * diff2;
505: }
506: sample_variance /= (test_size - 1);
507: sample_std = PetscSqrtReal(sample_variance);
509: /* Normalize skewness and kurtosis */
510: skewness = (skewness / test_size) / PetscPowReal(sample_std, 3.0);
511: kurtosis = (kurtosis / test_size) / PetscPowReal(sample_std, 4.0) - 3.0; /* Excess kurtosis */
513: PetscCall(VecRestoreArray(test_vec, &array));
515: /* Report results */
516: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nTarget parameters:\n"));
517: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean : %.6f\n", (double)mean_target));
518: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Std Dev : %.6f\n", (double)std_target));
520: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nSample statistics (n=%" PetscInt_FMT "):\n", (PetscInt)test_size));
521: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean : %.6f (error: %.6f)\n", (double)sample_mean, (double)PetscAbsReal(sample_mean - mean_target)));
522: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Std Dev : %.6f (error: %.6f)\n", (double)sample_std, (double)PetscAbsReal(sample_std - std_target)));
523: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Skewness : %.6f (expected ~0 for Gaussian)\n", (double)skewness));
524: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Kurtosis : %.6f (expected ~0 for Gaussian)\n", (double)kurtosis));
526: /* Statistical tests with reasonable tolerances for finite samples */
527: PetscReal mean_error = PetscAbsReal(sample_mean - mean_target);
528: PetscReal std_error = PetscAbsReal(sample_std - std_target);
529: PetscReal mean_tolerance = 3.0 * std_target / PetscSqrtReal((PetscReal)test_size); /* 3-sigma rule */
530: PetscReal std_tolerance = 0.1 * std_target; /* 10% tolerance for std dev */
531: PetscReal skew_tolerance = 0.1; /* Skewness should be near 0 */
532: PetscReal kurt_tolerance = 0.5; /* Excess kurtosis should be near 0 */
534: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nStatistical tests:\n"));
535: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean test : %s (error %.6f < tolerance %.6f)\n", mean_error < mean_tolerance ? "PASS" : "FAIL", (double)mean_error, (double)mean_tolerance));
536: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Std dev test : %s (error %.6f < tolerance %.6f)\n", std_error < std_tolerance ? "PASS" : "FAIL", (double)std_error, (double)std_tolerance));
537: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Skewness test : %s (|skewness| %.6f < tolerance %.6f)\n", PetscAbsReal(skewness) < skew_tolerance ? "PASS" : "FAIL", (double)PetscAbsReal(skewness), (double)skew_tolerance));
538: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Kurtosis test : %s (|kurtosis| %.6f < tolerance %.6f)\n", PetscAbsReal(kurtosis) < kurt_tolerance ? "PASS" : "FAIL", (double)PetscAbsReal(kurtosis), (double)kurt_tolerance));
540: /* Overall test result */
541: PetscBool all_pass = (mean_error < mean_tolerance) && (std_error < std_tolerance) && (PetscAbsReal(skewness) < skew_tolerance) && (PetscAbsReal(kurtosis) < kurt_tolerance);
542: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nOverall result: %s\n", all_pass ? "PASS - Distribution is Gaussian" : "FAIL - Distribution may not be Gaussian"));
543: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==============================================\n\n"));
545: PetscCall(VecDestroy(&test_vec));
546: }
547: }
549: /* Cleanup */
550: PetscCall(MatDestroy(&H));
551: PetscCall(VecDestroy(&x_forecast));
552: PetscCall(VecDestroy(&x_mean));
553: PetscCall(VecDestroy(&obs_error_var));
554: PetscCall(VecDestroy(&obs_noise));
555: PetscCall(VecDestroy(&observation));
556: PetscCall(VecDestroy(&rmse_work));
557: PetscCall(VecDestroy(&truth_state));
558: PetscCall(VecDestroy(&x0));
559: PetscCall(PetscDADestroy(&da));
560: PetscCall(DMDestroy(&da_state));
561: PetscCall(Lorenz96ContextDestroy(&l95_ctx));
562: PetscCall(Lorenz96ContextDestroy(&truth_ctx));
563: PetscCall(PetscRandomDestroy(&rng));
565: PetscCall(PetscFinalize());
566: return 0;
567: }
569: /*TEST
571: testset:
572: nsize: 1
573: args: -steps 112 -burn 10 -obs_freq 1 -obs_error 1 -petscda_view -petscda_ensemble_size 30
575: test:
576: requires: !complex
577: suffix: eigen
578: args: -petscda_ensemble_sqrt_type eigen
580: test:
581: suffix: chol
582: args: -petscda_ensemble_sqrt_type cholesky
584: TEST*/