File size: 3,892 Bytes
d2f661a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | # following https://pysteps.readthedocs.io/en/stable/auto_examples/plot_steps_nowcast.html
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:
# occasional PySTEPS errors that happen with little/no precip
# therefore returning all zeros makes sense
R_f = self.zero_prediction(R, zerovalue)
else:
raise
# Back-transform to rain rates
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
|