def nowcast_main_loop(
    precip,
    velocity,
    state,
    timesteps,
    extrap_method,
    func,
    extrap_kwargs=None,
    velocity_pert_gen=None,
    params=None,
    ensemble=False,
    num_ensemble_members=1,
    callback=None,
    return_output=True,
    num_workers=1,
    measure_time=False,
):
    """Utility method for advection-based nowcast models that are applied in
    the Lagrangian coordinates. In addition, this method allows the case, where
    one or more components of the model (e.g. an autoregressive process) require
    using regular integer time steps but the user-supplied values are irregular
    or non-integer.

    Parameters
    ----------
    precip : array_like
        Array of shape (m,n) containing the most recently observed precipitation
        field.
    velocity : array_like
        Array of shape (2,m,n) containing the x- and y-components of the
        advection field.
    state : object
        The initial state of the nowcast model.
    timesteps : int or list of floats
        Number of time steps to forecast or a list of time steps for which the
        forecasts are computed. The elements of the list are required to be in
        ascending order.
    extrap_method : str, optional
        Name of the extrapolation method to use. See the documentation of
        :py:mod:`pysteps.extrapolation.interface`.
    ensemble : bool
        Set to True to produce a nowcast ensemble.
    num_ensemble_members : int
        Number of ensemble members. Applicable if ensemble is set to True.
    func : function
        A function that takes the current state of the nowcast model and its
        parameters and returns a forecast field and the new state. The shape of
        the forecast field is expected to be (m,n) for a deterministic nowcast
        and (n_ens_members,m,n) for an ensemble.
    extrap_kwargs : dict, optional
        Optional dictionary containing keyword arguments for the extrapolation
        method. See the documentation of pysteps.extrapolation.
    velocity_pert_gen : list, optional
        Optional list of functions that generate velocity perturbations. The
        length of the list is expected to be n_ens_members. The functions
        are expected to take lead time (relative to timestep index) as input
        argument and return a perturbation field of shape (2,m,n).
    params : dict, optional
        Optional dictionary containing keyword arguments for func.
    callback : function, optional
        Optional function that is called after computation of each time step of
        the nowcast. The function takes one argument: the nowcast array. This
        can be used, for instance, writing output files.
    return_output : bool, optional
        Set to False to disable returning the output forecast fields and return
        None instead. This can save memory if the intermediate results are
        instead written to files using the callback function.
    num_workers : int, optional
        Number of parallel workers to use. Applicable if a nowcast ensemble is
        generated.
    measure_time : bool, optional
        If set to True, measure, print and return the computation time.

    Returns
    -------
    out : list
        List of forecast fields for the given time steps. If measure_time is
        True, return a pair, where the second element is the total computation
        time in the loop.
    """
    precip_forecast_out = None

    timesteps, original_timesteps, timestep_type = create_timestep_range(timesteps)

    state_cur = state
    if not ensemble:
        precip_forecast_prev = precip[np.newaxis, :]
    else:
        precip_forecast_prev = np.stack([precip for _ in range(num_ensemble_members)])
    displacement = None
    t_prev = 0.0
    t_total = 0.0

    # initialize the extrapolator
    extrapolator = extrapolation.get_method(extrap_method)

    x_values, y_values = np.meshgrid(
        np.arange(precip.shape[1]), np.arange(precip.shape[0])
    )

    xy_coords = np.stack([x_values, y_values])

    if extrap_kwargs is None:
        extrap_kwargs = dict()
    else:
        extrap_kwargs = extrap_kwargs.copy()
    extrap_kwargs["xy_coords"] = xy_coords
    extrap_kwargs["return_displacement"] = True

    if measure_time:
        starttime_total = time.time()

    # loop through the integer time steps or bins if non-integer time steps
    # were given
    for t, subtimestep_idx in enumerate(timesteps):
        if timestep_type == "list":
            subtimesteps = [original_timesteps[t_] for t_ in subtimestep_idx]
        else:
            subtimesteps = [t]

        if (timestep_type == "list" and subtimesteps) or (
            timestep_type == "int" and t > 0
        ):
            is_nowcast_time_step = True
        else:
            is_nowcast_time_step = False

        # print a message if nowcasts are computed for the current integer time
        # step (this is not necessarily the case, since the current bin might
        # not contain any time steps)
        if is_nowcast_time_step:
            print(
                f"Computing nowcast for time step {t}... ",
                end="",
                flush=True,
            )

            if measure_time:
                starttime = time.time()

        # call the function to iterate the integer-timestep part of the model
        # for one time step
        precip_forecast_new, state_new = func(state_cur, params)

        if not ensemble:
            precip_forecast_new = precip_forecast_new[np.newaxis, :]

        # advect the currect forecast field to the subtimesteps in the current
        # timestep bin and append the results to the output list
        # apply temporal interpolation to the forecasts made between the
        # previous and the next integer time steps
        for t_sub in subtimesteps:
            if t_sub > 0:
                t_diff_prev_int = t_sub - int(t_sub)
                if t_diff_prev_int > 0.0:
                    precip_forecast_ip = (
                        1.0 - t_diff_prev_int
                    ) * precip_forecast_prev + t_diff_prev_int * precip_forecast_new
                else:
                    precip_forecast_ip = precip_forecast_prev

                t_diff_prev = t_sub - t_prev
                t_total += t_diff_prev

                if displacement is None:
                    displacement = [None for _ in range(precip_forecast_ip.shape[0])]

                if precip_forecast_out is None and return_output:
                    precip_forecast_out = [
                        [] for _ in range(precip_forecast_ip.shape[0])
                    ]

                precip_forecast_out_cur = [
                    None for _ in range(precip_forecast_ip.shape[0])
                ]

                def worker1(i):
                    extrap_kwargs_ = extrap_kwargs.copy()
                    extrap_kwargs_["displacement_prev"] = displacement[i]
                    extrap_kwargs_["allow_nonfinite_values"] = (
                        True if np.any(~np.isfinite(precip_forecast_ip[i])) else False
                    )

                    if velocity_pert_gen is not None:
                        velocity_ = velocity + velocity_pert_gen[i](t_total)
                    else:
                        velocity_ = velocity

                    precip_forecast_ep, displacement[i] = extrapolator(
                        precip_forecast_ip[i],
                        velocity_,
                        [t_diff_prev],
                        **extrap_kwargs_,
                    )

                    precip_forecast_out_cur[i] = precip_forecast_ep[0]
                    if return_output:
                        precip_forecast_out[i].append(precip_forecast_ep[0])

                if DASK_IMPORTED and ensemble and num_ensemble_members > 1:
                    res = []
                    for i in range(precip_forecast_ip.shape[0]):
                        res.append(dask.delayed(worker1)(i))
                    dask.compute(*res, num_workers=num_workers)
                else:
                    for i in range(precip_forecast_ip.shape[0]):
                        worker1(i)

                if callback is not None:
                    precip_forecast_out_cur = np.stack(precip_forecast_out_cur)
                    callback(precip_forecast_out_cur)

                precip_forecast_out_cur = None
                t_prev = t_sub

        # advect the forecast field by one time step if no subtimesteps in the
        # current interval were found
        if not subtimesteps:
            t_diff_prev = t + 1 - t_prev
            t_total += t_diff_prev

            if displacement is None:
                displacement = [None for _ in range(precip_forecast_new.shape[0])]

            def worker2(i):
                extrap_kwargs_ = extrap_kwargs.copy()
                extrap_kwargs_["displacement_prev"] = displacement[i]

                if velocity_pert_gen is not None:
                    velocity_ = velocity + velocity_pert_gen[i](t_total)
                else:
                    velocity_ = velocity

                _, displacement[i] = extrapolator(
                    None,
                    velocity_,
                    [t_diff_prev],
                    **extrap_kwargs_,
                )

            if DASK_IMPORTED and ensemble and num_ensemble_members > 1:
                res = []
                for i in range(precip_forecast_new.shape[0]):
                    res.append(dask.delayed(worker2)(i))
                dask.compute(*res, num_workers=num_workers)
            else:
                for i in range(precip_forecast_new.shape[0]):
                    worker2(i)

            t_prev = t + 1

        precip_forecast_prev = precip_forecast_new
        state_cur = state_new

        if is_nowcast_time_step:
            if measure_time:
                print(f"{time.time() - starttime:.2f} seconds.")
            else:
                print("done.")

    if return_output:
        precip_forecast_out = np.stack(precip_forecast_out)
        if not ensemble:
            precip_forecast_out = precip_forecast_out[0, :]

    if measure_time:
        return precip_forecast_out, time.time() - starttime_total
    else:
        return precip_forecast_out
Utility method for advection-based nowcast models that are applied in
the Lagrangian coordinates. In addition, this method allows the case, where
one or more components of the model (e.g. an autoregressive process) require
using regular integer time steps but the user-supplied values are irregular
or non-integer.

Parameters
----------
precip : array_like
Array of shape (m,n) containing the most recently observed precipitation
field.
velocity : array_like
Array of shape (2,m,n) containing the x- and y-components of the
advection field.
state : object
The initial state of the nowcast model.
timesteps : int or list of floats
Number of time steps to forecast or a list of time steps for which the
forecasts are computed. The elements of the list are required to be in
ascending order.
extrap_method : str, optional
Name of the extrapolation method to use. See the documentation of
:py:mod:`pysteps.extrapolation.interface`.
ensemble : bool
Set to True to produce a nowcast ensemble.
num_ensemble_members : int
Number of ensemble members. Applicable if ensemble is set to True.
func : function
A function that takes the current state of the nowcast model and its
parameters and returns a forecast field and the new state. The shape of
the forecast field is expected to be (m,n) for a deterministic nowcast
and (n_ens_members,m,n) for an ensemble.
extrap_kwargs : dict, optional
Optional dictionary containing keyword arguments for the extrapolation
method. See the documentation of pysteps.extrapolation.
velocity_pert_gen : list, optional
Optional list of functions that generate velocity perturbations. The
length of the list is expected to be n_ens_members. The functions
are expected to take lead time (relative to timestep index) as input
argument and return a perturbation field of shape (2,m,n).
params : dict, optional
Optional dictionary containing keyword arguments for func.
callback : function, optional
Optional function that is called after computation of each time step of
the nowcast. The function takes one argument: the nowcast array. This
can be used, for instance, writing output files.
return_output : bool, optional
Set to False to disable returning the output forecast fields and return
None instead. This can save memory if the intermediate results are
instead written to files using the callback function.
num_workers : int, optional
Number of parallel workers to use. Applicable if a nowcast ensemble is
generated.
measure_time : bool, optional
If set to True, measure, print and return the computation time.

Returns
-------
out : list
List of forecast fields for the given time steps. If measure_time is
True, return a pair, where the second element is the total computation
time in the loop.
pysteps/nowcasts/utils.py