diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 030e5b58..f74d8aeb 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -29,6 +29,7 @@ jobs: python-version: '3.12' - name: Build + uses: extractions/setup-just@v1 run: | sudo apt-get update sudo apt-get install -y python3-pip @@ -36,7 +37,7 @@ jobs: python3 -m venv .venv-gh-pages source .venv-gh-pages/bin/activate python3 -m pip install "$PWD"[doc,test] - make html + just html - name: Deploy uses: peaceiris/actions-gh-pages@v4 if: ${{ startsWith(github.ref, 'refs/tags/') }} diff --git a/Makefile b/Makefile deleted file mode 100644 index 471c91a1..00000000 --- a/Makefile +++ /dev/null @@ -1,39 +0,0 @@ -.SHELLFLAGS += -u -.ONESHELL: -MAKEFLAGS += --no-builtin-rules -VERSION = $(shell python -m setuptools_scm -f plain) - -# NOTE: Keep this at the top! Bare `make` should just build the sdist + bdist. -dist: - python -m build - -release: dist - python -m twine upload dist/* - -test: - mkdir out - pytest -v --outdir=out - -# WARNING: --math fetches .js code from a CDN, be careful where it comes from: -# https://github.com/mitmproxy/pdoc/security/advisories/GHSA-5vgj-ggm4-fg62 -html: - 2>pdoc.log pdoc -t docs/template -o html pydrex !pydrex.mesh !pydrex.distributed tests \ - --favicon "https://raw.githubusercontent.com/seismic-anisotropy/PyDRex/main/docs/assets/favicon32.png" \ - --footer-text "PyDRex $(VERSION)" \ - --math - -# WARNING: --math fetches .js code from a CDN, be careful where it comes from: -# https://github.com/mitmproxy/pdoc/security/advisories/GHSA-5vgj-ggm4-fg62 -live_docs: - pdoc -t docs/template pydrex !pydrex.mesh !pydrex.distributed tests \ - --favicon "https://raw.githubusercontent.com/seismic-anisotropy/PyDRex/main/docs/assets/favicon32.png" \ - --footer-text "PyDRex $(VERSION)" \ - --math - -clean: - rm -rf dist - rm -rf out - rm -rf html - rm -rf pdoc.log - -.PHONY: release test live_docs clean diff --git a/README.md b/README.md index 00a108c1..616890d8 100644 --- a/README.md +++ b/README.md @@ -6,24 +6,25 @@ #### Simulate crystallographic preferred orientation evolution in polycrystals -This repository contains a Python 3 reimplementation of the D-Rex model +PyDRex is a new Python 3 implementation of the D-Rex model for the evolution of crystallographic preferred orientation in polycrystals. The code is available for use under the [GNU GPL v3](https://www.gnu.org/licenses/gpl-3.0.en.html) license. Documentation is accessible via Python's REPL `help()` and also [online](https://seismic-anisotropy.github.io/PyDRex/). ## Installing -Check `requires-python` in `pyproject.toml` for the minimum required Python -version. The software is tested on Linux, MacOS and Windows, however large -simulations require substantial computational resources usually afforded by HPC -Linux clusters. Linux shell scripts for setting up a [Ray](https://www.ray.io/) cluster +Tagged package versions can be installed from [PyPi](https://pypi.org/project/pydrex/). +The minimum required Python version is displayed in the sidebar (search for `Requires: Python`). +PyDRex is tested on Linux, MacOS and Windows, +however large simulations require substantial computational resources +usually only afforded by HPC Linux clusters. +Linux shell scripts for setting up a [Ray](https://www.ray.io/) cluster on distributed memory PBS systems are provided in the `tools` folder. Optional mesh generation using [gmsh](https://pypi.org/project/gmsh/) is available, however the `gmsh` module requires the [`glu`](https://gitlab.freedesktop.org/mesa/glu) library (which may require manual installation on some systems). -Tagged package versions can be installed from [PyPi](https://pypi.org/project/pydrex/). To install the latest version, execute: pip install pydrex @@ -56,13 +57,15 @@ They have their own README file as well. ## Building offline documentation -The documentation can be built offline from a git clone or source distribution. +The documentation can be built offline from a git clone or unpacked source distribution. Install documentation builder dependencies with `pip install '.[doc]'`. +Developers are also recommended to download the [just](https://github.com/casey/just) command runner. +Otherwise, build commands can be found in the provided `justfile`. -Run `make html` from the terminal to generate PyDRex's documentation +Run `just html` from the terminal to generate PyDRex's documentation (available in the `html` directory), including the API reference. The homepage will be `html/index.html`. -Alternatively, run `make live_docs` to build and serve the documentation on a `localhost` port. +Alternatively, run `just live_docs` to build and serve the documentation on a `localhost` port. Follow the displayed prompts to open the live documentation in a browser. It should automatically reload after changes to the source code. @@ -72,3 +75,8 @@ For a Linux or MacOS development environment, clone the [source code](https://gi and execute the Bash script `tools/venv_install.sh`. This will set up a local Python virtual environment with an [editable install](https://setuptools.pypa.io/en/latest/userguide/development_mode.html) of PyDRex that can be activated/deactivated by following the displayed prompts. + +For documentation contributions, please note that +only `h2` headings will be rendered within HTML anchor elements. +Avoid using markdown headings with 3 or more leading hashes in most cases, +as these will not appear in the documentation sidebar nor be linkable. diff --git a/docs/assets/cornerflow2d_simple_example.png b/docs/assets/cornerflow2d_simple_example.png new file mode 100644 index 00000000..ec712e8f Binary files /dev/null and b/docs/assets/cornerflow2d_simple_example.png differ diff --git a/docs/template/index.html.jinja2 b/docs/template/index.html.jinja2 index c1ab9959..131e0d7b 100644 --- a/docs/template/index.html.jinja2 +++ b/docs/template/index.html.jinja2 @@ -4,9 +4,11 @@ {% block nav %} -

{{ footer_text }}

-
+
+
+
{{ footer_text }}
+
{% endblock %} {% block content %} @@ -15,13 +17,37 @@ {% filter to_html %} # Welcome to PyDRex! -> This is the documentation website for [PyDRex](https://github.com/seismic-anisotropy/PyDRex), +> This is the documentation website for [PyDRex](https://pypi.org/project/pydrex/), a Python implementation of the D-Rex model for kinematic texture simulations. -Use the links below to browse the API documentation or the tests and examples: +Click the buttons below to browse the documentation or read about the tests and examples. + {% endfilter %} + +
+ +
+ + +
+ +
+ + {% filter to_html %} +> The `pydrex` package is published on PyPI (the Python Package Index). +The source code is hosted on GitHub, with an issue tracker and discussion page. {% endfilter %} -
-
+
+ +
+ + +
{% include "search.html.jinja2" %} diff --git a/docs/template/theme.css b/docs/template/theme.css index 03b59811..702307a4 100644 --- a/docs/template/theme.css +++ b/docs/template/theme.css @@ -38,8 +38,17 @@ } } -.homediv { padding-top: 56px; max-width: 1080px; } +/* Center search box and version number in nav bars. */ +/* Then fix up internal div display with some margins. */ +input { display: block; margin: auto; } +nav h6 { text-align: center; } +nav div { padding-right: 8%; } + +/* Add padding for nav bar footer in case of long version numbers (dev docs). */ +nav footer { padding-right: 5px; } +/* Styles for greeting page and buttons. */ +.homediv { padding-top: 56px; max-width: 1080px; } .button { background-color: var(--muted); color: var(--pdoc-background); @@ -55,10 +64,11 @@ cursor: pointer; width: 48%; } - .button:hover { background-color: var(--text); color: var(--active); } +/* Render display equations in horizontally scrollable box. */ +/* Also tries to hide the scroll bar (sometimes it doesn't work). */ .katex-display { overflow: auto hidden } diff --git a/examples/standalone/cornerflow_simple.py b/examples/standalone/cornerflow_simple.py new file mode 100755 index 00000000..f7b9564c --- /dev/null +++ b/examples/standalone/cornerflow_simple.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python +import doctest +import functools as ft +import logging +from dataclasses import dataclass + +import matplotlib as mpl +import numpy as np + +import pydrex +from pydrex import pathlines +from pydrex import visualisation as vis + + +# Define simulation parameters. Check to ensure serializability with a simple doctest. +@dataclass(frozen=True) +class Parameters(pydrex.DefaultParams): # Inherit default PyDRex parameters. + """Parameters for a simple cornerflow example. + + >>> params = Parameters() + >>> # Evaluating hash() will raise an error if there are mutable attributes. + >>> isinstance(hash(params), int) + True + + """ + + # Add parameters that are specific to this simulation. + # You MUST use type annotations, otherwise these will be set as class attributes! + # + plate_speed: float = 2.0 / (100 * 365 * 86400) + domain_height: float = 2e5 + domain_width: float = 1e6 + n_timesteps: int = 50 + pathline_ends: tuple = (-0.1, -0.3, -0.54, -0.78) # % of domain height, negative z. + domain_coords: tuple = ("X", "Z") + max_strain: float = 10.0 + out_figure: str = "cornerflow2d_simple_example.png" + out_logfile: str = "cornerflow2d_simple_example.log" + + +doctest.testmod() # Run doctests for serializability check. + +# Initialise parameter store. Optionally modify default PyDRex parameters. +params = Parameters( + phase_assemblage=(pydrex.MineralPhase.olivine, pydrex.MineralPhase.enstatite), + phase_fractions=(0.7, 0.3), + gbm_mobility=10, + number_of_grains=5000, +) + +# Get velocity and velocity gradient callables. +f_velocity, f_velocity_grad = pydrex.velocity.corner_2d( + *params.domain_coords, params.plate_speed +) + +# Turn params.pathline_ends into an array of 3D coordinates. +final_locations = [ + np.array([params.domain_width, 0.0, z * params.domain_height]) + for z in params.pathline_ends +] + + +# Set up CPO computation for a pathline. +def run(params, f_velocity, f_velocity_grad, min_coords, max_coords, final_location): + # Get pydrex logger, prefer logging messages instead of using `print()`. + logger = logging.getLogger("pydrex") + # Create mineral phase. Each phase is tracked separately. + olA = pydrex.Mineral( + pydrex.MineralPhase.olivine, + pydrex.MineralFabric.olivine_A, + pydrex.DeformationRegime.matrix_dislocation, + n_grains=params.number_of_grains, + ) + ens = pydrex.Mineral( + pydrex.MineralPhase.enstatite, + pydrex.MineralFabric.enstatite_AB, + pydrex.DeformationRegime.matrix_dislocation, + n_grains=params.number_of_grains, + ) + # Calculate pathline by solving backwards from the final location. + timestamps, f_position = pathlines.get_pathline( + final_location, + f_velocity, + f_velocity_grad, + min_coords, + max_coords, + max_strain=params.max_strain, + regular_steps=params.n_timesteps, + ) + positions = [f_position(t) for t in timestamps] + velocity_gradients = [f_velocity_grad(np.nan, np.asarray(x)) for x in positions] + + # All minerals start undeformed, at zero strain. + deformation_gradient = np.identity(3) + strains = np.empty_like(timestamps) + strains[0] = 0 + + # Main CPO computation loop. + for i, time in enumerate(timestamps[:-1], start=1): + strains[i] = strains[i - 1] + pydrex.utils.strain_increment( + timestamps[i] - time, velocity_gradients[i] + ) + logger.info( + "path_id = %d; step = %d/%d (ε = %.2f)", + hash(abs(final_location[-1])), + i, + params.n_timesteps, + strains[i], + ) + deformation_gradient = pydrex.update_all( + (olA, ens), + params.as_dict(), + deformation_gradient, + f_velocity_grad, + pathline=(time, timestamps[i], f_position), + ) + return timestamps, positions, strains, (olA, ens), deformation_gradient + + +# Set up storage for results. +results: list[dict] = [] +all_strains = [] +all_positions = [] + +# Get a multiprocessing Pool and record which backend is being used. +Pool, HAS_RAY = pydrex.utils.import_proc_pool() + +# Solve CPO on multiple pathlines in parallel using multiprocessing. +with pydrex.io.logfile_enable(params.out_logfile): + # Using more than 4 cores is not needed and will introduce wasted overhead. + n_cpus = min(4, pydrex.utils.default_ncpus()) + with Pool(processes=n_cpus) as pool: + _run = ft.partial( + run, + params, + f_velocity, + f_velocity_grad, + np.array([0.0, 0.0, -params.domain_height]), + np.array([params.domain_width, 0.0, 0.0]), + # Different final_location depending on pathline. + ) + for i, out in enumerate(pool.imap(_run, final_locations)): + times, positions, strains, (olA, ens), deformation_gradient = out + # Only 1 phase here, 100% olivine. + avg_stiffnesses = pydrex.minerals.voigt_averages( + (olA, ens), (olA.phase, ens.phase), params.phase_fractions + ) + result = pydrex.diagnostics.elasticity_components(avg_stiffnesses) + all_strains.append(strains) + all_positions.append(positions) + results.append(result) + +# Simple visualisation using provided plotting function. +fig = vis.figure(figsize=(10, 3)) +ax = fig.add_subplot() +markers = ("s", "o", "v", "*") +cmaps = ["cmc.batlow_r"] * len(markers) + +for i, z_exit in enumerate(np.array(params.pathline_ends) * params.domain_height): + velocity_grid_res = None # Resolution of grid used to compute streamlines. + density = None # Density of actual plotted streamlines. + if i == 0: # Only plot velocity streamlines once. + velocity_grid_res = [50, 10] + density = [0.05 * r for r in velocity_grid_res] + + # Set up plotting function arguments. + velocity = (f_velocity, velocity_grid_res) + geometry = ( + all_positions[i], + [0, -params.domain_height], + [params.domain_width + 1e5, 0], + ) + ref_axes = "".join(params.domain_coords) + + # The M-index (`pydrex.diagnostics.misoientation_index`) is a much better measure of + # texture strength, but is very expensive to compute. For this simple example, we + # just use the percentage of the average elastic tensor that deviates from the + # closest isotropic elastic tensor. + cpo = (results[i]["percent_anisotropy"], results[i]["hexagonal_axis"]) + + norm = mpl.colors.Normalize(vmin=all_strains[0][0], vmax=all_strains[0][-1]) + fig, ax, q, s = vis.steady_box2d( + ax, + velocity, + geometry, + ref_axes, + cpo, + colors=all_strains[i], + marker=markers[i], + cmap=cmaps[i], + norm=norm, + density=density, + ) +fig.colorbar( + mpl.cm.ScalarMappable(norm=norm, cmap=cmaps[0]), + ax=ax, + aspect=75, + location="top", + label="Strain (ε)", +) +fig.savefig(params.out_figure) diff --git a/justfile b/justfile new file mode 100644 index 00000000..3a2f12cb --- /dev/null +++ b/justfile @@ -0,0 +1,33 @@ +VERSION := `python -m setuptools_scm -f plain` + +build: + python -m build + +release: build + python -m twine upload dist/* + +test: + mkdir out + pytest -v --outdir=out + +html: + # WARNING: --math fetches .js code from a CDN, be careful where it comes from: + # https://github.com/mitmproxy/pdoc/security/advisories/GHSA-5vgj-ggm4-fg62 + 2>pdoc.log pdoc -t docs/template -o html pydrex !pydrex.mesh !pydrex.distributed tests \ + --favicon "https://raw.githubusercontent.com/seismic-anisotropy/PyDRex/main/docs/assets/favicon32.png" \ + --footer-text "PyDRex {{VERSION}}" \ + --math + +live_docs: + # WARNING: --math fetches .js code from a CDN, be careful where it comes from: + # https://github.com/mitmproxy/pdoc/security/advisories/GHSA-5vgj-ggm4-fg62 + pdoc -t docs/template pydrex !pydrex.mesh !pydrex.distributed tests \ + --favicon "https://raw.githubusercontent.com/seismic-anisotropy/PyDRex/main/docs/assets/favicon32.png" \ + --footer-text "PyDRex {{VERSION}}" \ + --math + +clean: + rm -rf dist + rm -rf out + rm -rf html + rm -rf pdoc.log diff --git a/src/pydrex/__init__.py b/src/pydrex/__init__.py index 15fa6a67..f0dd9baf 100644 --- a/src/pydrex/__init__.py +++ b/src/pydrex/__init__.py @@ -7,44 +7,57 @@ **This software is currently in early development (alpha) and therefore subject to breaking changes without notice.** -## About - -Viscoplastic deformation of minerals, e.g. in Earth's mantle, leads to distinct -signatures in the mineral texture. Many minerals naturally occur in -polycrystalline form, which means that they are composed of many grains with -different volumes and lattice orientations. Preferential alignment of the -average lattice orientation is called crystallographic preferred orientation -(CPO). PyDRex simulates the development and evolution of CPO in deforming -polycrystals, as well as tracking macroscopic finite strain measures. -Currently, the code supports olivine and enstatite mineral phases. The -following features are provided: -- JIT-compiled CPO solver, based on the D-Rex model, which updates the - polycrystal orientation distribution depending on the macroscopic velocity - gradients -- `Mineral` class which stores attributes of a distinct mineral phase in the - polycrystal and its texture snapshots -- Voigt averaging to calculate the average elastic tensor of a textured, - multiphase polycrystal -- Decomposition of average elastic tensors into components attributed to +## Introduction + +Viscoplastic deformation of minerals, e.g. in Earth's mantle, leads to +distinct signatures in the mineral texture. Many minerals naturally +occur in polycrystalline form, which means that they are composed of +many grains with different volumes and lattice orientations. Preferential +alignment of the average lattice orientation is called crystallographic +preferred orientation (CPO). PyDRex simulates the development +and evolution of CPO in deforming polycrystals, as well as tracking +macroscopic finite strain measures. Currently, the code supports olivine +and enstatite mineral phases. **These are some of the main features of PyDRex:** + +- **JIT-compiled CPO solver**, based on the D-Rex model[^1], which updates + the polycrystal orientation distribution depending on the macroscopic + velocity gradients + +- **Voigt averaging** to calculate the average elastic tensor of + a textured, multiphase polycrystal + +- **Decomposition of elastic tensors** into components attributed to minerals with distinct lattice symmetries -- Crystallographic pole figure visualisation (contouring is a work in progress) -- [work in progress] Texture diagnostics: M-index, bingham average, - Point-Girdle-Random symmetry, coaxial a.k.a "BA" index, etc. -- [work in progress] Seismic anisotropy diagnostics: % tensorial anisotropy, + +- **Texture diagnostics:** M-index, Bingham average, Point-Girdle-Random + symmetry, coaxial a.k.a "BA" index, etc. + +- **Seismic anisotropy diagnostics**: percent tensorial anisotropy, hexagonal symmetry a.k.a transverse isotropy direction, etc. -The core CPO solver is based on the original Fortran 90 implementation by Édouard Kaminski, -which can be [downloaded from this link (~90KB)](http://www.ipgp.fr/~kaminski/web_doudoud/DRex.tar.gz). -The reference papers are [Kaminski & Ribe (2001)](https://doi.org/10.1016/s0012-821x(01)00356-9) -and [Kaminski & Ribe (2004)](https://doi.org/10.1111%2Fj.1365-246x.2004.02308.x), -and an open-access paper which discusses the model is [Fraters & Billen (2021)](https://doi.org/10.1029/2021gc009846). +- Crystallographic **pole figures** for visualisation[⁰](#wip) + +- Finite element **mesh generator** for simple rectangular meshes using + [gmsh](https://pypi.org/project/gmsh/) + +- Steady-state **analytical velocity fields** for various idealised flows[⁰](#wip) + +[^1]: The core CPO solver is based on the original Fortran 90 implementation by Édouard Kaminski, + which can be [downloaded from this link (~90KB)](http://www.ipgp.fr/~kaminski/web_doudoud/DRex.tar.gz). + The reference papers are [Kaminski & Ribe (2001)](https://doi.org/10.1016/s0012-821x(01)00356-9) + and [Kaminski & Ribe (2004)](https://doi.org/10.1111%2Fj.1365-246x.2004.02308.x), + and an open-access paper which discusses the model is [Fraters & Billen (2021)](https://doi.org/10.1029/2021gc009846). ## Installation -The minimum required Python version is set using `requires-python` in the -[`pyproject.toml`](https://github.com/seismic-anisotropy/PyDRex/blob/main/pyproject.toml) file. -For installation instructions, -see [the README](https://github.com/seismic-anisotropy/PyDRex/blob/main/README.md) file. +The minimum required Python version is displayed in the +[PyPI package](https://pypi.org/project/pydrex/) metadata +(search that page for `Requires: Python`). +To install the latest version of PyDRex for a compatible Python version, execute: + + pip install pydrex + +See the PyPI page for more installation options. ## Documentation @@ -82,7 +95,7 @@ where $b$ is the length of the Burgers' vector, $σ$ is the stress and $μ$ is the shear modulus. The value of the exponent $p$ is given by the `stress_exponent` input parameter. For an overview of available parameters, -see [the `pydrex.mock` source code, for now...] +see `pydrex.core.DefaultParams`. The effects of dynamic recrystallization are twofold. Grains with a higher than average dislocation density may be affected by either grain nucleation, which is @@ -103,16 +116,35 @@ boundary migration in nucleated grains. It also manifests as an upper bound on texture strength. -## Parameter reference +## Simple examples -Model parameters will eventually be provided in a `.toml` file. -For now just pass a dictionary to `config` in the `Mineral.update_orientations` method. -A draft of the input file spec is shown below: +These simple examples demonstrate the basics of using PyDRex +and provide a starting point for integrating PyDRex into more complex simulations. -```toml -.. include:: data/specs/spec.toml +### Steady, viscous 2D flow in a corner + +This example simulates CPO development during flow along four pathlines in a domain with +two impermeable boundaries (top and left sides) and two permeable boundaries (bottom and +right sides). The velocity is prescribed using a built-in analytical solution, which +approximates the flow field for material involved in mid-ocean ridge upwellings. + +```python +.. include:: ../../examples/standalone/cornerflow_simple.py ``` +![Simple 2D corner flow](https://raw.githubusercontent.com/seismic-anisotropy/PyDRex/main/docs/assets/cornerflow2d_simple_example.png) + +--- + + +
    +
  1. These features are only partially implemented so far. Check the + open issues + to see details about known bugs and planned features. +
  2. +
+ + """ # Set up the top-level pydrex namespace for convenient usage. diff --git a/src/pydrex/minerals.py b/src/pydrex/minerals.py index 034abd68..9e811dc2 100644 --- a/src/pydrex/minerals.py +++ b/src/pydrex/minerals.py @@ -697,6 +697,10 @@ def voigt_averages( - `phase_assemblage` — collection of unique mineral phases in the aggregate - `phase_fractions` — collection of volume fractions for each phase in `phase_assemblage` (values should sum to 1). + - `elastic_tensors` — optional elastic tensors for mineral phases, in 6x6 Voigt + matrix format, ordered according to `phase_assemblage` (to implement custom + stiffness tensors, either modify the attributes of a `StiffnessTensors` instance + or use a subclass) Raises a ValueError if the minerals contain an unequal number of grains or stored texture results. @@ -717,29 +721,20 @@ def voigt_averages( # TODO: Perform rotation directly on the 6x6 matrices, see Carcione 2007. # This trick is implemented in cpo_elastic_tensor.cc in Aspect. + + # Generate 4th order tensors. + phase_tensors = [_tensors.voigt_to_elastic_tensor(S) for S in elastic_tensors] + average_tensors = np.zeros((n_steps, 6, 6)) for i in range(n_steps): for mineral in minerals: for n in range(n_grains): - match mineral.phase: - case _core.MineralPhase.olivine: - average_tensors[i] += _tensors.elastic_tensor_to_voigt( - _tensors.rotate( - elastic_tensors.olivine, - mineral.orientations[i][n, ...].transpose(), - ) - * mineral.fractions[i][n] - * phase_fractions[phase_assemblage.index(mineral.phase)] - ) - case _core.MineralPhase.enstatite: - average_tensors[i] += _tensors.elastic_tensor_to_voigt( - _tensors.rotate( - elastic_tensors.enstatite, - mineral.orientations[i][n, ...].transpose(), - ) - * mineral.fractions[i][n] - * phase_fractions[phase_assemblage.index(mineral.phase)] - ) - case _: - raise ValueError(f"unsupported mineral phase: {mineral.phase}") + average_tensors[i] += _tensors.elastic_tensor_to_voigt( + _tensors.rotate( + phase_tensors[phase_assemblage.index(mineral.phase)], + mineral.orientations[i][n, ...].transpose(), + ) + * mineral.fractions[i][n] + * phase_fractions[phase_assemblage.index(mineral.phase)] + ) return average_tensors diff --git a/src/pydrex/visualisation.py b/src/pydrex/visualisation.py index c5542860..fc2cad00 100644 --- a/src/pydrex/visualisation.py +++ b/src/pydrex/visualisation.py @@ -186,33 +186,41 @@ def steady_box2d( ¹with signature `f(t, x)` where `t` is not used and `x` is a 3D position vector - Additional keyword arguments are passed to the `matplotlib.axes.Axes.quiver` call - used to plot the velocity vectors. + Additional keyword arguments are passed to the `matplotlib.axes.Axes.streamplot` + call used to plot the velocity field, except `'norm'` and `'width'`, which are + passed to `matplotlib.axes.Axes.quiver` (used to plot the CPO bars). Returns the figure handle, the axes handle, the quiver collection (velocities) and the scatter collection (pathline). """ + # Set up base figure and main axis labels. fig, ax = figure_unless(ax) ax.set_xlabel(f"{ref_axes[0]} {label_suffix}") ax.set_ylabel(f"{ref_axes[1]} {label_suffix}") + # Unpack callable + data input tuples. get_velocity, resolution = velocity positions, min_coords, max_coords = geometry x_min, y_min = min_coords x_max, y_max = max_coords + # Apply axis settings. ax.set_xlim((x_min, x_max)) ax.set_ylim((y_min, y_max)) ax.set_aspect(aspect) ax.xaxis.set_major_formatter(tick_formatter) ax.yaxis.set_major_formatter(tick_formatter) - _ref_axes = ref_axes.lower() - axes_map = {"x": 0, "y": 1, "z": 2} - horizontal = axes_map[_ref_axes[0]] - vertical = axes_map[_ref_axes[1]] + # Unpack arguments relevant to velocity streamplot (rest of kwargs also go there). + horizontal, vertical = _geo.to_indices2d(*ref_axes) + zorder = kwargs.pop("zorder", 10) + # Unpack arguments relevant to CPO bars. + norm = kwargs.pop("norm", None) + width = kwargs.pop("width", 3e-3) + + # Plot velocity streamlines. velocities = None if resolution is not None: x_res, y_res = resolution @@ -230,19 +238,21 @@ def steady_box2d( U[i] = v3d[horizontal] V[i] = v3d[vertical] - velocities = ax.quiver( + velocities = ax.streamplot( X_grid, Y_grid, U.reshape(X_grid.shape), V.reshape(Y_grid.shape), - pivot=kwargs.pop("pivot", "mid"), - alpha=kwargs.pop("alpha", 0.25), + color=kwargs.pop("color", "#ff002244"), + density=kwargs.pop("density", tuple([r * 0.1 for r in resolution])), + zorder=zorder, **kwargs, ) - dummy_dim = ({0, 1, 2} - set(_geo.to_indices2d(*ref_axes))).pop() + # Plot CPO bars. + dummy_dim = ({0, 1, 2} - {horizontal, vertical}).pop() xi_2D = np.asarray([_utils.remove_dim(p, dummy_dim) for p in positions]) - qcoll: plt.Quiver | plt.PathCollection + qcoll: plt.Quiver | plt.PathCollection # This variable can be either type. if cpo is None: qcoll = ax.scatter(xi_2D[:, 0], xi_2D[:, 1], marker=marker, c=colors, cmap=cmap) else: @@ -260,11 +270,12 @@ def steady_box2d( cpo_2D[:, 1], colors, cmap=cmap, + norm=norm, pivot="mid", - width=kwargs.pop("width", 3e-3), + width=width, headaxislength=0, headlength=0, - zorder=kwargs.pop("zorder", 10) + 1, # Always above velocity vectors. + zorder=zorder + 1, # Always above velocity vectors. ) return fig, ax, velocities, qcoll diff --git a/tests/README.md b/tests/README.md index ea0e8e76..ab0d68f8 100644 --- a/tests/README.md +++ b/tests/README.md @@ -1,7 +1,8 @@ # PyDRex tests Running the base test suite requires [pytest](https://docs.pytest.org) and up to ~16GB RAM. -By default, `make test` will attempt to write persistent logs and output to `./out/`. +If the [just](https://github.com/casey/just) command runner is installed, +`just test` will run the test suite and attempt to write persistent logs and output to `./out/`. Alternatively, run `pytest` from the root of the source tree. To print more verbose test progress and information, use the flag `pytest -v`. diff --git a/tests/test_vortex_2d.py b/tests/test_vortex_2d.py index c72aaf19..55a04d2e 100644 --- a/tests/test_vortex_2d.py +++ b/tests/test_vortex_2d.py @@ -121,7 +121,6 @@ def run_singlephase(params: dict, seed: int, assert_each=None, **kwargs) -> tupl strains, cmap="cmc.batlow_r", aspect="equal", - alpha=1, ) fig_path.colorbar(s, ax=ax_path, aspect=25, label="Strain (ε)") @@ -272,7 +271,6 @@ def test_xz(self, outdir, seed, n_grains): cmap="cmc.batlow_r", tick_formatter=lambda x, pos: str(x), aspect="equal", - alpha=1, ) fig_path.colorbar(s, ax=ax_path, aspect=25, label="Strain (ε)") fig_path.savefig(_io.resolve_path(f"{out_basepath}_path.pdf"))