# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] toc=true
# <h1>Table of Contents<span class="tocSkip"></span></h1>
# <div class="toc" style="margin-top: 1em;"><ul class="toc-item"><li><span><a href="#Create-the-dataset" data-toc-modified-id="Create-the-dataset-1"><span class="toc-item-num">1&nbsp;&nbsp;</span>Create the dataset</a></span></li><li><span><a href="#Create-neural-network" data-toc-modified-id="Create-neural-network-2"><span class="toc-item-num">2&nbsp;&nbsp;</span>Create neural network</a></span><ul class="toc-item"><li><span><a href="#Now-with-constraints" data-toc-modified-id="Now-with-constraints-2.1"><span class="toc-item-num">2.1&nbsp;&nbsp;</span>Now with constraints</a></span></li></ul></li></ul></div>
# -

import numpy as np
import matplotlib.pyplot as plt
# %matplotlib inline

# ## Create the dataset

z = np.linspace(0, 1, 20)
z_plot = np.linspace(0, 1, 21)


def get_y(x1, x2, x3):
    x1 = np.atleast_1d(x1); x2 = np.atleast_1d(x2); x3 = np.atleast_1d(x3)
    tmp =  (x1.reshape(x1.shape[0], 1) * np.sin(1 * np.pi * z.reshape(1, 20)) + 
            x2.reshape(x1.shape[0], 1) * np.sin(4 * np.pi * z.reshape(1, 20)) + 
            x3.reshape(x1.shape[0], 1) * np.sin(5 * np.pi * z.reshape(1, 20)))
    tmp = (tmp.T - np.mean(tmp, axis=1)).T
    rain = np.maximum(0, np.sum(tmp[:, 10:], axis=1).reshape((-1, 1))) * 10
    tmp[:, 10:] =  tmp[:, 10:] - (rain / 10.)
    return np.concatenate([tmp, rain], axis=1)


# create sample data
n_samples = 10000
x1 = np.random.normal(size=n_samples)
x2 = np.random.normal(size=n_samples)
x3 = np.random.normal(size=n_samples)
y = get_y(x1, x2, x3)

y.shape

i = np.random.randint(0, n_samples)
plt.plot(y[i], z_plot);
np.sum(y[i]), y[i, -1]

plt.hist(y[:, -1])

# Create test and valid data
x = np.stack([x1, x2, x3], axis=1)
# Add noise to xs
x += (np.random.rand(x.shape[0], x.shape[1]) - 0.2) * 0.6

i_split = int(0.8*n_samples)

x_train = x[:i_split]
y_train = y[:i_split]
x_valid = x[i_split:]
y_valid = y[i_split:]

# ## Create neural network 

import keras
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import Adam

import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth = True   # Allocates as much memory as needed.
keras.backend.tensorflow_backend.set_session(tf.Session(config=config))

model = Sequential([
    Dense(200, input_shape=(3,), activation='relu'),
    Dense(200, activation='relu'),
    Dense(21)
])

model.summary()

model.compile(Adam(0.001), 'mse')

model.fit(x_train, y_train, batch_size=512, epochs=10, validation_data=(x_valid, y_valid))

y_pred = model.predict(x_valid)

i = np.random.randint(0, n_samples*0.2)
plt.plot(y_valid[i], z_plot, label='true')
plt.plot(y_pred[i], z_plot, label='pred')

np.mean(np.abs(np.mean(y_pred[:, :], axis=1))), np.mean(np.abs(np.mean(y_valid[:], axis=1)))

plt.hist(y_pred[:, -1], bins=50);

# ### Now with constraints

import keras.backend as K
def transform(x):
    """x[sample, z-1]"""
    x_rain = K.reshape(K.maximum(0., x[:, -1]), (-1, 1))
    x_first = - K.reshape(K.sum(x[:, :-1], axis=1), (-1, 1)) - x_rain
    return K.concatenate([x_first, x[:, :-1], x_rain], axis=1)


from keras.losses import mean_squared_error
def transformed_mse(y_true, y_pred):
    return mean_squared_error(y_true, transform(y_pred))


model_transform = Sequential([
    Dense(200, input_shape=(3,), activation='relu'),
    Dense(200, activation='relu'),
    Dense(21-1)
])

model_transform.compile(Adam(0.001), transformed_mse)

model_transform.fit(x_train, y_train, batch_size=512, epochs=10, 
                    validation_data=(x_valid, y_valid))

# +
y_pred_trans = model_transform.predict(x_valid)

y_rain = np.maximum(0, y_pred_trans[:, -1]).reshape((-1, 1))
y_first = -y_pred_trans[:, :-1].sum(axis=1).reshape((-1, 1)) - y_rain

y_pred_trans = np.concatenate([y_first, y_pred_trans[:, :-1], y_rain], axis=1)
# -

i = np.random.randint(0, 200)
i = 1948
plt.plot(y_valid[i], z_plot, label='true')
plt.plot(y_pred_trans[i], z_plot, label='pred')
plt.legend();
np.sum(y_pred_trans[i]), np.sum(y_valid[i])

max_error = np.argmax(np.mean(np.abs((y_valid - y_pred_trans)), axis=1))
max_error

y_pred_trans.shape

np.mean(np.abs(np.mean(y_pred_trans, axis=1))), np.mean(np.abs(np.mean(y_valid, axis=1)))

plt.hist(y_pred_trans[:, -1], bins=50);

plt.hist(y_valid[:, -1], bins=50);



StarcoderdataJupyterScriptsDedupFiltered