# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:tf]
#     language: python
#     name: conda-env-tf-py
# ---

# +
from numpy.random import seed
from tensorflow.compat.v1 import set_random_seed
set_random_seed(2)
seed(1)

from tcn import TCN
from tensorflow.keras.layers import Dense
from tensorflow.keras import Input, Model
from pathlib import Path

import dask.array as da
import zarr
import pandas as pd
import numpy as np
import tensorflow as tf
from tensorflow import keras
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import robust_scale
from sklearn.metrics import fbeta_score
from tensorflow.keras import backend as K


# -

#Using fbeta because we want more False Positives than False Negatives
def fbeta(y_true, y_pred, threshold_shift=0):
    beta = 2

    # just in case of hipster activation at the final layer
    y_pred = K.clip(y_pred, 0, 1)

    # shifting the prediction threshold from .5 if needed
    y_pred_bin = K.round(y_pred + threshold_shift)

    tp = K.sum(K.round(y_true * y_pred_bin)) + K.epsilon()
    fp = K.sum(K.round(K.clip(y_pred_bin - y_true, 0, 1)))
    fn = K.sum(K.round(K.clip(y_true - y_pred, 0, 1)))

    precision = tp / (tp + fp)
    recall = tp / (tp + fn)

    beta_squared = beta ** 2
    return (beta_squared + 1) * (precision * recall) / (beta_squared * precision + recall)


data_folder = Path("..", "data", "interim")
train_dask = da.from_zarr(str(data_folder.joinpath("breath_data_train.zarr")))
train = train_dask.compute()

# +
#The TCN package expected 3 axes, even if it's 2D, so I added an extra
X = train[:,:-2][:, :, np.newaxis]

y = train[:, -2]

# +
model = keras.models.Sequential(
    [
        TCN(
            input_shape=[7500, 1],
            kernel_size=2,
            activation="relu",
            dilations=[rate for rate in (1, 2, 4, 8) * 2],
            return_sequences=False,
        ),
        Dense(1, activation="sigmoid"),
    ]
)

model.compile(optimizer="adam", loss="mse", metrics=["mse", "mae", "mape"])

history = model.fit(X, y, epochs=10, batch_size=2, validation_split=0.2,)

StarcoderdataJupyterScriptsDedupFiltered