test: Update calibration test for the paper

tests/test_simple_shear_2d.py

@@ -376,7 +376,7 @@ def test_dvdx_GBM(self, outdir, seeds_nearX45, ncpus):
         _seeds = seeds_nearX45
         n_seeds = len(_seeds)
         angles = np.empty((len(gbm_mobilities), n_seeds, len(timestamps)))
-        θ_fse = np.empty_like(timestamps)
+        θ_fse = np.zeros_like(timestamps)
         strains = timestamps * strain_rate
         M0_drexF90, M50_drexF90, M200_drexF90 = self.interp_GBM_FortranDRex(strains)
 
@@ -459,7 +459,7 @@ def test_dvdx_GBM(self, outdir, seeds_nearX45, ncpus):
                 np.savez(
                     f"{out_basepath}_angles.npz",
                     angles=result_angles,
-                    err=result_angles_err
+                    err=result_angles_err,
                 )
                 fig, ax, colors = _vis.alignment(
                     None,
@@ -481,8 +481,8 @@ def test_dvdx_GBM(self, outdir, seeds_nearX45, ncpus):
                     ["k", "k"],
                     ["none", None],
                     [
-                        "Zhang & Karato, 1995\n(1200°C)",
-                        "Zhang & Karato, 1995\n(1300°C)",
+                        "Zhang & Karato, 1995\n(1473 K)",
+                        "Zhang & Karato, 1995\n(1573 K)",
                     ],
                     _core.MineralFabric.olivine_A,
                 )
@@ -549,15 +549,16 @@ def test_GBM_calibration(self, outdir, seeds, ncpus):
         shear_direction = [0, 1, 0]  # Used to calculate the angular diagnostics.
         strain_rate = 1
         _, get_velocity_gradient = _velocity.simple_shear_2d("Y", "X", strain_rate)
-        timestamps = np.linspace(0, 5.5, 251)  # Solve until D₀t=5.5 ('shear' γ=11).
+        timestamps = np.linspace(0, 3.2, 65)  # Solve until D₀t=3.2 ('shear' γ=6.4).
         params = _io.DEFAULT_PARAMS
         params["number_of_grains"] = 5000
-        gbm_mobilities = (5, 10, 15, 20)  # Must be in ascending order.
-        markers = (".", "*", "d", "s")
+        gbm_mobilities = (0, 10, 50, 125)  # Must be in ascending order.
+        markers = ("x", "*", "1", ".")
         # Uses 100 seeds by default; use all 1000 if you have more RAM and CPU time.
         _seeds = seeds[:100]
         n_seeds = len(_seeds)
         angles = np.empty((len(gbm_mobilities), n_seeds, len(timestamps)))
+        θ_fse = np.zeros_like(timestamps)
         strains = timestamps * strain_rate
 
         optional_logging = cl.nullcontext()
@@ -569,6 +570,10 @@ def test_GBM_calibration(self, outdir, seeds, ncpus):
         with optional_logging:
             clock_start = process_time()
             for m, gbm_mobility in enumerate(gbm_mobilities):
+                if m == 0:
+                    return_fse = True
+                else:
+                    return_fse = False
                 params["gbm_mobility"] = gbm_mobility
                 _run = ft.partial(
                     self.run,
@@ -577,7 +582,7 @@ def test_GBM_calibration(self, outdir, seeds, ncpus):
                     strain_rate,
                     get_velocity_gradient,
                     shear_direction,
-                    return_fse=False,
+                    return_fse=return_fse,
                 )
                 with Pool(processes=ncpus) as pool:
                     for s, out in enumerate(pool.imap_unordered(_run, _seeds)):
@@ -594,6 +599,11 @@ def test_GBM_calibration(self, outdir, seeds, ncpus):
                                 f"{out_basepath}.npz",
                                 postfix=f"M{_io.stringify(gbm_mobility)}",
                             )
+                        if return_fse:
+                            θ_fse += fse_angles
+
+                if return_fse:
+                    θ_fse /= n_seeds
 
                 if outdir is not None:
                     labels.append(f"$M^∗$ = {params['gbm_mobility']}")
@@ -626,14 +636,23 @@ def test_GBM_calibration(self, outdir, seeds, ncpus):
                     schema,
                     [[int(D * 200) for D in strains]],  # Shear strain % is 200 * D₀.
                 )
+                np.savez(
+                    _io.resolve_path(f"{out_basepath}_ensemble_means.npz"),
+                    angles=result_angles,
+                    err=result_angles_err,
+                )
+                fig = _vis.figure(
+                    figsize=(_vis.DEFAULT_FIG_WIDTH * 3, _vis.DEFAULT_FIG_HEIGHT)
+                )
                 fig, ax, colors = _vis.alignment(
-                    None,
+                    fig.add_subplot(),
                     strains,
                     result_angles,
                     markers,
                     labels,
                     err=result_angles_err,
                     θ_max=80,
+                    θ_fse=θ_fse,
                 )
                 (
                     _,
@@ -666,9 +685,7 @@ def test_GBM_calibration(self, outdir, seeds, ncpus):
                     ],
                     fabric=_core.MineralFabric.olivine_A,
                 )
-                # There is a lot of stuff on this legend, so put it outside the axes.
-                # These values might need to be tweaked depending on the font size, etc.
-                _legend = _utils.redraw_legend(ax, fig=fig, bbox_to_anchor=(1.77, 0.99))
+                _legend = _utils.redraw_legend(ax, loc="upper right", ncols=3)
                 fig.savefig(
                     _io.resolve_path(f"{out_basepath}.pdf"),
                     bbox_extra_artists=(_legend,),
@@ -684,7 +701,7 @@ def test_GBM_calibration(self, outdir, seeds, ncpus):
                                 _angles(
                                     np.take(data_Skemer2016.shear_strain, indices) / 200
                                 ),
-                                np.take(data_Skemer2016.angle, indices)
+                                np.take(data_Skemer2016.angle, indices),
                             )
                         ]
                     )