PyMPDATA_examples.Jarecka_et_al_2015.formulae

  1# adapted from:
  2# https://github.com/igfuw/shallow-water-elliptic-drop/blob/master/analytical/analytic_equations.py
  3# Code used in the paper of Jarecka, Jaruga, Smolarkiewicz -
  4# "A Spreading Drop of Shallow Water" (JCP 289, doi:10.1016/j.jcp.2015.02.003).
  5
  6import numba
  7import numpy as np
  8from scipy.integrate import odeint
  9
 10from PyMPDATA.impl.enumerations import ARG_DATA, ARG_FOCUS, MAX_DIM_NUM
 11
 12
 13def amplitude(x, y, lx, ly):
 14    A = 1 / lx / ly
 15    h = A * (1 - (x / lx) ** 2 - (y / ly) ** 2)
 16    return np.where(h > 0, h, 0)
 17
 18
 19@numba.njit()
 20def deriv(y, _):
 21    """
 22    return derivatives of [lambda_x, dlambda_x/dt, lambda_y, dlambda_y/dt
 23    four first-order ODEs based on  eq. 7  (Jarecka, Jaruga, Smolarkiewicz)
 24    """
 25    return np.array((y[1], 2.0 / y[0] ** 2 / y[2], y[3], 2.0 / y[0] / y[2] ** 2))
 26
 27
 28def d2_el_lamb_lamb_t_evol(times, lamb_x0, lamb_y0):
 29    """
 30    solving coupled nonlinear second-order ODEs - eq. 7  (Jarecka, Jaruga, Smolarkiewicz)
 31    returning array with first dim denoting time, second dim:
 32    [lambda_x, dot{lambda_x}, lambda_y, dot{lambda_y}
 33    """
 34    assert times[0] == 0
 35    yinit = np.array([lamb_x0, 0.0, lamb_y0, 0.0])  # initial values (dot_lamb = 0.)
 36    result, info = odeint(deriv, yinit, times, full_output=True)
 37    assert info["message"] == "Integration successful."
 38    return result
 39
 40
 41def make_rhs_indexers(ats, grid_step, time_step, options):
 42    @numba.njit(**options.jit_flags)
 43    def rhs(m, _0, h, _1, _2, _3):
 44        retval = (
 45            m
 46            - ((ats(*h, +1) - ats(*h, -1)) / 2) / 2 * ats(*h, 0) * time_step / grid_step
 47        )
 48        return retval
 49
 50    return rhs
 51
 52
 53def make_rhs(grid_step, time_step, axis, options, traversals):
 54    indexers = traversals.indexers[traversals.n_dims]
 55    apply_scalar = traversals.apply_scalar(loop=False)
 56
 57    formulae_rhs = tuple(
 58        (
 59            make_rhs_indexers(
 60                ats=indexers.ats[axis],
 61                grid_step=grid_step[axis],
 62                time_step=time_step,
 63                options=options,
 64            ),
 65            None,
 66            None,
 67        )
 68    )
 69
 70    @numba.njit(**options.jit_flags)
 71    def apply(traversals_data, momentum, h):
 72        null_scalarfield, null_scalarfield_bc = traversals_data.null_scalar_field
 73        null_vectorfield, null_vectorfield_bc = traversals_data.null_vector_field
 74        return apply_scalar(
 75            *formulae_rhs,
 76            *momentum.field,
 77            *null_vectorfield,
 78            null_vectorfield_bc,
 79            *h.field,
 80            h.bc,
 81            *null_scalarfield,
 82            null_scalarfield_bc,
 83            *null_scalarfield,
 84            null_scalarfield_bc,
 85            *null_scalarfield,
 86            null_scalarfield_bc,
 87            traversals_data.buffer
 88        )
 89
 90    return apply
 91
 92
 93def make_interpolate_indexers(ati, options):
 94    @numba.njit(**options.jit_flags)
 95    def interpolate(momentum_x, _, momentum_y):
 96        momenta = (momentum_x[ARG_FOCUS], (momentum_x[ARG_DATA], momentum_y[ARG_DATA]))
 97        return ati(*momenta, 0.5)
 98
 99    return interpolate
100
101
102def make_interpolate(options, traversals):
103    indexers = traversals.indexers[traversals.n_dims]
104    apply_vector = traversals.apply_vector()
105
106    formulae_interpolate = tuple(
107        (
108            make_interpolate_indexers(ati=indexers.ati[i], options=options)
109            if indexers.ati[i] is not None
110            else None
111        )
112        for i in range(MAX_DIM_NUM)
113    )
114
115    @numba.njit(**options.jit_flags)
116    def apply(traversals_data, momentum_x, momentum_y, advector):
117        null_vectorfield, null_vectorfield_bc = traversals_data.null_vector_field
118        return apply_vector(
119            *formulae_interpolate,
120            *advector.field,
121            *momentum_x.field,
122            momentum_x.bc,
123            *null_vectorfield,
124            null_vectorfield_bc,
125            *momentum_y.field,
126            momentum_y.bc,
127            traversals_data.buffer
128        )
129
130    return apply