def _linda_forecast(
    precip,
    precip_lagr_diff,
    timesteps,
    forecast_gen,
    precip_pert_gen,
    velocity_pert_gen,
    n_ensemble_members,
    seed,
    measure_time,
    print_info,
    return_output,
    callback,
):
    """Compute LINDA nowcast."""
    # compute convolved difference fields
    precip_lagr_diff = precip_lagr_diff.copy()

    for i in range(precip_lagr_diff.shape[0]):
        for _ in range(forecast_gen["ari_order"] - i):
            precip_lagr_diff[i] = _composite_convolution(
                precip_lagr_diff[i],
                forecast_gen["kernels_1"],
                forecast_gen["interp_weights"],
            )

    # initialize the random generators
    if precip_pert_gen is not None:
        rs_precip_pert = []
        np.random.seed(seed)
        for _ in range(n_ensemble_members):
            rs = np.random.RandomState(seed)
            rs_precip_pert.append(rs)
            seed = rs.randint(0, high=1e9)
    else:
        rs_precip_pert = None

    if velocity_pert_gen is not None:
        velocity_perturbators = []
        np.random.seed(seed)
        for _ in range(n_ensemble_members):
            vp = velocity_pert_gen["init_func"](seed)
            velocity_perturbators.append(
                lambda t, vp=vp: velocity_pert_gen["gen_func"](
                    vp, t * velocity_pert_gen["timestep"]
                )
            )
            seed = np.random.RandomState(seed).randint(0, high=1e9)
    else:
        velocity_perturbators = None

    state = {
        "precip_forecast": [precip[-1].copy() for _ in range(n_ensemble_members)],
        "precip_lagr_diff": [
            precip_lagr_diff.copy() for _ in range(n_ensemble_members)
        ],
        "rs_precip_pert": rs_precip_pert,
    }
    params = {
        "interp_weights": forecast_gen["interp_weights"],
        "kernels_1": forecast_gen["kernels_1"],
        "kernels_2": forecast_gen["kernels_2"],
        "mask_adv": forecast_gen["mask_adv"],
        "num_ens_members": n_ensemble_members,
        "num_workers": forecast_gen["num_workers"],
        "num_ensemble_workers": min(n_ensemble_members, forecast_gen["num_workers"]),
        "precip_pert_gen": precip_pert_gen,
        "psi": forecast_gen["psi"],
    }

    precip_forecast = nowcast_main_loop(
        precip[-1],
        forecast_gen["velocity"],
        state,
        timesteps,
        forecast_gen["extrap_method"],
        _update,
        extrap_kwargs=forecast_gen["extrap_kwargs"],
        velocity_pert_gen=velocity_perturbators,
        params=params,
        ensemble=True,
        num_ensemble_members=n_ensemble_members,
        callback=callback,
        return_output=return_output,
        num_workers=forecast_gen["num_workers"],
        measure_time=measure_time,
    )
    if measure_time:
        precip_forecast, mainloop_time = precip_forecast

    if return_output:
        if not forecast_gen["add_perturbations"]:
            precip_forecast = precip_forecast[0]
        if measure_time:
            return precip_forecast, mainloop_time
        else:
            return precip_forecast
    else:
        return None
Compute LINDA nowcast.
pysteps/nowcasts/linda.py