Updated var names for 4dvar. xb = bg, xa = analysis, x = temporary during loops, 0 denotes first timestep

Author: kysolvik

dabench/dacycler/_var4d.py

@@ -126,7 +126,7 @@ def _calc_J_term(self,
     @partial(jax.jit, static_argnums=[0, 1])
     def _innerloop_4d(self,
                       system_dim: int,
-                      Xb_ds: XarrayDatasetLike,
+                      X_ds: XarrayDatasetLike,
                       xb0_ds: XarrayDatasetLike,
                       obs_vals: ArrayLike,
                       Hs: ArrayLike,
@@ -137,8 +137,8 @@ def _innerloop_4d(self,
                       obs_time_mask: ArrayLike
                       ) -> XarrayDatasetLike:
         """4DVar innerloop"""
-        x0_prev_ds = Xb_ds.isel(time=0)
-        Xb_ar = Xb_ds.to_stacked_array('system',['time'])
+        x0_ds = X_ds.isel(time=0)
+        X_ar = X_ds.to_stacked_array('system',['time'])
 
         # Set up Variables
         SumMtHtRinvHM = jnp.zeros_like(B)             # A input
@@ -151,22 +151,22 @@ def _innerloop_4d(self,
                     lambda: self._calc_J_term(
                         Hs.at[i].get(mode='clip'),
                         M.data[j],
-                        Rinv, obs_vals[i], Xb_ar.data[j]),
+                        Rinv, obs_vals[i], X_ar.data[j]),
                     lambda: (jnp.zeros_like(SumMtHtRinvHM),
                              jnp.zeros_like(SumMtHtRinvD))
                     )
             SumMtHtRinvHM += Jb
             SumMtHtRinvD += Jo
         # Compute initial departure
-        db0 = (xb0_ds - x0_prev_ds).to_stacked_array('system',[]).data
+        db0 = (xb0_ds - x0_ds).to_stacked_array('system',[]).data
 
         # Solve Ax=b for the initial perturbation
         dx0 = self._solve(db0, SumMtHtRinvHM, SumMtHtRinvD, B)
 
         # New x0 guess is the last guess plus the analyzed delta
-        x0_new_ds = x0_prev_ds + dx0.ravel()
+        xa0_ds = x0_ds + dx0.ravel()
 
-        return x0_new_ds
+        return xa0_ds
 
     def _make_outerloop_4d(self,
                            xb0_ds: XarrayDatasetLike,
@@ -185,18 +185,18 @@ def _outerloop_4d(x0_ds: XarrayDatasetLike,
             # Get TLM and current forecast trajectory
             # Based on current best guess for x0
             x0_ds = x0_ds.to_xarray()
-            xb_ds, M = self.model_obj.compute_tlm(
+            X_ds, M = self.model_obj.compute_tlm(
                 n_steps=n_steps,
                 state_vec=x0_ds
             )
 
             # 4D-Var inner loop
-            x0_new_ds = self._innerloop_4d(
-                self.system_dim, xb_ds, xb0_ds, obs_values,
+            xa0_ds = self._innerloop_4d(
+                self.system_dim, X_ds, xb0_ds, obs_values,
                 Hs, B, Rinv, M, obs_window_indices, obs_time_mask
                 )
 
-            return xj.from_xarray(x0_new_ds.assign_coords(x0_ds.coords)), x0_ds
+            return xj.from_xarray(xa0_ds.assign_coords(x0_ds.coords)), x0_ds
 
         return _outerloop_4d
 
@@ -236,7 +236,7 @@ def _solve(self,
         return dx0
 
     def _cycle_obsop(self,
-                     x0_ds: XarrayDatasetLike,
+                     xb0_ds: XarrayDatasetLike,
                      obs_values: ArrayLike,
                      obs_loc_indices: ArrayLike,
                      obs_time_mask: ArrayLike,
@@ -284,14 +284,12 @@ def _cycle_obsop(self,
         # Static Variables
         Rinv = jscipy.linalg.inv(R)
 
-        # Best guess for x0 starts as background
-        x0_new_ds = deepcopy(x0_ds)
-
         outerloop_4d_func = self._make_outerloop_4d(
-                x0_ds, Hs, B, Rinv, obs_values, obs_window_indices,
+                xb0_ds, Hs, B, Rinv, obs_values, obs_window_indices,
                 obs_time_mask, self.steps_per_window)
 
-        x0_new_ds, all_x0s = jax.lax.scan(outerloop_4d_func, init=xj.from_xarray(x0_new_ds),
+        xa0_ds, all_x0s = jax.lax.scan(outerloop_4d_func,
+                                          init=xj.from_xarray(xb0_ds),
                 xs=None, length=self.n_outer_loops)
 
-        return x0_new_ds.to_xarray()
+        return xa0_ds.to_xarray()

dabench/dacycler/_var4d_backprop.py

@@ -134,13 +134,13 @@ def _callback_raise_error(self,
 
     # @partial(jax.jit, static_argnums=[0])
     def _calc_obs_term(self,
-                       Xb: ArrayLike,
+                       X: ArrayLike,
                        obs_vals: ArrayLike,
                        Ht: ArrayLike,
                        Rinv: ArrayLike
                        ) -> jax.Array:
-        Yb = Xb @ Ht
-        resid = Yb.ravel() - obs_vals.ravel()
+        Y = X @ Ht
+        resid = Y.ravel() - obs_vals.ravel()
 
         return jnp.sum(resid.T @ Rinv @ resid)
 
@@ -163,15 +163,15 @@ def loss_4dvarcost(x0: XarrayDatasetLike) -> jax.Array:
 
             # Make new prediction
             # NOTE: [1] selects the full forecast instead of last timestep only
-            Xb = self._step_forecast(
+            X = self._step_forecast(
                 x0, n_steps)[1].to_stacked_array('system',['time']).data
 
             # Calculate observation term of J_0
             obs_term = 0
             for i, j in enumerate(obs_window_indices):
                 obs_term += jax.lax.cond(
                         obs_time_mask.at[i].get(mode='fill', fill_value=0),
-                        lambda: self._calc_obs_term(Xb[j], obs_vals[i],
+                        lambda: self._calc_obs_term(X[j], obs_vals[i],
                                                     Hs.at[i].get(mode='clip').T,
                                                     Rinv),
                         lambda: 0.0
@@ -220,15 +220,15 @@ def _backprop_epoch(
             updates, opt_state = optimizer.update(dx0_hess, opt_state)
             x0_ar.data = optax.apply_updates(
                 x0_ar.data, updates)
-            x0_new_ds = x0_ar.to_unstacked_dataset('system').assign_attrs(
+            xa0_ds = x0_ar.to_unstacked_dataset('system').assign_attrs(
                 x0_ds.attrs
             )
-            return (xj.from_xarray(x0_new_ds), init_loss, opt_state), loss_val
+            return (xj.from_xarray(xa0_ds), init_loss, opt_state), loss_val
 
         return _backprop_epoch
 
     def _cycle_obsop(self,
-                     x0_ds: XarrayDatasetLike,
+                     xb0_ds: XarrayDatasetLike,
                      obs_values: ArrayLike,
                      obs_loc_indices: ArrayLike,
                      obs_time_mask: ArrayLike,
@@ -281,7 +281,7 @@ def _cycle_obsop(self,
                 Binv + Hs.at[0].get().T @ Rinv @ Hs.at[0].get())
 
         loss_func = self._make_loss(
-                x0_ds,
+                xb0_ds,
                 obs_values,
                 Hs,
                 Binv,
@@ -295,15 +295,15 @@ def _cycle_obsop(self,
                 1,
                 self.lr_decay)
         optimizer = optax.sgd(lr)
-        opt_state = optimizer.init(x0_ds.to_stacked_array('system',[]).data)
+        opt_state = optimizer.init(xb0_ds.to_stacked_array('system',[]).data)
 
         # Make initial forecast and calculate loss
         backprop_epoch_func = self._make_backprop_epoch(loss_func, optimizer,
                                                         hessian_inv)
         epoch_state_tuple, loss_vals = jax.lax.scan(
-                backprop_epoch_func, init=(xj.from_xarray(x0_ds), 0., opt_state),
+                backprop_epoch_func, init=(xj.from_xarray(xb0_ds), 0., opt_state),
                 xs=jnp.arange(self.num_iters))
 
-        x0_new_ds = epoch_state_tuple[0].to_xarray()
+        xa0_ds = epoch_state_tuple[0].to_xarray()
 
-        return x0_new_ds
+        return xa0_ds