fit_kriging_model.py

"""Fit a spatiotemporal Kriging model to the given points containing radar- and
gauge-measured rainfall accumulations. This implementation uses 3d ordinary
Kriging from PyKrige (https://geostat-framework.readthedocs.io/projects/pykrige/en/stable)
so that the third dimension is reserved for time.

Input
-----
- radar-gauge pair file generated by running collect_radar_gauge_pairs.py with
  attribute "gauge_location" included

Output
------
Pickle dump containing the fitted model.

Configuration files (in the config/<profile> directory)
-------------------------------------------------------
- fit_kriging_model.cfg
"""

import argparse
import configparser
import os
import pickle

import numpy as np
from pykrige.ok3d import OrdinaryKriging3D

# parse command-line arguments
argparser = argparse.ArgumentParser()
argparser.add_argument("rgpairfile", type=str, help="radar-gauge pair file")
argparser.add_argument("outfile", type=str, help="output file")
argparser.add_argument("profile", type=str, help="configuration profile to use")
args = argparser.parse_args()

# read configuration file
config = configparser.ConfigParser()
config.read(os.path.join("config", args.profile, "fit_kriging_model.cfg"))

radar_gauge_pairs = pickle.load(open(args.rgpairfile, "rb"))

# collect data points for Kriging
x = []
y = []
z = []
val = []

for timestamp in radar_gauge_pairs.keys():
    for fmisid in radar_gauge_pairs[timestamp].keys():
        p = radar_gauge_pairs[timestamp][fmisid]
        x_, y_ = p[2]["gauge_location"]

        x.append(x_)
        y.append(y_)
        z.append(timestamp.timestamp())
        val.append(np.log10(p[1] / p[0]))

if config["kriging"]["time_scaling_factor"] == "auto":
    anisotropy_scaling_z = 0.5 * (np.std(x) + np.std(y)) / np.std(z)
else:
    anisotropy_scaling_z = float(config["kriging"]["time_scaling_factor"])

model = OrdinaryKriging3D(
    x,
    y,
    z,
    val,
    variogram_model="exponential",
    anisotropy_scaling_z=anisotropy_scaling_z,
    verbose=True,
)

pickle.dump(model, open(args.outfile, "wb"))