Merge pull request #152 from seismic-anisotropy/2d_cornerflow

test: Add reworked 2D cornerflow example

src/pydrex/diagnostics.py

@@ -25,6 +25,7 @@
 from pydrex import stats as _stats
 from pydrex import tensors as _tensors
 from pydrex import geometry as _geo
+from pydrex import logger as _log
 
 
 def elasticity_components(voigt_matrices):
@@ -264,7 +265,7 @@ def symmetry(orientations, axis="a"):
 
 
 def misorientation_indices(
-    orientation_stack, system: _geo.LatticeSystem, bins=None, ncpus=None
+    orientation_stack, system: _geo.LatticeSystem, bins=None, ncpus=None, pool=None,
 ):
     """Calculate M-indices for a series of polycrystal textures.
 
@@ -275,21 +276,28 @@ def misorientation_indices(
     Uses the multiprocessing library to calculate texture indices for multiple snapshots
     simultaneously. If `ncpus` is `None` the number of CPU cores to use is chosen
     automatically based on the maximum number available to the Python interpreter,
-    otherwise the specified number of cores is requested.
+    otherwise the specified number of cores is requested. Alternatively, an existing
+    instance of `multiprocessing.Pool` can be provided.
 
     See `misorientation_index` for documentation of the remaining arguments.
 
     """
+    if ncpus is not None and pool is not None:
+        _log.warning("ignoring `ncpus` argument because a Pool was provided")
     m_indices = np.empty(len(orientation_stack))
     _run = ft.partial(
         misorientation_index,
         system=system,
         bins=bins,
     )
-    if ncpus is None:
-        ncpus = len(os.sched_getaffinity(0)) - 1
-    with Pool(processes=ncpus) as pool:
-        for i, out in enumerate(pool.imap_unordered(_run, orientation_stack)):
+    if pool is None:
+        if ncpus is None:
+            ncpus = len(os.sched_getaffinity(0)) - 1
+        with Pool(processes=ncpus) as pool:
+            for i, out in enumerate(pool.imap(_run, orientation_stack)):
+                m_indices[i] = out
+    else:
+        for i, out in enumerate(pool.imap(_run, orientation_stack)):
             m_indices[i] = out
     return m_indices
 

src/pydrex/geometry.py

@@ -188,7 +188,8 @@ def poles(orientations, ref_axes="xz", hkl=[1, 0, 0]):
     horizontal and vertical axes of pole figures used to plot the directions.
 
     """
-    upward_axes = (set("xyz") - set(ref_axes)).pop()
+    _ref_axes = ref_axes.lower()
+    upward_axes = (set("xyz") - set(_ref_axes)).pop()
     axes_map = {"x": 0, "y": 1, "z": 2}
 
     # Get directions in the right-handed frame.
@@ -198,8 +199,8 @@ def poles(orientations, ref_axes="xz", hkl=[1, 0, 0]):
 
     # Rotate into the chosen reference frame.
     zvals = directions[:, axes_map[upward_axes]]
-    yvals = directions[:, axes_map[ref_axes[1]]]
-    xvals = directions[:, axes_map[ref_axes[0]]]
+    yvals = directions[:, axes_map[_ref_axes[1]]]
+    xvals = directions[:, axes_map[_ref_axes[0]]]
     return xvals, yvals, zvals
 
 

src/pydrex/velocity.py

@@ -1,4 +1,10 @@
-"""> PyDRex: Steady-state solutions of velocity (gradients) for various flows."""
+"""> PyDRex: Steady-state solutions of velocity (gradients) for various flows.
+
+For the sake of consistency, all callables returned from methods in this module expect a
+3D position vector as input. They also return 3D tensors in all cases. This means they
+can be directly used as arguments to e.g. `pydrex.minerals.Mineral.update_orientations`.
+
+"""
 import functools as ft
 
 import numba as nb
@@ -45,10 +51,38 @@ def _cell_2d(x, horizontal=0, vertical=2, velocity_edge=1):
     return v
 
 
+@nb.njit(fastmath=True)
+def _corner_2d(x, horizontal=0, vertical=2, plate_speed=1):
+    h = x[horizontal]
+    v = x[vertical]
+    if np.abs(h) < 1e-15 and np.abs(v) < 1e-15:
+        return np.full(3, np.nan)
+    out = np.zeros(3)
+    prefactor = 2 * plate_speed / np.pi
+    out[horizontal] = prefactor * (np.arctan2(h, -v) + h * v / (h**2 + v**2))
+    out[vertical] = prefactor * v**2 / (h**2 + v**2)
+    return out
+
+
+@nb.njit(fastmath=True)
+def _corner_2d_grad(x, horizontal=0, vertical=2, plate_speed=1):
+    h = x[horizontal]
+    v = x[vertical]
+    if np.abs(h) < 1e-15 and np.abs(v) < 1e-15:
+        return np.full((3, 3), np.nan)
+    grad_v = np.zeros((3, 3))
+    prefactor = 4 * plate_speed / (np.pi * (h**2 + v**2)**2)
+    grad_v[horizontal, horizontal] = -(h**2) * v
+    grad_v[horizontal, vertical] = h**3
+    grad_v[vertical, horizontal] = -h * v**2
+    grad_v[vertical, vertical] = h**2 * v
+    return prefactor * grad_v
+
+
 def simple_shear_2d(direction, deformation_plane, strain_rate):
     """Return simple shear velocity and velocity gradient callables.
 
-    The returned callables have signature f(x) where x is a position vector.
+    The returned callables have signature f(x) where x is a 3D position vector.
 
     Args:
     - `direction` (one of {"X", "Y", "Z"}) — velocity vector direction
@@ -93,9 +127,16 @@ def simple_shear_2d(direction, deformation_plane, strain_rate):
 
 
 def cell_2d(horizontal, vertical, velocity_edge):
-    """Return velocity and velocity gradient callable for a steady-state 2D Stokes cell.
+    r"""Get velocity and velocity gradient callables for a steady-state 2D Stokes cell.
+
+    The velocity field is defined by:
+    $$
+    \bm{u} = \cos(π x/2)\sin(π x/2) \bm{\hat{h}} - \sin(π x/2)\cos(π x/2) \bm{\hat{v}}
+    $$
+    where $\bm{\hat{h}}$ and $\bm{\hat{v}}$ are unit vectors in the chosen horizontal
+    and vertical directions, respectively.
 
-    The returned callables have signature f(x) where x is a position vector.
+    The returned callables have signature f(x) where x is a 3D position vector.
 
     Args:
     - `horizontal` (one of {"X", "Y", "Z"}) — horizontal direction
@@ -137,3 +178,87 @@ def cell_2d(horizontal, vertical, velocity_edge):
             velocity_edge=velocity_edge,
         ),
     )
+
+
+def corner_2d(horizontal, vertical, plate_speed):
+    r"""Get velocity and velocity gradient callables for a steady-state 2D corner flow.
+
+    The velocity field is defined by:
+    $$
+    \bm{u} = \frac{dr}{dt} \bm{\hat{r}} + r \frac{dθ}{dt} \bm{\hat{θ}}
+    = \frac{2 U}{π}(θ\sinθ - \cosθ) ⋅ \bm{\hat{r}} + \frac{2 U}{π}θ\cosθ ⋅ \bm{\hat{θ}}
+    $$
+    where $θ = 0$ points vertically downwards along the ridge axis
+    and $θ = π/2$ points along the surface. $U$ is the half spreading velocity.
+    Streamlines for the flow obey:
+    $$
+    ψ = \frac{2 U r}{π}θ\cosθ
+    $$
+    and are related to the velocity through:
+    $$
+    \bm{u} = -\frac{1}{r} ⋅ \frac{dψ}{dθ} ⋅ \bm{\hat{r}} + \frac{dψ}{dr}\bm{\hat{θ}}
+    $$
+    Conversion to Cartesian ($x,y,z$) coordinates yields:
+    $$
+    \bm{u} = \frac{2U}{π} \left[
+    \tan^{-1}\left(\frac{x}{-z}\right) + \frac{xz}{x^{2} + z^{2}} \right] \bm{\hat{x}} +
+    \frac{2U}{π} \frac{z^{2}}{x^{2} + z^{2}} \bm{\hat{z}}
+    $$
+    where
+    \begin{align\*}
+    x &= r \sinθ \cr
+    z &= -r \cosθ
+    \end{align\*}
+    and the velocity gradient is:
+    $$
+    L = \frac{4 U}{π{(x^{2}+z^{2})}^{2}} ⋅
+    \begin{bmatrix}
+        -x^{2}z & 0 & x^{3} \cr
+        0 & 0 & 0 \cr
+        -xz^{2} & 0 & x^{2}z
+    \end{bmatrix}
+    $$
+    See also Fig. 5 in [Kaminski & Ribe, 2002](https://doi.org/10.1029/2001GC000222).
+
+    The returned callables have signature f(x) where x is a 3D position vector.
+
+    Args:
+    - `horizontal` (one of {"X", "Y", "Z"}) — horizontal direction
+    - `vertical` (one of {"X", "Y", "Z"}) — vertical direction
+    - `plate_speed` (float) — speed of the “plate” i.e. upper boundary
+
+    """
+    geometry = (horizontal, vertical)
+    match geometry:
+        case ("X", "Y"):
+            _geometry = (0, 1)
+        case ("X", "Z"):
+            _geometry = (0, 2)
+        case ("Y", "X"):
+            _geometry = (1, 0)
+        case ("Y", "Z"):
+            _geometry = (1, 2)
+        case ("Z", "X"):
+            _geometry = (2, 0)
+        case ("Z", "Y"):
+            _geometry = (2, 1)
+        case _:
+            raise ValueError(
+                "unsupported convection cell geometry with"
+                + f" horizontal = {horizontal}, vertical = {vertical}"
+            )
+
+    return (
+        ft.partial(
+            _corner_2d,
+            horizontal=_geometry[0],
+            vertical=_geometry[1],
+            plate_speed=plate_speed,
+        ),
+        ft.partial(
+            _corner_2d_grad,
+            horizontal=_geometry[0],
+            vertical=_geometry[1],
+            plate_speed=plate_speed,
+        ),
+    )