pydrex.velocity

  1"""> PyDRex: Steady-state solutions of velocity (gradients) for various flows.
  2
  3For the sake of consistency, all callables returned from methods in this module expect a
  4time scalar and 3D position vector as inputs. They also return 3D tensors in all cases.
  5This means they can be directly used as arguments to e.g.
  6`pydrex.minerals.Mineral.update_orientations`.
  7
  8"""
  9
 10import functools as ft
 11
 12import numba as nb
 13import numpy as np
 14
 15from pydrex import geometry as _geo
 16
 17
 18@nb.njit(fastmath=True)
 19def _simple_shear_2d_grad(t, x, direction, deformation_plane, strain_rate):
 20    grad_v = np.zeros((3, 3))
 21    grad_v[direction, deformation_plane] = 2 * strain_rate
 22    return grad_v
 23
 24
 25@nb.njit(fastmath=True)
 26def _simple_shear_2d(t, x, direction, deformation_plane, strain_rate):
 27    v = np.zeros(3)
 28    v[direction] = x[deformation_plane] * strain_rate
 29    return v
 30
 31
 32@nb.njit(fastmath=True)
 33def _cell_2d_grad(t, x, horizontal, vertical, velocity_edge, edge_length):
 34    _lim = edge_length / 2
 35    if np.abs(x[horizontal]) > _lim or np.abs(x[vertical]) > _lim:
 36        # NOTE: At the moment, this prints type info rather than values. Probably a
 37        # numba bug, version 0.59 shouldn't require these to be compile-time constants.
 38        raise ValueError(
 39            f"position ({x[horizontal]}, {x[vertical]}) is outside domain with xᵢ ∈ [-{_lim}, {_lim}]"
 40        )
 41    cos_πx_divd = np.cos(np.pi * x[horizontal] / edge_length)
 42    cos_πz_divd = np.cos(np.pi * x[vertical] / edge_length)
 43    sin_πx_divd = np.sin(np.pi * x[horizontal] / edge_length)
 44    sin_πz_divd = np.sin(np.pi * x[vertical] / edge_length)
 45    grad_v = np.zeros((3, 3))
 46    grad_v[horizontal, horizontal] = -np.pi / edge_length * sin_πz_divd * sin_πx_divd
 47    grad_v[horizontal, vertical] = np.pi / edge_length * cos_πz_divd * cos_πx_divd
 48    grad_v[vertical, vertical] = -np.pi / edge_length * cos_πz_divd * cos_πx_divd
 49    grad_v[vertical, horizontal] = np.pi / edge_length * sin_πz_divd * sin_πx_divd
 50    return velocity_edge * grad_v
 51
 52
 53@nb.njit(fastmath=True)
 54def _cell_2d(t, x, horizontal, vertical, velocity_edge, edge_length):
 55    _lim = edge_length / 2
 56    if np.abs(x[horizontal]) > _lim or np.abs(x[vertical]) > _lim:
 57        # NOTE: At the moment, this prints type info rather than values. Probably a
 58        # numba bug, version 0.59 shouldn't require these to be compile-time constants.
 59        raise ValueError(
 60            f"position ({x[horizontal]}, {x[vertical]}) is outside domain with xᵢ ∈ [-{_lim}, {_lim}]"
 61        )
 62    cos_πx_divd = np.cos(np.pi * x[horizontal] / edge_length)
 63    cos_πz_divd = np.cos(np.pi * x[vertical] / edge_length)
 64    sin_πx_divd = np.sin(np.pi * x[horizontal] / edge_length)
 65    sin_πz_divd = np.sin(np.pi * x[vertical] / edge_length)
 66    v = np.zeros(3)
 67    v[horizontal] = velocity_edge * cos_πx_divd * sin_πz_divd
 68    v[vertical] = -velocity_edge * sin_πx_divd * cos_πz_divd
 69    return v
 70
 71
 72@nb.njit(fastmath=True)
 73def _corner_2d(t, x, horizontal, vertical, plate_speed):
 74    h = x[horizontal]
 75    v = x[vertical]
 76    if np.abs(h) < 1e-15 and np.abs(v) < 1e-15:
 77        return np.full(3, np.nan)
 78    out = np.zeros(3)
 79    prefactor = 2 * plate_speed / np.pi
 80    out[horizontal] = prefactor * (np.arctan2(h, -v) + h * v / (h**2 + v**2))
 81    out[vertical] = prefactor * v**2 / (h**2 + v**2)
 82    return out
 83
 84
 85@nb.njit(fastmath=True)
 86def _corner_2d_grad(t, x, horizontal, vertical, plate_speed):
 87    h = x[horizontal]
 88    v = x[vertical]
 89    if np.abs(h) < 1e-15 and np.abs(v) < 1e-15:
 90        return np.full((3, 3), np.nan)
 91    grad_v = np.zeros((3, 3))
 92    prefactor = 4 * plate_speed / (np.pi * (h**2 + v**2) ** 2)
 93    grad_v[horizontal, horizontal] = -(h**2) * v
 94    grad_v[horizontal, vertical] = h**3
 95    grad_v[vertical, horizontal] = -h * v**2
 96    grad_v[vertical, vertical] = h**2 * v
 97    return prefactor * grad_v
 98
 99
100def simple_shear_2d(
101    direction: str, deformation_plane: str, strain_rate: float
102) -> tuple:
103    """Return simple shear velocity and velocity gradient callables.
104
105    The returned callables have signature f(t, x) where x is a 3D position vector.
106
107    - `direction` (one of {"X", "Y", "Z"}) — velocity vector direction
108    - `deformation_plane` (one of {"X", "Y", "Z"}) — direction of velocity gradient
109    - `strain_rate` — 1/2 × magnitude of the largest eigenvalue of the velocity
110      gradient
111
112    .. note::
113        Input arrays to the returned callables must have homogeneous element types.
114        Arrays with e.g. both floating point and integer values are not supported.
115
116    Examples:
117
118    >>> import numpy as np
119    >>> u, L = simple_shear_2d("X", "Z", 1e-4)
120    >>> u(np.nan, np.array([0, 0, 0]))  # Time is not used here.
121    array([0., 0., 0.])
122    >>> u(np.nan, np.array([0, 0, 1]))
123    array([0.0001, 0.    , 0.    ])
124    >>> u(np.nan, np.array([0.0, 0.0, 2.0]))
125    array([0.0002, 0.    , 0.    ])
126    >>> L(np.nan, np.array([0, 0, 0]))
127    array([[0.    , 0.    , 0.0002],
128           [0.    , 0.    , 0.    ],
129           [0.    , 0.    , 0.    ]])
130    >>> L(np.nan, np.array([0.0, 0.0, 1.0]))
131    array([[0.    , 0.    , 0.0002],
132           [0.    , 0.    , 0.    ],
133           [0.    , 0.    , 0.    ]])
134
135    """
136    try:
137        indices = _geo.to_indices2d(direction, deformation_plane)
138    except ValueError:
139        raise ValueError(
140            "unsupported shear type with"
141            + f" direction = {direction}, deformation_plane = {deformation_plane}"
142        )
143
144    return (
145        ft.partial(
146            _simple_shear_2d,
147            direction=indices[0],
148            deformation_plane=indices[1],
149            strain_rate=strain_rate,
150        ),
151        ft.partial(
152            _simple_shear_2d_grad,
153            direction=indices[0],
154            deformation_plane=indices[1],
155            strain_rate=strain_rate,
156        ),
157    )
158
159
160def cell_2d(
161    horizontal: str, vertical: str, velocity_edge: float, edge_length: float = 2.0
162) -> tuple:
163    r"""Get velocity and velocity gradient callables for a steady-state 2D Stokes cell.
164
165    The cell is centered at (0,0) and the velocity field is defined by:
166    $$
167    \bm{u} = U\cos(π x/d)\sin(π z/d) \bm{\hat{h}} - U\sin(π x/d)\cos(π z/d) \bm{\hat{v}}
168    $$
169    where $\bm{\hat{h}}$ and $\bm{\hat{v}}$ are unit vectors in the chosen horizontal
170    and vertical directions, respectively. The velocity at the cell edge has a magnitude
171    of $U$ and $d$ is the length of a cell edge.
172
173    The returned callables have signature f(t, x) where x is a 3D position vector.
174
175    - `horizontal` (one of {"X", "Y", "Z"}) — horizontal direction
176    - `vertical` (one of {"X", "Y", "Z"}) — vertical direction
177    - `velocity_edge` — velocity magnitude at the center of the cell edge
178    - `edge_length` (optional) — the edge length of the cell (= 2 by default)
179
180    Examples:
181
182    >>> import numpy as np
183    >>> u, L = cell_2d("X", "Z", 1)
184    >>> u(np.nan, np.array([0, 0, 0]))  # Time value is not used for steady flows.
185    array([ 0.,  0., -0.])
186    >>> u(np.nan, np.array([0, 0, 1]))
187    array([ 1.,  0., -0.])
188    >>> u(np.nan, np.array([0, 1, 0]))  # Y-value is not used.
189    array([ 0.,  0., -0.])
190    >>> u(np.nan, np.array([0, 0, -1]))
191    array([-1.,  0., -0.])
192    >>> u(np.nan, np.array([1, 0, 0]))
193    array([ 0.,  0., -1.])
194    >>> u(np.nan, np.array([-0.5, 0.0, 0.0]))
195    array([0.        , 0.        , 0.70710678])
196    >>> L(np.nan, np.array([0, 0, 0]))
197    array([[-0.        ,  0.        ,  1.57079633],
198           [ 0.        ,  0.        ,  0.        ],
199           [ 0.        ,  0.        , -1.57079633]])
200    >>> L(np.nan, np.array([0.5, 0.0, 0.0]))
201    array([[-0.        ,  0.        ,  1.11072073],
202           [ 0.        ,  0.        ,  0.        ],
203           [ 0.        ,  0.        , -1.11072073]])
204    >>> L(np.nan, np.array([0, 0, 0])) == L(np.nan, np.array([0, 1, 0]))  # Y-value is not used.
205    array([[ True,  True,  True],
206           [ True,  True,  True],
207           [ True,  True,  True]])
208    >>> L(np.nan, np.array([1, 0, 0])) == L(np.nan, np.array([0, 0, 1]))
209    array([[ True,  True,  True],
210           [ True,  True,  True],
211           [ True,  True,  True]])
212    >>> L(np.nan, np.array([1, 0, 0])) == L(np.nan, np.array([-1, 0, 0]))
213    array([[ True,  True,  True],
214           [ True,  True,  True],
215           [ True,  True,  True]])
216    >>> L(np.nan, np.array([1, 0, 0])) == L(np.nan, np.array([0, 0, -1]))
217    array([[ True,  True,  True],
218           [ True,  True,  True],
219           [ True,  True,  True]])
220    >>> L(np.nan, np.array([0.5, 0.0, 0.5]))
221    array([[-0.78539816,  0.        ,  0.78539816],
222           [ 0.        ,  0.        ,  0.        ],
223           [ 0.78539816,  0.        , -0.78539816]])
224
225    >>> u, L = cell_2d("X", "Z", 6.3e-10, 1e5)
226    >>> u(np.nan, np.array([0, 0, 0]))
227    array([ 0.,  0., -0.])
228    >>> u(np.nan, np.array([0.0, 0.0, -5e4]))
229    array([-6.3e-10,  0.0e+00, -0.0e+00])
230    >>> u(np.nan, np.array([2e2, 0e0, 0e0]))
231    array([ 0.0000000e+00,  0.0000000e+00, -3.9583807e-12])
232
233    """
234    if edge_length < 0:
235        raise ValueError(f"edge length of 2D cell must be positive, not {edge_length}")
236    try:
237        indices = _geo.to_indices2d(horizontal, vertical)
238    except ValueError:
239        raise ValueError(
240            "unsupported convection cell geometry with"
241            + f" horizontal = {horizontal}, vertical = {vertical}"
242        )
243
244    return (
245        ft.partial(
246            _cell_2d,
247            horizontal=indices[0],
248            vertical=indices[1],
249            velocity_edge=velocity_edge,
250            edge_length=edge_length,
251        ),
252        ft.partial(
253            _cell_2d_grad,
254            horizontal=indices[0],
255            vertical=indices[1],
256            velocity_edge=velocity_edge,
257            edge_length=edge_length,
258        ),
259    )
260
261
262def corner_2d(horizontal: str, vertical: str, plate_speed: float) -> tuple:
263    r"""Get velocity and velocity gradient callables for a steady-state 2D corner flow.
264
265    The velocity field is defined by:
266    $$
267    \bm{u} = \frac{dr}{dt} \bm{\hat{r}} + r \frac{dθ}{dt} \bm{\hat{θ}}
268    = \frac{2 U}{π}(θ\sinθ - \cosθ) ⋅ \bm{\hat{r}} + \frac{2 U}{π}θ\cosθ ⋅ \bm{\hat{θ}}
269    $$
270    where $θ = 0$ points vertically downwards along the ridge axis
271    and $θ = π/2$ points along the surface. $U$ is the half spreading velocity.
272    Streamlines for the flow obey:
273    $$
274    ψ = \frac{2 U r}{π}θ\cosθ
275    $$
276    and are related to the velocity through:
277    $$
278    \bm{u} = -\frac{1}{r} ⋅ \frac{dψ}{dθ} ⋅ \bm{\hat{r}} + \frac{dψ}{dr}\bm{\hat{θ}}
279    $$
280    Conversion to Cartesian ($x,y,z$) coordinates yields:
281    $$
282    \bm{u} = \frac{2U}{π} \left[
283    \tan^{-1}\left(\frac{x}{-z}\right) + \frac{xz}{x^{2} + z^{2}} \right] \bm{\hat{x}} +
284    \frac{2U}{π} \frac{z^{2}}{x^{2} + z^{2}} \bm{\hat{z}}
285    $$
286    where
287    \begin{align\*}
288    x &= r \sinθ \cr
289    z &= -r \cosθ
290    \end{align\*}
291    and the velocity gradient is:
292    $$
293    L = \frac{4 U}{π{(x^{2}+z^{2})}^{2}} ⋅
294    \begin{bmatrix}
295        -x^{2}z & 0 & x^{3} \cr
296        0 & 0 & 0 \cr
297        -xz^{2} & 0 & x^{2}z
298    \end{bmatrix}
299    $$
300    See also Fig. 5 in [Kaminski & Ribe, 2002](https://doi.org/10.1029/2001GC000222).
301
302    The returned callables have signature f(t, x) where x is a 3D position vector.
303
304    - `horizontal` (one of {"X", "Y", "Z"}) — horizontal direction
305    - `vertical` (one of {"X", "Y", "Z"}) — vertical direction
306    - `plate_speed` (float) — speed of the “plate” i.e. upper boundary
307
308    """
309    try:
310        indices = _geo.to_indices2d(horizontal, vertical)
311    except ValueError:
312        raise ValueError(
313            "unsupported convection cell geometry with"
314            + f" horizontal = {horizontal}, vertical = {vertical}"
315        )
316
317    return (
318        ft.partial(
319            _corner_2d,
320            horizontal=indices[0],
321            vertical=indices[1],
322            plate_speed=plate_speed,
323        ),
324        ft.partial(
325            _corner_2d_grad,
326            horizontal=indices[0],
327            vertical=indices[1],
328            plate_speed=plate_speed,
329        ),
330    )