[WIP] update groups, i.e. domain and observation localization

[odunbar]

Purpose

Closes #377
Indirectly, Closes #288

To-do

• example with a coupled-system exhibiting some conditional independence of parameters
• Functionality for: ETKI,UKI,EKS,SEKI
• Fix localizers, currently requiring the dimensionality in advance of use, rather than at the time of use - can this be adjusted, or will localizers also have to be part of the update groups

Content

• UpdateGroup object
• unit tests for UpdateGroup
• implemented an interface with all EKP methods, in particular this necessitated lifting saving of "g", and verbose flag statements out of the individual updates thus reducing a lot of duplicated code
• implemented functionality for EKI
• Constructed an example based on learning "fast" and "slow" parameters of a coupled L96 model, hopefully improved by the grouping process.
• Updated to work with new Observations and ObservationSeries framework: In particular, a nice constructor that takes in prior, observation, dict_in_to_out and produces a list of UpdateGroups that can be passed to EKP.
• It also now works with minibatching, by repeating index blocks for each observation in the minibatch

Example: Multiscale lorenz 96 with j=k

Parameters (Slow)
F: 3-vector of amplitude, frequency and mean of a periodic forcing on slow variables
G: Coupling term in slow equation to the fast variables
Parameters (Fast)
h: Coupling term in the fast variables
c: Timescale separation parameter of the fast variables
b: nonlinearity scaling in the fast variables
With labelled observations observed "Independently" I.e. we assume independent with regards to observation (probably untrue)
<X>, <Y>, <X^2>, <Y^2>, <XY>
one can define two update groups as

group_identifiers = Dict(
    ["F"] => ["<X>", "<X^2>"], 
    ["G", "h", "c", "b"] => ["<Y>", "<Y^2>", "<XY>"],
)
update_groups = create_update_groups(priors, observation, group_identifiers) # handles all the indexing 

By allowing more correlations in the observations e.g. create <X><X^2> and <Y><Y^2><XY>one can also define groups by these identifiers

group_identifiers = Dict(
    ["F"] => ["<X><X^2>"], 
    ["G", "h", "c", "b"] => ["<Y><Y^2><XY>"],
)
update_groups = create_update_groups(priors, observation, group_identifiers)

and pass to the EKP object with update_groups = update_groups flag
Spin up over a year.

samples of the trajectories used for data

forward map used for data (noise covaraince and true data <X><Y><X^2><Y^2><XY> nomalized by the noise)

Shown below is the initial and final 50-ensemble (normalized by the noise) pushed through the forward map, over the true sample +/- 2sd (grey), after 5 iterations.

• True parameters [8.0, 2.0, 0.05555555555555555], 10.0, 10.0, 10.0, 10.0 ]
• Left: Full calibration, {F,G,h,c,b} with all data {<X><Y><X^2><Y^2><XY>}.

[7.864338120823016, 1.1500565821985165, 0.05883690301359528, 2.7165402443001034, 2.8197381618086217, 6.8072979930231945, 7.978078567957204]

• Right: calibration with grouped {F}, {G,h,c,b} and fast/slow data batches {<X><X^2>}, {<Y><Y^2><XY>}.

[5.093003061750248, 2.1573021824213985, 0.0060389967710004905, 1.0526243678897405, 0.8561707192010413, 8.331178239855038, 3.594166654372452]

---

• I have read and checked the items on the review checklist.