MILES-GUESS Regression Example Notebook

John Schreck, David John Gagne, Charlie Becker, Gabrielle Gantos, Dhamma Kimpara, Thomas Martin

Objective

This notebook is meant to demonstrate a functional example of how to use the miles-guess repository for training an evidential regression model.

Steps to get this notebook running: 1) Follow package installation steps in the miles-guess ReadMe. 2) Run the cells in this notebook in order.

Quick Intro to Surface Layer (SL)

The Surface Layer (SL) problem attempts to predict energy fluxes from Earth’s surface given near-surface atmospheric conditions. The dataset used in this example problem originates from the Royal Netherlands Meteorological Institute (KNMI) Cabauw Experimental Site for Atmospheric Research. The outgoing longwave radiation, converted to skin temperature, is derived from a device located approximately 200 meters from the flux tower.

Key Points

• Evidential deep learning is capable of training a single model to simultaneously learn the problem and estimate the problem’s aleatoric and epistemic uncertainties. See this paper here().

• This notebook represents an example of evidential deep learning for a regression problem. Specifically, that regression problem

• Data? (Make sure to mention the subsetting of the dataset for ptype)

• What will user be prepared to do with miles-guess after running this notebook?

import os, tqdm, yaml
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

from sklearn.model_selection import GroupShuffleSplit
from sklearn.preprocessing import MinMaxScaler, RobustScaler

from mlguess.keras.models import EvidentialRegressorDNN
from mlguess.keras.callbacks import get_callbacks
from mlguess.regression_uq import compute_results

2024-03-12 17:54:32.060087: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-03-12 17:54:32.110346: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-03-12 17:54:32.949343: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT

Config File

Parameters of interest

#TODO: Highlight parameters of interest

config = "../config/surface_layer/mlp.yml"

with open(config) as cf:
    conf = yaml.load(cf, Loader=yaml.FullLoader)

Data

data = pd.read_csv("../data/sample_cabauw_surface_layer.csv")

Train-Valid-Test Splits

This is a three-step process: 1. Group data by day. 2. Split all of the data on the day column between train (90%) and test (10%). The test data will be consisten accross all trained models and all data and model ensembles. 3. Split the 90% training data from Step 1 into training and validation.

data["day"] = data["Time"].apply(lambda x: str(x).split(" ")[0])

data_seed = 0
flat_seed = 1000

# Need the same test_data for all trained models (data and model ensembles)
gsp = GroupShuffleSplit(n_splits=1, random_state=flat_seed, train_size=0.9)
splits = list(gsp.split(data, groups=data["day"]))
train_index, test_index = splits[0]
train_data, test_data = data.iloc[train_index].copy(), data.iloc[test_index].copy()

# Make N train-valid splits using day as grouping variable
gsp = GroupShuffleSplit(n_splits=1,  random_state=flat_seed, train_size=0.885)
splits = list(gsp.split(train_data, groups=train_data["day"]))
train_index, valid_index = splits[data_seed]
train_data, valid_data = train_data.iloc[train_index].copy(), train_data.iloc[valid_index].copy()

Data Scaling

input_cols = conf["data"]["input_cols"]
#TODO: Should we include the other two output variables as potential options?
output_cols = ["friction_velocity:surface:m_s-1"]

x_scaler, y_scaler = RobustScaler(), MinMaxScaler((0, 1))
x_train = x_scaler.fit_transform(train_data[input_cols])
x_valid = x_scaler.transform(valid_data[input_cols])
x_test = x_scaler.transform(test_data[input_cols])

y_train = y_scaler.fit_transform(train_data[output_cols])
y_valid = y_scaler.transform(valid_data[output_cols])
y_test = y_scaler.transform(test_data[output_cols])

Train the model

# TODO: Why overwrite these variables in the config?
conf["model"]["epochs"] = 1
conf["model"]["verbose"] = 1

model = EvidentialRegressorDNN(**conf["model"])
model.compile()

2024-03-12 17:54:35.228747: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:282] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected

model.fit(x_train,
          y_train,
          validation_data=(x_valid, y_valid))

  86/1795 ━━━━━━━━━━━━━━━━━━━━ 1s 588us/step - loss: 3.0557e-05

/glade/work/ggantos/conda-envs/guess/lib/python3.10/site-packages/keras/src/optimizers/base_optimizer.py:576: UserWarning: Gradients do not exist for variables ['bias', 'kernel', 'bias'] when minimizing the loss. If using `model.compile()`, did you forget to provide a `loss` argument?
  warnings.warn(

1795/1795 ━━━━━━━━━━━━━━━━━━━━ 1s 669us/step - loss: 2.9939e-05 - val_loss: 2.8053e-05

<keras.src.callbacks.history.History at 0x1460bc91e200>

Predict with the model

y_pred = model.predict(x_test, return_uncertainties=False)

8/8 ━━━━━━━━━━━━━━━━━━━━ 0s 1ms/step

y_pred

array([[ 0.06768488,  0.61546266,  1.7284483 ,  0.65492064],
       [-0.02590589,  0.7062391 ,  1.6261563 ,  0.66753787],
       [ 0.12891376,  0.7776666 ,  1.5624065 ,  0.73487395],
       ...,
       [-0.01874506,  0.7131814 ,  1.6728607 ,  0.7000923 ],
       [ 0.00279877,  0.73572874,  1.6917179 ,  0.7401652 ],
       [-0.07375085,  0.6940977 ,  1.7539929 ,  0.6484979 ]],
      dtype=float32)

mae = np.mean(np.abs(y_pred[:, 0] - test_data[output_cols[0]]))
mae

0.2545100680072519

Create a Monte Carlo ensemble

ensemble:
    n_models: 1
    n_splits: 1
    monte_carlo_passes: 100

Create an ensemble of models (random initialization) using cross-validation splits. If MC passes > 0, an ensemble is created after each model finishes training on the test holdout. The LOTV may then be applied to the ensemble created from cross-validation. Otherwise a single ensemble is created but the LOTV is not applied.

# TODO: will have a separate ensemble class implemented in the future
results = model.predict_monte_carlo(x_test,
                                    conf["monte_carlo_passes"],
                                    y_scaler)

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[13], line 1
----> 1 results = model.predict_monte_carlo(x_test,
      2                                     conf["monte_carlo_passes"],
      3                                     y_scaler)

AttributeError: 'EvidentialRegressorDNN' object has no attribute 'predict_monte_carlo'

mu_ensemble = np.mean(results, axis=0)
var_ensemble = np.var(results, axis=0)

config = "../config/surface_layer/gaussian.yml"
with open(config) as cf:
    conf = yaml.load(cf, Loader=yaml.FullLoader)

conf["model"]["epochs"] = 1
conf["model"]["verbose"] = 1

gauss_model = GaussianRegressorDNN(**conf["model"])
gauss_model.build_neural_network(x_train.shape[-1], y_train.shape[-1])

gauss_model.fit(
    x_train,
    y_train,
    validation_data=(x_valid, y_valid),
    callbacks=get_callbacks(conf)
)

mu, var = gauss_model.predict_uncertainty(x_test, y_scaler)

# compute variance and std from learned parameters
#mu, var = gauss_model.calc_uncertainties(y_pred, y_scaler)

mae = np.mean(np.abs(mu[:, 0]-test_data[output_cols[0]]))
print(mae, np.mean(var) ** (1/2))

sns.jointplot(x=test_data[output_cols[0]], y=mu[:, 0], kind='hex')
plt.xlabel('Target')
plt.ylabel('Predicted Target')
plt.show()

sns.jointplot(x=mu[:, 0], y=np.sqrt(var)[:, 0], kind='hex')
plt.xlabel('Predicted mu')
plt.ylabel('Predicted sigma')
plt.show()

config = "../config/surface_layer/evidential.yml"

with open(config) as cf:
    conf = yaml.load(cf, Loader=yaml.FullLoader)

conf["model"]["epochs"] = 5
conf["model"]["verbose"] = 1

ev_model = EvidentialRegressorDNN(**conf["model"])
ev_model.build_neural_network(x_train.shape[-1], y_train.shape[-1])

ev_model.fit(
    x_train,
    y_train,
    validation_data=(x_valid, y_valid),
    callbacks=get_callbacks(conf))

result = ev_model.predict_uncertainty(x_test, scaler=y_scaler)
mu, aleatoric, epistemic = result

mae = np.mean(np.abs(mu[:, 0] - test_data[output_cols[0]]))
print(mae, np.mean(aleatoric)**(1/2), np.mean(epistemic)**(1/2))

compute_results(test_data,
                output_cols,
                mu,
                np.sqrt(aleatoric),
                np.sqrt(epistemic))

config = "../config/surface_layer/gaussian.yml"
with open(config) as cf:
    conf = yaml.load(cf, Loader=yaml.FullLoader)

conf["save_loc"] = "./"
conf["model"]["epochs"] = 1
conf["model"]["verbose"] = 0
n_splits = conf["ensemble"]["n_splits"]

# make save directory for model weights
os.makedirs(os.path.join(conf["save_loc"], "cv_ensemble", "models"), exist_ok=True)

data_seed = 0
gsp = GroupShuffleSplit(n_splits=1, random_state=flat_seed, train_size=0.9)
splits = list(gsp.split(data, groups=data["day"]))
train_index, test_index = splits[0]

ensemble_mu = np.zeros((n_splits, test_data.shape[0], 1))
ensemble_var = np.zeros((n_splits, test_data.shape[0], 1))

for data_seed in tqdm.tqdm(range(n_splits)):
    data = pd.read_csv(fn)
    data["day"] = data["Time"].apply(lambda x: str(x).split(" ")[0])

    # Need the same test_data for all trained models (data and model ensembles)
    flat_seed = 1000
    gsp = GroupShuffleSplit(n_splits=1,
                            random_state=flat_seed,
                            train_size=0.9)
    splits = list(gsp.split(data, groups=data["day"]))
    train_index, test_index = splits[0]
    train_data, test_data = data.iloc[train_index].copy(), data.iloc[test_index].copy()

    # Make N train-valid splits using day as grouping variable
    gsp = GroupShuffleSplit(n_splits=n_splits, random_state=flat_seed, train_size=0.885)
    splits = list(gsp.split(train_data, groups=train_data["day"]))
    train_index, valid_index = splits[data_seed]
    train_data, valid_data = train_data.iloc[train_index].copy(), train_data.iloc[valid_index].copy()

    x_scaler, y_scaler = RobustScaler(), MinMaxScaler((0, 1))
    x_train = x_scaler.fit_transform(train_data[input_cols])
    x_valid = x_scaler.transform(valid_data[input_cols])
    x_test = x_scaler.transform(test_data[input_cols])

    y_train = y_scaler.fit_transform(train_data[output_cols])
    y_valid = y_scaler.transform(valid_data[output_cols])
    y_test = y_scaler.transform(test_data[output_cols])

    model = GaussianRegressorDNN(**conf["model"])
    model.build_neural_network(x_train.shape[-1], y_train.shape[-1])

    model.fit(
        x_train,
        y_train,
        validation_data=(x_valid, y_valid),
        callbacks=get_callbacks(conf))

    model.model_name = f"cv_ensemble/models/model_seed0_split{data_seed}.h5"
    model.save_model()

    # Save the best model
    model.model_name = "cv_ensemble/models/best.h5"
    model.save_model()

    mu, var = model.predict_uncertainty(x_test, y_scaler)
    mae = np.mean(np.abs(mu[:, 0]-test_data[output_cols[0]]))

    ensemble_mu[data_seed] = mu
    ensemble_var[data_seed] = var

model = GaussianRegressorDNN().load_model(conf)

ensemble_mu, ensemble_var = model.predict_ensemble(x_test, scaler=y_scaler)

epistemic = np.var(ensemble_mu, axis=0)
aleatoric = np.mean(ensemble_var, axis=0)

print(epistemic.mean()**(1/2), aleatoric.mean()**(1/2))

compute_results(test_data,
                output_cols,
                np.mean(ensemble_mu, axis=0),
                np.sqrt(aleatoric),
                np.sqrt(epistemic))

monte_carlo_steps = 10

ensemble_mu, ensemble_var = model.predict_monte_carlo(x_test,
                                                      monte_carlo_steps,
                                                      scaler=y_scaler)

ensemble_epistemic = np.var(ensemble_mu, axis=0)
ensemble_aleatoric = np.mean(ensemble_var, axis=0)
ensemble_mean = np.mean(ensemble_mu, axis=0)