|
|
from datetime import datetime, timedelta
|
|
|
import glob
|
|
|
import os
|
|
|
|
|
|
from fire import Fire
|
|
|
import h5py
|
|
|
from matplotlib import pyplot as plt
|
|
|
import numpy as np
|
|
|
|
|
|
from ldcast import forecast
|
|
|
from ldcast.visualization import plots
|
|
|
|
|
|
|
|
|
def read_data(
|
|
|
data_dir="../data/demo/20210622",
|
|
|
t0=datetime(2021,6,22,18,35),
|
|
|
interval=timedelta(minutes=5),
|
|
|
past_timesteps=4,
|
|
|
crop_box=((128,480), (160,608))
|
|
|
):
|
|
|
cb = crop_box
|
|
|
R_past = []
|
|
|
t = t0 - (past_timesteps-1) * interval
|
|
|
for i in range(past_timesteps):
|
|
|
timestamp = t.strftime("%y%j%H%M")
|
|
|
fn = f"RZC{timestamp}VL.[80]01.h5"
|
|
|
fn = os.path.join(data_dir, fn)
|
|
|
found_files = glob.glob(fn)
|
|
|
if found_files:
|
|
|
fn = found_files[0]
|
|
|
else:
|
|
|
raise FileNotFoundError(f"Unable to find data file {fn}.")
|
|
|
with h5py.File(fn, 'r') as f:
|
|
|
R = f["dataset1"]["data1"]["data"][:]
|
|
|
R = R[cb[0][0]:cb[0][1], cb[1][0]:cb[1][1]]
|
|
|
R_past.append(R)
|
|
|
t += interval
|
|
|
|
|
|
R_past = np.stack(R_past, axis=0)
|
|
|
return R_past
|
|
|
|
|
|
|
|
|
def plot_border(ax, crop_box=((128,480), (160,608))):
|
|
|
import shapefile
|
|
|
border = shapefile.Reader("../data/Border_CH.shp")
|
|
|
shapes = list(border.shapeRecords())
|
|
|
for shape in shapes:
|
|
|
x = np.array([i[0]/1000. for i in shape.shape.points[:]])
|
|
|
y = np.array([i[1]/1000. for i in shape.shape.points[:]])
|
|
|
ax.plot(
|
|
|
x-crop_box[1][0]-255, 480-y-crop_box[0][0],
|
|
|
'k', linewidth=1.0
|
|
|
)
|
|
|
|
|
|
|
|
|
def plot_frame(R, fn, draw_border=True, t=None, label=None):
|
|
|
fig = plt.figure(dpi=150)
|
|
|
ax = fig.add_subplot()
|
|
|
plots.plot_precip_image(ax, R)
|
|
|
if draw_border:
|
|
|
plot_border(ax)
|
|
|
if t is not None:
|
|
|
timestamp = "%Y-%m-%d %H:%M UTC"
|
|
|
if label is not None:
|
|
|
timestamp += f" ({label})"
|
|
|
ax.text(
|
|
|
0.02, 0.98, t.strftime(timestamp),
|
|
|
horizontalalignment='left', verticalalignment='top',
|
|
|
transform=ax.transAxes
|
|
|
)
|
|
|
|
|
|
fig.savefig(fn, bbox_inches='tight')
|
|
|
plt.close(fig)
|
|
|
|
|
|
|
|
|
def forecast_demo(
|
|
|
ldm_weights_fn="../models/genforecast/genforecast-radaronly-256x256-20step.pt",
|
|
|
autoenc_weights_fn="../models/autoenc/autoenc-32-0.01.pt",
|
|
|
num_diffusion_iters=50,
|
|
|
out_dir="../figures/demo/",
|
|
|
data_dir="../data/demo/20210622",
|
|
|
t0=datetime(2021,6,22,18,35),
|
|
|
interval=timedelta(minutes=5),
|
|
|
past_timesteps=4,
|
|
|
crop_box=((128,480), (160,608)),
|
|
|
draw_border=True,
|
|
|
ensemble_members=1,
|
|
|
):
|
|
|
R_past = read_data(
|
|
|
data_dir=data_dir, t0=t0, interval=interval,
|
|
|
past_timesteps=past_timesteps, crop_box=crop_box
|
|
|
)
|
|
|
if ensemble_members == 1:
|
|
|
fc = forecast.Forecast(
|
|
|
ldm_weights_fn=ldm_weights_fn,
|
|
|
autoenc_weights_fn=autoenc_weights_fn
|
|
|
)
|
|
|
R_pred = fc(
|
|
|
R_past,
|
|
|
num_diffusion_iters=num_diffusion_iters
|
|
|
)
|
|
|
elif ensemble_members > 1:
|
|
|
fc = forecast.ForecastDistributed(
|
|
|
ldm_weights_fn=ldm_weights_fn,
|
|
|
autoenc_weights_fn=autoenc_weights_fn,
|
|
|
)
|
|
|
R_past = R_past.reshape((1,) + R_past.shape)
|
|
|
R_pred = fc(
|
|
|
R_past,
|
|
|
num_diffusion_iters=num_diffusion_iters,
|
|
|
ensemble_members=ensemble_members
|
|
|
)
|
|
|
R_past = R_past[0,...]
|
|
|
R_pred = R_pred[0,...].mean(axis=-1)
|
|
|
else:
|
|
|
raise ValueError("ensemble_members must be > 0")
|
|
|
|
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
for k in range(R_past.shape[0]):
|
|
|
fn = os.path.join(out_dir, f"R_past-{k:02d}.png")
|
|
|
t = t0 - (R_past.shape[0]-k-1) * interval
|
|
|
plot_frame(R_past[k,:,:], fn, draw_border=draw_border,
|
|
|
t=t, label="Real")
|
|
|
for k in range(R_pred.shape[0]):
|
|
|
fn = os.path.join(out_dir, f"R_pred-{k:02d}.png")
|
|
|
t = t0 + (k+1)*interval
|
|
|
plot_frame(R_pred[k,:,:], fn, draw_border=draw_border,
|
|
|
t=t, label="Predicted")
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
Fire(forecast_demo)
|
|
|
|
