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 LETKF (NONE-localization) 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: /* Work with a local (ghosted) vector so the Lorenz-96 stencil has the
94: * required neighbors for periodic boundary conditions. */
95: PetscCall(DMDAGetCorners(l95->da, &xs, NULL, NULL, &xm, NULL, NULL));
96: PetscCall(DMGetLocalVector(l95->da, &X_local));
97: PetscCall(DMGlobalToLocalBegin(l95->da, X, INSERT_VALUES, X_local));
98: PetscCall(DMGlobalToLocalEnd(l95->da, X, INSERT_VALUES, X_local));
99: PetscCall(DMDAVecGetArrayRead(l95->da, X_local, &x));
100: PetscCall(DMDAVecGetArray(l95->da, F_vec, &f));
102: /* Standard Lorenz-96 tendency: (x_{i+1} - x_{i-2}) * x_{i-1} - x_i + F. */
103: for (i = xs; i < xs + xm; i++) f[i] = (x[i + 1] - x[i - 2]) * x[i - 1] - x[i] + l95->F;
105: PetscCall(DMDAVecRestoreArrayRead(l95->da, X_local, &x));
106: PetscCall(DMDAVecRestoreArray(l95->da, F_vec, &f));
107: PetscCall(DMRestoreLocalVector(l95->da, &X_local));
108: PetscFunctionReturn(PETSC_SUCCESS);
109: }
111: /*
112: Lorenz96ContextCreate - Create and initialize a Lorenz96 context with reusable TS object
114: Input Parameters:
115: + da - DM for state space
116: . n - State dimension
117: . F - Forcing parameter
118: . dt - Time step size
120: Output Parameter:
121: . ctx - Initialized Lorenz96 context
122: */
123: static PetscErrorCode Lorenz96ContextCreate(DM da, PetscInt n, PetscReal F, PetscReal dt, Lorenz96Ctx **ctx)
124: {
125: Lorenz96Ctx *l95;
127: PetscFunctionBeginUser;
128: PetscCall(PetscNew(&l95));
129: l95->da = da;
130: l95->n = n;
131: l95->F = F;
132: l95->dt = dt;
134: /* Create and configure a reusable time stepper to avoid repeated allocation/deallocation */
135: PetscCall(TSCreate(PetscObjectComm((PetscObject)da), &l95->ts));
136: PetscCall(TSSetProblemType(l95->ts, TS_NONLINEAR));
137: PetscCall(TSSetRHSFunction(l95->ts, NULL, Lorenz96RHS, l95));
138: PetscCall(TSSetType(l95->ts, TSRK));
139: PetscCall(TSRKSetType(l95->ts, TSRK4));
140: PetscCall(TSSetTimeStep(l95->ts, dt));
141: PetscCall(TSSetMaxTime(l95->ts, dt));
142: PetscCall(TSSetExactFinalTime(l95->ts, TS_EXACTFINALTIME_MATCHSTEP));
143: PetscCall(TSSetFromOptions(l95->ts));
145: *ctx = l95;
146: PetscFunctionReturn(PETSC_SUCCESS);
147: }
149: /*
150: Lorenz96ContextDestroy - Destroy a Lorenz96 context and its TS object
151: */
152: static PetscErrorCode Lorenz96ContextDestroy(Lorenz96Ctx **ctx)
153: {
154: PetscFunctionBeginUser;
155: if (!ctx || !*ctx) PetscFunctionReturn(PETSC_SUCCESS);
156: PetscCall(TSDestroy(&(*ctx)->ts));
157: PetscCall(PetscFree(*ctx));
158: PetscFunctionReturn(PETSC_SUCCESS);
159: }
161: /* Advance a single state vector one TS step. Used by the truth trajectory and as the per-column kernel of Lorenz96Step(). */
162: static PetscErrorCode Lorenz96StepVec(Lorenz96Ctx *l95, Vec x)
163: {
164: PetscFunctionBeginUser;
165: PetscCall(TSSetStepNumber(l95->ts, 0));
166: PetscCall(TSSetTime(l95->ts, 0.0));
167: PetscCall(TSSolve(l95->ts, x));
168: PetscFunctionReturn(PETSC_SUCCESS);
169: }
171: /*
172: Lorenz96Step - Advance every column of the ensemble matrix one time step using Lorenz-96 dynamics.
174: Input Parameters:
175: + ensemble - ensemble matrix whose columns are advanced in place
176: - ctx - Lorenz96 context (contains reusable TS)
178: Notes:
179: TS only advances one state at a time, so loop over the columns here. Uses a single explicit RK4 step
180: with the pre-configured TS object for efficiency.
181: */
182: static PetscErrorCode Lorenz96Step(Mat ensemble, PetscCtx ctx)
183: {
184: Lorenz96Ctx *l95 = (Lorenz96Ctx *)ctx;
185: PetscInt n;
187: PetscFunctionBeginUser;
188: PetscCall(MatGetSize(ensemble, NULL, &n));
189: /* Collective: dense ensemble Mat is row-distributed, so every rank visits every global column j and
190: MatDenseGetColumnVec returns the parallel column-Vec that all ranks step together. */
191: for (PetscInt j = 0; j < n; j++) {
192: Vec col;
194: PetscCall(MatDenseGetColumnVec(ensemble, j, &col));
195: PetscCall(Lorenz96StepVec(l95, col));
196: PetscCall(MatDenseRestoreColumnVec(ensemble, j, &col));
197: }
198: PetscFunctionReturn(PETSC_SUCCESS);
199: }
201: /*
202: CreateIdentityObservationMatrix - Create identity observation matrix H for Lorenz-96
204: For the fully observed case, H is an nxn identity matrix representing y = H*x where
205: each observation corresponds directly to a state variable.
207: Input Parameter:
208: . n - State dimension (number of grid points)
210: Output Parameter:
211: . H - Identity observation matrix (n x n), sparse AIJ format
212: */
213: static PetscErrorCode CreateIdentityObservationMatrix(PetscInt n, Mat *H)
214: {
215: PetscFunctionBeginUser;
216: /* Create identity observation matrix H (n x n) */
217: PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, n, n, 1, NULL, 0, NULL, H));
218: PetscCall(MatSetFromOptions(*H));
220: /* Set diagonal entries to 1.0 for identity mapping */
221: for (PetscInt i = 0; i < n; i++) PetscCall(MatSetValue(*H, i, i, 1.0, INSERT_VALUES));
223: PetscCall(MatAssemblyBegin(*H, MAT_FINAL_ASSEMBLY));
224: PetscCall(MatAssemblyEnd(*H, MAT_FINAL_ASSEMBLY));
225: PetscFunctionReturn(PETSC_SUCCESS);
226: }
228: /*
229: ValidateParameters - Validate input parameters and apply constraints
231: Input/Output Parameters:
232: + n - State dimension
233: . steps - Number of time steps
234: . burn - Burn-in steps
235: . obs_freq - Observation frequency
236: . ensemble_size - Ensemble size
237: . dt - Time step
238: . F - Forcing parameter
239: - obs_error_std - Observation error standard deviation
241: This routine should be reduced/eliminated. It is the job of PETSc functions to validate input, not the tutorials code
242: */
243: static PetscErrorCode ValidateParameters(PetscInt *n, PetscInt *steps, PetscInt *burn, PetscInt *obs_freq, PetscInt *ensemble_size, PetscReal *dt, PetscReal *F, PetscReal *obs_error_std)
244: {
245: PetscFunctionBeginUser;
246: /* Validate and constrain integer parameters */
247: PetscCheck(*n > 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "State dimension n must be positive, got %" PetscInt_FMT, *n);
248: PetscCheck(*steps >= 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number of steps must be non-negative, got %" PetscInt_FMT, *steps);
249: 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);
251: /* Apply constraints */
252: if (*obs_freq < MIN_OBS_FREQ) {
253: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Observation frequency adjusted from %" PetscInt_FMT " to %" PetscInt_FMT "\n", *obs_freq, (PetscInt)MIN_OBS_FREQ));
254: *obs_freq = MIN_OBS_FREQ;
255: }
256: 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));
257: if (*burn > *steps) {
258: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: Burn-in steps (%" PetscInt_FMT ") exceeds total steps (%" PetscInt_FMT "), setting burn = steps\n", *burn, *steps));
259: *burn = *steps;
260: }
262: /* Validate real-valued parameters */
263: PetscCheck(*dt > 0.0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Time step dt must be positive, got %g", (double)*dt);
264: 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);
265: PetscCheck(PetscIsNormalReal(*F), PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Forcing parameter F must be a normal real number");
266: PetscFunctionReturn(PETSC_SUCCESS);
267: }
269: /*
270: ComputeRMSE - Compute root mean square error between two vectors
272: Input Parameters:
273: + v1 - First vector
274: . v2 - Second vector
275: - n - Vector dimension (for normalization)
277: Output Parameter:
278: . rmse - Root mean square error
279: */
280: static PetscErrorCode ComputeRMSE(Vec v1, Vec v2, Vec work, PetscInt n, PetscReal *rmse)
281: {
282: PetscReal norm;
284: PetscFunctionBeginUser;
285: PetscCall(VecWAXPY(work, -1.0, v2, v1));
286: PetscCall(VecNorm(work, NORM_2, &norm));
287: *rmse = norm / PetscSqrtReal((PetscReal)n);
288: PetscFunctionReturn(PETSC_SUCCESS);
289: }
291: int main(int argc, char **argv)
292: {
293: Lorenz96Ctx *l95_ctx = NULL, *truth_ctx = NULL;
294: DM da_state;
295: PetscDA da;
296: Vec x0, x_mean, x_forecast;
297: Vec truth_state, rmse_work;
298: Vec observation, obs_noise, obs_error_var;
299: Mat H = NULL; /* Observation operator matrix */
300: PetscRandom rng;
301: PetscInt n = DEFAULT_N, steps = DEFAULT_STEPS, burn = DEFAULT_BURN, obs_freq = DEFAULT_OBS_FREQ;
302: PetscInt random_seed = DEFAULT_RANDOM_SEED, ensemble_size = DEFAULT_ENSEMBLE_SIZE;
303: PetscInt n_stat_steps = 0, obs_count = 0, step, progress_interval;
304: PetscReal F = DEFAULT_F, dt = DEFAULT_DT, obs_error_std = DEFAULT_OBS_ERROR_STD;
305: PetscReal ensemble_init_std = 1; /* Initial ensemble spread */
306: PetscReal rmse_forecast = 0.0, rmse_analysis = 0.0;
307: PetscReal sum_rmse_forecast = 0.0, sum_rmse_analysis = 0.0;
309: PetscFunctionBeginUser;
310: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
312: /* Parse command-line options */
313: PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Lorenz-96 LETKF Quick Example", NULL);
314: PetscCall(PetscOptionsInt("-n", "State dimension", "", n, &n, NULL));
315: PetscCall(PetscOptionsInt("-steps", "Number of time steps", "", steps, &steps, NULL));
316: PetscCall(PetscOptionsInt("-burn", "Burn-in steps excluded from statistics", "", burn, &burn, NULL));
317: PetscCall(PetscOptionsInt("-obs_freq", "Observation frequency", "", obs_freq, &obs_freq, NULL));
318: PetscCall(PetscOptionsReal("-F", "Forcing parameter", "", F, &F, NULL));
319: PetscCall(PetscOptionsReal("-dt", "Time step size", "", dt, &dt, NULL));
320: PetscCall(PetscOptionsReal("-obs_error", "Observation error standard deviation", "", obs_error_std, &obs_error_std, NULL));
321: PetscCall(PetscOptionsReal("-ensemble_init_std", "Initial ensemble spread standard deviation", "", ensemble_init_std, &ensemble_init_std, NULL));
322: PetscCall(PetscOptionsInt("-random_seed", "Random seed for ensemble perturbations", "", random_seed, &random_seed, NULL));
323: PetscOptionsEnd();
325: /* Validate and constrain parameters */
326: PetscCall(ValidateParameters(&n, &steps, &burn, &obs_freq, &ensemble_size, &dt, &F, &obs_error_std));
328: /* Calculate progress reporting interval (avoid division by zero) */
329: progress_interval = (steps >= PROGRESS_INTERVALS) ? (steps / PROGRESS_INTERVALS) : 1;
331: /* Create 1D periodic DM for state space */
332: PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, n, 1, 2, NULL, &da_state));
333: PetscCall(DMSetFromOptions(da_state));
334: PetscCall(DMSetUp(da_state));
335: PetscCall(DMDASetUniformCoordinates(da_state, 0.0, (PetscReal)n, 0.0, 0.0, 0.0, 0.0));
337: /* Create Lorenz96 context with reusable TS object */
338: PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &l95_ctx));
339: PetscCall(Lorenz96ContextCreate(da_state, n, F, dt, &truth_ctx));
341: /* Initialize random number generator */
342: PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rng));
343: PetscCall(PetscRandomSetSeed(rng, (unsigned long)random_seed));
344: PetscCall(PetscRandomSetFromOptions(rng));
345: PetscCall(PetscRandomSeed(rng));
347: /* Initialize state vectors */
348: PetscCall(DMCreateGlobalVector(da_state, &x0));
349: PetscCall(PetscRandomSetInterval(rng, -.1 * F, .1 * F));
350: PetscCall(VecSetRandom(x0, rng));
351: PetscCall(PetscRandomSetInterval(rng, 0, 1));
353: /* Initialize truth trajectory */
354: PetscCall(VecDuplicate(x0, &truth_state));
355: PetscCall(VecCopy(x0, truth_state));
356: PetscCall(VecDuplicate(x0, &rmse_work));
358: /* Spin up truth to get onto attractor */
359: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Spinning up truth for %" PetscInt_FMT " steps...\n", (PetscInt)SPINUP_STEPS));
360: for (PetscInt k = 0; k < SPINUP_STEPS; k++) PetscCall(Lorenz96StepVec(truth_ctx, truth_state));
362: /* Initialize observation vectors */
363: PetscCall(VecDuplicate(x0, &observation));
364: PetscCall(VecDuplicate(x0, &obs_noise));
365: PetscCall(VecDuplicate(x0, &obs_error_var));
366: PetscCall(VecSet(obs_error_var, obs_error_std * obs_error_std));
368: /* Initialize ensemble statistics vectors */
369: PetscCall(VecDuplicate(x0, &x_mean));
370: PetscCall(VecDuplicate(x0, &x_forecast));
372: /* Create identity observation matrix H */
373: PetscCall(CreateIdentityObservationMatrix(n, &H));
375: /* Create and configure PetscDA for ensemble data assimilation */
376: PetscCall(PetscDACreate(PETSC_COMM_WORLD, &da));
377: PetscCall(PetscDALETKFSetLocalizationType(da, PETSCDA_LETKF_LOC_NONE));
378: PetscCall(PetscDASetSizes(da, n, n));
379: PetscCall(PetscDAEnsembleSetSize(da, ensemble_size));
380: PetscCall(PetscDASetFromOptions(da));
381: PetscCall(PetscDAEnsembleGetSize(da, &ensemble_size));
382: PetscCall(PetscDASetUp(da));
383: PetscCall(PetscDASetObsErrorVariance(da, obs_error_var));
385: /* Initialize ensemble members from spun-up truth state with appropriate spread */
386: PetscCall(PetscDAEnsembleInitialize(da, truth_state, ensemble_init_std, rng));
388: /* Print configuration summary */
389: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Lorenz-96 LETKF Example\n"));
390: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "======================\n"));
391: PetscCall(PetscPrintf(PETSC_COMM_WORLD,
392: " State dimension : %" PetscInt_FMT "\n"
393: " Ensemble size : %" PetscInt_FMT "\n"
394: " Forcing parameter (F) : %.4f\n"
395: " Time step (dt) : %.4f\n"
396: " Total steps : %" PetscInt_FMT "\n"
397: " Burn-in steps : %" PetscInt_FMT "\n"
398: " Observation frequency : %" PetscInt_FMT "\n"
399: " Observation noise std : %.3f\n"
400: " Ensemble init std : %.3f\n"
401: " Random seed : %" PetscInt_FMT "\n\n",
402: n, ensemble_size, (double)F, (double)dt, steps, burn, obs_freq, (double)obs_error_std, (double)ensemble_init_std, random_seed));
404: /* Main assimilation cycle: forecast and analysis steps */
405: for (step = 0; step <= steps; step++) {
406: PetscReal time = step * dt;
408: /* Forecast step: compute ensemble mean and forecast RMSE */
409: PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
410: PetscCall(VecCopy(x_mean, x_forecast));
411: PetscCall(ComputeRMSE(x_forecast, truth_state, rmse_work, n, &rmse_forecast));
412: rmse_analysis = rmse_forecast; /* Default to forecast RMSE if no analysis */
414: /* Analysis step: assimilate observations when available */
415: if (step % obs_freq == 0 && step > 0) {
416: /* Generate synthetic noisy observations from truth */
417: PetscCall(VecSetRandomGaussian(obs_noise, rng, 0.0, obs_error_std));
418: PetscCall(VecWAXPY(observation, 1.0, obs_noise, truth_state));
420: /* Perform ETKF analysis with observation matrix H */
421: PetscCall(PetscDAEnsembleAnalysis(da, observation, H));
423: /* Compute analysis RMSE */
424: PetscCall(PetscDAEnsembleComputeMean(da, x_mean));
425: PetscCall(ComputeRMSE(x_mean, truth_state, rmse_work, n, &rmse_analysis));
426: obs_count++;
427: }
429: /* Accumulate statistics after burn-in period */
430: if (step >= burn) {
431: sum_rmse_forecast += rmse_forecast;
432: sum_rmse_analysis += rmse_analysis;
433: n_stat_steps++;
434: }
436: /* Progress reporting */
437: if ((step % progress_interval == 0) || (step == steps) || (step == 0)) {
438: 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]" : ""));
439: }
441: /* Propagate ensemble and truth trajectory */
442: if (step < steps) {
443: PetscCall(PetscDAEnsembleForecast(da, Lorenz96Step, l95_ctx));
444: PetscCall(Lorenz96StepVec(truth_ctx, truth_state));
445: }
446: }
448: /* Report final statistics */
449: if (n_stat_steps > 0) {
450: PetscReal avg_rmse_forecast = sum_rmse_forecast / n_stat_steps;
451: PetscReal avg_rmse_analysis = sum_rmse_analysis / n_stat_steps;
452: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nStatistics (%" PetscInt_FMT " post-burn-in steps):\n", n_stat_steps));
453: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==================================================\n"));
454: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean RMSE (forecast) : %.5f\n", (double)avg_rmse_forecast));
455: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean RMSE (analysis) : %.5f\n", (double)avg_rmse_analysis));
456: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Observations used : %" PetscInt_FMT "\n\n", obs_count));
457: } else {
458: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nWarning: No post-burn-in statistics collected (burn >= steps)\n\n"));
459: }
461: /* Test VecSetRandomGaussian to verify Gaussian distribution */
462: {
463: Vec test_vec;
464: PetscInt test_size = 10000; /* Large sample for statistical testing */
465: PetscScalar *array;
466: PetscReal mean_target = 2.0, std_target = 1.5;
467: PetscReal sample_mean = 0.0, sample_variance = 0.0, sample_std;
468: PetscReal skewness = 0.0, kurtosis = 0.0;
469: PetscBool test_gaussian = PETSC_FALSE;
471: PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_gaussian", &test_gaussian, NULL));
473: if (test_gaussian) {
474: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n==============================================\n"));
475: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Testing VecSetRandomGaussian\n"));
476: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==============================================\n"));
478: /* Create test vector */
479: PetscCall(VecCreate(PETSC_COMM_WORLD, &test_vec));
480: PetscCall(VecSetSizes(test_vec, PETSC_DECIDE, test_size));
481: PetscCall(VecSetFromOptions(test_vec));
483: /* Generate Gaussian random numbers */
484: PetscCall(VecSetRandomGaussian(test_vec, rng, mean_target, std_target));
486: /* Get array for statistical analysis */
487: PetscCall(VecGetArray(test_vec, &array));
489: /* Compute sample mean */
490: for (PetscInt i = 0; i < test_size; i++) sample_mean += PetscRealPart(array[i]);
491: sample_mean /= test_size;
493: /* Compute sample variance and higher moments */
494: for (PetscInt i = 0; i < test_size; i++) {
495: PetscReal diff = PetscRealPart(array[i]) - sample_mean;
496: PetscReal diff2 = diff * diff;
497: sample_variance += diff2;
498: skewness += diff * diff2;
499: kurtosis += diff2 * diff2;
500: }
501: sample_variance /= (test_size - 1);
502: sample_std = PetscSqrtReal(sample_variance);
504: /* Normalize skewness and kurtosis */
505: skewness = (skewness / test_size) / PetscPowReal(sample_std, 3.0);
506: kurtosis = (kurtosis / test_size) / PetscPowReal(sample_std, 4.0) - 3.0; /* Excess kurtosis */
508: PetscCall(VecRestoreArray(test_vec, &array));
510: /* Report results */
511: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nTarget parameters:\n"));
512: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean : %.6f\n", (double)mean_target));
513: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Std Dev : %.6f\n", (double)std_target));
515: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nSample statistics (n=%" PetscInt_FMT "):\n", (PetscInt)test_size));
516: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Mean : %.6f (error: %.6f)\n", (double)sample_mean, (double)PetscAbsReal(sample_mean - mean_target)));
517: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Std Dev : %.6f (error: %.6f)\n", (double)sample_std, (double)PetscAbsReal(sample_std - std_target)));
518: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Skewness : %.6f (expected ~0 for Gaussian)\n", (double)skewness));
519: PetscCall(PetscPrintf(PETSC_COMM_WORLD, " Kurtosis : %.6f (expected ~0 for Gaussian)\n", (double)kurtosis));
521: /* Statistical tests with reasonable tolerances for finite samples */
522: PetscReal mean_error = PetscAbsReal(sample_mean - mean_target);
523: PetscReal std_error = PetscAbsReal(sample_std - std_target);
524: PetscReal mean_tolerance = 3.0 * std_target / PetscSqrtReal((PetscReal)test_size); /* 3-sigma rule */
525: PetscReal std_tolerance = 0.1 * std_target; /* 10% tolerance for std dev */
526: PetscReal skew_tolerance = 0.1; /* Skewness should be near 0 */
527: PetscReal kurt_tolerance = 0.5; /* Excess kurtosis should be near 0 */
529: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nStatistical tests:\n"));
530: 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));
531: 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));
532: 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));
533: 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));
535: /* Overall test result */
536: PetscBool all_pass = (mean_error < mean_tolerance) && (std_error < std_tolerance) && (PetscAbsReal(skewness) < skew_tolerance) && (PetscAbsReal(kurtosis) < kurt_tolerance);
537: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nOverall result: %s\n", all_pass ? "PASS - Distribution is Gaussian" : "FAIL - Distribution may not be Gaussian"));
538: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "==============================================\n\n"));
540: PetscCall(VecDestroy(&test_vec));
541: }
542: }
544: PetscCall(PetscDAView(da, PETSC_VIEWER_STDOUT_WORLD));
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_ensemble_size 30
572: test:
573: requires: !complex
574: suffix: eigen
576: TEST*/