Create a Forward map wrapper to more easily test pipelines without using an emulator (#402)

* moved MLTs into subdirectory

* added ForwardMapWrapper, made compatible with MCMC and encoder utilities

* featured in sinusoid example

* remove anonymous function and have G in space so sample.jl doesnt crash

* add FM into sinusoid

* format

* Update src/Emulator.jl

Co-authored-by: ArneBouillon <45404227+ArneBouillon@users.noreply.github.com>

* Update docs/src/examples/sinusoid_example.md

Co-authored-by: ArneBouillon <45404227+ArneBouillon@users.noreply.github.com>

* add a warning message that the input encoder is not used

* some starting tests

* remove extra end

* scope "sample"

* improve const

* update test_throws

* qualify sample and predict

* coverage

* format

---------

Co-authored-by: ArneBouillon <45404227+ArneBouillon@users.noreply.github.com>

Author: odunbar

docs/src/assets/sinusoid_MCMC_hist_FM.png

@@ -31,7 +31,6 @@ y_{obs} = G(\theta) + \gamma, \qquad \gamma \sim \mathcal{N}(0, \Gamma)
 where $\Gamma$ is the observational covariance matrix. We will assume the noise to be independent for each observable, giving us a diagonal covariance matrix.
 
 
-
 # Walkthrough of code
 
 You can find the full scripts to reproduce this tutorial in `examples/Sinusoid/`. The code is split into four sections:
@@ -253,7 +252,11 @@ encoder_schedule = [(decorrelate_sample_cov(), "in"), (decorrelate_structure_mat
 encoder_kwargs = (; obs_noise_cov = Γ,)
 
 # Build emulator with data
-emulator_gp = Emulator(gauss_proc, input_output_pairs; encoder_schedule = encoder_schedule, encoder_kwargs = encoder_kwargs)
+emulator_gp = Emulator(gauss_proc,
+    input_output_pairs;
+    encoder_schedule = encoder_schedule,
+    encoder_kwargs = encoder_kwargs,
+)
 optimize_hyperparameters!(emulator_gp)
 ```
 For this simple example, we already know the observational noise `Γ=0.2*I`, so we set `noise_learn = false`. 
@@ -291,8 +294,12 @@ random_features = VectorRandomFeatureInterface(
     optimizer_options = optimizer_options,
 )
 
-emulator_random_features =
-    Emulator(random_features, input_output_pairs; encoder_schedule = deepcopy(encoder_schedule), encoder_kwargs=deepcopy(encoder_kwargs))
+emulator_random_features = Emulator(
+    random_features,
+    input_output_pairs;
+    encoder_schedule = deepcopy(encoder_schedule),
+    encoder_kwargs=deepcopy(encoder_kwargs),
+)
 optimize_hyperparameters!(emulator_random_features)
 ```
 ### Emulator Validation
@@ -342,7 +349,22 @@ and the random features absolute errors are here:
 Both these error maps look similar. Importantly, we want the emulator to show the low errors in the region around the true parameter values near $\theta=(3, 6)$ (i.e, near where the EKI points converge, shown by the scatter points in
 the previous plot). This the region that we will be sampling in the next step. 
 We see low errors near here for both outputs and for both emulators. Now we have validated these emulators, 
-we will proceed the last step of CES: Sampling of the posterior distribution. 
+we will proceed the last step of CES: Sampling of the posterior distribution.
+
+## An aside, Forward map wrapper
+
+What if we run a simple case, and wish to see the emulator performance against the actual forward map sampling in the pipeline? We allow for an exact substitution by using our `ForwardMapWrapper`
+
+```julia
+emulator_fm = forward_map_wrapper(
+    G,
+    prior,
+    input_output_pairs;
+    encoder_schedule = deepcopy(encoder_schedule),
+    encoder_kwargs = deepcopy(encoder_kwargs),
+)
+```
+In this way we can test against the posterior found with the true model, which in this setting is cheap enough to run.
 
 ## Sample
 
@@ -441,14 +463,26 @@ process and we find similar results:
 | parameters        | mean   | std  |
 | :---------------- | :----: | :--: |
 | amplitude         |   2.99 | 0.21 |
-| vert_shift        |   7.06 | 4.89 |
+| vert_shift        |   7.06 | 0.48 |
 
 ![RF_2d_posterior](../assets/sinusoid_MCMC_hist_RF.png)
 
+
+### Sample with the true forward map
+
+| parameters        | mean   | std  |
+| :---------------- | :----: | :--: |
+| amplitude         |   3.00 | 0.22 |
+| vert_shift        |   7.08 | 0.44 |
+
+![FM_2d_posterior](../assets/sinusoid_MCMC_hist_FM.png)
+
+
 It is reassuring to see that our uncertainty quantification methods are robust to the different emulator
 choices here. This is because our particular GP and RF emulators showed similar accuracy during validation. 
 However, this result is highly sensitive to the choices of GP kernel and RF kernel structure. If you find very 
 different posterior distributions for different emulators, it is likely that the kernel choices need be refined.
 The kernel choices must be flexible enough to accurately capture the relationships between the inputs and outputs. 
 We recommend trying a variety of different emulator configurations and carefully considering emulator validation 
 on samples that the emulator has not been trained on. 
+

docs/src/examples/sinusoid_example.md

@@ -77,6 +77,7 @@ gppackage = Emulators.GPJL()
 gauss_proc = Emulators.GaussianProcess(gppackage, noise_learn = false)
 
 
+
 # Build emulator with data
 emulator_gp = Emulator(
     gauss_proc,
@@ -120,13 +121,24 @@ random_features = VectorRandomFeatureInterface(
     optimizer_options = optimizer_options,
 )
 
-emulator_random_features = Emulator(
+emulator_rf = Emulator(
     random_features,
     input_output_pairs;
     encoder_schedule = deepcopy(encoder_schedule),
     encoder_kwargs = deepcopy(encoder_kwargs),
 )
-optimize_hyperparameters!(emulator_random_features)
+optimize_hyperparameters!(emulator_rf)
+
+
+# We also create a forward map wrapper instead of using an emulator.
+# In this way we can test our encoded sampling pipeline with the exact model.
+emulator_fm = forward_map_wrapper(
+    G,
+    prior,
+    input_output_pairs;
+    encoder_schedule = deepcopy(encoder_schedule),
+    encoder_kwargs = deepcopy(encoder_kwargs),
+)
 
 
 ## Emulator Validation
@@ -155,20 +167,28 @@ g_true_grid = reshape(g_true, (output_dim, N_grid, N_grid))
 
 # Next, predict with emulators. We first need to transform to the unconstrained space.
 input_grid_unconstrained = Emulators.transform_constrained_to_unconstrained(prior, input_grid)
+
 gp_mean, gp_cov = Emulators.predict(emulator_gp, input_grid_unconstrained, transform_to_real = true)
-rf_mean, rf_cov = Emulators.predict(emulator_random_features, input_grid_unconstrained, transform_to_real = true)
+rf_mean, rf_cov = Emulators.predict(emulator_rf, input_grid_unconstrained, transform_to_real = true)
+fm_mean, fm_cov = Emulators.predict(emulator_fm, input_grid_unconstrained; transform_to_real = true)
+
 
 # Reshape into (output_dim x 50 x 50) grid
 output_dim = 2
 gp_grid = reshape(gp_mean, (output_dim, N_grid, N_grid))
 rf_grid = reshape(rf_mean, (output_dim, N_grid, N_grid))
+fm_grid = reshape(fm_mean, (output_dim, N_grid, N_grid))
 
 # Convert cov matrix into std and reshape
 gp_std = [sqrt.(diag(gp_cov[j])) for j in 1:length(gp_cov)]
 gp_std_grid = reshape(permutedims(reduce(vcat, [x' for x in gp_std]), (2, 1)), (output_dim, N_grid, N_grid))
 rf_std = [sqrt.(diag(rf_cov[j])) for j in 1:length(rf_cov)]
 rf_std_grid = reshape(permutedims(reduce(vcat, [x' for x in rf_std]), (2, 1)), (output_dim, N_grid, N_grid))
 
+fm_std = [sqrt.(diag(fm_cov[j])) for j in 1:length(fm_cov)]
+fm_std_grid = reshape(permutedims(reduce(vcat, [x' for x in fm_std]), (2, 1)), (output_dim, N_grid, N_grid))
+
+
 ## Plot
 # First, we will plot the ground truth. We have 2 parameters and 2 outputs, so we will create a contour plot 
 # for each output to show how they vary against the two inputs. 
@@ -277,6 +297,43 @@ p = plot(
 )
 savefig(p, joinpath(data_save_directory, "sinusoid_RF_emulator_contours.png"))
 
+# Plot FM emulator contours
+p1 = contour(
+    amp_range,
+    vshift_range,
+    fm_grid[1, :, :];
+    fill = true,
+    clims = range_clims,
+    xlabel = "Amplitude",
+    ylabel = "Vertical Shift",
+    title = "FM Sinusoid Range",
+)
+plot!(inputs[1, :], inputs[2, :]; seriestype = :scatter, zcolor = outputs[1, :], label = :false)
+p2 = contour(
+    amp_range,
+    vshift_range,
+    fm_grid[2, :, :];
+    fill = true,
+    clims = mean_clims,
+    xlabel = "Amplitude",
+    ylabel = "Vertical Shift",
+    title = "FM Sinusoid Mean",
+)
+plot!(inputs[1, :], inputs[2, :]; seriestype = :scatter, zcolor = outputs[2, :], label = :false)
+p = plot(
+    p1,
+    p2,
+    right_margin = 3mm,
+    bottom_margin = 3mm,
+    size = (600, 300),
+    layout = (1, 2),
+    guidefontsize = 12,
+    tickfontsize = 10,
+    legendfontsize = 10,
+)
+savefig(p, joinpath(data_save_directory, "sinusoid_FM_emulator_contours.png"))
+
+
 # Both the GP and RF emulator give similar results to the ground truth G(θ), indicating they are correctly
 # learning the relationships between the parameters and the outputs. We also see the contours agree with the 
 # colors of the training data points. 
@@ -347,6 +404,34 @@ p2 = contour(
 )
 p = plot(p1, p2, size = (600, 300), layout = (1, 2), guidefontsize = 12, tickfontsize = 10, legendfontsize = 10)
 savefig(p, joinpath(data_save_directory, "sinusoid_RF_emulator_std_contours.png"))
+
+# Plot FM std estimates
+p1 = contour(
+    amp_range,
+    vshift_range,
+    fm_std_grid[1, :, :];
+    c = :cividis,
+    fill = true,
+    clims = range_std_clims,
+    xlabel = "Amplitude",
+    ylabel = "Vertical Shift",
+    title = "FM 1σ in Sinusoid Range",
+)
+p2 = contour(
+    amp_range,
+    vshift_range,
+    fm_std_grid[2, :, :];
+    c = :cividis,
+    fill = true,
+    clims = mean_std_clims,
+    xlabel = "Amplitude",
+    ylabel = "Vertical Shift",
+    title = "FM 1σ in Sinusoid Mean",
+)
+p = plot(p1, p2, size = (600, 300), layout = (1, 2), guidefontsize = 12, tickfontsize = 10, legendfontsize = 10)
+savefig(p, joinpath(data_save_directory, "sinusoid_FM_emulator_std_contours.png"))
+
+
 # The GP and RF uncertainty predictions are similar and show lower uncertainties around the region of interest
 # where we have more training points.
 
@@ -380,6 +465,7 @@ p2 = contour(
 p = plot(p1, p2, size = (600, 300), layout = (1, 2), guidefontsize = 12, tickfontsize = 10, legendfontsize = 10)
 savefig(p, joinpath(data_save_directory, "sinusoid_GP_errors_contours.png"))
 
+# Now RF
 rf_diff_grid = abs.(rf_grid - g_true_grid)
 p1 = contour(
     amp_range,
@@ -406,6 +492,33 @@ p2 = contour(
 p = plot(p1, p2, size = (600, 300), layout = (1, 2), guidefontsize = 12, tickfontsize = 10, legendfontsize = 10)
 savefig(p, joinpath(data_save_directory, "sinusoid_RF_errors_contours.png"))
 
+# Now FM
+fm_diff_grid = abs.(fm_grid - g_true_grid)
+p1 = contour(
+    amp_range,
+    vshift_range,
+    fm_diff_grid[1, :, :];
+    c = :cividis,
+    fill = true,
+    clims = range_diff_clims,
+    xlabel = "Amplitude",
+    ylabel = "Vertical Shift",
+    title = "FM error in Sinusoid Range",
+)
+p2 = contour(
+    amp_range,
+    vshift_range,
+    fm_diff_grid[2, :, :];
+    c = :cividis,
+    fill = true,
+    clims = mean_diff_clims,
+    xlabel = "Amplitude",
+    ylabel = "Vertical Shift",
+    title = "FM error in Sinusoid Mean",
+)
+p = plot(p1, p2, size = (600, 300), layout = (1, 2), guidefontsize = 12, tickfontsize = 10, legendfontsize = 10)
+savefig(p, joinpath(data_save_directory, "sinusoid_FM_errors_contours.png"))
+
 # Here, we want the emulator to show the low errors in the region around the true parameter values near θ = (3, 6),
 # as this is the region that we will be sampling in the next step. This appears to the be case for both 
 # outputs and both emulators. 
@@ -415,12 +528,18 @@ println(mean(gp_diff_grid, dims = (2, 3)))
 println(mean(rf_diff_grid, dims = (2, 3)))
 
 
+
+
+
+
 save(
     joinpath(data_save_directory, "emulators.jld2"),
     "emulator_gp",
     emulator_gp,
-    "emulator_random_features",
-    emulator_random_features,
+    "emulator_rf",
+    emulator_rf,
+    "emulator_fm",
+    emulator_fm,
     "prior",
     prior,
     "rng",