Minibatching compatability with CES, internalize flattening operations #436

Issue

When using EnsembleKalmanProcesses.jl with several observations $$(y_1,\dots,y_N)$$ one might use minibatching to reduce costs. This creates $$(B_1,\dots,B_n)$$ for $$n \leq N$$ and subdivides the forward map $$G = [G_{B_1}(\theta),\dots,G_{B_n}(\theta)] $$, then performs the $$k$$-th EKP_update based on $$G_{B_k}$$ only. While this can be used as a time-saving/complexity reduction assumption to help use EKP optimize with expensive models, it generates pairs of the form:

$$ \Big( \big(\theta_{1}, G_{B_1}(\theta_1) \big), ..., \big(\theta_{k}, G_{B_k}(\theta_k) \big), \Big) $$ therefore we do not have pairs $$\theta_i, G(\theta_i)$$

CES will currently handle the limited setting of multiple minibatches but only if they are size 1

There is a small modifier to handle most cases we need though. By assuming that the minibatches are from different evaluations of a stochastic simulator, This is the most likely use case for most users data, as often the different data $y_i$ come from this underlying assumption. This would mean that really our data is "y", and not a stacked object. Then by construction of $$B_i$$, this is just sampling from $$y_1,\dots,y_N$$, so just multiple copies of the "y". We in theory have to just break down the batches into their individual components $$B_1 = {i_1,i_2,i_3...}$$ that each look like "y", giving pairs :

$$ \Big( \big(\theta_{1}, G_{i_1}(\theta_1) \big), ..., \big(\theta_{K}, G_{i_K}(\theta_k) \big), ...\Big) $$

for $$G_{i_1}$$ being the forward map component fitting only to $$y_{i_1}$$. We can then train an emulator that treats all the $$G_{i}$$ as random realizations of and underlying "$$G$$", again, this is true for most users.

In this case really we just require a flattening operation to be done by the user as follows:

• get_u and get_g eki iterations
• get the minibatch epochs.
• for each minibatch (size e.g. 3), duplicate any parameter "u" 3 times, and split the output "g" into it's un-batched form"g=[g_1,g_2,g_3]" , then flatten the list (u,g_1), (u,g_2),(u, g_3)
• create the training/test data from here and continue with "CES"
• Also modify encoder_kwargs to be consistent.

Note this is nice though, as it still allows for

• batches of different sizes (i.e. different numbers of y_i's)
• batches to overlap (i.e. a batches [1,2,3] and [2,3,4])

Improvement

As stated in the most broad sense, most users do (and all users should ) use batching in this way where we assume batches just stack observational "atoms" that are all statistically the same. In this way we can modify internal operations and default behaviour to capture this:

• the minibatcher should internalize this flattening so the user does not need to perform it externally
• We also directly extract encoder_kwargs (that may include e.g. output samples) from EKI, and these do not flatten the batches, this sould be updated to do so by default