Source code for fluidsim.base.time_stepping.pseudo_spect

"""Time stepping (:mod:`fluidsim.base.time_stepping.pseudo_spect`)
========================================================================

Provides:

.. autoclass:: TimeSteppingPseudoSpectral
   :members:
   :private-members:

.. todo::

  It would be interesting to also implement the Adams-Bashforth (leapfrog)
  scheme with phase-shifting. It is very close to
  :func:`_time_step_RK2_phaseshift` with 2 evaluations of the non-linear terms
  per time step (but with 2 symmetrical and equivalent steps).

.. note::

  For a theoretical presentation of phase-shifting, see the book Numerical
  Experiments in Homogeneous Turbulence (Robert S. Rogallo,
  https://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/19810022965.pdf).

"""

from random import randint

import numpy as np

from transonic import Transonic, Type, NDim, Array, boost, Union

from .base import TimeSteppingBase

ts = Transonic()

N = NDim(2, 3, 4)
A = Array[np.complex128, N, "C"]
Am1 = Array[np.complex128, N - 1, "C"]

N123 = NDim(1, 2, 3)
A123c = Array[np.complex128, N123, "C"]
A123f = Array[np.float64, N123, "C"]

T = Type(np.float64, np.complex128)
A1 = Array[T, N, "C"]
A2 = Array[T, N - 1, "C"]
ArrayDiss = Union[A1, A2]

uniform = np.random.default_rng().uniform


@boost
def step_Euler(
    state_spect: A, dt: float, tendencies: A, diss: ArrayDiss, output: A
):
    output[:] = (state_spect + dt * tendencies) * diss
    return output


@boost
def step_Euler_inplace(state_spect: A, dt: float, tendencies: A, diss: ArrayDiss):
    step_Euler(state_spect, dt, tendencies, diss, state_spect)


@boost
def step_like_RK2(
    state_spect: A, dt: float, tendencies: A, diss: ArrayDiss, diss2: ArrayDiss
):
    state_spect[:] = state_spect * diss + dt * diss2 * tendencies


@boost
def mean_with_phaseshift(
    tendencies_0: A, tendencies_1_shift: A, phaseshift: Am1, output: A
):
    output[:] = 0.5 * (tendencies_0 + tendencies_1_shift / phaseshift)
    return output


@boost
def mul(phaseshift: Am1, state_spect: A, output: A):
    output[:] = phaseshift * state_spect
    return output


@boost
def div_inplace(arr: A, phaseshift: Am1):
    arr /= phaseshift
    return arr


@boost
def compute_phaseshift_terms(
    phase_alpha: A123f,
    phase_beta: A123f,
    phaseshift_alpha: A123c,
    phaseshift_beta: A123c,
):
    phaseshift_alpha[:] = np.exp(1j * phase_alpha)
    phaseshift_beta[:] = np.exp(1j * phase_beta)
    return phaseshift_alpha, phaseshift_beta


class ExactLinearCoefs:
    """Handle the computation of the exact coefficient for the RK4."""

    def __init__(self, time_stepping):
        self.time_stepping = time_stepping
        sim = time_stepping.sim
        self.shapeK_loc = sim.oper.shapeK_loc
        self.freq_lin = time_stepping.freq_lin

        self.exact = np.empty_like(self.freq_lin)
        self.exact2 = np.empty_like(self.freq_lin)

        if sim.params.time_stepping.USE_CFL:
            self.get_updated_coefs = self.get_updated_coefs_CLF
            self.dt_old = 0.0
        else:
            self.compute(time_stepping.deltat)
            self.get_updated_coefs = self.get_coefs

    def compute(self, dt):
        """Compute the exact coefficients."""
        f_lin = self.freq_lin
        exact = self.exact
        exact2 = self.exact2

        if ts.is_transpiled:
            ts.use_block("exact_lin_compute")
        else:
            # transonic block (
            #     A1 f_lin, exact, exact2;
            #     float dt
            # )
            # transonic block (
            #     A2 f_lin, exact, exact2;
            #     float dt
            # )
            exact[:] = np.exp(-dt * f_lin)
            exact2[:] = np.exp(-dt / 2 * f_lin)
        self.dt_old = dt

    def get_updated_coefs_CLF(self):
        """Get the exact coefficient updated if needed."""
        dt = self.time_stepping.deltat
        if self.dt_old != dt:
            self.compute(dt)
        return self.exact, self.exact2

    def get_coefs(self):
        """Get the exact coefficients as stored."""
        return self.exact, self.exact2


[docs] class TimeSteppingPseudoSpectral(TimeSteppingBase): """Time stepping class for pseudo-spectral solvers."""
[docs] @staticmethod def _complete_params_with_default(params): """This static method is used to complete the *params* container.""" TimeSteppingBase._complete_params_with_default(params) params.time_stepping.USE_CFL = True params.time_stepping._set_child( "phaseshift_random", attribs=dict(nb_pairs=1, nb_steps_compute_new_pair=None), )
def __init__(self, sim): super().__init__(sim) self.init_from_params() def init_from_params(self): self._init_freq_lin() self._init_compute_time_step() self._init_exact_linear_coef() self._init_time_scheme() def _init_freq_lin(self): f_d, f_d_hypo = self.sim.compute_freq_diss() freq_dissip = f_d + f_d_hypo if hasattr(self.sim, "compute_freq_complex"): freq_complex = self._compute_freq_complex() self.freq_lin = freq_dissip + freq_complex freq_max = freq_complex.imag.max() self.deltat_max = 0.78 * np.pi / freq_max else: self.freq_lin = freq_dissip def _init_time_scheme(self): type_time_scheme = self.params.time_stepping.type_time_scheme if type_time_scheme.startswith("RK"): self._state_spect_tmp = np.empty_like(self.sim.state.state_spect) if type_time_scheme.endswith("_random"): self._init_phaseshift_random() if not hasattr(self.sim.oper, "get_phases_random"): raise NotImplementedError if type_time_scheme == "Euler": time_step_RK = self._time_step_Euler elif type_time_scheme == "Euler_phaseshift": time_step_RK = self._time_step_Euler_phaseshift elif type_time_scheme == "Euler_phaseshift_random": time_step_RK = self._time_step_Euler_phaseshift_random elif type_time_scheme == "RK2": time_step_RK = self._time_step_RK2 elif type_time_scheme == "RK2_trapezoid": time_step_RK = self._time_step_RK2_trapezoid elif type_time_scheme == "RK2_phaseshift": time_step_RK = self._time_step_RK2_phaseshift elif type_time_scheme == "RK2_phaseshift_random": time_step_RK = self._time_step_RK2_phaseshift_random elif type_time_scheme == "RK2_phaseshift_exact": time_step_RK = self._time_step_RK2_phaseshift_exact elif type_time_scheme == "RK4": self._state_spect_tmp1 = np.empty_like(self.sim.state.state_spect) time_step_RK = self._time_step_RK4 else: raise ValueError(f'Problem name time_scheme ("{type_time_scheme}")') self._time_step_RK = time_step_RK def _compute_freq_complex(self): state_spect = self.sim.state.state_spect freq_complex = np.empty_like(state_spect) for ik, key in enumerate(state_spect.keys): freq_complex[ik] = self.sim.compute_freq_complex(key) return freq_complex def _init_exact_linear_coef(self): self.exact_linear_coefs = ExactLinearCoefs(self)
[docs] def one_time_step_computation(self): """One time step.""" self._time_step_RK() self.sim.oper.dealiasing(self.sim.state.state_spect) self.sim.state.statephys_from_statespect() # np.isnan(np.sum seems to be really fast if np.isnan(np.sum(self.sim.state.state_spect[0])): raise ValueError(f"nan at it = {self.it}, t = {self.t:.4f}")
[docs] def _time_step_Euler(self): r"""Forward Euler method. Notes ----- .. |p| mathmacro:: \partial .. |dt| mathmacro:: \mathop{dt} .. |dx| mathmacro:: \mathop{dx} We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), The forward Euler method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time :math:`t = 0`. Euler approximation : .. math:: \p_t \log S = \sigma + \frac{N_0}{S_0}, Integrating from :math:`t` to :math:`t+\dt`, it gives: .. math:: S_{\dt} = (S_0 + N_0 \dt) e^{\sigma \dt}. """ dt = self.deltat diss = self.exact_linear_coefs.get_updated_coefs()[0] compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect tendencies_0 = compute_tendencies() step_Euler_inplace(state_spect, dt, tendencies_0, diss)
[docs] def _get_phaseshift(self): """Compute the phase shift term.""" if hasattr(self, "_phaseshift") and self._phaseshift is not None: return self._phaseshift oper = self.sim.oper ndim = len(oper.axes) if ndim == 1: phase = 0.5 * oper.deltax * oper.kx elif ndim == 2: phase = 0.5 * (oper.deltax * oper.KX + oper.deltay * oper.KY) elif ndim == 3: phase = 0.5 * ( oper.deltax * oper.Kx + oper.deltay * oper.Ky + oper.deltaz * oper.Kz ) else: raise NotImplementedError self._phaseshift = np.exp(1j * phase) return self._phaseshift
def _init_phaseshift_random(self): params_phaseshift = self.params.time_stepping.phaseshift_random if params_phaseshift.nb_steps_compute_new_pair is None: if params_phaseshift.nb_pairs == 1: params_phaseshift.nb_steps_compute_new_pair = 2 else: params_phaseshift.nb_steps_compute_new_pair = ( 4 * params_phaseshift.nb_pairs ) self._index_phaseshift = 1 self._previous_index_pair = 0 self._previous_index_flip = 0 self._pairs_phaseshift = [] for _ in range(params_phaseshift.nb_pairs): phaseshift_alpha = np.empty( self.sim.oper.shapeK_loc, dtype=np.complex128 ) phaseshift_beta = np.empty_like(phaseshift_alpha) phase_alpha, phase_beta = self.sim.oper.get_phases_random() self._pairs_phaseshift.append( compute_phaseshift_terms( phase_alpha, phase_beta, phaseshift_alpha, phaseshift_beta ) )
[docs] def _get_phaseshift_random(self): """Compute two random phase shift terms.""" params_phaseshift = self.params.time_stepping.phaseshift_random nb_pairs = params_phaseshift.nb_pairs nb_steps_compute_new_pair = params_phaseshift.nb_steps_compute_new_pair if nb_pairs == 1 and nb_steps_compute_new_pair == 1: phaseshift_alpha, phaseshift_beta = self._pairs_phaseshift[0] elif nb_pairs == 1 and nb_steps_compute_new_pair == 2: pair = self._pairs_phaseshift[0] if self._index_phaseshift == 1: phaseshift_alpha, phaseshift_beta = pair else: phaseshift_beta, phaseshift_alpha = pair else: index_pair = randint(0, nb_pairs - 1) pair = self._pairs_phaseshift[index_pair] index_flip = randint(0, 1) if ( index_pair == self._previous_index_pair and index_flip == self._previous_index_flip ): index_flip = 0 if index_flip else 1 self._previous_index_pair = index_pair self._previous_index_flip = index_flip if index_flip: phaseshift_alpha, phaseshift_beta = pair else: phaseshift_beta, phaseshift_alpha = pair if self._index_phaseshift == nb_steps_compute_new_pair: self._index_phaseshift = 1 phase_alpha, phase_beta = self.sim.oper.get_phases_random() phaseshift_alpha, phaseshift_beta = self._pairs_phaseshift.pop(0) self._pairs_phaseshift.append( compute_phaseshift_terms( phase_alpha, phase_beta, phaseshift_alpha, phaseshift_beta ) ) else: self._index_phaseshift += 1 return phaseshift_alpha, phaseshift_beta
[docs] def _time_step_Euler_phaseshift(self): r"""Forward Euler method, dealiasing with phase-shifting. Notes ----- We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), The forward Euler method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time :math:`t = 0`. - Euler approximation: .. math:: \p_t \log S = \sigma + \frac{N_\mathrm{dealias}}{S_0}, Integrating from :math:`t` to :math:`t+\dt`, it gives: .. math:: S_{\dt} = (S_0 + N_\mathrm{dealias} \dt) e^{\sigma \dt}. where the dealiased non-linear term is :math:`N_\mathrm{dealias} = (N_0 + \tilde N_0)/2` and the phase-shifted nonlinear term :math:`\tilde N_0` is given by .. math:: \tilde N_0 = e^{-\frac{1}{2}k\dx}N\left(e^{\frac{1}{2}k\dx}S_0\right). """ dt = self.deltat diss = self.exact_linear_coefs.get_updated_coefs()[0] compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect # regular tendencies tendencies_0 = compute_tendencies() # shifted tendencies phaseshift = self._get_phaseshift() tendencies_shifted = ( compute_tendencies(phaseshift * state_spect) / phaseshift ) # dealiased tendencies tendencies_dealiased = 0.5 * (tendencies_0 + tendencies_shifted) step_Euler_inplace(state_spect, dt, tendencies_dealiased, diss)
[docs] def _time_step_Euler_phaseshift_random(self): r"""Forward Euler method, dealiasing with phase-shifting. Notes ----- We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), The forward Euler method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time :math:`t = 0`. - Euler approximation: .. math:: \p_t \log S = \sigma + \frac{N_\mathrm{dealias}}{S_0}, Integrating from :math:`t` to :math:`t+\dt`, it gives: .. math:: S_{\dt} = (S_0 + N_\mathrm{dealias} \dt) e^{\sigma \dt}. where the dealiased non-linear term :math:`N_\mathrm{dealias} = (\tilde N_{0\alpha} + \tilde N_{0\beta})/2` is computed as the average of two terms shifted with dependant phases :math:`\phi_\alpha = \alpha \dx k` and :math:`\phi_\beta = \beta \dx k` with :math:`\alpha` taken randomly between -1 and 1 and :math:`|\alpha - \beta| = 0.5`. """ dt = self.deltat diss = self.exact_linear_coefs.get_updated_coefs()[0] compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect # shifted tendencies phaseshift_alpha, phaseshift_beta = self._get_phaseshift_random() tendencies_alpha = ( compute_tendencies(phaseshift_alpha * state_spect) / phaseshift_alpha ) tendencies_beta = ( compute_tendencies(phaseshift_beta * state_spect) / phaseshift_beta ) # dealiased tendencies tendencies_dealiased = 0.5 * (tendencies_alpha + tendencies_beta) step_Euler_inplace(state_spect, dt, tendencies_dealiased, diss)
[docs] def _time_step_RK2(self): r"""Runge-Kutta 2 method. Notes ----- We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), The Runge-Kutta 2 method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time :math:`t = 0`. - Approximation 1: .. math:: \p_t \log S = \sigma + \frac{N_0}{S_0}, Integrating from :math:`t` to :math:`t+\dt/2`, it gives: .. math:: S_1 = (S_0 + \frac{\dt}{2} N_0) e^{\sigma \frac{\dt}{2}}. - Approximation 2: .. math:: \p_t \log S = \sigma + \frac{N_1}{S_1}, Integrating from :math:`t` to :math:`t+\dt` and retaining only the terms in :math:`(N\dt/S)^1` gives: .. math:: S_2 = S_0 e^{\sigma \dt} + \dt N_1 e^{\sigma \frac{\dt}{2}}. """ dt = self.deltat diss, diss2 = self.exact_linear_coefs.get_updated_coefs() compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect tendencies_0 = compute_tendencies() state_spect_12 = self._state_spect_tmp step_Euler( state_spect, dt / 2, tendencies_0, diss2, output=state_spect_12 ) tendencies_12 = compute_tendencies(state_spect_12, old=tendencies_0) step_like_RK2(state_spect, dt, tendencies_12, diss, diss2)
[docs] def _time_step_RK2_trapezoid(self): r"""Runge-Kutta 2 method with trapezoidal rule (Heun's method). Notes ----- We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), Heun's method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time :math:`t = 0`. - Approximation 1: .. math:: \p_t \log S = \sigma + \frac{N(S_0)}{S_0}, Integrating from :math:`t` to :math:`t+\dt`, it gives: .. math:: S_1 = (S_0 + N_0 \dt) e^{\sigma \dt}. - Approximation 2: .. math:: \p_t \log S = \sigma + \frac{1}{2}\left( \frac{N_0}{S_0} + \frac{N_1}{S_1}\right), Integrating from :math:`t` to :math:`t+\dt` and retaining only the terms in :math:`(N\dt/S)^1` gives: .. math:: S_2 = S_0 e^{\sigma \dt} + \frac{\dt}{2} (N_0 e^{\sigma \dt} + N_1). """ dt = self.deltat diss, diss2 = self.exact_linear_coefs.get_updated_coefs() compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect tendencies_0 = compute_tendencies() state_spect_1 = self._state_spect_tmp step_Euler(state_spect, dt, tendencies_0, diss, output=state_spect_1) tendencies_1 = compute_tendencies(state_spect_1) state_spect[:] = ( state_spect + dt / 2 * tendencies_0 ) * diss + dt / 2 * tendencies_1
[docs] def _time_step_RK2_phaseshift(self): r"""Runge-Kutta 2 method with phase-shifting. Notes ----- We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), Heun's method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time :math:`t = 0`. - Approximation 1: .. math:: \p_t \log S = \sigma + \frac{N_0}{S_0}, Integrating from :math:`t` to :math:`t+\dt`, it gives: .. math:: S_1 = (S_0 + N_0 \dt) e^{\sigma \dt}. - Approximation 2: .. math:: \p_t \log S = \sigma + \frac{N_d}{S_0 e^{\sigma \frac{\dt}{2}}}, where the dealiased non-linear term is :math:`N_\mathrm{dealias} = (N_0 + \tilde N_1)/2` and the phase-shifted nonlinear term :math:`\tilde N_1` is given by .. math:: \tilde N_1 = e^{-\frac{1}{2}k\dx}N\left(e^{\frac{1}{2}k\dx}S_1\right). Integrating from :math:`t` to :math:`t+\dt` and retaining only the terms in :math:`(N\dt/S)^1` gives: .. math:: S_2 = S_0 e^{\sigma \dt} + \dt N_d e^{\sigma \frac{\dt}{2}}. """ dt = self.deltat diss, diss2 = self.exact_linear_coefs.get_updated_coefs() compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect tendencies_0 = compute_tendencies() state_spect_1 = self._state_spect_tmp step_Euler(state_spect, dt, tendencies_0, diss, output=state_spect_1) phaseshift = self._get_phaseshift() tendencies_1_shift = compute_tendencies(phaseshift * state_spect_1) tendencies_d = self._state_spect_tmp if ts.is_transpiled: ts.use_block("rk2_tendencies_d") else: # transonic block ( # A tendencies_d, tendencies_0, tendencies_1_shift; # Am1 phaseshift; # ) tendencies_d[:] = 0.5 * ( tendencies_0 + tendencies_1_shift / phaseshift ) step_like_RK2(state_spect, dt, tendencies_d, diss, diss2)
[docs] def _time_step_RK2_phaseshift_random(self): r"""Runge-Kutta 2 method with phase-shifting (random). Notes ----- We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), Heun's method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time :math:`t = 0`. - Approximation 1: .. math:: \p_t \log S = \sigma + \frac{\tilde N_{0\alpha}}{S_0}, Integrating from :math:`t` to :math:`t+\dt`, it gives: .. math:: S_1 = (S_0 + \tilde N_{0\alpha} \dt) e^{\sigma \dt}. - Approximation 2: .. math:: \p_t \log S = \sigma + \frac{N_d}{S_0 e^{\sigma \frac{\dt}{2}}}, where the dealiased non-linear term is :math:`N_d = (\tilde N_{0\alpha} + \tilde N_{1\beta})/2`. Integrating from :math:`t` to :math:`t+\dt` and retaining only the terms in :math:`(N\dt/S)^1` gives: .. math:: S_2 = S_0 e^{\sigma \dt} + \dt N_d e^{\sigma \frac{\dt}{2}}. """ dt = self.deltat diss, diss2 = self.exact_linear_coefs.get_updated_coefs() phaseshift_alpha, phaseshift_beta = self._get_phaseshift_random() compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect tmp0 = np.empty_like(state_spect) state_spect_shift = mul(phaseshift_alpha, state_spect, output=tmp0) tendencies_0_shift = compute_tendencies( state_spect_shift, old=state_spect_shift ) tendencies_0 = div_inplace(tendencies_0_shift, phaseshift_alpha) state_spect_1 = step_Euler( state_spect, dt, tendencies_0, diss, output=self._state_spect_tmp ) tmp1 = np.empty_like(state_spect) state_spect_1_shift = mul(phaseshift_beta, state_spect_1, output=tmp1) tendencies_1_shift = compute_tendencies( state_spect_1_shift, old=state_spect_1_shift ) tendencies_d = mean_with_phaseshift( tendencies_0, tendencies_1_shift, phaseshift_beta, output=self._state_spect_tmp, ) step_like_RK2(state_spect, dt, tendencies_d, diss, diss2)
[docs] def _time_step_RK2_phaseshift_exact(self): r"""Runge-Kutta 2 method with phase-shifting for exact dealiasing. Notes ----- It requires 4 evaluations of the nonlinear terms (as RK4). We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), Heun's method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time :math:`t = 0`. - Approximation 1: .. math:: \p_t \log S = \sigma + \frac{N_{d0}}{S_0}, Integrating from :math:`t` to :math:`t+\dt`, it gives: .. math:: S_1 = (S_0 + N_{d0} \dt) e^{\sigma \dt}. - Approximation 2: .. math:: \p_t \log S = \sigma + \frac{N_d}{S_0 e^{\sigma \frac{\dt}{2}}}, where the dealiased non-linear term is :math:`N_d = (N_{d0} + N_{d1})/2`. Integrating from :math:`t` to :math:`t+\dt` and retaining only the terms in :math:`(N\dt/S)^1` gives: .. math:: S_2 = S_0 e^{\sigma \dt} + \dt N_d e^{\sigma \frac{\dt}{2}}. """ dt = self.deltat diss, diss2 = self.exact_linear_coefs.get_updated_coefs() phaseshift = self._get_phaseshift() compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect tmp0 = np.empty_like(state_spect) tmp1 = np.empty_like(state_spect) tmp2 = np.empty_like(state_spect) tmp3 = np.empty_like(state_spect) tendencies_0 = compute_tendencies(state_spect, old=tmp0) state_spect_shift = mul(phaseshift, state_spect, output=tmp1) tendencies_0_shift = compute_tendencies( state_spect_shift, old=state_spect_shift ) tendencies_d0 = mean_with_phaseshift( tendencies_0, tendencies_0_shift, phaseshift, output=tmp2 ) state_spect_1 = step_Euler( state_spect, dt, tendencies_d0, diss, output=self._state_spect_tmp ) tendencies_1 = compute_tendencies(state_spect_1, old=tmp0) state_spect_shift = mul(phaseshift, state_spect_1, output=tmp1) tendencies_1_shift = compute_tendencies( state_spect_shift, old=state_spect_shift ) tendencies_d = tmp3 if ts.is_transpiled: ts.use_block("rk2_exact") else: # based on approximation 1 # transonic block ( # A tendencies_d, tendencies_d0, tendencies_1, tendencies_1_shift; # Am1 phaseshift # ) tendencies_d[:] = 0.5 * ( tendencies_d0 + 0.5 * (tendencies_1 + tendencies_1_shift / phaseshift) ) step_like_RK2(state_spect, dt, tendencies_d, diss, diss2)
[docs] def _time_step_RK4(self): r"""Runge-Kutta 4 method. Notes ----- .. |SA1dt2| mathmacro:: S_{A1 \mathop{dt}/2} We consider an equation of the form .. math:: \p_t S = \sigma S + N(S), The Runge-Kutta 4 method computes an approximation of the solution after a time increment :math:`\dt`. We denote the initial time as :math:`t = 0`. This time scheme uses 4 approximations. Only the terms in :math:`\dt^1` are retained. - Approximation 1: .. math:: \p_t \log S = \sigma + \frac{N_0}{S_0}, Integrating from :math:`t` to :math:`t+\dt/2` gives: .. math:: \SA1dt2 = (S_0 + N_0 \dt/2) e^{\sigma \frac{\dt}{2}}. Integrating from :math:`t` to :math:`t+\dt` gives: .. math:: S_{A1\dt} = (S_0 + N_0 \dt) e^{\sigma \dt}. - Approximation 2: .. math:: \p_t \log S = \sigma + \frac{N(\SA1dt2)}{ \SA1dt2 }, Integrating from :math:`t` to :math:`t+\dt/2` gives: .. |SA2dt2| mathmacro:: S_{A2 \mathop{dt}/2} .. math:: \SA2dt2 = S_0 e^{\sigma \frac{\dt}{2}} + N(\SA1dt2) \frac{\dt}{2}. Integrating from :math:`t` to :math:`t+\dt` gives: .. math:: S_{A2\dt} = S_0 e^{\sigma \dt} + N(\SA1dt2) e^{\sigma \frac{\dt}{2}} \dt. - Approximation 3: .. math:: \p_t \log S = \sigma + \frac{N(\SA2dt2)}{ \SA2dt2 }, Integrating from :math:`t` to :math:`t+\dt` gives: .. math:: S_{A3\dt} = S_0 e^{\sigma \dt} + N(\SA2dt2) e^{\sigma \frac{\dt}{2}} \dt. - Approximation 4: .. math:: \p_t \log S = \sigma + \frac{N(S_{A3\dt})}{ S_{A3\dt} }, Integrating from :math:`t` to :math:`t+\dt` gives: .. math:: S_{A4\dt} = S_0 e^{\sigma \dt} + N(S_{A3\dt}) \dt. The final result is a pondered average of the results of 4 approximations for the time :math:`t+\dt`: .. math:: \frac{1}{3} \left[ \frac{1}{2} S_{A1\dt} + S_{A2\dt} + S_{A3\dt} + \frac{1}{2} S_{A4\dt} \right], which is equal to: .. math:: S_0 e^{\sigma \dt} + \frac{\dt}{3} \left[ \frac{1}{2} N_0 e^{\sigma \dt} + N(\SA1dt2) e^{\sigma \frac{\dt}{2}} + N(\SA2dt2) e^{\sigma \frac{\dt}{2}} + \frac{1}{2} N(S_{A3\dt})\right]. """ dt = self.deltat diss, diss2 = self.exact_linear_coefs.get_updated_coefs() compute_tendencies = self.sim.tendencies_nonlin state_spect = self.sim.state.state_spect tendencies_0 = compute_tendencies() state_spect_tmp1 = self._state_spect_tmp1 # rk4_step0 state_spect_tmp = step_Euler( state_spect, dt / 6, tendencies_0, diss, output=self._state_spect_tmp ) state_spect_12_approx1 = step_Euler( state_spect, dt / 2, tendencies_0, diss2, output=state_spect_tmp1, ) tendencies_1 = compute_tendencies( state_spect_12_approx1, old=tendencies_0 ) del state_spect_12_approx1 state_spect_12_approx2 = state_spect_tmp1 if ts.is_transpiled: ts.use_block("rk4_step1") else: # based on approximation 1 # transonic block ( # A state_spect, state_spect_tmp, # state_spect_12_approx2, tendencies_1; # A1 diss2; # float dt # ) # transonic block ( # A state_spect, state_spect_tmp, # state_spect_12_approx2, tendencies_1; # A2 diss2; # float dt # ) state_spect_tmp[:] += dt / 3 * diss2 * tendencies_1 state_spect_12_approx2[:] = ( state_spect * diss2 + dt / 2 * tendencies_1 ) tendencies_2 = compute_tendencies( state_spect_12_approx2, old=tendencies_1 ) del state_spect_12_approx2 state_spect_1_approx = state_spect_tmp1 if ts.is_transpiled: ts.use_block("rk4_step2") else: # based on approximation 2 # transonic block ( # A state_spect, state_spect_tmp, # state_spect_1_approx, tendencies_2; # A1 diss, diss2; # float dt # ) # transonic block ( # A state_spect, state_spect_tmp, # state_spect_1_approx, tendencies_2; # A2 diss, diss2; # float dt # ) state_spect_tmp[:] += dt / 3 * diss2 * tendencies_2 state_spect_1_approx[:] = ( state_spect * diss + dt * diss2 * tendencies_2 ) tendencies_3 = compute_tendencies(state_spect_1_approx, old=tendencies_2) del state_spect_1_approx if ts.is_transpiled: ts.use_block("rk4_step3") else: # result using the 4 approximations # transonic block ( # A state_spect, state_spect_tmp, tendencies_3; # float dt # ) state_spect[:] = state_spect_tmp + dt / 6 * tendencies_3