|
|
|
|
|
from datetime import timedelta
|
|
|
|
|
|
import dask
|
|
|
import numpy as np
|
|
|
from pysteps import nowcasts
|
|
|
from pysteps.motion.lucaskanade import dense_lucaskanade
|
|
|
from pysteps.utils import transformation
|
|
|
|
|
|
|
|
|
class PySTEPSModel:
|
|
|
def __init__(
|
|
|
self,
|
|
|
data_format='channels_first',
|
|
|
future_timesteps=20,
|
|
|
ensemble_size=32,
|
|
|
km_per_pixel=1.0,
|
|
|
interval=timedelta(minutes=5),
|
|
|
transform_to_rainrate=None,
|
|
|
transform_from_rainrate=None,
|
|
|
):
|
|
|
self.transform_to_rainrate = transform_to_rainrate
|
|
|
self.transform_from_rainrate = transform_from_rainrate
|
|
|
self.data_format = data_format
|
|
|
self.nowcast_method = nowcasts.get_method("steps")
|
|
|
self.future_timesteps = future_timesteps
|
|
|
self.ensemble_size = ensemble_size
|
|
|
self.km_per_pixel = km_per_pixel
|
|
|
self.interval = interval
|
|
|
|
|
|
def zero_prediction(self, R, zerovalue):
|
|
|
out_shape = (self.future_timesteps,) + R.shape[1:] + \
|
|
|
(self.ensemble_size,)
|
|
|
return np.full(out_shape, zerovalue, dtype=R.dtype)
|
|
|
|
|
|
def predict_sample(self, x, threshold=-10.0, zerovalue=-15.0):
|
|
|
R = self.transform_to_rainrate(x)
|
|
|
(R, _) = transformation.dB_transform(
|
|
|
R, threshold=0.1, zerovalue=zerovalue
|
|
|
)
|
|
|
R[~np.isfinite(R)] = zerovalue
|
|
|
if (R == zerovalue).all():
|
|
|
R_f = self.zero_prediction(R, zerovalue)
|
|
|
else:
|
|
|
V = dense_lucaskanade(R)
|
|
|
try:
|
|
|
R_f = self.nowcast_method(
|
|
|
R,
|
|
|
V,
|
|
|
self.future_timesteps,
|
|
|
n_ens_members=self.ensemble_size,
|
|
|
n_cascade_levels=6,
|
|
|
precip_thr=threshold,
|
|
|
kmperpixel=self.km_per_pixel,
|
|
|
timestep=self.interval.total_seconds()/60,
|
|
|
noise_method="nonparametric",
|
|
|
vel_pert_method="bps",
|
|
|
mask_method="incremental",
|
|
|
num_workers=2
|
|
|
)
|
|
|
R_f = R_f.transpose(1,2,3,0)
|
|
|
except (ValueError, RuntimeError) as e:
|
|
|
zero_error = str(e).endswith("contains non-finite values") or \
|
|
|
str(e).startswith("zero-size array to reduction operation") or \
|
|
|
str(e).endswith("nonstationary AR(p) process")
|
|
|
if zero_error:
|
|
|
|
|
|
|
|
|
R_f = self.zero_prediction(R, zerovalue)
|
|
|
else:
|
|
|
raise
|
|
|
|
|
|
|
|
|
R_f = transformation.dB_transform(
|
|
|
R_f, threshold=threshold, inverse=True
|
|
|
)[0]
|
|
|
|
|
|
if self.transform_from_rainrate is not None:
|
|
|
R_f = self.transform_from_rainrate(R_f)
|
|
|
|
|
|
return R_f
|
|
|
|
|
|
def __call__(self, x, parallel=True):
|
|
|
while isinstance(x, list) or isinstance(x, tuple):
|
|
|
x = x[0]
|
|
|
x = np.array(x, copy=False)
|
|
|
if self.data_format == "channels_first":
|
|
|
x = x.transpose(0,2,3,4,1)
|
|
|
|
|
|
pred = self.predict_sample
|
|
|
if parallel:
|
|
|
pred = dask.delayed(pred)
|
|
|
y = [
|
|
|
pred(x[i,:,:,:,0])
|
|
|
for i in range(x.shape[0])
|
|
|
]
|
|
|
if parallel:
|
|
|
y = dask.compute(y, scheduler="threads", num_workers=len(y))[0]
|
|
|
y = np.stack(y, axis=0)
|
|
|
|
|
|
if self.data_format == "channels_first":
|
|
|
y = np.expand_dims(y, 1)
|
|
|
else:
|
|
|
y = np.expand_dims(y, -2)
|
|
|
|
|
|
return y
|
|
|
|
