<jupyter_start><jupyter_text>This notebook is intended as a simplified and straightforward to use implementation of the main neural network-based model considered in the paper (the `nn_aux_emb` model). There are minor changes in the setup compared to the original implementation for the paper (detailed below), see the folder `nn_postprocessing` for the original code).<jupyter_code>import pandas as pd
import numpy as np
# read data (can be downloaded from https://doi.org/10.6084/m9.figshare.13516301.v1)

data = pd.read_feather(path = "data_RL18.feather")
data.station = pd.to_numeric(data.station, downcast = 'integer')

# drop soil moisture predictions due to missing values
# Note that this is a minor change compared to the paper, but does not have a significant effect
data = data.drop(['sm_mean', 'sm_var'], axis=1)

# split into train and test data
eval_start = 1626724
train_end = 1626723

train_features_raw = data.iloc[:train_end,3:42].to_numpy()
train_targets = data.iloc[:train_end,2].to_numpy()
train_IDs = data.iloc[:train_end,1].to_numpy()

test_features_raw = data.iloc[eval_start:,3:42].to_numpy()
test_targets = data.iloc[eval_start:,2].to_numpy()
test_IDs = data.iloc[eval_start:,1].to_numpy()
# normalize data

def normalize(data, method=None, shift=None, scale=None):
    result = np.zeros(data.shape)
    if method == "MAX":
        scale = np.max(data, axis=0)
        shift = np.zeros(scale.shape)
    for index in range(len(data[0])):
        result[:,index] = (data[:,index] - shift[index]) / scale[index]
    return result, shift, scale

train_features, train_shift, train_scale = normalize(train_features_raw[:,:], method="MAX")

test_features = normalize(test_features_raw[:,:], shift=train_shift, scale=train_scale)[0]
# helper functions for NN models

import tensorflow as tf

from tensorflow.keras.layers import Input, Dense, Embedding, Flatten, Concatenate
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam


def crps_cost_function(y_true, y_pred, theano=False):
    """Compute the CRPS cost function for a normal distribution defined by
    the mean and standard deviation.

    Code inspired by <NAME> (HITS).

    Args:
        y_true: True values
        y_pred: Tensor containing predictions: [mean, std]
        theano: Set to true if using this with pure theano.

    Returns:
        mean_crps: Scalar with mean CRPS over batch
    """

    # Split input
    mu = y_pred[:, 0]
    sigma = y_pred[:, 1]
    # Ugly workaround for different tensor allocation in keras and theano
    if not theano:
        y_true = y_true[:, 0]   # Need to also get rid of axis 1 to match!

    # To stop sigma from becoming negative we first have to 
    # convert it the the variance and then take the square
    # root again. 
    var = sigma ** 2
    # The following three variables are just for convenience
    loc = (y_true - mu) / tf.sqrt(var)
    phi = 1.0 / tf.sqrt(2.0 * np.pi) * tf.exp(-loc ** 2 / 2.0)
    Phi = 0.5 * (1.0 + tf.math.erf(loc / tf.sqrt(2.0)))
    # First we will compute the crps for each input/target pair
    crps =  tf.sqrt(var) * (loc * (2. * Phi - 1.) + 2 * phi - 1. / tf.sqrt(np.pi))
    # Then we take the mean. The cost is now a scalar
    return tf.math.reduce_mean(crps)

def build_emb_model(n_features, n_outputs, hidden_nodes, emb_size, max_id,
                    compile=False, lr=0.01,
                    loss=crps_cost_function,
                    activation='relu', reg=None):
    """
    Args:
        n_features: Number of features
        n_outputs: Number of outputs
        hidden_nodes: int or list of hidden nodes
        emb_size: Embedding size
        max_id: Max embedding ID
        compile: If true, compile model
        lr: learning rate
        loss: loss function
        activation: Activation function for hidden layer

    Returns:
        model: Keras model
    """
    if type(hidden_nodes) is not list:
        hidden_nodes = [hidden_nodes]

    features_in = Input(shape=(n_features,))
    id_in = Input(shape=(1,))
    emb = Embedding(max_id + 1, emb_size)(id_in)
    emb = Flatten()(emb)
    x = Concatenate()([features_in, emb])
    for h in hidden_nodes:
        x = Dense(h, activation=activation, kernel_regularizer=reg)(x)
    x = Dense(n_outputs, activation='linear', kernel_regularizer=reg)(x)
    model = Model(inputs=[features_in, id_in], outputs=x)

    if compile:
        opt = Adam(learning_rate=lr)
        model.compile(optimizer=opt, loss=loss)
    return model
from tensorflow.keras.backend import clear_session

# training multiple models in a loop

emb_size = 2
max_id = int(tf.math.reduce_max([train_IDs.max(), test_IDs.max()]))
n_features = train_features.shape[1]
n_outputs = 2

nreps = 10
trn_scores = []
test_scores = []
preds = []

from tqdm.notebook import tqdm

for i in tqdm(range(nreps)):
    clear_session()
    
    features_in = Input(shape=(n_features,))
    id_in = Input(shape=(1,))
    emb = Embedding(max_id + 1, emb_size)(id_in)
    emb = Flatten()(emb)
    x = Concatenate()([features_in, emb])
    x = Dense(512, activation='relu')(x)
    x = Dense(n_outputs, activation='linear')(x)
    nn_aux_emb = Model(inputs=[features_in, id_in], outputs=x)

    opt = Adam(learning_rate=0.002)
    nn_aux_emb.compile(optimizer=opt, loss=crps_cost_function)
    
    nn_aux_emb.fit([train_features, train_IDs], train_targets, epochs=15, batch_size=4096, verbose=0)   
    
    trn_scores.append(nn_aux_emb.evaluate([train_features, train_IDs], train_targets, 4096, verbose=0))
    test_scores.append(nn_aux_emb.evaluate([test_features, test_IDs], test_targets, 4096, verbose=0))
    preds.append(nn_aux_emb.predict([test_features, test_IDs], 4096, verbose=0))
test_scores
# evaluate ensemble of models

from scipy.stats import norm

def crps_normal(mu, sigma, y):
    """
    Compute CRPS for a Gaussian distribution. 
    """
    # Make sure sigma is positive
    sigma = np.abs(sigma)
    loc = (y - mu) / sigma
    crps = sigma * (loc * (2 * norm.cdf(loc) - 1) + 
                    2 * norm.pdf(loc) - 1. / np.sqrt(np.pi))
    return crps

preds = np.array(preds)
preds[:, :, 1] = np.abs(preds[:, :, 1]) # Make sure std is positive
mean_preds = np.mean(preds, 0)
ens_score = crps_normal(mean_preds[:, 0], mean_preds[:, 1], test_targets).mean()
print(f'Ensemble test score = {ens_score}')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1156015037593985, 'fraction_numerical': 0.018483709273182956, 'mean_word_length': 3.5251594613749115, 'pattern_counts': {'":': 1, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '366935', 'license': 'MIT', 'n_tokens_mistral': 2239, 'n_tokens_neox': 2126, 'n_words': 600, 'path': 'RL18_replication.ipynb', 'repo_name': 'XuetongWang-DRR/ppnn', 'size': '6391'}