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: PetscFunctionBeginUser;
204: /* Create identity observation matrix H (n x n) */
205: PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, n, n, 1, NULL, 0, NULL, H));
206: PetscCall(MatSetFromOptions(*H));
208: /* Set diagonal entries to 1.0 for identity mapping */
209: for (PetscInt i = 0; i < n; i++) PetscCall(MatSetValue(*H, i, i, 1.0, INSERT_VALUES));
211: PetscCall(MatAssemblyBegin(*H, MAT_FINAL_ASSEMBLY));
212: PetscCall(MatAssemblyEnd(*H, MAT_FINAL_ASSEMBLY));
213: PetscFunctionReturn(PETSC_SUCCESS);
214: }
216: /*
217: ValidateParameters - Validate input parameters and apply constraints
219: Input/Output Parameters:
220: + n - State dimension
221: . steps - Number of time steps
222: . burn - Burn-in steps
223: . obs_freq - Observation frequency
224: . ensemble_size - Ensemble size
225: . dt - Time step
226: . F - Forcing parameter
227: - obs_error_std - Observation error standard deviation
229: This routine should be reduced/eliminated. It is the job of PETSc functions to validate input, not the tutorials code
230: */
231: static PetscErrorCode ValidateParameters(PetscInt *n, PetscInt *steps, PetscInt *burn, PetscInt *obs_freq, PetscInt *ensemble_size, PetscReal *dt, PetscReal *F, PetscReal *obs_error_std)
232: {
233: PetscFunctionBeginUser;
234: /* Validate and constrain integer parameters */
235: PetscCheck(*n > 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "State dimension n must be positive, got %" PetscInt_FMT, *n);
236: PetscCheck(*steps >= 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number of steps must be non-negative, got %" PetscInt_FMT, *steps);
237: 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);
239: /* Apply constraints */
240: if (*obs_freq < MIN_OBS_FREQ) {
241: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Observation frequency adjusted from %" PetscInt_FMT " to %" PetscInt_FMT "\n", *obs_freq, (PetscInt)MIN_OBS_FREQ));
242: *obs_freq = MIN_OBS_FREQ;
243: }
244: 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));
245: if (*burn > *steps) {
246: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Burn-in steps (%" PetscInt_FMT ") exceeds total steps (%" PetscInt_FMT "), setting burn = steps\n", *burn, *steps));
247: *burn = *steps;
248: }
250: /* Validate real-valued parameters */
251: PetscCheck(*dt > 0.0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Time step dt must be positive, got %g", (double)*dt);
252: 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);
253: PetscCheck(PetscIsNormalReal(*F), PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Forcing parameter F must be a normal real number");
254: PetscFunctionReturn(PETSC_SUCCESS);
255: }
257: /*
258: ComputeRMSE - Compute root mean square error between two vectors
260: Input Parameters:
261: + v1 - First vector
262: . v2 - Second vector
263: - n - Vector dimension (for normalization)
265: Output Parameter:
266: . rmse - Root mean square error
267: */
268: static PetscErrorCode ComputeRMSE(Vec v1, Vec v2, Vec work, PetscInt n, PetscReal *rmse)
269: {
270: PetscReal norm;
272: PetscFunctionBeginUser;
273: PetscCall(VecWAXPY(work, -1.0, v2, v1));
274: PetscCall(VecNorm(work, NORM_2, &norm));
275: *rmse = norm / PetscSqrtReal((PetscReal)n);
276: PetscFunctionReturn(PETSC_SUCCESS);
277: }
279: int main(int argc, char **argv)
280: {
281: /* Configuration parameters */
282: PetscInt n = DEFAULT_N;
283: PetscInt steps = DEFAULT_STEPS;
284: PetscInt burn = DEFAULT_BURN;
285: PetscInt obs_freq = DEFAULT_OBS_FREQ;
286: PetscInt random_seed = DEFAULT_RANDOM_SEED;
287: PetscInt ensemble_size = DEFAULT_ENSEMBLE_SIZE;
288: PetscReal F = DEFAULT_F;
289: PetscReal dt = DEFAULT_DT;
290: PetscReal obs_error_std = DEFAULT_OBS_ERROR_STD;
291: PetscReal ensemble_init_std = 1; /* Initial ensemble spread */
293: /* PETSc objects */
294: Lorenz96Ctx *l95_ctx = NULL, *truth_ctx = NULL;
295: DM da_state;
296: PetscDA da;
297: Vec x0, x_mean, x_forecast;
298: Vec truth_state, rmse_work;
299: Vec observation, obs_noise, obs_error_var;
300: PetscRandom rng;
301: Mat H = NULL; /* Observation operator matrix */
303: /* Statistics tracking */
304: PetscReal rmse_forecast = 0.0, rmse_analysis = 0.0;
305: PetscReal sum_rmse_forecast = 0.0, sum_rmse_analysis = 0.0;
306: PetscInt n_stat_steps = 0;
307: PetscInt obs_count = 0;
308: PetscInt step, progress_interval;
310: PetscFunctionBeginUser;
311: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
313: /* Parse command-line options */
314: PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Lorenz-96 ETKF Quick Example", NULL);
315: PetscCall(PetscOptionsInt("-n", "State dimension", "", n, &n, NULL));
316: PetscCall(PetscOptionsInt("-steps", "Number of time steps", "", steps, &steps, NULL));
317: PetscCall(PetscOptionsInt("-burn", "Burn-in steps excluded from statistics", "", burn, &burn, NULL));
318: PetscCall(PetscOptionsInt("-obs_freq", "Observation frequency", "", obs_freq, &obs_freq, NULL));
319: PetscCall(PetscOptionsReal("-F", "Forcing parameter", "", F, &F, NULL));
320: PetscCall(PetscOptionsReal("-dt", "Time step size", "", dt, &dt, NULL));
321: PetscCall(PetscOptionsReal("-obs_error", "Observation error standard deviation", "", obs_error_std, &obs_error_std, NULL));
322: PetscCall(PetscOptionsReal("-ensemble_init_std", "Initial ensemble spread standard deviation", "", ensemble_init_std, &ensemble_init_std, NULL));
323: PetscCall(PetscOptionsInt("-random_seed", "Random seed for ensemble perturbations", "", random_seed, &random_seed, NULL));
324: PetscOptionsEnd();
326: /* Validate and constrain parameters */
327: PetscCall(ValidateParameters(&n, &steps, &burn, &obs_freq, &ensemble_size, &dt, &F, &obs_error_std));
329: /* Calculate progress reporting interval (avoid division by zero) */
330: progress_interval = (steps >= PROGRESS_INTERVALS) ? (steps / PROGRESS_INTERVALS) : 1;
332: /* Create 1D periodic DM for state space */
333: PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, n, 1, 2, NULL, &da_state));
334: PetscCall(DMSetFromOptions(da_state));
335: PetscCall(DMSetUp(da_state));
336: PetscCall(DMDASetUniformCoordinates(da_state, 0.0, (PetscReal)n, 0.0, 0.0, 0.0, 0.0));
338: /* Create Lorenz96 context with reusable TS object */
339: PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &l95_ctx));
340: PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &truth_ctx));
342: /* Initialize random number generator */
343: PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rng));
344: PetscCall(PetscRandomSetSeed(rng, (unsigned long)random_seed));
345: PetscCall(PetscRandomSetFromOptions(rng));
346: PetscCall(PetscRandomSeed(rng));
348: /* Initialize state vectors */
349: PetscCall(DMCreateGlobalVector(da_state, &x0));
350: PetscCall(PetscRandomSetInterval(rng, -.1 * F, .1 * F));
351: PetscCall(VecSetRandom(x0, rng));
352: PetscCall(PetscRandomSetInterval(rng, 0, 1));
354: /* Initialize truth trajectory */
355: PetscCall(VecDuplicate(x0, &truth_state));
356: PetscCall(VecCopy(x0, truth_state));
357: PetscCall(VecDuplicate(x0, &rmse_work));
359: /* Spin up truth to get onto attractor */
360: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Spinning up truth for %" PetscInt_FMT " steps...\n", (PetscInt)SPINUP_STEPS));
361: for (PetscInt k = 0; k < SPINUP_STEPS; k++) PetscCall(Lorenz96Step(truth_state, truth_state, truth_ctx));
363: /* Initialize observation vectors */
364: PetscCall(VecDuplicate(x0, &observation));
365: PetscCall(VecDuplicate(x0, &obs_noise));
366: PetscCall(VecDuplicate(x0, &obs_error_var));
367: PetscCall(VecSet(obs_error_var, obs_error_std * obs_error_std));
369: /* Initialize ensemble statistics vectors */
370: PetscCall(VecDuplicate(x0, &x_mean));
371: PetscCall(VecDuplicate(x0, &x_forecast));
373: /* Create identity observation matrix H */
374: PetscCall(CreateIdentityObservationMatrix(n, &H));
376: /* Create and configure PetscDA for ensemble data assimilation */
377: PetscCall(PetscDACreate(PETSC_COMM_WORLD, &da));
378: PetscCall(PetscDASetType(da, PETSCDAETKF)); /* Set ETKF type */
379: PetscCall(PetscDASetSizes(da, n, n));
380: PetscCall(PetscDAEnsembleSetSize(da, ensemble_size));
381: PetscCall(PetscDASetFromOptions(da));
382: PetscCall(PetscDAEnsembleGetSize(da, &ensemble_size));
383: PetscCall(PetscDASetUp(da));
384: PetscCall(PetscDAViewFromOptions(da, NULL, "-petscda_view"));
385: PetscCall(PetscDASetObsErrorVariance(da, obs_error_var));
387: /* Initialize ensemble members from spun-up truth state with appropriate spread */
388: PetscCall(PetscDAEnsembleInitialize(da, truth_state, ensemble_init_std, rng));
390: /* Print configuration summary */
391: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Lorenz-96 ETKF Example\n"));
392: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "======================\n"));
393: PetscCall(PetscPrintf(PETSC_COMM_WORLD,
394: " State dimension : %" PetscInt_FMT "\n"
395: " Ensemble size : %" PetscInt_FMT "\n"
396: " Forcing parameter (F) : %.4f\n"
397: " Time step (dt) : %.4f\n"
398: " Total steps : %" PetscInt_FMT "\n"
399: " Burn-in steps : %" PetscInt_FMT "\n"
400: " Observation frequency : %" PetscInt_FMT "\n"
401: " Observation noise std : %.3f\n"
402: " Ensemble init std : %.3f\n"
403: " Random seed : %" PetscInt_FMT "\n\n",
404: n, ensemble_size, (double)F, (double)dt, steps, burn, obs_freq, (double)obs_error_std, (double)ensemble_init_std, random_seed));
406: /* Main assimilation cycle: forecast and analysis steps */
407: for (step = 0; step <= steps; step++) {
408: PetscReal time = step * dt;
410: /* Forecast step: compute ensemble mean and forecast RMSE */
411: PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
412: PetscCall(VecCopy(x_mean, x_forecast));
413: PetscCall(ComputeRMSE(x_forecast, truth_state, rmse_work, n, &rmse_forecast));
414: rmse_analysis = rmse_forecast; /* Default to forecast RMSE if no analysis */
416: /* Analysis step: assimilate observations when available */
417: if (step % obs_freq == 0 && step > 0) {
418: /* Generate synthetic noisy observations from truth */
419: PetscCall(VecSetRandomGaussian(obs_noise, rng, 0.0, obs_error_std));
420: PetscCall(VecWAXPY(observation, 1.0, obs_noise, truth_state));
422: /* Perform ETKF analysis with observation matrix H */
423: PetscCall(PetscDAEnsembleAnalysis(da, observation, H));
425: /* Compute analysis RMSE */
426: PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
427: PetscCall(ComputeRMSE(x_mean, truth_state, rmse_work, n, &rmse_analysis));
428: obs_count++;
429: }
431: /* Accumulate statistics after burn-in period */
432: if (step >= burn) {
433: sum_rmse_forecast += rmse_forecast;
434: sum_rmse_analysis += rmse_analysis;
435: n_stat_steps++;
436: }
438: /* Progress reporting */
439: if ((step % progress_interval == 0) || (step == steps) || (step == 0)) {
440: 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]" : ""));
441: }
443: /* Propagate ensemble and truth trajectory */
444: if (step < steps) {
445: PetscCall(PetscDAEnsembleForecast(da, Lorenz96Step, l95_ctx));
446: PetscCall(Lorenz96Step(truth_state, truth_state, truth_ctx));
447: }
448: }
450: /* Report final statistics */
451: if (n_stat_steps > 0) {
452: PetscReal avg_rmse_forecast = sum_rmse_forecast / n_stat_steps;
453: PetscReal avg_rmse_analysis = sum_rmse_analysis / n_stat_steps;
454: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nStatistics (%" PetscInt_FMT " post-burn-in steps):\n", n_stat_steps));
455: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==================================================\n"));
456: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean RMSE (forecast) : %.5f\n", (double)avg_rmse_forecast));
457: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean RMSE (analysis) : %.5f\n", (double)avg_rmse_analysis));
458: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Observations used : %" PetscInt_FMT "\n\n", obs_count));
459: } else {
460: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nWarning: No post-burn-in statistics collected (burn >= steps)\n\n"));
461: }
463: /* Test VecSetRandomGaussian to verify Gaussian distribution */
464: {
465: Vec test_vec;
466: PetscInt test_size = 10000; /* Large sample for statistical testing */
467: PetscScalar *array;
468: PetscReal mean_target = 2.0, std_target = 1.5;
469: PetscReal sample_mean = 0.0, sample_variance = 0.0, sample_std;
470: PetscReal skewness = 0.0, kurtosis = 0.0;
471: PetscBool test_gaussian = PETSC_FALSE;
473: PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_gaussian", &test_gaussian, NULL));
475: if (test_gaussian) {
476: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n==============================================\n"));
477: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Testing VecSetRandomGaussian\n"));
478: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==============================================\n"));
480: /* Create test vector */
481: PetscCall(VecCreate(PETSC_COMM_WORLD, &test_vec));
482: PetscCall(VecSetSizes(test_vec, PETSC_DECIDE, test_size));
483: PetscCall(VecSetFromOptions(test_vec));
485: /* Generate Gaussian random numbers */
486: PetscCall(VecSetRandomGaussian(test_vec, rng, mean_target, std_target));
488: /* Get array for statistical analysis */
489: PetscCall(VecGetArray(test_vec, &array));
491: /* Compute sample mean */
492: for (PetscInt i = 0; i < test_size; i++) sample_mean += PetscRealPart(array[i]);
493: sample_mean /= test_size;
495: /* Compute sample variance and higher moments */
496: for (PetscInt i = 0; i < test_size; i++) {
497: PetscReal diff = PetscRealPart(array[i]) - sample_mean;
498: PetscReal diff2 = diff * diff;
499: sample_variance += diff2;
500: skewness += diff * diff2;
501: kurtosis += diff2 * diff2;
502: }
503: sample_variance /= (test_size - 1);
504: sample_std = PetscSqrtReal(sample_variance);
506: /* Normalize skewness and kurtosis */
507: skewness = (skewness / test_size) / PetscPowReal(sample_std, 3.0);
508: kurtosis = (kurtosis / test_size) / PetscPowReal(sample_std, 4.0) - 3.0; /* Excess kurtosis */
510: PetscCall(VecRestoreArray(test_vec, &array));
512: /* Report results */
513: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nTarget parameters:\n"));
514: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean : %.6f\n", (double)mean_target));
515: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Std Dev : %.6f\n", (double)std_target));
517: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nSample statistics (n=%" PetscInt_FMT "):\n", (PetscInt)test_size));
518: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean : %.6f (error: %.6f)\n", (double)sample_mean, (double)PetscAbsReal(sample_mean - mean_target)));
519: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Std Dev : %.6f (error: %.6f)\n", (double)sample_std, (double)PetscAbsReal(sample_std - std_target)));
520: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Skewness : %.6f (expected ~0 for Gaussian)\n", (double)skewness));
521: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Kurtosis : %.6f (expected ~0 for Gaussian)\n", (double)kurtosis));
523: /* Statistical tests with reasonable tolerances for finite samples */
524: PetscReal mean_error = PetscAbsReal(sample_mean - mean_target);
525: PetscReal std_error = PetscAbsReal(sample_std - std_target);
526: PetscReal mean_tolerance = 3.0 * std_target / PetscSqrtReal((PetscReal)test_size); /* 3-sigma rule */
527: PetscReal std_tolerance = 0.1 * std_target; /* 10% tolerance for std dev */
528: PetscReal skew_tolerance = 0.1; /* Skewness should be near 0 */
529: PetscReal kurt_tolerance = 0.5; /* Excess kurtosis should be near 0 */
531: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nStatistical tests:\n"));
532: 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));
533: 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));
534: 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));
535: 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));
537: /* Overall test result */
538: PetscBool all_pass = (mean_error < mean_tolerance) && (std_error < std_tolerance) && (PetscAbsReal(skewness) < skew_tolerance) && (PetscAbsReal(kurtosis) < kurt_tolerance);
539: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nOverall result: %s\n", all_pass ? "PASS - Distribution is Gaussian" : "FAIL - Distribution may not be Gaussian"));
540: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==============================================\n\n"));
542: PetscCall(VecDestroy(&test_vec));
543: }
544: }
546: /* Cleanup */
547: PetscCall(MatDestroy(&H));
548: PetscCall(VecDestroy(&x_forecast));
549: PetscCall(VecDestroy(&x_mean));
550: PetscCall(VecDestroy(&obs_error_var));
551: PetscCall(VecDestroy(&obs_noise));
552: PetscCall(VecDestroy(&observation));
553: PetscCall(VecDestroy(&rmse_work));
554: PetscCall(VecDestroy(&truth_state));
555: PetscCall(VecDestroy(&x0));
556: PetscCall(PetscDADestroy(&da));
557: PetscCall(DMDestroy(&da_state));
558: PetscCall(Lorenz96ContextDestroy(&l95_ctx));
559: PetscCall(Lorenz96ContextDestroy(&truth_ctx));
560: PetscCall(PetscRandomDestroy(&rng));
562: PetscCall(PetscFinalize());
563: return 0;
564: }
566: /*TEST
568: testset:
569: nsize: 1
570: args: -steps 112 -burn 10 -obs_freq 1 -obs_error 1 -petscda_view -petscda_ensemble_size 30
572: test:
573: requires: !complex
574: suffix: eigen
575: args: -petscda_ensemble_sqrt_type eigen
577: test:
578: suffix: chol
579: args: -petscda_ensemble_sqrt_type cholesky
581: TEST*/