tests.test_vortex_2d

  1"""> PyDRex: tests for CPO stability in 2D vortex and Stokes cell flows."""
  2
  3import contextlib as cl
  4import functools as ft
  5
  6import numpy as np
  7import pytest
  8from numpy import testing as nt
  9
 10from pydrex import core as _core
 11from pydrex import diagnostics as _diagnostics
 12from pydrex import io as _io
 13from pydrex import logger as _log
 14from pydrex import minerals as _minerals
 15from pydrex import pathlines as _path
 16from pydrex import stats as _stats
 17from pydrex import utils as _utils
 18from pydrex import velocity as _velocity
 19from pydrex import visualisation as _vis
 20from pydrex import geometry as _geo
 21
 22Pool, HAS_RAY = _utils.import_proc_pool()
 23
 24SUBDIR = "2d_vortex"
 25"""Subdirectory of `outdir` used to store outputs from these tests."""
 26
 27
 28def run_singlephase(params: dict, seed: int, assert_each=None, **kwargs) -> tuple:
 29    """Run 2D convection cell simulation for a single mineral phase.
 30
 31    Uses A-type olivine by default.
 32
 33    Args:
 34    - `params` — see `pydrex.core.DefaultParams`
 35    - `seed` — seed for random number generation
 36    - `assert_each` — optional callable with signature `f(mineral,
 37      deformation_gradient)` that performs assertions at each step
 38
 39    Optional keyword args are consumed by:
 40    1. `pydrex.velocity.cell_2d` and
 41    2. the `pydrex.minerals.Mineral` constructor
 42
 43    Returns a tuple containing:
 44    1. the `mineral` (instance of `pydrex.minerals.Mineral`)
 45    2. the resampled texture (a tuple of the orientations and fractions)
 46    3. a tuple of `fig, ax, q, s` as returned from `pydrex.visualisation.pathline_box2d`
 47
 48    """
 49    horizontal: str = kwargs.pop("horizontal", "X")
 50    vertical: str = kwargs.pop("vertical", "Z")
 51    velocity_edge: float = kwargs.pop("velocity_edge", 6.342e-10)
 52    edge_length: float = kwargs.pop("edge_length", 2e5)
 53
 54    max_strain = kwargs.pop(
 55        # This should be enough to go around the cell one time.
 56        "max_strain",
 57        int(np.ceil(velocity_edge * (edge_length / 2) ** 2)),
 58    )
 59    get_velocity, get_velocity_gradient = _velocity.cell_2d(
 60        horizontal,
 61        vertical,
 62        velocity_edge,
 63        edge_length,
 64    )
 65    mineral = _minerals.Mineral(
 66        phase=kwargs.pop("phase", _core.MineralPhase.olivine),
 67        fabric=kwargs.pop("fabric", _core.MineralFabric.olivine_A),
 68        regime=kwargs.pop("regime", _core.DeformationRegime.matrix_dislocation),
 69        n_grains=params["number_of_grains"],
 70        seed=seed,
 71        **kwargs,
 72    )
 73
 74    size = edge_length / 2
 75    dummy_dim = ({0, 1, 2} - set(_geo.to_indices2d(horizontal, vertical))).pop()
 76
 77    timestamps, get_position = _path.get_pathline(
 78        _utils.add_dim([0.5, -0.75], dummy_dim) * size,
 79        get_velocity,
 80        get_velocity_gradient,
 81        _utils.add_dim([-size, -size], dummy_dim),
 82        _utils.add_dim([size, size], dummy_dim),
 83        max_strain,
 84        regular_steps=max_strain * 10,
 85    )
 86    positions = [get_position(t) for t in timestamps]
 87    velocity_gradients = [  # Steady flow, time variable is np.nan.
 88        get_velocity_gradient(np.nan, np.asarray(x)) for x in positions
 89    ]
 90
 91    strains = np.empty_like(timestamps)
 92    strains[0] = 0
 93    deformation_gradient = np.eye(3)
 94
 95    for t, time in enumerate(timestamps[:-1], start=1):
 96        strains[t] = strains[t - 1] + (
 97            _utils.strain_increment(timestamps[t] - time, velocity_gradients[t])
 98        )
 99        _log.info("step %d/%d (ε = %.2f)", t, len(timestamps) - 1, strains[t])
100
101        deformation_gradient = mineral.update_orientations(
102            params,
103            deformation_gradient,
104            get_velocity_gradient,
105            pathline=(time, timestamps[t], get_position),
106        )
107        if assert_each is not None:
108            assert_each(mineral, deformation_gradient)
109
110    orientations, fractions = _stats.resample_orientations(
111        mineral.orientations, mineral.fractions, seed=seed
112    )
113    cpo_strengths = np.full(len(orientations), 1.0)
114    cpo_vectors = [_diagnostics.bingham_average(o) for o in orientations]
115    fig_path, ax_path, q, s = _vis.steady_box2d(
116        None,
117        (get_velocity, [20, 20]),
118        (positions, [-size, -size], [size, size]),
119        horizontal + vertical,
120        (cpo_strengths, cpo_vectors),
121        strains,
122        cmap="cmc.batlow_r",
123        aspect="equal",
124        alpha=1,
125    )
126    fig_path.colorbar(s, ax=ax_path, aspect=25, label="Strain (ε)")
127
128    return mineral, (orientations, fractions), (fig_path, ax_path, q, s)
129
130
131class TestCellOlivineA:
132    """Tests for A-type olivine polycrystals in a 2D Stokes cell."""
133
134    class_id = "cell_olivineA"
135    _ensemble_n_grains = [100, 500, 1000, 5000, 10000]
136
137    @classmethod
138    def _make_ensemble_figure(cls, outdir):
139        # Create the combined figure from outputs of the parametrized ensemble test.
140        data = []
141        out_basepath = f"{outdir}/{SUBDIR}/{cls.class_id}"
142        try:
143            for n_grains in cls._ensemble_n_grains:
144                data.append(np.load(f"{out_basepath}_xz_ensemble_N{n_grains}_data.npz"))
145        except FileNotFoundError:
146            _log.debug(
147                "skipping visualisation of 2D cell ensemble results (missing datafiles)"
148            )
149            return
150
151        fig = _vis.figure()
152        mosaic = fig.subplot_mosaic([["a)"], ["b)"]], sharex=True)
153        fig, mosaic["a)"], colors = _vis.alignment(
154            mosaic["a)"],
155            np.asarray([d["strains"] for d in data]),
156            np.asarray([d["angles_mean"] for d in data]),
157            ("s", "o", "v", "*", "."),
158            list(map(str, cls._ensemble_n_grains)),
159            err=np.asarray([d["angles_err"] for d in data]),
160        )
161        fig, mosaic["b)"], colors = _vis.alignment(
162            mosaic["b)"],
163            np.asarray([d["strains"] for d in data]),
164            np.asarray([d["max_sizes_mean"] for d in data]),
165            ("s", "o", "v", "*", "."),
166            list(map(str, cls._ensemble_n_grains)),
167            err=np.asarray([d["max_sizes_err"] for d in data]),
168            θ_max=4,
169        )
170        mosaic["b)"].set_ylabel(r"$\log_{10}(\overline{S}_{\mathrm{max}})$")
171
172        mosaic["a)"].label_outer()
173        mosaic["b)"].get_legend().remove()
174        fig.savefig(
175            _io.resolve_path(
176                f"{outdir}/{SUBDIR}/{cls.class_id}_xz_ensemble_combined.pdf"
177            )
178        )
179
180    @classmethod
181    def run(
182        cls,
183        params,
184        final_location,
185        get_velocity,
186        get_velocity_gradient,
187        min_coords,
188        max_coords,
189        max_strain,
190        seed=None,
191    ):
192        """Run 2D Stokes cell A-type olivine simulation."""
193        mineral = _minerals.Mineral(
194            phase=_core.MineralPhase.olivine,
195            fabric=_core.MineralFabric.olivine_A,
196            regime=_core.DeformationRegime.matrix_dislocation,
197            n_grains=params["number_of_grains"],
198            seed=seed,
199        )
200        deformation_gradient = np.eye(3)
201
202        timestamps, get_position = _path.get_pathline(
203            final_location,
204            get_velocity,
205            get_velocity_gradient,
206            min_coords,
207            max_coords,
208            max_strain,
209            regular_steps=int(max_strain * 10),
210        )
211        positions = [get_position(t) for t in timestamps]
212        velocity_gradients = [
213            get_velocity_gradient(np.nan, np.asarray(x)) for x in positions
214        ]
215        strains = np.empty_like(timestamps)
216        strains[0] = 0
217        for t, time in enumerate(timestamps[:-1], start=1):
218            strains[t] = strains[t - 1] + (
219                _utils.strain_increment(timestamps[t] - time, velocity_gradients[t])
220            )
221            _log.info("step %d/%d (ε = %.2f)", t, len(timestamps) - 1, strains[t])
222
223            deformation_gradient = mineral.update_orientations(
224                params,
225                deformation_gradient,
226                get_velocity_gradient,
227                pathline=(time, timestamps[t], get_position),
228            )
229        return timestamps, positions, strains, mineral, deformation_gradient
230
231    @pytest.mark.big
232    def test_xz_10k(self, outdir, seed):
233        """Run 2D cell test with 10000 grains (~14GiB RAM requirement)."""
234        self.test_xz(outdir, seed, 10000)
235
236    @pytest.mark.skipif(_utils.in_ci("win32"), reason="Unable to allocate memory")
237    @pytest.mark.parametrize("n_grains", [100, 500, 1000, 5000])
238    def test_xz(self, outdir, seed, n_grains):
239        """Test to check that 5000 grains is "enough" to resolve transient features."""
240        if outdir is not None:
241            out_basepath = f"{outdir}/{SUBDIR}/{self.class_id}_xz_N{n_grains}"
242
243        params = _core.DefaultParams().as_dict()
244        params["number_of_grains"] = n_grains
245        get_velocity, get_velocity_gradient = _velocity.cell_2d("X", "Z", 1)
246
247        timestamps, positions, strains, mineral, deformation_gradient = self.run(
248            params,
249            np.asarray([0.5, 0.0, -0.75]),
250            get_velocity,
251            get_velocity_gradient,
252            np.asarray([-1, 0, -1]),
253            np.asarray([1, 0, 1]),
254            7,
255            seed=seed,
256        )
257        angles = [
258            _diagnostics.smallest_angle(
259                _diagnostics.bingham_average(a, axis="a"), get_velocity(np.nan, x)
260            )
261            for a, x in zip(mineral.orientations, positions, strict=True)
262        ]
263        if outdir is not None:
264            # First figure with the domain and pathline.
265            fig_path, ax_path, q, s = _vis.steady_box2d(
266                None,
267                (get_velocity, [20, 20]),
268                (positions, [-1, -1], [1, 1]),
269                "XZ",
270                None,
271                strains,
272                cmap="cmc.batlow_r",
273                tick_formatter=lambda x, pos: str(x),
274                aspect="equal",
275                alpha=1,
276            )
277            fig_path.colorbar(s, ax=ax_path, aspect=25, label="Strain (ε)")
278            fig_path.savefig(_io.resolve_path(f"{out_basepath}_path.pdf"))
279            # Second figure with the angles and grain sizes at every 10 strain values.
280            fig = _vis.figure()
281            axθ = fig.add_subplot(2, 1, 1)
282            fig, axθ, colors = _vis.alignment(
283                axθ,
284                strains,
285                angles,
286                (".",),
287                (None,),
288                colors=[strains],
289                cmaps=["cmc.batlow_r"],
290            )
291            ax_sizes = fig.add_subplot(2, 1, 2, sharex=axθ)
292            fig, ax_sizes, parts = _vis.grainsizes(
293                ax_sizes, strains[::10], mineral.fractions[::10]
294            )
295            axθ.label_outer()
296            fig.savefig(_io.resolve_path(f"{out_basepath}.pdf"))
297
298        # Some checks for when we should have "enough" grains.
299        # Based on empirical model outputs, it seems like the dip at ε ≈ 3.75 is the
300        # least sensitive feature to the random state (seed) so we will use that.
301        if n_grains >= 5000:
302            # Can we resolve the temporary alignment to below 20° at ε ≈ 3.75?
303            mean_θ_in_dip = np.mean(angles[34:43])
304            assert mean_θ_in_dip < 12, mean_θ_in_dip
305            # Can we resolve corresponding dip in max grain size (normalized, log_10)?
306            mean_size_in_dip = np.log10(
307                np.mean([np.max(f) for f in mineral.fractions[34:43]]) * n_grains
308            )
309            assert 2 < mean_size_in_dip < 3, mean_size_in_dip
310            # Can we resolve subsequent peak in max grain size (normalized, log_10)?
311            max_size_post_dip = np.log10(
312                np.max([np.max(f) for f in mineral.fractions[43:]]) * n_grains
313            )
314            assert max_size_post_dip > 3, max_size_post_dip
315
316    @pytest.mark.slow
317    @pytest.mark.parametrize("n_grains", [100, 500, 1000, 5000, 10000])
318    def test_xz_ensemble(self, outdir, seeds_nearX45, ncpus, n_grains):
319        """Test to demonstrate stability of the dip at ε ≈ 3.75 for 5000+ grains."""
320        _seeds = seeds_nearX45
321        n_seeds = len(_seeds)
322        if outdir is not None:
323            out_basepath = f"{outdir}/{SUBDIR}/{self.class_id}_xz_ensemble_N{n_grains}"
324
325        params = _core.DefaultParams().as_dict()
326        params["number_of_grains"] = n_grains
327        get_velocity, get_velocity_gradient = _velocity.cell_2d("X", "Z", 1)
328
329        _run = ft.partial(
330            self.run,
331            params,
332            np.asarray([0.5, 0.0, -0.75]),
333            get_velocity,
334            get_velocity_gradient,
335            np.asarray([-1, 0, -1]),
336            np.asarray([1, 0, 1]),
337            7,
338        )
339        angles = np.empty((n_seeds, 70))
340        max_sizes = np.empty_like(angles)
341        with Pool(processes=ncpus) as pool:
342            for s, out in enumerate(pool.imap_unordered(_run, _seeds)):
343                timestamps, positions, strains, mineral, deformation_gradient = out
344                angles[s] = [
345                    _diagnostics.smallest_angle(
346                        _diagnostics.bingham_average(a, axis="a"),
347                        get_velocity(np.nan, x),
348                    )
349                    for a, x in zip(mineral.orientations, positions, strict=True)
350                ]
351                max_sizes[s] = np.log10(np.max(mineral.fractions, axis=1) * n_grains)
352
353        if outdir is not None:
354            # Figure with the angles and max grain sizes (ensemble averages).
355            fig = _vis.figure()
356            axθ = fig.add_subplot(2, 1, 1)
357            angles_mean = np.mean(angles, axis=0)
358            angles_err = np.std(angles, axis=0)
359            fig, axθ, colors = _vis.alignment(
360                axθ,
361                strains,
362                angles_mean,
363                (".",),
364                (None,),
365                err=angles_err,
366            )
367            ax_maxsize = fig.add_subplot(2, 1, 2, sharex=axθ)
368            ax_maxsize.set_ylabel(r"$\log_{10}(\overline{S}_{\mathrm{max}})$")
369            max_sizes_mean = np.mean(max_sizes, axis=0)
370            ax_maxsize.plot(strains, max_sizes_mean, color=colors[0])
371            max_sizes_err = np.std(max_sizes, axis=0)
372            ax_maxsize.fill_between(
373                strains,
374                max_sizes_mean - max_sizes_err,
375                max_sizes_mean + max_sizes_err,
376                alpha=0.22,
377                color=colors[0],
378            )
379            axθ.label_outer()
380            fig.savefig(_io.resolve_path(f"{out_basepath}.pdf"))
381            np.savez(
382                _io.resolve_path(f"{out_basepath}_data.npz"),
383                strains=strains,
384                max_sizes_mean=max_sizes_mean,
385                max_sizes_err=max_sizes_err,
386                angles_mean=angles_mean,
387                angles_err=angles_err,
388            )
389
390
391class TestDiffusionCreep:
392    """Tests for diffusion creep regime."""
393
394    class_id = "diff_creep"
395
396    @pytest.mark.skipif(_utils.in_ci("win32"), reason="Unable to allocate memory")
397    def test_cell_olA(self, outdir, seed, ncpus, orientations_init_y):
398        params = _core.DefaultParams().as_dict()
399        params["gbm_mobility"] = 10
400
401        def get_assert_each(i):  # The order of orientations_init_y is significant.
402            @_utils.serializable
403            def assert_each(mineral, deformation_gradient):
404                # Check that surrogate grain sizes are not changing.
405                nt.assert_allclose(
406                    mineral.fractions[-1], mineral.fractions[-2], atol=1e-16, rtol=0
407                )
408                p, g, r = _diagnostics.symmetry_pgr(mineral.orientations[-1])
409                nt.assert_allclose(
410                    np.array([p, g, r]),
411                    _diagnostics.symmetry_pgr(mineral.orientations[-2]),
412                    atol=0.25,
413                    rtol=0,
414                )
415                match i:
416                    case 0:
417                        # Check that symmetry remains mostly random.
418                        assert r > 0.9, f"{r}"
419                    case 1:
420                        # Check that symmetry remains mostly girdled.
421                        assert g > 0.9, f"{g}"
422                    case 2:
423                        # Check that symmetry remains mostly clustered.
424                        assert p > 0.9, f"{g}"
425
426            return assert_each
427
428        @_utils.serializable
429        def _run(assert_each, orientations_init):
430            return run_singlephase(
431                params,
432                seed,
433                regime=_core.DeformationRegime.matrix_diffusion,
434                orientations_init=orientations_init,
435            )
436
437        optional_logging = cl.nullcontext()
438        if outdir is not None:
439            out_basepath = f"{outdir}/{SUBDIR}/{self.class_id}_olA"
440            optional_logging = _io.logfile_enable(f"{out_basepath}.log")
441
442        assert_each_list = [
443            get_assert_each(i) for i, _ in enumerate(orientations_init_y)
444        ]
445        orientations_init_list = [
446            f(params["number_of_grains"]) for f in orientations_init_y
447        ]
448        with optional_logging:
449            with Pool(processes=ncpus) as pool:
450                for i, out in enumerate(
451                    pool.starmap(_run, zip(assert_each_list, orientations_init_list))
452                ):
453                    mineral, resampled_texture, fig_objects = out
454                    if outdir is not None:
455                        fig_objects[0].savefig(
456                            _io.resolve_path(f"{out_basepath}_path_{i}.pdf")
457                        )