tests.test_simple_shear_2d

  1"""> PyDRex: 2D simple shear tests."""
  2
  3import contextlib as cl
  4import functools as ft
  5import sys
  6from time import process_time
  7
  8import numpy as np
  9import pytest
 10from numpy import asarray as Ŋ
 11from numpy import testing as nt
 12from scipy.interpolate import PchipInterpolator
 13
 14from pydrex import core as _core
 15from pydrex import diagnostics as _diagnostics
 16from pydrex import geometry as _geo
 17from pydrex import io as _io
 18from pydrex import logger as _log
 19from pydrex import minerals as _minerals
 20from pydrex import pathlines as _paths
 21from pydrex import stats as _stats
 22from pydrex import utils as _utils
 23from pydrex import velocity as _velocity
 24from pydrex import visualisation as _vis
 25
 26Pool, HAS_RAY = _utils.import_proc_pool()
 27
 28# Subdirectory of `outdir` used to store outputs from these tests.
 29SUBDIR = "2d_simple_shear"
 30
 31
 32class TestPreliminaries:
 33    """Preliminary tests to check that various auxiliary routines are working."""
 34
 35    def test_strain_increment(self):
 36        """Test for accumulating strain via strain increment calculations."""
 37        _, get_velocity_gradient = _velocity.simple_shear_2d("X", "Z", 1)
 38        timestamps = np.linspace(0, 1, 10)  # Solve until D₀t=1 (tensorial strain).
 39        strains_inc = np.zeros_like(timestamps)
 40        L = get_velocity_gradient(np.nan, Ŋ([0e0, 0e0, 0e0]))
 41        for i, ε in enumerate(strains_inc[1:]):
 42            strains_inc[i + 1] = strains_inc[i] + _utils.strain_increment(
 43                timestamps[1] - timestamps[0],
 44                L,
 45            )
 46        # For constant timesteps, check strains == positive_timestamps * strain_rate.
 47        nt.assert_allclose(strains_inc, timestamps, atol=6e-16, rtol=0)
 48
 49        # Same thing, but for strain rate similar to experiments.
 50        _, get_velocity_gradient = _velocity.simple_shear_2d("Y", "X", 1e-5)
 51        timestamps = np.linspace(0, 1e6, 10)  # Solve until D₀t=10 (tensorial strain).
 52        strains_inc = np.zeros_like(timestamps)
 53        L = get_velocity_gradient(np.nan, Ŋ([0e0, 0e0, 0e0]))
 54        for i, ε in enumerate(strains_inc[1:]):
 55            strains_inc[i + 1] = strains_inc[i] + _utils.strain_increment(
 56                timestamps[1] - timestamps[0],
 57                L,
 58            )
 59        nt.assert_allclose(strains_inc, timestamps * 1e-5, atol=5e-15, rtol=0)
 60
 61        # Again, but this time the particle will move (using get_pathline).
 62        # We use a 400km x 400km box and a strain rate of 1e-15 s⁻¹.
 63        get_velocity, get_velocity_gradient = _velocity.simple_shear_2d("X", "Z", 1e-15)
 64        timestamps, get_position = _paths.get_pathline(
 65            Ŋ([1e5, 0e0, 1e5]),
 66            get_velocity,
 67            get_velocity_gradient,
 68            Ŋ([-2e5, 0e0, -2e5]),
 69            Ŋ([2e5, 0e0, 2e5]),
 70            2,
 71            regular_steps=10,
 72        )
 73        positions = [get_position(t) for t in timestamps]
 74        velocity_gradients = [get_velocity_gradient(np.nan, Ŋ(x)) for x in positions]
 75
 76        # Check that polycrystal is experiencing steady velocity gradient.
 77        nt.assert_array_equal(
 78            velocity_gradients, np.full_like(velocity_gradients, velocity_gradients[0])
 79        )
 80        # Check that positions are changing as expected.
 81        xdiff = np.diff(Ŋ([x[0] for x in positions]))
 82        zdiff = np.diff(Ŋ([x[2] for x in positions]))
 83        assert xdiff[0] > 0
 84        assert zdiff[0] == 0
 85        nt.assert_allclose(xdiff, np.full_like(xdiff, xdiff[0]), rtol=0, atol=1e-10)
 86        nt.assert_allclose(zdiff, np.full_like(zdiff, zdiff[0]), rtol=0, atol=1e-10)
 87        strains_inc = np.zeros_like(timestamps)
 88        for t, time in enumerate(timestamps[:-1], start=1):
 89            strains_inc[t] = strains_inc[t - 1] + (
 90                _utils.strain_increment(timestamps[t] - time, velocity_gradients[t])
 91            )
 92        # fig, ax, _, _ = _vis.pathline_box2d(
 93        #     None,
 94        #     get_velocity,
 95        #     "xz",
 96        #     strains_inc,
 97        #     positions,
 98        #     ".",
 99        #     Ŋ([-2e5, -2e5]),
100        #     Ŋ([2e5, 2e5]),
101        #     [20, 20],
102        # )
103        # fig.savefig("/tmp/fig.png")
104        nt.assert_allclose(
105            strains_inc,
106            (timestamps - timestamps[0]) * 1e-15,
107            atol=5e-15,
108            rtol=0,
109        )
110
111
112class TestOlivineA:
113    """Tests for stationary A-type olivine polycrystals in 2D simple shear."""
114
115    class_id = "olivineA"
116
117    @classmethod
118    def run(
119        cls,
120        params,
121        timestamps,
122        strain_rate,
123        get_velocity_gradient,
124        shear_direction,
125        seed=None,
126        return_fse=None,
127        get_position=lambda t: np.zeros(3),  # Stationary particles by default.
128    ):
129        """Reusable logic for 2D olivine (A-type) simple shear tests.
130
131        Returns a tuple with the mineral and the FSE angle (or `None` if `return_fse` is
132        `None`).
133
134        """
135        mineral = _minerals.Mineral(
136            phase=_core.MineralPhase.olivine,
137            fabric=_core.MineralFabric.olivine_A,
138            regime=_core.DeformationRegime.matrix_dislocation,
139            n_grains=params["number_of_grains"],
140            seed=seed,
141        )
142        deformation_gradient = np.eye(3)  # Undeformed initial state.
143        θ_fse = np.empty_like(timestamps)
144        θ_fse[0] = 45
145
146        for t, time in enumerate(timestamps[:-1], start=1):
147            # Set up logging message depending on dynamic parameter and seeds.
148            msg_start = (
149                f"N = {params['number_of_grains']}; "
150                + f"λ∗ = {params['nucleation_efficiency']}; "
151                + f"X = {params['gbs_threshold']}; "
152                + f"M∗ = {params['gbm_mobility']}; "
153            )
154            if seed is not None:
155                msg_start += f"# {seed}; "
156
157            _log.info(msg_start + "step %s/%s (t = %s)", t, len(timestamps) - 1, time)
158
159            deformation_gradient = mineral.update_orientations(
160                params,
161                deformation_gradient,
162                get_velocity_gradient,
163                pathline=(time, timestamps[t], get_position),
164            )
165            _log.debug(
166                "› velocity gradient = %s",
167                get_velocity_gradient(np.nan, np.full(3, np.nan)).flatten(),
168            )
169            _log.debug("› strain D₀t = %.2f", strain_rate * timestamps[t])
170            _log.debug(
171                "› grain fractions: median = %s, max = %s, min = %s",
172                np.median(mineral.fractions[-1]),
173                np.max(mineral.fractions[-1]),
174                np.min(mineral.fractions[-1]),
175            )
176            if return_fse:
177                _, fse_v = _diagnostics.finite_strain(deformation_gradient)
178                θ_fse[t] = _diagnostics.smallest_angle(fse_v, shear_direction)
179            else:
180                θ_fse = None
181
182        return mineral, θ_fse
183
184    @classmethod
185    def interp_GBM_Kaminski2001(cls, strains):
186        """Interpolate Kaminski & Ribe, 2001 data to get target angles at `strains`."""
187        _log.info("interpolating target CPO angles...")
188        data = _io.read_scsv(_io.data("thirdparty") / "Kaminski2001_GBMshear.scsv")
189        cs_M0 = PchipInterpolator(
190            _utils.remove_nans(data.equivalent_strain_M0) / 200,
191            _utils.remove_nans(data.angle_M0),
192        )
193        cs_M50 = PchipInterpolator(
194            _utils.remove_nans(data.equivalent_strain_M50) / 200,
195            _utils.remove_nans(data.angle_M50),
196        )
197        cs_M200 = PchipInterpolator(
198            _utils.remove_nans(data.equivalent_strain_M200) / 200,
199            _utils.remove_nans(data.angle_M200),
200        )
201        return [cs_M0(strains), cs_M50(strains), cs_M200(strains)]
202
203    @classmethod
204    def interp_GBM_FortranDRex(cls, strains):
205        """Interpolate angles produced using 'tools/drex_forward_simpleshear.f90'."""
206        _log.info("interpolating target CPO  angles...")
207        data = _io.read_scsv(_io.data("drexF90") / "olA_D1E4_dt50_X0_L5.scsv")
208        data_strains = np.linspace(0, 1, 200)
209        cs_M0 = PchipInterpolator(data_strains, _utils.remove_nans(data.M0_angle))
210        cs_M50 = PchipInterpolator(data_strains, _utils.remove_nans(data.M50_angle))
211        cs_M200 = PchipInterpolator(data_strains, _utils.remove_nans(data.M200_angle))
212        return [cs_M0(strains), cs_M50(strains), cs_M200(strains)]
213
214    @classmethod
215    def interp_GBS_FortranDRex(cls, strains):
216        """Interpolate angles produced using 'tools/drex_forward_simpleshear.f90'."""
217        _log.info("interpolating target CPO  angles...")
218        data = _io.read_scsv(_io.data("thirdparty") / "a_axis_GBS_fortran.scsv")
219        data_strains = np.linspace(0, 1, 200)
220        cs_X0 = PchipInterpolator(data_strains, _utils.remove_nans(data.a_mean_X0))
221        cs_X20 = PchipInterpolator(data_strains, _utils.remove_nans(data.a_mean_X20))
222        cs_X40 = PchipInterpolator(data_strains, _utils.remove_nans(data.a_mean_X40))
223        return [cs_X0(strains), cs_X20(strains), cs_X40(strains)]
224
225    @classmethod
226    def interp_GBS_long_FortranDRex(cls, strains):
227        """Interpolate angles produced using 'tools/drex_forward_simpleshear.f90'."""
228        _log.info("interpolating target CPO  angles...")
229        data = _io.read_scsv(_io.data("thirdparty") / "a_axis_GBS_long_fortran.scsv")
230        data_strains = np.linspace(0, 2.5, 500)
231        cs_X0 = PchipInterpolator(data_strains, _utils.remove_nans(data.a_mean_X0))
232        cs_X20 = PchipInterpolator(data_strains, _utils.remove_nans(data.a_mean_X20))
233        cs_X40 = PchipInterpolator(data_strains, _utils.remove_nans(data.a_mean_X40))
234        return [cs_X0(strains), cs_X20(strains), cs_X40(strains)]
235
236    @classmethod
237    def interp_GBS_Kaminski2004(cls, strains):
238        """Interpolate Kaminski & Ribe, 2001 data to get target angles at `strains`."""
239        _log.info("interpolating target CPO angles...")
240        data = _io.read_scsv(_io.data("thirdparty") / "Kaminski2004_GBSshear.scsv")
241        cs_X0 = PchipInterpolator(
242            _utils.remove_nans(data.dimensionless_time_X0),
243            45 + _utils.remove_nans(data.angle_X0),
244        )
245        cs_X0d2 = PchipInterpolator(
246            _utils.remove_nans(data.dimensionless_time_X0d2),
247            45 + _utils.remove_nans(data.angle_X0d2),
248        )
249        cs_X0d4 = PchipInterpolator(
250            _utils.remove_nans(data.dimensionless_time_X0d4),
251            45 + _utils.remove_nans(data.angle_X0d4),
252        )
253        return [cs_X0(strains), cs_X0d2(strains), cs_X0d4(strains)]
254
255    @pytest.mark.skipif(_utils.in_ci("win32"), reason="Unable to allocate memory")
256    def test_zero_recrystallisation(self, seed):
257        """Check that M*=0 is a reliable switch to turn off recrystallisation."""
258        params = _core.DefaultParams().as_dict()
259        params["gbm_mobility"] = 0
260        strain_rate = 1
261        timestamps = np.linspace(0, 1, 25)  # Solve until D₀t=1 (tensorial strain).
262        shear_direction = Ŋ([0, 1, 0], dtype=np.float64)
263        _, get_velocity_gradient = _velocity.simple_shear_2d("Y", "X", strain_rate)
264        mineral, _ = self.run(
265            params,
266            timestamps,
267            strain_rate,
268            get_velocity_gradient,
269            shear_direction,
270            seed=seed,
271        )
272        for fractions in mineral.fractions[1:]:
273            nt.assert_allclose(fractions, mineral.fractions[0], atol=1e-15, rtol=0)
274
275    @pytest.mark.skipif(_utils.in_ci("win32"), reason="Unable to allocate memory")
276    @pytest.mark.parametrize("gbm_mobility", [50, 100, 150])
277    def test_grainsize_median(self, seed, gbm_mobility):
278        """Check that M*={50,100,150}, λ*=5 causes decreasing grain size median."""
279        params = _core.DefaultParams().as_dict()
280        params["gbm_mobility"] = gbm_mobility
281        params["nucleation_efficiency"] = 5
282        strain_rate = 1
283        timestamps = np.linspace(0, 1, 25)  # Solve until D₀t=1 (tensorial strain).
284        n_timestamps = len(timestamps)
285        shear_direction = Ŋ([0, 1, 0], dtype=np.float64)
286        _, get_velocity_gradient = _velocity.simple_shear_2d("Y", "X", strain_rate)
287        mineral, _ = self.run(
288            params,
289            timestamps,
290            strain_rate,
291            get_velocity_gradient,
292            shear_direction,
293            seed=seed,
294        )
295        medians = np.empty(n_timestamps)
296        for i, fractions in enumerate(mineral.fractions):
297            medians[i] = np.median(fractions)
298
299        # The first diff is positive (~1e-6) before nucleation sets in.
300        nt.assert_array_less(np.diff(medians)[1:], np.full(n_timestamps - 2, 0))
301
302    @pytest.mark.slow
303    @pytest.mark.parametrize("gbs_threshold", [0, 0.2, 0.4])
304    @pytest.mark.parametrize("nucleation_efficiency", [3, 5, 10])
305    def test_dvdx_ensemble(
306        self, outdir, seeds_nearX45, ncpus, gbs_threshold, nucleation_efficiency
307    ):
308        r"""Test a-axis alignment to shear in Y direction (init. SCCS near 45° from X).
309
310        Velocity gradient:
311        $$\bm{L} = \begin{bmatrix} 0 & 0 & 0 \cr 2 & 0 & 0 \cr 0 & 0 & 0 \end{bmatrix}$$
312
313        """
314        strain_rate = 1
315        timestamps = np.linspace(0, 1, 201)  # Solve until D₀t=1 ('shear' γ=2).
316        n_timestamps = len(timestamps)
317        # Use `seeds` instead of `seeds_nearX45` if you have even more RAM and CPU time.
318        _seeds = seeds_nearX45
319        n_seeds = len(_seeds)
320
321        shear_direction = Ŋ([0, 1, 0], dtype=np.float64)
322        _, get_velocity_gradient = _velocity.simple_shear_2d("Y", "X", strain_rate)
323
324        gbm_mobilities = [0, 50, 125, 200]
325        markers = ("x", "*", "d", "s")
326
327        _id = f"X{_io.stringify(gbs_threshold)}_L{_io.stringify(nucleation_efficiency)}"
328        # Output setup with optional logging and data series labels.
329        θ_fse = np.empty_like(timestamps)
330        angles = np.empty((len(gbm_mobilities), n_seeds, n_timestamps))
331        optional_logging = cl.nullcontext()
332        if outdir is not None:
333            out_basepath = f"{outdir}/{SUBDIR}/{self.class_id}_dvdx_ensemble_{_id}"
334            optional_logging = _io.logfile_enable(f"{out_basepath}.log")
335            labels = []
336
337        with optional_logging:
338            clock_start = process_time()
339            for m, gbm_mobility in enumerate(gbm_mobilities):
340                if m == 0:
341                    return_fse = True
342                else:
343                    return_fse = False
344
345                params = {
346                    "phase_assemblage": (_core.MineralPhase.olivine,),
347                    "phase_fractions": (1.0,),
348                    "enstatite_fraction": 0.0,
349                    "stress_exponent": 1.5,
350                    "deformation_exponent": 3.5,
351                    "gbm_mobility": gbm_mobility,
352                    "gbs_threshold": gbs_threshold,
353                    "nucleation_efficiency": nucleation_efficiency,
354                    "number_of_grains": 5000,
355                    "initial_olivine_fabric": "A",
356                }
357
358                _run = ft.partial(
359                    self.run,
360                    params,
361                    timestamps,
362                    strain_rate,
363                    get_velocity_gradient,
364                    shear_direction,
365                    return_fse=return_fse,
366                )
367                with Pool(processes=ncpus) as pool:
368                    for s, out in enumerate(pool.imap_unordered(_run, _seeds)):
369                        mineral, fse_angles = out
370                        angles[m, s, :] = [
371                            _diagnostics.smallest_angle(v, shear_direction)
372                            for v in _diagnostics.elasticity_components(
373                                _minerals.voigt_averages([mineral], params)
374                            )["hexagonal_axis"]
375                        ]
376                        # Save the whole mineral for the first seed only.
377                        if outdir is not None and s == 0:
378                            postfix = (
379                                f"M{_io.stringify(gbm_mobility)}"
380                                + f"_X{_io.stringify(gbs_threshold)}"
381                                + f"_L{_io.stringify(nucleation_efficiency)}"
382                            )
383                            mineral.save(f"{out_basepath}.npz", postfix=postfix)
384                        if return_fse:
385                            θ_fse += fse_angles
386
387                if return_fse:
388                    θ_fse /= n_seeds
389
390                if outdir is not None:
391                    labels.append(f"$M^∗$ = {gbm_mobility}")
392
393            _log.info("elapsed CPU time: %s", np.abs(process_time() - clock_start))
394
395        # Take ensemble means and optionally plot figure.
396        strains = timestamps * strain_rate
397        _log.info("postprocessing results for %s", _id)
398        result_angles = angles.mean(axis=1)
399        result_angles_err = angles.std(axis=1)
400
401        if outdir is not None:
402            schema = {
403                "delimiter": ",",
404                "missing": "-",
405                "fields": [
406                    {
407                        "name": "strain",
408                        "type": "integer",
409                        "unit": "percent",
410                        "fill": 999999,
411                    }
412                ],
413            }
414            np.savez(
415                f"{out_basepath}.npz",
416                angles=result_angles,
417                angles_err=result_angles_err,
418            )
419            _io.save_scsv(
420                f"{out_basepath}_strains.scsv",
421                schema,
422                [[int(D * 200) for D in strains]],  # Shear strain % is 200 * D₀.
423            )
424            fig, ax, colors = _vis.alignment(
425                None,
426                strains,
427                result_angles,
428                markers,
429                labels,
430                err=result_angles_err,
431                θ_max=60,
432                θ_fse=θ_fse,
433            )
434            fig.savefig(_io.resolve_path(f"{out_basepath}.pdf"))
435
436    @pytest.mark.slow
437    def test_dvdx_GBM(self, outdir, seeds_nearX45, ncpus):
438        r"""Test a-axis alignment to shear in Y direction (init. SCCS near 45° from X).
439
440        Velocity gradient:
441        $$
442        \bm{L} = 10^{-4} ×
443            \begin{bmatrix} 0 & 0 & 0 \cr 2 & 0 & 0 \cr 0 & 0 & 0 \end{bmatrix}
444        $$
445
446        Results are compared to the Fortran 90 output.
447
448        """
449        shear_direction = Ŋ([0, 1, 0], dtype=np.float64)
450        strain_rate = 1e-4
451        _, get_velocity_gradient = _velocity.simple_shear_2d("Y", "X", strain_rate)
452        timestamps = np.linspace(0, 1e4, 51)  # Solve until D₀t=1 ('shear' γ=2).
453        i_strain_40p = 10  # Index of 40% strain, lower strains are not relevant here.
454        i_strain_100p = 25  # Index of 100% strain, when M*=0 matches FSE.
455        params = _core.DefaultParams().as_dict()
456        params["gbs_threshold"] = 0  # No GBS, to match the Fortran parameters.
457        gbm_mobilities = (0, 10, 50, 125, 200)  # Must be in ascending order.
458        markers = ("x", ".", "*", "d", "s")
459        # Use `seeds` instead of `seeds_nearX45` if you have even more RAM and CPU time.
460        _seeds = seeds_nearX45
461        n_seeds = len(_seeds)
462        angles = np.empty((len(gbm_mobilities), n_seeds, len(timestamps)))
463        θ_fse = np.zeros_like(timestamps)
464        strains = timestamps * strain_rate
465        M0_drexF90, M50_drexF90, M200_drexF90 = self.interp_GBM_FortranDRex(strains)
466
467        optional_logging = cl.nullcontext()
468        if outdir is not None:
469            out_basepath = f"{outdir}/{SUBDIR}/{self.class_id}_mobility"
470            optional_logging = _io.logfile_enable(f"{out_basepath}.log")
471            labels = []
472
473        with optional_logging:
474            clock_start = process_time()
475            for m, gbm_mobility in enumerate(gbm_mobilities):
476                if m == 0:
477                    return_fse = True
478                else:
479                    return_fse = False
480                params["gbm_mobility"] = gbm_mobility
481
482                _run = ft.partial(
483                    self.run,
484                    params,
485                    timestamps,
486                    strain_rate,
487                    get_velocity_gradient,
488                    shear_direction,
489                    return_fse=True,
490                )
491                with Pool(processes=ncpus) as pool:
492                    for s, out in enumerate(pool.imap_unordered(_run, _seeds)):
493                        mineral, fse_angles = out
494                        angles[m, s, :] = [
495                            _diagnostics.smallest_angle(v, shear_direction)
496                            for v in _diagnostics.elasticity_components(
497                                _minerals.voigt_averages([mineral], params)
498                            )["hexagonal_axis"]
499                        ]
500                        # Save the whole mineral for the first seed only.
501                        if outdir is not None and s == 0:
502                            mineral.save(
503                                f"{out_basepath}.npz",
504                                postfix=f"M{_io.stringify(gbm_mobility)}",
505                            )
506                        if return_fse:
507                            θ_fse += fse_angles
508
509                if return_fse:
510                    θ_fse /= n_seeds
511
512                if outdir is not None:
513                    labels.append(f"$M^∗$ = {params['gbm_mobility']}")
514
515            _log.info("elapsed CPU time: %s", np.abs(process_time() - clock_start))
516
517            # Take ensemble means and optionally plot figure.
518            _log.info("postprocessing results...")
519            result_angles = angles.mean(axis=1)
520            result_angles_err = angles.std(axis=1)
521
522            if outdir is not None:
523                schema = {
524                    "delimiter": ",",
525                    "missing": "-",
526                    "fields": [
527                        {
528                            "name": "strain",
529                            "type": "integer",
530                            "unit": "percent",
531                            "fill": 999999,
532                        }
533                    ],
534                }
535                _io.save_scsv(
536                    f"{out_basepath}_strains.scsv",
537                    schema,
538                    [[int(D * 200) for D in strains]],  # Shear strain % is 200 * D₀.
539                )
540                np.savez(
541                    f"{out_basepath}_angles.npz",
542                    angles=result_angles,
543                    err=result_angles_err,
544                )
545                fig, ax, colors = _vis.alignment(
546                    None,
547                    strains,
548                    result_angles,
549                    markers,
550                    labels,
551                    err=result_angles_err,
552                    θ_max=60,
553                    θ_fse=θ_fse,
554                )
555                ax.plot(strains, M0_drexF90, c=colors[0])
556                ax.plot(strains, M50_drexF90, c=colors[2])
557                ax.plot(strains, M200_drexF90, c=colors[4])
558                _vis.show_Skemer2016_ShearStrainAngles(
559                    ax,
560                    ["Z&K 1200 C", "Z&K 1300 C"],
561                    ["v", "^"],
562                    ["k", "k"],
563                    ["none", None],
564                    [
565                        "Zhang & Karato, 1995\n(1473 K)",
566                        "Zhang & Karato, 1995\n(1573 K)",
567                    ],
568                    _core.MineralFabric.olivine_A,
569                )
570                # There is a lot of stuff on this legend, so put it outside the axes.
571                # These values might need to be tweaked depending on the font size, etc.
572                _legend = _utils.redraw_legend(ax, fig=fig, bbox_to_anchor=(1.66, 0.99))
573                fig.savefig(
574                    _io.resolve_path(f"{out_basepath}.pdf"),
575                    bbox_extra_artists=(_legend,),
576                    bbox_inches="tight",
577                )
578
579            # Check that GBM speeds up the alignment between 40% and 100% strain.
580            _log.info("checking grain orientations...")
581            for i, θ in enumerate(result_angles[:-1], start=1):
582                nt.assert_array_less(
583                    result_angles[i][i_strain_40p:i_strain_100p],
584                    θ[i_strain_40p:i_strain_100p],
585                )
586
587            # Check that M*=0 matches FSE (±1°) past 100% strain.
588            nt.assert_allclose(
589                result_angles[0][i_strain_100p:],
590                θ_fse[i_strain_100p:],
591                atol=1,
592                rtol=0,
593            )
594
595            # Check that results match Fortran output past 40% strain.
596            nt.assert_allclose(
597                result_angles[0][i_strain_40p:],
598                M0_drexF90[i_strain_40p:],
599                atol=0,
600                rtol=0.1,  # At 40% strain the match is worse than at higher strain.
601            )
602            nt.assert_allclose(
603                result_angles[2][i_strain_40p:],
604                M50_drexF90[i_strain_40p:],
605                atol=1,
606                rtol=0,
607            )
608            nt.assert_allclose(
609                result_angles[4][i_strain_40p:],
610                M200_drexF90[i_strain_40p:],
611                atol=1.5,
612                rtol=0,
613            )
614
615    @pytest.mark.slow
616    def test_GBM_calibration(self, outdir, seeds, ncpus):
617        r"""Compare results for various values of $$M^∗$$ to A-type olivine data.
618
619        Velocity gradient:
620        $$
621        \bm{L} = 10^{-4} ×
622            \begin{bmatrix} 0 & 0 & 0 \cr 2 & 0 & 0 \cr 0 & 0 & 0 \end{bmatrix}
623        $$
624
625        Unlike `test_dvdx_GBM`,
626        grain boudary sliding is enabled here (see `_core.DefaultParams`).
627        Data are provided by [Skemer & Hansen, 2016](http://dx.doi.org/10.1016/j.tecto.2015.12.003).
628
629        """
630        shear_direction = Ŋ([0, 1, 0], dtype=np.float64)
631        strain_rate = 1
632        _, get_velocity_gradient = _velocity.simple_shear_2d("Y", "X", strain_rate)
633        timestamps = np.linspace(0, 3.2, 65)  # Solve until D₀t=3.2 ('shear' γ=6.4).
634        params = _core.DefaultParams().as_dict()
635        params["number_of_grains"] = 5000
636        gbm_mobilities = (0, 10, 50, 125)  # Must be in ascending order.
637        markers = ("x", "*", "1", ".")
638        # Uses 100 seeds by default; use all 1000 if you have more RAM and CPU time.
639        _seeds = seeds[:100]
640        n_seeds = len(_seeds)
641        angles = np.empty((len(gbm_mobilities), n_seeds, len(timestamps)))
642        θ_fse = np.zeros_like(timestamps)
643        strains = timestamps * strain_rate
644
645        optional_logging = cl.nullcontext()
646        if outdir is not None:
647            out_basepath = f"{outdir}/{SUBDIR}/{self.class_id}_calibration"
648            optional_logging = _io.logfile_enable(f"{out_basepath}.log")
649            labels = []
650
651        with optional_logging:
652            clock_start = process_time()
653            for m, gbm_mobility in enumerate(gbm_mobilities):
654                return_fse = True if m == 0 else False
655                params["gbm_mobility"] = gbm_mobility
656                _run = ft.partial(
657                    self.run,
658                    params,
659                    timestamps,
660                    strain_rate,
661                    get_velocity_gradient,
662                    shear_direction,
663                    return_fse=return_fse,
664                )
665                with Pool(processes=ncpus) as pool:
666                    for s, out in enumerate(pool.imap_unordered(_run, _seeds)):
667                        mineral, fse_angles = out
668                        angles[m, s, :] = [
669                            _diagnostics.smallest_angle(v, shear_direction)
670                            for v in _diagnostics.elasticity_components(
671                                _minerals.voigt_averages([mineral], params)
672                            )["hexagonal_axis"]
673                        ]
674                        # Save the whole mineral for the first seed only.
675                        if outdir is not None and s == 0:
676                            mineral.save(
677                                f"{out_basepath}.npz",
678                                postfix=f"M{_io.stringify(gbm_mobility)}",
679                            )
680                        if return_fse:
681                            θ_fse += fse_angles
682
683                if return_fse:
684                    θ_fse /= n_seeds
685
686                if outdir is not None:
687                    labels.append(f"$M^∗$ = {params['gbm_mobility']}")
688
689            _log.info("elapsed CPU time: %s", np.abs(process_time() - clock_start))
690
691            # Take ensemble means and optionally plot figure.
692            _log.info("postprocessing results...")
693            result_angles = angles.mean(axis=1)
694            result_angles_err = angles.std(axis=1)
695
696            if outdir is not None:
697                schema = {
698                    "delimiter": ",",
699                    "missing": "-",
700                    "fields": [
701                        {
702                            "name": "strain",
703                            "type": "integer",
704                            "unit": "percent",
705                            "fill": 999999,
706                        }
707                    ],
708                }
709                _io.save_scsv(
710                    f"{out_basepath}_strains.scsv",
711                    schema,
712                    [[int(D * 200) for D in strains]],  # Shear strain % is 200 * D₀.
713                )
714                np.savez(
715                    _io.resolve_path(f"{out_basepath}_ensemble_means.npz"),
716                    angles=result_angles,
717                    err=result_angles_err,
718                )
719                fig = _vis.figure(
720                    figsize=(_vis.DEFAULT_FIG_WIDTH * 3, _vis.DEFAULT_FIG_HEIGHT)
721                )
722                fig, ax, colors = _vis.alignment(
723                    fig.add_subplot(),
724                    strains,
725                    result_angles,
726                    markers,
727                    labels,
728                    err=result_angles_err,
729                    θ_max=80,
730                    θ_fse=θ_fse,
731                )
732                (
733                    _,
734                    _,
735                    _,
736                    data_Skemer2016,
737                    indices,
738                ) = _vis.show_Skemer2016_ShearStrainAngles(
739                    ax,
740                    [
741                        "Z&K 1200 C",
742                        "Z&K 1300 C",
743                        "Skemer 2011",
744                        "Hansen 2014",
745                        "Warren 2008",
746                        "Webber 2010",
747                        "H&W 2015",
748                    ],
749                    ["v", "^", "o", "s", "v", "o", "s"],
750                    ["k", "k", "k", "k", "k", "k", "k"],
751                    ["none", "none", "none", "none", None, None, None],
752                    [
753                        "Zhang & Karato, 1995 (1473 K)",
754                        "Zhang & Karato, 1995 (1573 K)",
755                        "Skemer et al., 2011 (1500 K)",
756                        "Hansen et al., 2014 (1473 K)",
757                        "Warren et al., 2008",
758                        "Webber et al., 2010",
759                        "Hansen & Warren, 2015",
760                    ],
761                    fabric=_core.MineralFabric.olivine_A,
762                )
763                _legend = _utils.redraw_legend(ax, loc="upper right", ncols=3)
764                fig.savefig(
765                    _io.resolve_path(f"{out_basepath}.pdf"),
766                    bbox_extra_artists=(_legend,),
767                    bbox_inches="tight",
768                )
769                r2vals = []
770                for angles in result_angles:
771                    _angles = PchipInterpolator(strains, angles)
772                    r2 = np.sum(
773                        [
774                            (a - b) ** 2
775                            for a, b in zip(
776                                _angles(
777                                    np.take(data_Skemer2016.shear_strain, indices) / 200
778                                ),
779                                np.take(data_Skemer2016.angle, indices),
780                            )
781                        ]
782                    )
783                    r2vals.append(r2)
784                _log.info(
785                    "Sums of squared residuals (r-values) for each M∗: %s", r2vals
786                )
787
788    @pytest.mark.big
789    def test_dudz_pathline(self, outdir, seed):
790        """Test alignment of olivine a-axis for a polycrystal advected on a pathline."""
791        test_id = "dudz_pathline"
792        optional_logging = cl.nullcontext()
793        if outdir is not None:
794            out_basepath = f"{outdir}/{SUBDIR}/{self.class_id}_{test_id}"
795            optional_logging = _io.logfile_enable(f"{out_basepath}.log")
796
797        with optional_logging:
798            shear_direction = Ŋ([1, 0, 0], dtype=np.float64)
799            strain_rate = 1e-15  # Moderate, realistic shear in the upper mantle.
800            get_velocity, get_velocity_gradient = _velocity.simple_shear_2d(
801                "X", "Z", strain_rate
802            )
803            n_timesteps = 10
804            timestamps, get_position = _paths.get_pathline(
805                Ŋ([1e5, 0e0, 1e5]),
806                get_velocity,
807                get_velocity_gradient,
808                Ŋ([-2e5, 0e0, -2e5]),
809                Ŋ([2e5, 0e0, 2e5]),
810                2,
811                regular_steps=n_timesteps,
812            )
813            positions = [get_position(t) for t in timestamps]
814            velocity_gradients = [
815                get_velocity_gradient(np.nan, Ŋ(x)) for x in positions
816            ]
817
818            params = _core.DefaultParams().as_dict()
819            params["gbm_mobility"] = 10
820            params["number_of_grains"] = 5000
821            olA = _minerals.Mineral(n_grains=params["number_of_grains"], seed=seed)
822            deformation_gradient = np.eye(3)
823            strains = np.zeros_like(timestamps)
824            for t, time in enumerate(timestamps[:-1], start=1):
825                strains[t] = strains[t - 1] + (
826                    _utils.strain_increment(timestamps[t] - time, velocity_gradients[t])
827                )
828                _log.info("step %d/%d (ε = %.2f)", t, len(timestamps) - 1, strains[t])
829                deformation_gradient = olA.update_orientations(
830                    params,
831                    deformation_gradient,
832                    get_velocity_gradient,
833                    pathline=(time, timestamps[t], get_position),
834                )
835
836            if outdir is not None:
837                olA.save(f"{out_basepath}.npz")
838
839            orient_resampled, fractions_resampled = _stats.resample_orientations(
840                olA.orientations, olA.fractions, seed=seed
841            )
842            # About 36GB, 26 min needed with float64. GitHub macos runner has 14GB.
843            misorient_indices = _diagnostics.misorientation_indices(
844                orient_resampled,
845                _geo.LatticeSystem.orthorhombic,
846                ncpus=3,
847            )
848            cpo_vectors = np.zeros((n_timesteps + 1, 3))
849            cpo_angles = np.zeros(n_timesteps + 1)
850            for i, matrices in enumerate(orient_resampled):
851                cpo_vectors[i] = _diagnostics.bingham_average(
852                    matrices,
853                    axis=_minerals.OLIVINE_PRIMARY_AXIS[olA.fabric],
854                )
855                cpo_angles[i] = _diagnostics.smallest_angle(
856                    cpo_vectors[i], shear_direction
857                )
858
859            # Check for mostly decreasing CPO angles (exclude initial condition).
860            _log.debug("cpo angles: %s", cpo_angles)
861            nt.assert_array_less(np.diff(cpo_angles[1:]), np.ones(n_timesteps - 1))
862            # Check for increasing CPO strength (M-index).
863            _log.debug("cpo strengths: %s", misorient_indices)
864            nt.assert_array_less(
865                np.full(n_timesteps, -0.01), np.diff(misorient_indices)
866            )
867            # Check that last angle is <5° (M*=125) or <10° (M*=10).
868            assert cpo_angles[-1] < 10
869
870        if outdir is not None:
871            fig, ax, _, _ = _vis.steady_box2d(
872                None,
873                (get_velocity, [20, 20]),
874                (positions, Ŋ([-2e5, -2e5]), Ŋ([2e5, 2e5])),
875                "xz",
876                (cpo_vectors, misorient_indices),
877                strains,
878            )
879            fig.savefig(f"{out_basepath}.pdf")