fix: Fix velocity gradient definition in Fluidity example and use M*=10

examples/fluidity/advection2d/advection2d.flml

@@ -94,16 +94,12 @@
                 <string_value type="code" language="python" lines="20">def val(X, t, dt, fields, n):
     import numpy as np
     import pydrex
-
-    from advection2d import velocity_gradient
+    from advection2d import STRAIN_RATE, _velocity_grad
 
     # Construct 3D velocity gradient matrix and position vector.
-    velocity_gradient_init = velocity_gradient(X, t)  # Takes a 2D position vector.
     position = np.asarray(X) if len(X) == 3 else np.insert(X, 1, 0)
-    if velocity_gradient_init.shape == (2, 2):
-        velocity_gradient_init = np.insert(
-            np.insert(velocity_gradient_init, [1], 0, axis=0), [1], 0, axis=1
-        )
+    velocity_gradient_init = _velocity_grad(position)
+
     # Create new mineral with random initial orientations and undeformed initial state.
     # The extra 22 values are stored after the actual CPO data and they are:
     # a scalar that records the current accumulated strain,
@@ -132,7 +128,7 @@
     from pydrex import utils
     from pydrex import logger as _log
 
-    from advection2d import gen_interp_position, gen_interp_velocity_grad, PARAMS
+    from advection2d import cb_interp_position, cb_interp_velocity_grad, PARAMS
 
     # Unpack values from previous step and create Mineral.
     n_grains = int((n - 22) / 10)
@@ -156,9 +152,7 @@
     position = np.asarray(X) if len(X) == 3 else np.insert(X, 1, 0)
     velocity_gradient = fields["VelocityGradient"]
     if velocity_gradient.shape == (2, 2):
-        velocity_gradient = np.insert(
-            np.insert(velocity_gradient, [1], 0, axis=0), [1], 0, axis=1
-        )
+        velocity_gradient = utils.add_dim(velocity_gradient, 1)
     _log.info("deformation gradient: %s", deformation_gradient_prev.flatten())
     _log.info("position: %s", gen_interp_position(position_prev, position, t, dt)(t))
     _log.info("velocity gradient: %s",

examples/fluidity/advection2d/advection2d.py

@@ -1,42 +1,49 @@
 import configparser
 import pathlib
 
-import numpy as np
 import numba as nb
+import numpy as np
+
+from pydrex.utils import add_dim, remove_dim
+from pydrex.velocity import simple_shear_2d
 
 config = configparser.ConfigParser()
 config.read(pathlib.PurePath(__file__).with_suffix(".ini"))
 init = config["initial conditions"]
 
 
-def gen_init_positions():
-    """Generate initial particle positions."""
+def init_positions():
+    """Return initial particle positions."""
     init_horiz = np.array([float(x) for x in init["INIT_HORIZ"].split()])
     init_vert = float(init["INIT_VERT"])
     return np.array([[x, init_vert] for x in init_horiz])
 
 
-def gen_interp_position(position_prev, position, t, dt):
-    """Generate position interpolator for CPO integration from `position_prev` to `position`."""
+def cb_interp_position(position_prev, position, t, dt):
+    """Return position interpolator for CPO integration from `position_prev` to `position`."""
+
     @nb.njit(fastmath=True)
     def _interp_position(int_time):
         return position_prev + (int_time - t + dt) / dt * (position - position_prev)
+
     return _interp_position
 
 
-def gen_interp_velocity_grad(velocity_grad_prev, velocity_grad, t, dt):
-    """Generate ∇v interpolator for CPO integration between `velocity_grad_prev` and `velocity_grad`."""
+def cb_interp_velocity_grad(velocity_grad_prev, velocity_grad, t, dt):
+    """Return ∇v interpolator for CPO integration between `velocity_grad_prev` and `velocity_grad`."""
+
     @nb.njit(fastmath=True)
     def _interp_vel_grad(int_time):
         return velocity_grad_prev + (int_time - t + dt) / dt * (
             velocity_grad - velocity_grad_prev
         )
+
     return _interp_vel_grad
 
 
 STRAIN_RATE = float(init["STRAIN_RATE"])
 # Initial position of the particle.
-INITIAL_POSITIONS = gen_init_positions()
+INITIAL_POSITIONS = init_positions()
 # Recrystallisation parameter configuration.
 PARAMS = {
     "stress_exponent": float(init["STRESS_EXPONENT"]),
@@ -48,15 +55,12 @@ def _interp_vel_grad(int_time):
     "olivine_fraction": 1,
 }
 
-@nb.njit(fastmath=True)
-def velocity(coords, time):
-    """Return prescribed velocity for 2D simple shear flow."""
-    return np.array([coords[1] * STRAIN_RATE, 0.0])
+_velocity, _velocity_grad = simple_shear_2d("X", "Z", STRAIN_RATE)
+
+
+def velocity(X, t):
+    return remove_dim(_velocity(add_dim(X, 1)), 1)
 
 
-@nb.njit(fastmath=True)
-def velocity_gradient(coords, time):
-    """Return prescribed velocity for 2D simple shear flow."""
-    grad_v = np.zeros((2, 2))
-    grad_v[0, 2] = 2 * STRAIN_RATE
-    return grad_v
+def velocity_gradient(X, t):
+    return remove_dim(_velocity_grad(add_dim(X, 1)), 1)

src/pydrex/utils.py

@@ -68,6 +68,78 @@ def remove_nans(a):
     return a[~np.isnan(a)]
 
 
+def remove_dim(a, dim):
+    """Remove all values corresponding to dimension `dim` from an array.
+
+    Note that a `dim` of 0 refers to the “x” values.
+
+    Examples:
+
+    >>> a = [1, 2, 3]
+    >>> remove_dim(a, 1)
+    array([2, 3])
+    >>> remove_dim(a, 2)
+    array([1, 3])
+    >>> remove_dim(a, 3)
+    array([1, 2])
+
+    >>> a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+    >>> remove_dim(a, 1)
+    array([[5, 6],
+           [8, 9]])
+    >>> remove_dim(a, 2)
+    array([[1, 3],
+           [7, 9]])
+    >>> remove_dim(a, 3)
+    array([[1, 2],
+           [4, 5]])
+
+    """
+    _a = np.asarray(a)
+    for i, _ in enumerate(_a.shape):
+        _a = np.delete(_a, [dim], axis=i)
+    return _a
+
+
+def add_dim(a, dim, val=0):
+    """Add entries of `val` corresponding to dimension `dim` to an array.
+
+    Note that a `dim` of 0 refers to the “x” values.
+
+    Examples:
+
+    >>> a = [1, 2]
+    >>> add_dim(a, 1)
+    array([0, 1, 2])
+    >>> add_dim(a, 2)
+    array([1, 0, 2])
+    >>> add_dim(a, 3)
+    array([1, 2, 0])
+
+    >>> add_dim([1.0, 2.0], 3)
+    array([1., 2., 0.])
+
+    >>> a = [[1, 2], [3, 4]]
+    >>> add_dim(a, 1)
+    array([[0, 0, 0],
+           [0, 1, 2],
+           [0, 3, 4]])
+    >>> add_dim(a, 2)
+    array([[1, 0, 2],
+           [0, 0, 0],
+           [3, 0, 4]])
+    >>> add_dim(a, 3)
+    array([[1, 2, 0],
+           [3, 4, 0],
+           [0, 0, 0]])
+
+    """
+    _a = np.asarray(a)
+    for i, _ in enumerate(_a.shape):
+        _a = np.insert(_a, [dim], 0, axis=i)
+    return _a
+
+
 def default_ncpus():
     """Get a safe default number of CPUs available for multiprocessing.
 

tests/test_simple_shear_2d.py

@@ -809,6 +809,7 @@ def test_dudz_pathline(self, outdir, seed):
             velocity_gradients = [get_velocity_gradient(Ŋ(x)) for x in positions]
 
             params = _io.DEFAULT_PARAMS
+            params["gbm_mobility"] = 10
             params["number_of_grains"] = 5000
             olA = _minerals.Mineral(n_grains=params["number_of_grains"], seed=seed)
             deformation_gradient = np.eye(3)
@@ -857,7 +858,7 @@ def test_dudz_pathline(self, outdir, seed):
                 np.full(n_timesteps, -0.01), np.diff(misorient_indices)
             )
             # Check that last angle is <5° (M*=125) or <10° (M*=10).
-            assert cpo_angles[-1] < 5
+            assert cpo_angles[-1] < 10
 
         if outdir is not None:
             fig, ax, _, _ = _vis.pathline_box2d(