GAN_IDL_Channel_sim.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1a4eead3",
   "metadata": {},
   "source": [
    "# DeepGANnel\n",
    "\n",
    "A GAN made to generate synthetic ion channel data from previously recorded data, using a CNN model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import os, sys, socket\n",
    "import random\n",
    "import glob\n",
    "import re\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from statistics import mean\n",
    "from copy import deepcopy\n",
    "from tqdm import tqdm\n",
    "\n",
    "import time\n",
    "from datetime import datetime\n",
    "\n",
    "from tensorflow.keras import layers\n",
    "from sklearn.preprocessing import minmax_scale, RobustScaler,robust_scale\n",
    "from IPython.display import clear_output\n",
    "\n",
    "# Experimental - adding 1/f noise instead of gaussian noise\n",
    "\n",
    "#from colorednoise import powerlaw_psd_gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90d4887d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Datasource\n",
    "DATAFOLDER='Lina2'\n",
    "# Constants\n",
    "# Have only verified 50, 100, 250\n",
    "BATCH_SIZE = 100\n",
    "# Have only verified 512, 1024, 2048, 4096\n",
    "WINDOW_SIZE = 4096\n",
    "\n",
    "#Add a little noise to facilitate CNN \"throughout\"\n",
    "NOISE_AMP = 0.01\n",
    "\n",
    "#0 is accept records that are all one state; 0.5 is,accept only records with more than 50% of each state\n",
    "#Not sure what happens with 0.5 or more so recommend range of 0 to 0.25\n",
    "#Note if too high preprocess data will loop indefinitely at some stage.\n",
    "THRESHOLD =0.05\n",
    "CROP=100\n",
    "\n",
    "#Try a critic approach rather than strictly discriminator\n",
    "Wasserstein=False\n",
    "ADAMO=True\n",
    "\n",
    "scripts = glob.glob('*.ipynb')\n",
    "myPC=socket.gethostname()\n",
    "with open (\"runs.txt\",\"a\") as runs:\n",
    "    runs.write(f\"This PC is {myPC}\\nDate is {datetime.now()}\\nTF= {tf.__version__}\\n\")\n",
    "    runs.write(f'Data is in folder {DATAFOLDER}\\n')\n",
    "    runs.write(f\"Script name: {scripts}\\n\")\n",
    "    runs.write(f\"Wasserstein was used? {Wasserstein}\\n\")\n",
    "    runs.write(f'Adam used (SGD if false) {ADAMO}')\n",
    "    runs.write(f\"BATCH_SIZE = {BATCH_SIZE}, WINDOW_SIZE = {WINDOW_SIZE}, NOISE_AMP = {NOISE_AMP}, THRESHOLD = {THRESHOLD}\\n\\n\")\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "985a88c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print (tf.__version__)\n",
    "print(tf.config.list_physical_devices())\n",
    "print(tf.test.is_built_with_cuda())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b150206",
   "metadata": {},
   "source": [
    "## Step 1: Preprocessing data\n",
    "\n",
    "First job is to sort the data to obtain batches. Discard any window that has no events in as the GAN will just output noisy \"flat\" data that isn't useful! We'll then save this data to use later on. We know the data is scaled so we don't need to do that here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7084b0a4",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def check_window(channel_record: np.ndarray) -> bool:\n",
    "    \n",
    "    \"\"\" Returns True if there is an event in the window \"\"\"\n",
    "    \n",
    "    \"\"\"return True if (1 in channel_record[:,1:]) and (0 in channel_record[:,1:]) else False\"\"\"\n",
    "    \"\"\"toughen this up, want more than one single event :-)\"\"\"\n",
    "    threshold = THRESHOLD\n",
    "    total=len(channel_record[:,1])\n",
    "    opens=len(channel_record[channel_record[:,1]==1])\n",
    "    val=True if ((opens/total >threshold) and (opens/total <(1-threshold))) else False\n",
    "    \"\"\"plt.plot(channel_record[:,0])\n",
    "    plt.plot(channel_record[:,1])\n",
    "    plt.show()\n",
    "    \n",
    "    print('                                                                   ',end='\\r')   \n",
    "    print(f' length = {total}, open={opens}, frac={opens/total:.2f} so return={val}', end='\\r')\n",
    "    if val == False:\n",
    "        time.sleep(1)\"\"\"\n",
    "    return val\n",
    "    \n",
    "def preprocess_data() -> ():\n",
    "    \n",
    "    \"\"\" \n",
    "    Preprocessing data into one large numpy array - along with a list of\n",
    "    marker start indexes for randomly sampling batches.\n",
    "    \n",
    "    If an index < marker AND index + WINDOW_SIZE > marker, they'll be a change\n",
    "    of file in the window, therefore it's invalid.\n",
    "    \"\"\"\n",
    "    \n",
    "    cwd = os.getcwd()\n",
    "    \"\"\"alternatives\n",
    "    files = os.listdir(f'{cwd}/synth_data')\n",
    "    files = glob.glob('synth_data/*.csv')\"\"\"\n",
    "    files = glob.glob(f'{DATAFOLDER}/*.csv')\n",
    "    print(files)\n",
    "    outer_index = 1\n",
    "    markers = []\n",
    "    output_data = np.array([[0,0]])\n",
    "    for filename in files:\n",
    "        print(f'Processing {filename}')\n",
    "        \"\"\"currently starting from row 2 incase headers etc\"\"\"\n",
    "        inner_index = 0\n",
    "        active = False\n",
    "        \"\"\"for the funky file\"\"\"\n",
    "        funky=False\n",
    "        if funky:\n",
    "            print(\"funky\")\n",
    "            data1 = pd.read_csv(filename).values[2:,5].reshape(-1,1)\n",
    "            \"\"\"channel\"\"\"\n",
    "            data2 = pd.read_csv(filename).values[2:,2] .reshape(-1,1)\n",
    "            data=np.concatenate((data1,data2),axis=1).reshape(-1,2).astype(np.float)\n",
    "        \"\"\"print (np.amin(data[:,1]),np.amax(data[:,1]))\"\"\"\n",
    "        Nat=False\n",
    "        if Nat:\n",
    "            data = pd.read_csv(filename).values[2:,1:]\n",
    "            \"\"\"if accidently flipped channel and current!\"\"\"\n",
    "            data[:,[0,1]]=data[:,[1,0]]\n",
    "            print(data)\n",
    "            simplify=False\n",
    "            if simplify:\n",
    "                for row in data:\n",
    "                    \"\"\"print(row)\"\"\"\n",
    "                    if row[1]==3:\n",
    "                        row[1]=1\n",
    "                    elif row[1]==2:\n",
    "                        row[1]=0\n",
    "                    else:\n",
    "                        row[1]=0\n",
    "            simplify2=False\n",
    "            if simplify2:\n",
    "                for row in data:\n",
    "                    \"\"\"print(row)\"\"\"\n",
    "                    if row[1]==2:\n",
    "                        row[1]=1\n",
    "                    elif row[1]==1:\n",
    "                        row[1]=0\n",
    "                    elif row[1]==0:\n",
    "                        row[1]=0\n",
    "                    else:\n",
    "                        row[1]=1\n",
    "        timecol=False\n",
    "        if timecol==True:\n",
    "            data = pd.read_csv(filename).values[2:,1:]\n",
    "        else:\n",
    "            data = pd.read_csv(filename).values[2:,:]\n",
    "        '''round idealisation to 0 and 1\n",
    "        this is typically, but not always necessary depending on record source'''\n",
    "        data[:,1]=np.round(data[:,1])\n",
    "        data[:,1]=data[:,1]/np.max(data[:,1])\n",
    "        data[:,1]=np.round(data[:,1])\n",
    "        \n",
    "        print('data.shape',data.shape)\n",
    "        print('data type', type(data))\n",
    "        print('data snapshot',data)\n",
    "        savechange=False\n",
    "        if savechange:\n",
    "            ddict = {'1cur': data[:, 0],\n",
    "            '2chan': data[:, 1]}\n",
    "            df = pd.DataFrame (ddict, columns = ['1cur','2chan'])\n",
    "            print (df)\n",
    "            df.to_csv('filename.csv', index =False, header=False)           \n",
    "            \n",
    "        \"\"\"data=data[~(data[:,1] > 1)]    \n",
    "        print (np.amin(data[:,1]),np.amax(data[:,1]))\"\"\"\n",
    "\n",
    "        for i in range(len(data)):\n",
    "            window_check = check_window(data[i:i+WINDOW_SIZE, :])\n",
    "            markers.append(window_check)\n",
    "            if window_check and not active:\n",
    "                # Start recording - make a note of the inner index.\n",
    "                inner_index = i\n",
    "                active = True\n",
    "            elif not window_check and active:\n",
    "                \n",
    "                # Stop \"recording\" and save the file to the output data.\n",
    "                # Also make a marker for safe indexing later.\n",
    "     \n",
    "                window_to_save = data[inner_index: i + WINDOW_SIZE - 1, :]\n",
    "                end_index = outer_index + i + WINDOW_SIZE - 1\n",
    "                \n",
    "                output_data = np.concatenate((output_data, window_to_save))\n",
    "                \"\"\"markers.append(end_index)\"\"\"\n",
    "                               \n",
    "                active = False\n",
    "        outer_index = len(output_data)\n",
    "    print(len(data),len(markers))    \n",
    "    \"\"\"return (output_data, markers)\"\"\"\n",
    "    print(markers[:10])\n",
    "    print(f'Number of good windows is {len([marker for marker in markers if marker==True])}')\n",
    "    return (data, markers)\n",
    "    \n",
    "preprocessed_data, markers = preprocess_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(markers[:10])\n",
    "print(f'Number of good windows is {len([marker for marker in markers if markers!=True])}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Roughly how many valid windows do we have?\n",
    "print(len(preprocessed_data)/len(markers))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(1,2,1)\n",
    "n, bins, patches =plt.hist(preprocessed_data[:,0],bins=100)\n",
    "plt.subplot(1,2,2)\n",
    "n, bins, patches =plt.hist(preprocessed_data[:,1],bins=100)\n",
    "#plt.xlim([-0.6,0.4])\n",
    "#plt.xlim([-0.6,0.4])\n",
    "data = {'1cur': preprocessed_data[:,0],\n",
    "                '2chan': preprocessed_data[:,1]}\n",
    "df = pd.DataFrame (data, columns = ['1cur','2chan'])\n",
    "#print (df)\n",
    "df.to_csv(f'{WINDOW_SIZE}PYRraw.csv', index =False, header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ad5c27f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Minmax scale the current on the whole bloody lot\"\"\"\n",
    "\"\"\"preprocessed_data[:,0]=minmax_scale(preprocessed_data[:, 0])\"\"\"\n",
    "scaler=RobustScaler().fit(np.asarray(preprocessed_data[:,0]).reshape(-1,1))\n",
    "x=scaler.transform(np.asarray(preprocessed_data[:,0]).reshape(-1,1))/2.0\n",
    "\"\"\"rather late spotted that doing it with the below leaves no chance to inverse it!\n",
    "so now creating a scale option that can be used at the end to unscale it back to meet the original\"\"\"\n",
    "preprocessed_data[:,0]=robust_scale(preprocessed_data[:,0])/2.0\n",
    "print (np.amin(preprocessed_data[:,1]),np.amax(preprocessed_data[:,1]))\n",
    "print (np.amin(preprocessed_data[:,0]),np.amax(preprocessed_data[:,0]))\n",
    "print(preprocessed_data.shape)\n",
    "print(type(preprocessed_data))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a547e86",
   "metadata": {},
   "source": [
    "## Step 2: Building Models\n",
    "\n",
    "Here we are using two CNN models for the generator and discriminator - we also need to build the loss functions and optimizers for the two models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fde93e8d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def make_generator_model() -> tf.keras.Model:\n",
    "    \"\"\"\n",
    "        Makes the generator model.\n",
    "    \"\"\"\n",
    "    model = tf.keras.Sequential(name=\"Generator\")\n",
    "\n",
    "    model.add(layers.Dense(2*16*64, use_bias=False, input_shape=(WINDOW_SIZE,), dtype='float32'))\n",
    "    model.add(layers.BatchNormalization())\n",
    "    model.add(layers.LeakyReLU())\n",
    "\n",
    "    model.add(layers.Reshape((2, 16, 64)))\n",
    "    assert model.output_shape == (None, 2, 16, 64) # Note: None is the batch size\n",
    "\n",
    "    size = 16\n",
    "    while size < WINDOW_SIZE/2:\n",
    "\n",
    "        num_filters = 128 if size < 512 else 256\n",
    "\n",
    "        model.add(layers.Conv2DTranspose(num_filters, (2, 5), strides=(1, 2), padding='same', use_bias=False))\n",
    "        model.add(layers.BatchNormalization())\n",
    "        model.add(layers.LeakyReLU())\n",
    "\n",
    "        size *= 2\n",
    "\n",
    "    model.add(layers.Conv2DTranspose(1, (2, 5), strides=(1, 2), padding='same', use_bias=False))\n",
    "    model.add(layers.BatchNormalization())\n",
    "    model.add(layers.LeakyReLU())\n",
    "\n",
    "    assert model.output_shape == (None, 2, WINDOW_SIZE, 1)\n",
    "\n",
    "    return model\n",
    "\n",
    "def make_discriminator_model() -> tf.keras.Model:\n",
    "    \"\"\"\n",
    "        Makes the discriminator model.\n",
    "    \"\"\"\n",
    "    model = tf.keras.Sequential(name=\"Discriminator\")\n",
    "\n",
    "    model.add(layers.Conv2D(32, (2, 100), strides=(1, 2), padding='same', input_shape=[2, WINDOW_SIZE, 1],dtype='float32'))\n",
    "    model.add(layers.LeakyReLU())\n",
    "    model.add(layers.Dropout(0.3))\n",
    "\n",
    "    model.add(layers.Conv2D(64, (2, 50), strides=(1, 2), padding='same'))\n",
    "    model.add(layers.LeakyReLU())\n",
    "    model.add(layers.Dropout(0.3))\n",
    "\n",
    "    model.add(layers.Conv2D(64, (2, 25), strides=(1, 2), padding='same'))\n",
    "    model.add(layers.LeakyReLU())\n",
    "    model.add(layers.Dropout(0.3))\n",
    "\n",
    "    model.add(layers.Conv2D(64, (2, 5), strides=(1, 2), padding='same'))\n",
    "    model.add(layers.LeakyReLU())\n",
    "    model.add(layers.Dropout(0.3))\n",
    "\n",
    "    model.add(layers.Conv2D(128, (2, 5), strides=(1, 2), padding='same'))\n",
    "    model.add(layers.LeakyReLU())\n",
    "    model.add(layers.Dropout(0.3))\n",
    "\n",
    "    model.add(layers.Conv2D(256, (2, 5), strides=(1, 2), padding='same'))\n",
    "    model.add(layers.LeakyReLU(0.2))\n",
    "    model.add(layers.Dropout(0.3))\n",
    "\n",
    "    model.add(layers.Flatten())\n",
    "    model.add(layers.Dense(1))\n",
    "\n",
    "    return model\n",
    "\n",
    "# Define loss functions\n",
    "cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
    "\n",
    "def discriminator_loss(real_output, fake_output) -> float:        \n",
    "    real_loss = cross_entropy(tf.ones_like(real_output), real_output)   \n",
    "    if Wasserstein==True:\n",
    "        fake_loss = cross_entropy(tf.ones_like(fake_output), fake_output)\n",
    "        total_loss =real_loss - fake_loss\n",
    "    else:\n",
    "        fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)\n",
    "        total_loss = real_loss + fake_loss\n",
    "    \n",
    "    return total_loss\n",
    "\n",
    "def generator_loss(fake_output) -> float:\n",
    "    return cross_entropy(tf.ones_like(fake_output), fake_output)\n",
    "\n",
    "\n",
    "if ADAMO==True:\n",
    "    # Optimizers\n",
    "    generator_optimizer = tf.keras.optimizers.Adam(1e-3, amsgrad=True)\n",
    "    discriminator_optimizer = tf.keras.optimizers.Adam(1e-5, amsgrad=True)\n",
    "else:\n",
    "    #Otherwise just used SGD\n",
    "    generator_optimizer = tf.keras.optimizers.SGD(1e-6)\n",
    "    discriminator_optimizer = tf.keras.optimizers.SGD(1e-6)\n",
    "\n",
    "\n",
    "# Now build everything.\n",
    "generator = make_generator_model()\n",
    "discriminator = make_discriminator_model()    \n",
    "    \n",
    "    \n",
    "    \n",
    "#it appears that the change_lrs() is currently impotent.\n",
    "def change_lrs(old_gen_lr: float, \n",
    "               old_disc_lr: float, \n",
    "               gen_loss_history: list, \n",
    "               disc_loss_history: list,\n",
    "               reduction_factor: int = 2, \n",
    "               threshold_factor: int = 5, \n",
    "               history_length: int = 10):\n",
    "    # Changing the learning rates dynamically based on the last 100 epochs \n",
    "    # If either model has a loss 5x the other over hte last 10 training steps,\n",
    "    # multiply the \"bad\" model's learning rate by a factor (default 2) and\n",
    "    # divide the \"good\" model's learning rate by the same amount\n",
    "    \n",
    "    # Also allow switching from \"automatic\" to \"manual\" using a text file \"lr.txt\"\n",
    "    \n",
    "    with open(\"lr.txt\", \"r+\") as my_file:\n",
    "        data = my_file.read()\n",
    "\n",
    "        split = data.split('\\n')\n",
    "        automatic = True if \"True\" in split[0] else False\n",
    "        new_gen_lr = 0\n",
    "        new_disc_lr = 0\n",
    "\n",
    "        if automatic:\n",
    "            # Automatic method ONLY LIKELY TO WORK IF POSITIVE LOSSES ALL AROUND\n",
    "            if len(gen_loss_history) == 0:\n",
    "                # If no history, keep it how it is\n",
    "                return old_gen_lr, old_disc_lr\n",
    "            \n",
    "            elif history_length > len(gen_loss_history):\n",
    "                mean_recent_gen = mean(gen_loss_history[-history_length:])\n",
    "                mean_recent_disc = mean(disc_loss_history[-history_length:])\n",
    "            else:\n",
    "                mean_recent_gen = mean(gen_loss_history)\n",
    "                mean_recent_disc = mean(disc_loss_history)\n",
    "                \n",
    "            if mean_recent_gen >= threshold_factor * mean_recent_disc:\n",
    "                new_gen_lr = old_gen_lr / 5\n",
    "                new_disc_lr = old_disc_lr * 5\n",
    "            elif mean_recent_disc >= threshold_factor * mean_recent_gen:\n",
    "                new_disc_lr = old_disc_lr / 5\n",
    "                new_gen_lr = old_gen_lr * 5\n",
    "                \n",
    "            else:\n",
    "                new_disc_lr = old_disc_lr\n",
    "                new_gen_lr = old_gen_lr\n",
    "                    \n",
    "        else:\n",
    "            # Manual method\n",
    "                \n",
    "            parse_lr_from_file = lambda string: float(string.split(\":\")[1])\n",
    "            new_gen_lr = parse_lr_from_file(split[1])\n",
    "            new_disc_lr = parse_lr_from_file(split[2])\n",
    "            \n",
    "        tf.keras.backend.set_value(generator_optimizer.learning_rate,new_gen_lr)\n",
    "        tf.keras.backend.set_value(discriminator_optimizer.learning_rate,new_disc_lr)\n",
    "        # Save the new data\n",
    "        new_data = f\"AUTOMATIC : {automatic}\\nGENERATOR_LEARNING_RATE : {new_gen_lr}\\nDISCRIMINATOR_LEARNING_RATE : {new_disc_lr}\"\n",
    "        my_file.seek(0)\n",
    "        my_file.truncate(0)\n",
    "        my_file.write(new_data)\n",
    "\n",
    "    return new_gen_lr, new_disc_lr\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Checkpoints and image generation\n",
    "checkpoint_dir = \"training_checkpoints\"\n",
    "checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer, discriminator_optimizer=discriminator_optimizer, generator=generator, discriminator=discriminator)\n",
    "\n",
    "def getREALepoch() -> int:\n",
    "    \"\"\"\n",
    "    Save the current image to the working directory of the program.\n",
    "    \"\"\"\n",
    "    currentImages = glob.glob(\"images/*.png\")\n",
    "   \n",
    "    numList = [0]\n",
    "    for img in currentImages:\n",
    "        i = os.path.splitext(img)[0]\n",
    "        try:\n",
    "            num = re.findall('[0-9]+$', i)[0]\n",
    "            numList.append(int(num))\n",
    "        except IndexError:\n",
    "            pass\n",
    "    numList = sorted(numList)\n",
    "    return numList[-1]+1\n",
    "\n",
    "def unscale_images(data: np.array) -> np.array:\n",
    "    \"\"\"assumes scaler sitting in the global space\"\"\"\n",
    "    for i in range(len(data[:,0,0,0])):\n",
    "        tempx=data[i,0,:,0].reshape(1,-1)*2\n",
    "        #tempx=scaler.inverse_transform(data[i,0,:,0].reshape(-1,1)*2).reshape(1,-1,1)\n",
    "        tempx=scaler.inverse_transform(tempx).reshape(1,-1,1)\n",
    "        tempy=data[i,1,:,0].reshape(1,-1,1)\n",
    "        tempxy=np.stack((tempx,tempy),axis=1)\n",
    "        if i==0:\n",
    "            newdata=tempxy\n",
    "        else:\n",
    "            newdata=np.append(newdata,tempxy,axis=0)   \n",
    "    newdata=np.asarray(newdata)  \n",
    "    return newdata\n",
    "\n",
    "\n",
    "def make_images(model, epoch, gen_loss_list, disc_loss_list,inlineplot=True, plot=True):\n",
    "    me=random.randrange(BATCH_SIZE)\n",
    "    noise = tf.random.normal([BATCH_SIZE, WINDOW_SIZE])\n",
    "    generated_output = np.asarray(model(noise, training = False))\n",
    "    \"\"\"print(generated_output.shape)\n",
    "    print(type(generated_output))\"\"\"\n",
    "    generated_output=np.copy(generated_output)\n",
    "    rawidl=generated_output[me, 1, :, 0].copy()\n",
    "    \"\"\"print(generated_output.flags)\"\"\"\n",
    "    twostate=False\n",
    "    if twostate:\n",
    "        if generated_output[me, 1, 0, 0] >= 0.5:\n",
    "            generated_output[me, 1, 0, 0]=1\n",
    "        else:\n",
    "            generated_output[me, 1, 0, 0]=0  \n",
    "    \n",
    "    thresh=0.6\n",
    "    if twostate:\n",
    "        for i in range(1,WINDOW_SIZE):\n",
    "            if generated_output[me, 1, i-1, 0]==0:\n",
    "                if generated_output[me, 1, i, 0]<thresh:\n",
    "                    generated_output[me, 1, i, 0]=0\n",
    "                else:\n",
    "                    generated_output[me, 1, i, 0]=1\n",
    "            if generated_output[me, 1, i-1, 0]==1:\n",
    "                if generated_output[me, 1, i, 0]>1-thresh:\n",
    "                    generated_output[me, 1, i, 0]=1\n",
    "                else:\n",
    "                    generated_output[me, 1, i, 0]=0            \n",
    "\n",
    "    generated_output[:, 1, :, 0]=np.round(generated_output[:, 1, :, 0])\n",
    "    generated_output[:, 1, :, 0]=np.clip(generated_output[:, 1, :, 0],0,1)\n",
    "    \n",
    "    if inlineplot==True:\n",
    "        %matplotlib inline\n",
    "\n",
    "        ax1_length = generated_output.shape[2]\n",
    "        plt.figure(figsize=[10,4])\n",
    "        plt.plot(range(ax1_length), generated_output[me, 0, :, :], label=\"Current\")\n",
    "        plt.plot(range(ax1_length), generated_output[me, 1, :, :],'r-', label=\"Channels\")\n",
    "        plt.plot(range(len(rawidl)), rawidl, 'g--',label=\"Raw Channels\")\n",
    "        plt.ylim([-1.3,1.3])\n",
    "        plt.legend()\n",
    "        plt.show()\n",
    "    \n",
    "    if plot==True:\n",
    "        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(20, 12))\n",
    "        ax1_length = generated_output.shape[2]\n",
    "        ax1.plot(range(ax1_length), generated_output[me, 0, :, :], label=\"Current\")\n",
    "        ax1.plot(range(ax1_length), generated_output[me, 1, :, :],'r-', label=\"Channels\")\n",
    "        ax1.plot(range(len(rawidl)), rawidl, 'g--',label=\"Raw Channels\")\n",
    "        ax1.set_ylim([-3,4])\n",
    "        ax1.legend()\n",
    "\n",
    "        ax2_length = len(gen_loss_list)\n",
    "        ax2.plot(range(ax2_length), [loss+1 for loss in disc_loss_list], label=\"Discriminator\")\n",
    "        ax2.plot(range(ax2_length), [loss+1 for loss in gen_loss_list], label=\"Generator\")\n",
    "        ax2.set_yscale(\"log\")\n",
    "        #ax2.set_ylim(ymax = 10)\n",
    "        ax2.legend()\n",
    "        preprocessed_data, markers\n",
    "\n",
    "        image_batch = make_batch(BATCH_SIZE, WINDOW_SIZE, preprocessed_data, markers)\n",
    "        me=random.randrange(len(image_batch[:,0,0,0]))\n",
    "        ax3.plot(image_batch[me,0,:,:])\n",
    "        ax3.plot(image_batch[me,1,:,:])\n",
    "        ax3.set_ylim([-3,4])\n",
    "\n",
    "        plt.savefig(f\"images/epoch_{getREALepoch()}.png\")\n",
    "        plt.close()\n",
    "    with open(\"save.txt\", \"r\") as my_file:\n",
    "        data = my_file.read()\n",
    "        split = data.split('\\n')\n",
    "        saveme = True if \"True\" in split[0] else False\n",
    "    #print (saveme)\n",
    "    if saveme:\n",
    "        generated_output = unscale_images(generated_output)\n",
    "        data = {'1cur': generated_output[me, 0, CROP:-CROP, 0],\n",
    "        '2chan': np.round(generated_output[me, 1, CROP:-CROP, 0])}\n",
    "        df = pd.DataFrame (data, columns = ['1cur','2chan'])\n",
    "        #print (df)\n",
    "        df.to_csv(f'{WINDOW_SIZE}gen_data{DATAFOLDER}.csv', mode='a', index =False, header=False)\n",
    "    return generated_output[:,0:2,:,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def train_step(real_data: np.ndarray, GenTRain: bool=True, DiscTrain: bool=True) -> (float, float):    \n",
    "    \"\"\" \n",
    "    One training step - generates noise, generates a output using the \n",
    "    generator model, then tests the discriminator on both the fake output and\n",
    "    real data.\n",
    "    \n",
    "    \"\"\"  \n",
    "    # Generate latent vector for the generator\n",
    "    noise = tf.random.normal([BATCH_SIZE, WINDOW_SIZE])\n",
    "\n",
    "    # Use generator to create fake output, and then test the discriminator\n",
    "    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
    "        generated_output = generator(noise, training = GenTRain)\n",
    "        generated_output = tf.dtypes.cast(generated_output, tf.float64)\n",
    "\n",
    "        real_output = discriminator(real_data, training = DiscTrain)\n",
    "        fake_output = discriminator(generated_output, training = DiscTrain)\n",
    "\n",
    "        gen_loss = generator_loss(fake_output)\n",
    "        disc_loss = discriminator_loss(real_output, fake_output)\n",
    "\n",
    "    # Calculate gradients\n",
    "    generator_gradients = gen_tape.gradient(gen_loss, generator.trainable_variables)\n",
    "    discriminator_gradients = disc_tape.gradient(disc_loss, discriminator.trainable_variables)\n",
    "\n",
    "    # Apply these gradients\n",
    "    generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables))\n",
    "    discriminator_optimizer.apply_gradients(zip(discriminator_gradients, discriminator.trainable_variables))\n",
    "    \n",
    "    return disc_loss, gen_loss\n",
    "\n",
    "\n",
    "\n",
    "def make_window(window_size: int, data: np.ndarray, markers: list) -> np.ndarray:\n",
    "    \n",
    "    \"\"\"\n",
    "        Grabs one window from the dataset, checking to see if it is valid or \n",
    "        not first\n",
    "    \"\"\"\n",
    "    invalid = True\n",
    "    \"\"\"while invalid:\n",
    "        index = random.randrange(len(data) - window_size)\n",
    "        crosses_marker = False\n",
    "        for marker in markers:\n",
    "            if index < marker and index + window_size > marker:\n",
    "                crosses_marker = True\n",
    "        if not crosses_marker:\n",
    "            invalid = False\"\"\"\n",
    "    loops=0\n",
    "    while invalid:\n",
    "        index = random.randrange(len(data) - window_size)\n",
    "        crosses_marker = False\n",
    "        if markers[index]==True:\n",
    "            invalid = False\n",
    "        loops+=1\n",
    "        if loops>100000:\n",
    "            #Eventually just give up.\n",
    "\n",
    "            print('HELP looks like threshold to strict, am about to crash!',end=\"\\r\")\n",
    "            return -1\n",
    "\n",
    "    \n",
    "    \"\"\"current = minmax_scale(data[index:index+window_size, 0])\"\"\"\n",
    "    current  = data[index:index+window_size, 0]+np.random.normal(0,NOISE_AMP,window_size)\n",
    "    channels = data[index:index+window_size, 1]+np.random.normal(0,NOISE_AMP,window_size)\n",
    "    return np.stack((current, channels), axis=1)\n",
    "\n",
    "def make_batch(batch_size: int, window_size: int, data: np.ndarray, markers: list) -> np.ndarray:\n",
    "    \"\"\"\n",
    "        Formats one of the files into the required size for the model.\n",
    "        Since we have already preprocessed the data, we know this data is \n",
    "        scaled, and contains events. All we need to do is shape.\n",
    "    \"\"\"\n",
    "    data = [make_window(window_size, data, markers) for _ in range(batch_size)]\n",
    "    # Force data to be correct shape.\n",
    "    reshaped = np.expand_dims(np.swapaxes(np.stack(data, axis=0), 1, 2), axis=-1)\n",
    "    return reshaped\n",
    "\n",
    "def train(dataset: np.ndarray, epochs: int, markers: list, start_epoch = 0) -> None:\n",
    "    \"\"\"\n",
    "        Function for training the model using the training steps as described\n",
    "        above. Includes plots for monitoring progress.\n",
    "    \"\"\"   \n",
    "    gen_loss_list, disc_loss_list = [], []\n",
    "    print(\"Starting Training\")\n",
    "    for epoch in range(epochs):\n",
    "        # Time it for the console\n",
    "        start = time.time()\n",
    "        GenTrain=False\n",
    "        DiscTrain=True\n",
    "        # The important part!\n",
    "        for _ in tqdm(range(BATCH_SIZE)):\n",
    "            GenTrain=True if GenTrain==False else False\n",
    "            DiscTrain=True if GenTrain==False else False\n",
    "            image_batch = make_batch(BATCH_SIZE, WINDOW_SIZE, dataset, markers)\n",
    "            disc_loss, gen_loss = train_step(image_batch, GenTRain=GenTrain, DiscTrain=DiscTrain)\n",
    "        \n",
    "        # Dynamic control of learning rates\n",
    "        gen_lr = tf.keras.backend.get_value(generator_optimizer.learning_rate)\n",
    "        disc_lr = tf.keras.backend.get_value(discriminator_optimizer.learning_rate)\n",
    "        \n",
    "        new_gen_lr, new_disc_lr = change_lrs(gen_lr, disc_lr, gen_loss_list, disc_loss_list)\n",
    "        # Console outputs        \n",
    "        gen_loss_list.append(gen_loss.numpy())\n",
    "        disc_loss_list.append(disc_loss.numpy())\n",
    "        clear_output(wait=True)\n",
    "        print(f\"Time for epoch {epoch + start_epoch} is {time.time() - start}.\")\n",
    "        print(f\"Discriminator loss: {mean(disc_loss_list)}.\")\n",
    "        print(f\"Generator loss: {mean(gen_loss_list)}.\")\n",
    "        print(f\"DLR : {tf.keras.backend.get_value(discriminator_optimizer.learning_rate)}.\")\n",
    "        print(f\"GLR: {tf.keras.backend.get_value(generator_optimizer.learning_rate)}.\")\n",
    "        make_images(generator, 0, gen_loss_list=[], disc_loss_list=[], inlineplot=True, plot=False)\n",
    "        if epoch%10 == 0:\n",
    "            checkpoint.save(file_prefix = 'checkpoints/checkpoint_')\n",
    "            make_images(generator, epoch + start_epoch, gen_loss_list, disc_loss_list, plot=True)\n",
    "        \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Test plot\n",
    "plt.figure(figsize=[10,4])\n",
    "image_batch = make_batch(BATCH_SIZE, WINDOW_SIZE, preprocessed_data, markers)\n",
    "print(image_batch.shape)\n",
    "plt.plot(image_batch[0,0,:,:])\n",
    "plt.plot(image_batch[0,1,:,:])\n",
    "plt.ylim(-1.3,1.3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca0f2c52",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#Each \"epoch\", of course a batch of 100 records.\n",
    "train(preprocessed_data, 10000, markers, 12664 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(gen_loss_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98efcc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_loss_list=[]\n",
    "disc_loss_list=[]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2074b811",
   "metadata": {},
   "source": [
    "## Various manipulations and outputs.\n",
    "Create a specific output file without needing to train.\n",
    "This cell does some monkey business now.  So assuming you have a live checkpoint and everything is running OK (you may have JUST restored a checkpoint).\n",
    "This will run off and save a strip of data.\n",
    "Then save the current model as a checkpoint in a different location to usual so we can then save it to GitHub. HOWEVER, it may be too large (>100MB) so it then zips the lot up into one archive... then splits this up so Git friendly.\n",
    "There is then a routine at the bottom to remerge these mini archives into one proper zip archive which you can then restore.\n",
    "In fact to do this there are a few steps.\n",
    "Run filesplits merge facility...  create that zip archive.  copy paste and rename those files into your checkpoint location and naming format.\n",
    "Then reload as per usual.  I have checked this works on another PC.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa21ebc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fsplit.filesplit import Filesplit\n",
    "import shutil, os\n",
    "\n",
    "epoch=0\n",
    "x,y=[],[]\n",
    "batches=2\n",
    "for _ in range(batches):\n",
    "    data=make_images(generator, epoch, gen_loss_list=[], disc_loss_list=[],plot=False)\n",
    "    #Oops post processing defo needs cropping, since our CNNs focus on window centres\n",
    "    #Added this now below...      \n",
    "    np.save(\"data.npy\", data)\n",
    "    '''data = unscale_images(data.reshape(BATCH_SIZE,2,WINDOW_SIZE,1)).reshape((BATCH_SIZE,2,WINDOW_SIZE))'''\n",
    "    for record in range(len(data[:,0,0])):\n",
    "        cropped=data[record,0,:]\n",
    "        cropped=cropped[CROP:-CROP]\n",
    "        x.extend([pt for pt in cropped])\n",
    "    for record in range(len(data[:,0,0])):\n",
    "        cropped=data[record,1,:]\n",
    "        cropped=cropped[CROP:-CROP]\n",
    "        y.extend([pt for pt in cropped])\n",
    "        \n",
    "print('length of data strip is',len(x))\n",
    "plt.subplot(2,2,1)\n",
    "n, bins, patches =plt.hist(x,bins=100)\n",
    "    \n",
    "plt.subplot(2,2,2)\n",
    "n, bins, patches =plt.hist(y,bins=100)\n",
    "plt.vlines(-0.01,0,15000,\"red\")\n",
    "plt.subplot(2,2,3)\n",
    "plt.plot(x[:20000])\n",
    "plt.subplot(2,2,4)\n",
    "plt.plot(y[:20000])\n",
    "#plt.xlim([0,0.4])\n",
    "\n",
    "#plt.ylim([0,11000])\n",
    "datadf = {'1cur':x,\n",
    "        '2chan': y}\n",
    "df = pd.DataFrame (datadf, columns = ['1cur','2chan'])\n",
    "df.to_csv(f'{DATAFOLDER}/{WINDOW_SIZE}{DATAFOLDER}gen_data.csv', index =False, header=False)\n",
    "print(np.mean(np.asarray(x)))\n",
    "print(max(y))\n",
    "print(y.count(2))\n",
    "'''ALSO critically; write the current checkpoint to the folder too so reproducible'''\n",
    "shutil.rmtree('checkpointsaved')\n",
    "file_prefix= 'checkpointsaved/checkpointsaved_'\n",
    "temploc='temp/chkpt.zip'\n",
    "checkpoint.save(file_prefix = file_prefix)\n",
    "'''zip it into one file'''\n",
    "import zipfile\n",
    "\n",
    "def zipdir(path, ziph):\n",
    "    # ziph is zipfile handle\n",
    "    for root, dirs, files in os.walk(path):\n",
    "        for file in files:\n",
    "            ziph.write(os.path.join(root, file), \n",
    "                       os.path.relpath(os.path.join(root, file), \n",
    "                                       os.path.join(path, '..')))\n",
    "      \n",
    "zipf = zipfile.ZipFile(temploc, 'w', zipfile.ZIP_DEFLATED)\n",
    "zipdir('checkpointsaved', zipf)\n",
    "zipf.close()\n",
    "\n",
    "\n",
    "'''but have to split it small or Git no cope'''\n",
    "fs = Filesplit()\n",
    "fs.split(file=temploc, split_size=50000000,output_dir=f'{DATAFOLDER}/chk')\n",
    "print(\"sorted\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65faa203",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "for epoch in range(10):  \n",
    "    x=make_images(generator, epoch, gen_loss_list=[], disc_loss_list=[],plot=False)\n",
    "    if epoch%1==0:\n",
    "        # Console outputs\n",
    "        clear_output(wait=True)\n",
    "        print(f\"Time for epoch {epoch} is {(time.time() - start)/BATCH_SIZE}s.\")\n",
    "        start = time.time()\n",
    "        \n",
    "np.save(\"numpy{DATAFOLDER}.NPY\",x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f32e7e2",
   "metadata": {},
   "source": [
    "## If you wish to reload a previous checkpoint and retrain or just simulate\n",
    "edit this cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c300df0",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint.restore('checkpoints/checkpoint_-1048').assert_existing_objects_matched()\n",
    "make_images(generator, 0, gen_loss_list=[], disc_loss_list=[],plot=False)\n",
    "tf.keras.backend.set_value(generator_optimizer.learning_rate,1e-8)\n",
    "tf.keras.backend.set_value(discriminator_optimizer.learning_rate,1e-8)\n",
    "gen_loss_list=[]\n",
    "disc_loss_list=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33095830",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.utils import plot_model\n",
    "plot_model(generator, to_file='generator.png')\n",
    "generator.summary()\n",
    "plot_model(discriminator, to_file='discriminator.png')\n",
    "discriminator.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20bf7c91",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''When restoring use this to restore those checkpoint files'''\n",
    "from fsplit.filesplit import Filesplit\n",
    "fs = Filesplit()\n",
    "fs.merge(input_dir=f'{DATAFOLDER}/chk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.utils import plot_model\n",
    "import tensorflow as tf\n",
    "print(tf.__version__)\n",
    "plot_model(\n",
    "    discriminator, to_file='discriminator.png',dpi=300, show_shapes=True,\n",
    ")\n",
    "plot_model(\n",
    "    generator, to_file='generator.png',dpi=300, show_shapes=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}