{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "43a1495a-017d-412f-bfb8-76ad9a9e5875",
   "metadata": {},
   "source": [
    "# Lorenz '63 DIRESA tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f159961e-b306-4428-be35-c0c124c64b9e",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/gdepaepe/diresa/blob/main/diresa.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a1aba31-5e47-4c4c-b2bb-7610e9d6c7b4",
   "metadata": {},
   "source": [
    "### 1. Install packages\n",
    "The `diresa` package depends on the `tensorflow` and `tensorflow_probability` packages. This tutorial also uses `numpy` and `matplotlib`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9fb5ebcc-7dae-4050-a4db-fcd340d59121",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Install needed packages\n",
    "!pip install numpy\n",
    "!pip install matplotlib\n",
    "!pip install tensorflow\n",
    "!pip install tensorflow_probability\n",
    "!pip install diresa"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dc542ee-8836-4074-9a91-9fb1ab6ab3ec",
   "metadata": {},
   "source": [
    "### 2. Load the dataset\n",
    "In this tutorial, we are going to compress the 3D lorenz '63 butterfly into a 2D latent space. The `lorenz.csv` contains a list of butterfly points, with three colums for the X, Y and Z coordinate. The DIRESA model has 2 inputs: the original dataset and a shuffled version of this dataset for the twin encoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a7ccbcdb-1231-4321-aaea-4ecee58962d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://gitlab.com/etrovub/ai4wcm/public/diresa/-/raw/master/docs/lorenz.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b3db4af4-7525-4dc1-bbbd-911f7a01cad0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape lorenz.csv : (40000, 3)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "data_file = \"lorenz.csv\"\n",
    "data = np.loadtxt(data_file, delimiter=\",\")\n",
    "print(\"Shape\", data_file, \":\", data.shape)\n",
    "train = data[:30000]\n",
    "val = data[30000:]\n",
    "id_train = np.argsort((np.random.random(train.shape[0])))\n",
    "id_val = np.argsort((np.random.random(val.shape[0])))\n",
    "train_twin = train[id_train]\n",
    "val_twin = val[id_val]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94e31f8b-fc13-450a-be41-25bc5542c3c6",
   "metadata": {},
   "source": [
    "### 3. Build the DIRESA model\n",
    "We can build a DIRESA model with convolutional and/or dense layers with the `build_diresa` function. We can also build a DIRESA model based on a custom encoder and decoder with the `diresa_model` function (see below). We build here a model with an input shape of `(3,)` for the 3D butterfly points. Our encoder model has 3 dense layers with 40, 20 and 2 units (the latter is the dimension of the latent space). The decoder is a reflection of the encoder. The DIRESA model has 3 loss functions, the reconstruction loss (usually the MSE is used here), the covariance loss and a distance loss (here the MSE distance loss is used). Also the weights for the diffenent loss functions are specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6508df1e-4f69-44c8-b195-dd770e58b4ca",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from diresa.models import build_diresa\n",
    "from diresa.loss import mse_dist_loss, LatentCovLoss\n",
    "\n",
    "diresa = build_diresa(input_shape=(3,), dense_units=(40, 20, 2))\n",
    "\n",
    "diresa.compile(loss=['MSE', LatentCovLoss(1.), mse_dist_loss], loss_weights=[1., 3., 1.5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2953e57c-1e64-4bd6-be9d-4b8801199834",
   "metadata": {},
   "source": [
    "In order to lower the loss weight tuning effort, we will use annealing for the covariance loss. In this case, the covariance weight starts from an initial value (here the keras backend variable `cov_weight` is initialized to 0.) and is increased until the covariance loss reaches a certain target."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "14fec39d-b033-4ffc-862f-aecb487c67d3",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                Output Shape                 Param #   Connected to                  \n",
      "==================================================================================================\n",
      " Input (InputLayer)          [(None, 3)]                  0         []                            \n",
      "                                                                                                  \n",
      " Shuffled_Input (InputLayer  [(None, 3)]                  0         []                            \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " Encoder (Functional)        (None, 2)                    1022      ['Input[0][0]',               \n",
      "                                                                     'Shuffled_Input[0][0]']      \n",
      "|¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|\n",
      "| Encoder_Input (InputLayer  [(None, 3)]                  0         []                           |\n",
      "| )                                                                                              |\n",
      "|                                                                                                |\n",
      "| Enc_1_Dense_0 (Dense)      (None, 40)                   160       []                           |\n",
      "|                                                                                                |\n",
      "| Enc_1_Dense_1 (Dense)      (None, 20)                   820       []                           |\n",
      "|                                                                                                |\n",
      "| Enc_1_Dense_2 (Dense)      (None, 2)                    42        []                           |\n",
      "¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯\n",
      " Decoder (Functional)        (None, 3)                    1023      ['Encoder[0][0]']             \n",
      "|¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|\n",
      "| Decoder_Input (InputLayer  [(None, 2)]                  0         []                           |\n",
      "| )                                                                                              |\n",
      "|                                                                                                |\n",
      "| Dec_1_Dense_2 (Dense)      (None, 20)                   60        []                           |\n",
      "|                                                                                                |\n",
      "| Dec_1_Dense_1 (Dense)      (None, 40)                   840       []                           |\n",
      "|                                                                                                |\n",
      "| Dec_1_Dense_0 (Dense)      (None, 3)                    123       []                           |\n",
      "¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯\n",
      " Latent (Flatten)            (None, 2)                    0         ['Encoder[0][0]']             \n",
      "                                                                                                  \n",
      " dist_layer (DistLayer)      (None, 2)                    0         ['Input[0][0]',               \n",
      "                                                                     'Shuffled_Input[0][0]',      \n",
      "                                                                     'Encoder[0][0]',             \n",
      "                                                                     'Encoder[1][0]']             \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 2045 (7.99 KB)\n",
      "Trainable params: 2045 (7.99 KB)\n",
      "Non-trainable params: 0 (0.00 Byte)\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "import keras.backend as K\n",
    "from diresa.callback import LossWeightAnnealing\n",
    "\n",
    "cov_weight = K.variable(0.)\n",
    "diresa.compile(loss=['MSE', LatentCovLoss(cov_weight), mse_dist_loss], loss_weights=[1., 1., 1.5])\n",
    "diresa.summary(expand_nested=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09311d40-9ab9-42a1-8cf5-c9612a8d7f47",
   "metadata": {},
   "source": [
    "### 4. Train the DIRESA model\n",
    "We train the DIRESA model in a standard way. The output of the decoder should fit the input of the encoder. The batch size should be large enough for the calculation of the covariance loss, which calculates the covariance matrix of the latent space components over the batch. In the `LossWeightAnnealing` callback, we specify the target (`target_loss`) for the mean squared covariance between the latent components. Also the step size by which the annealing weight factor is increased (`anneal_step`) and epoch from which annealing is started (`start_epoch`) is specified. If annealing is not used, the fit method is called without callback function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6843ab10-79ae-49a9-adf5-8edb16257fd3",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20\n",
      "Cov/KL Loss weight: 0, Validation Reconstruction loss: 0.0174634, Validation Cov/KL loss: 1.\n",
      "59/59 - 2s - loss: 0.1540 - Decoder_loss: 0.0916 - Latent_loss: 0.0000e+00 - dist_layer_loss: 0.0416 - val_loss: 0.0547 - val_Decoder_loss: 0.0175 - val_Latent_loss: 0.0000e+00 - val_dist_layer_loss: 0.0248 - 2s/epoch - 42ms/step\n",
      "Epoch 2/20\n",
      "Cov/KL Loss weight: 0, Validation Reconstruction loss: 0.00296876, Validation Cov/KL loss: 1.\n",
      "59/59 - 0s - loss: 0.0251 - Decoder_loss: 0.0080 - Latent_loss: 0.0000e+00 - dist_layer_loss: 0.0114 - val_loss: 0.0062 - val_Decoder_loss: 0.0030 - val_Latent_loss: 0.0000e+00 - val_dist_layer_loss: 0.0022 - 201ms/epoch - 3ms/step\n",
      "Epoch 3/20\n",
      "Cov/KL Loss weight: 0.2, Validation Reconstruction loss: 0.00156346, Validation Cov/KL loss: 0.000142588.\n",
      "59/59 - 0s - loss: 0.0026 - Decoder_loss: 0.0018 - Latent_loss: 3.7637e-06 - dist_layer_loss: 5.3061e-04 - val_loss: 0.0021 - val_Decoder_loss: 0.0016 - val_Latent_loss: 2.8518e-05 - val_dist_layer_loss: 3.1686e-04 - 214ms/epoch - 4ms/step\n",
      "Epoch 4/20\n",
      "Cov/KL Loss weight: 0.4, Validation Reconstruction loss: 0.00155519, Validation Cov/KL loss: 0.000135184.\n",
      "59/59 - 0s - loss: 0.0019 - Decoder_loss: 0.0014 - Latent_loss: 6.8829e-06 - dist_layer_loss: 2.7929e-04 - val_loss: 0.0021 - val_Decoder_loss: 0.0016 - val_Latent_loss: 5.4074e-05 - val_dist_layer_loss: 3.0934e-04 - 205ms/epoch - 3ms/step\n",
      "Epoch 5/20\n",
      "Cov/KL Loss weight: 0.6, Validation Reconstruction loss: 0.00152977, Validation Cov/KL loss: 0.000128108.\n",
      "59/59 - 0s - loss: 0.0016 - Decoder_loss: 0.0013 - Latent_loss: 7.5046e-06 - dist_layer_loss: 2.5048e-04 - val_loss: 0.0021 - val_Decoder_loss: 0.0015 - val_Latent_loss: 7.6865e-05 - val_dist_layer_loss: 3.3634e-04 - 213ms/epoch - 4ms/step\n",
      "Epoch 6/20\n",
      "Cov/KL Loss weight: 0.8, Validation Reconstruction loss: 0.00153266, Validation Cov/KL loss: 0.000132983.\n",
      "59/59 - 0s - loss: 0.0016 - Decoder_loss: 0.0012 - Latent_loss: 8.5114e-06 - dist_layer_loss: 2.4649e-04 - val_loss: 0.0021 - val_Decoder_loss: 0.0015 - val_Latent_loss: 1.0639e-04 - val_dist_layer_loss: 3.0153e-04 - 291ms/epoch - 5ms/step\n",
      "Epoch 7/20\n",
      "Cov/KL Loss weight: 1, Validation Reconstruction loss: 0.00107131, Validation Cov/KL loss: 0.000132622.\n",
      "59/59 - 0s - loss: 0.0014 - Decoder_loss: 0.0011 - Latent_loss: 1.1285e-05 - dist_layer_loss: 2.3219e-04 - val_loss: 0.0016 - val_Decoder_loss: 0.0011 - val_Latent_loss: 1.3262e-04 - val_dist_layer_loss: 2.6278e-04 - 289ms/epoch - 5ms/step\n",
      "Epoch 8/20\n",
      "Cov/KL Loss weight: 1.2, Validation Reconstruction loss: 0.000982827, Validation Cov/KL loss: 0.000136045.\n",
      "59/59 - 0s - loss: 0.0014 - Decoder_loss: 0.0010 - Latent_loss: 1.1870e-05 - dist_layer_loss: 2.3644e-04 - val_loss: 0.0016 - val_Decoder_loss: 9.8283e-04 - val_Latent_loss: 1.6325e-04 - val_dist_layer_loss: 2.7859e-04 - 270ms/epoch - 5ms/step\n",
      "Epoch 9/20\n",
      "Cov/KL Loss weight: 1.4, Validation Reconstruction loss: 0.00107117, Validation Cov/KL loss: 0.000126257.\n",
      "59/59 - 0s - loss: 0.0013 - Decoder_loss: 9.4779e-04 - Latent_loss: 1.4912e-05 - dist_layer_loss: 2.2565e-04 - val_loss: 0.0017 - val_Decoder_loss: 0.0011 - val_Latent_loss: 1.7676e-04 - val_dist_layer_loss: 2.9609e-04 - 265ms/epoch - 4ms/step\n",
      "Epoch 10/20\n",
      "Cov/KL Loss weight: 1.6, Validation Reconstruction loss: 0.000828701, Validation Cov/KL loss: 0.000127913.\n",
      "59/59 - 0s - loss: 0.0013 - Decoder_loss: 8.9542e-04 - Latent_loss: 1.9677e-05 - dist_layer_loss: 2.2696e-04 - val_loss: 0.0014 - val_Decoder_loss: 8.2870e-04 - val_Latent_loss: 2.0466e-04 - val_dist_layer_loss: 2.6512e-04 - 215ms/epoch - 4ms/step\n",
      "Epoch 11/20\n",
      "Cov/KL Loss weight: 1.8, Validation Reconstruction loss: 0.00115553, Validation Cov/KL loss: 0.00012305.\n",
      "59/59 - 0s - loss: 0.0012 - Decoder_loss: 8.5127e-04 - Latent_loss: 1.8856e-05 - dist_layer_loss: 2.2253e-04 - val_loss: 0.0019 - val_Decoder_loss: 0.0012 - val_Latent_loss: 2.2149e-04 - val_dist_layer_loss: 3.7617e-04 - 252ms/epoch - 4ms/step\n",
      "Epoch 12/20\n",
      "Cov/KL Loss weight: 2, Validation Reconstruction loss: 0.000802105, Validation Cov/KL loss: 0.000132949.\n",
      "59/59 - 0s - loss: 0.0012 - Decoder_loss: 8.2960e-04 - Latent_loss: 2.2853e-05 - dist_layer_loss: 2.1870e-04 - val_loss: 0.0015 - val_Decoder_loss: 8.0211e-04 - val_Latent_loss: 2.6590e-04 - val_dist_layer_loss: 2.7728e-04 - 232ms/epoch - 4ms/step\n",
      "Epoch 13/20\n",
      "Cov/KL Loss weight: 2.2, Validation Reconstruction loss: 0.00102775, Validation Cov/KL loss: 0.000123016.\n",
      "59/59 - 0s - loss: 0.0011 - Decoder_loss: 7.4609e-04 - Latent_loss: 1.5520e-05 - dist_layer_loss: 2.1476e-04 - val_loss: 0.0018 - val_Decoder_loss: 0.0010 - val_Latent_loss: 2.7063e-04 - val_dist_layer_loss: 3.2407e-04 - 229ms/epoch - 4ms/step\n",
      "Epoch 14/20\n",
      "Cov/KL Loss weight: 2.4, Validation Reconstruction loss: 0.00148417, Validation Cov/KL loss: 0.000121763.\n",
      "59/59 - 0s - loss: 0.0011 - Decoder_loss: 7.3687e-04 - Latent_loss: 2.2741e-05 - dist_layer_loss: 2.1383e-04 - val_loss: 0.0024 - val_Decoder_loss: 0.0015 - val_Latent_loss: 2.9223e-04 - val_dist_layer_loss: 4.4190e-04 - 280ms/epoch - 5ms/step\n",
      "Epoch 15/20\n",
      "Cov/KL Loss weight: 2.6, Validation Reconstruction loss: 0.00118743, Validation Cov/KL loss: 0.000124618.\n",
      "59/59 - 0s - loss: 0.0011 - Decoder_loss: 7.3420e-04 - Latent_loss: 2.6980e-05 - dist_layer_loss: 2.1725e-04 - val_loss: 0.0020 - val_Decoder_loss: 0.0012 - val_Latent_loss: 3.2401e-04 - val_dist_layer_loss: 2.9670e-04 - 266ms/epoch - 5ms/step\n",
      "Epoch 16/20\n",
      "Cov/KL Loss weight: 2.8, Validation Reconstruction loss: 0.000591062, Validation Cov/KL loss: 0.000130965.\n",
      "59/59 - 0s - loss: 0.0010 - Decoder_loss: 6.9970e-04 - Latent_loss: 3.1541e-05 - dist_layer_loss: 2.0955e-04 - val_loss: 0.0014 - val_Decoder_loss: 5.9106e-04 - val_Latent_loss: 3.6670e-04 - val_dist_layer_loss: 2.6848e-04 - 327ms/epoch - 6ms/step\n",
      "Epoch 17/20\n",
      "Cov/KL Loss weight: 3, Validation Reconstruction loss: 0.000614805, Validation Cov/KL loss: 0.000131801.\n",
      "59/59 - 0s - loss: 0.0010 - Decoder_loss: 6.6951e-04 - Latent_loss: 4.0362e-05 - dist_layer_loss: 2.1089e-04 - val_loss: 0.0014 - val_Decoder_loss: 6.1480e-04 - val_Latent_loss: 3.9540e-04 - val_dist_layer_loss: 2.7220e-04 - 233ms/epoch - 4ms/step\n",
      "Epoch 18/20\n",
      "Cov/KL Loss weight: 3.2, Validation Reconstruction loss: 0.000582299, Validation Cov/KL loss: 0.000131345.\n",
      "59/59 - 0s - loss: 0.0010 - Decoder_loss: 6.5059e-04 - Latent_loss: 4.1834e-05 - dist_layer_loss: 2.0835e-04 - val_loss: 0.0014 - val_Decoder_loss: 5.8230e-04 - val_Latent_loss: 4.2030e-04 - val_dist_layer_loss: 2.6075e-04 - 301ms/epoch - 5ms/step\n",
      "Epoch 19/20\n",
      "Cov/KL Loss weight: 3.4, Validation Reconstruction loss: 0.000864667, Validation Cov/KL loss: 0.000132848.\n",
      "59/59 - 0s - loss: 9.8771e-04 - Decoder_loss: 6.4097e-04 - Latent_loss: 2.9934e-05 - dist_layer_loss: 2.1121e-04 - val_loss: 0.0018 - val_Decoder_loss: 8.6467e-04 - val_Latent_loss: 4.5168e-04 - val_dist_layer_loss: 3.5043e-04 - 284ms/epoch - 5ms/step\n",
      "Epoch 20/20\n",
      "Cov/KL Loss weight: 3.6, Validation Reconstruction loss: 0.000546364, Validation Cov/KL loss: 0.000126098.\n",
      "59/59 - 0s - loss: 9.8688e-04 - Decoder_loss: 6.3903e-04 - Latent_loss: 3.4604e-05 - dist_layer_loss: 2.0884e-04 - val_loss: 0.0014 - val_Decoder_loss: 5.4636e-04 - val_Latent_loss: 4.5395e-04 - val_dist_layer_loss: 2.4780e-04 - 286ms/epoch - 5ms/step\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.src.callbacks.History at 0x1879ce35790>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "callback = [LossWeightAnnealing(cov_weight, loss_name=\"val_Latent_loss\", target_loss=0.0001, anneal_step=0.2, start_epoch=3)]\n",
    "diresa.fit((train, train_twin), train, \n",
    "           validation_data=((val, val_twin), val),\n",
    "           epochs=20, batch_size=512, shuffle=True, verbose=2, callbacks=callback)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a9d89b-e05f-405f-8c1f-62fc393bc312",
   "metadata": {},
   "source": [
    "### 5. Encoder and decoder submodel\n",
    "We cut out the encoder and decoder submodels with the `cut_sub_model` function. So we can make predictions for latent and decoded space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2ada4d84-c7ef-4a72-b098-d919c2aa01c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "313/313 [==============================] - 0s 962us/step\n",
      "313/313 [==============================] - 0s 1ms/step\n"
     ]
    }
   ],
   "source": [
    "from diresa.toolbox import cut_sub_model\n",
    "compress_model = cut_sub_model(diresa, 'Encoder')\n",
    "decode_model = cut_sub_model(diresa, 'Decoder')\n",
    "latent = compress_model.predict(val)\n",
    "predict = decode_model.predict(latent)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fe75237-029c-4ac6-acdd-7c8dddd4f933",
   "metadata": {},
   "source": [
    "### 6. Show latent space\n",
    "We plot the 2D latent space. There is no correlation between the 2 latent components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ffb9f2d8-463a-4aa8-a2e8-4ecf8f51c297",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 103804 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "plt.title(\"Latent space\")\n",
    "plt.scatter(latent[:, 0], latent[:, 1], marker='.', s=0.1, color='C2')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e5ed8e-3c9d-43c2-9bbd-1c6b022f984b",
   "metadata": {},
   "source": [
    "### 7. Original versus decoded datset\n",
    "We compair the origonal dataset with the decoded one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "59695d86-8e7a-4b9b-889f-e24adb232199",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 148660 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(projection='3d')\n",
    "ax.scatter(train[:, 0], train[:, 1], train[:, 2], marker='.', s=0.1)\n",
    "ax.scatter(predict[:, 0], predict[:, 1], predict[:, 2], marker='.', s=0.1, color='C1')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e0e6d67-da2a-40ea-9716-38eee5559c2d",
   "metadata": {},
   "source": [
    "### 8. A covolutional example\n",
    "If your dataset consists of a number of variables (e.g. temperature and pressure, so 2 variables) over a 2 dimensional grid, convolutional layers can be used in the encoder/decoder. Here is an example for a grid (y, x) = (32, 64). The dataset would then have a size of (nbr_of_samples, 32, 64, 2). We will use a stack of 4 convolutional/maxpooling blocks in the encoder (the decoder mirrors the encoder). The first block uses 3 Conv2D layers, the second bock 2 and the third block 1, followed by a MaxPooling2D layer (`stack=(3, 2, 1)`). The number of filters in the first block is 32, in the second 16 and in the third 8 (`stack_filters=(32, 16, 8)`). The number of filters in Latent space, before flattening, is 1 (`latent_filters=1`). This will result in a latent size (before flattening) of (8, 16, 1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9acf7f76-bc9d-4219-97c8-869fe00ae655",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Model: \"model_1\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                Output Shape                 Param #   Connected to                  \n",
      "==================================================================================================\n",
      " Input (InputLayer)          [(None, 32, 64, 2)]          0         []                            \n",
      "                                                                                                  \n",
      " Shuffled_Input (InputLayer  [(None, 32, 64, 2)]          0         []                            \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " Encoder (Functional)        (None, 4, 8, 1)              27281     ['Input[0][0]',               \n",
      "                                                                     'Shuffled_Input[0][0]']      \n",
      "|¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|\n",
      "| Encoder_Input (InputLayer  [(None, 32, 64, 2)]          0         []                           |\n",
      "| )                                                                                              |\n",
      "|                                                                                                |\n",
      "| Enc_1_Conv2D_0 (Conv2D)    (None, 32, 64, 32)           608       []                           |\n",
      "|                                                                                                |\n",
      "| Enc_1_Conv2D_1 (Conv2D)    (None, 32, 64, 32)           9248      []                           |\n",
      "|                                                                                                |\n",
      "| Enc_1_Conv2D_2 (Conv2D)    (None, 32, 64, 32)           9248      []                           |\n",
      "|                                                                                                |\n",
      "| Enc_1_MaxPooling2D (MaxPo  (None, 16, 32, 32)           0         []                           |\n",
      "| oling2D)                                                                                       |\n",
      "|                                                                                                |\n",
      "| Enc_2_Conv2D_0 (Conv2D)    (None, 16, 32, 16)           4624      []                           |\n",
      "|                                                                                                |\n",
      "| Enc_2_Conv2D_1 (Conv2D)    (None, 16, 32, 16)           2320      []                           |\n",
      "|                                                                                                |\n",
      "| Enc_2_MaxPooling2D (MaxPo  (None, 8, 16, 16)            0         []                           |\n",
      "| oling2D)                                                                                       |\n",
      "|                                                                                                |\n",
      "| Enc_3_Conv2D_0 (Conv2D)    (None, 8, 16, 8)             1160      []                           |\n",
      "|                                                                                                |\n",
      "| Enc_3_MaxPooling2D (MaxPo  (None, 4, 8, 8)              0         []                           |\n",
      "| oling2D)                                                                                       |\n",
      "|                                                                                                |\n",
      "| Enc_4_Conv2D_0 (Conv2D)    (None, 4, 8, 1)              73        []                           |\n",
      "¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯\n",
      " Decoder (Functional)        (None, 32, 64, 2)            27282     ['Encoder[0][0]']             \n",
      "|¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|\n",
      "| Decoder_Input (InputLayer  [(None, 4, 8, 1)]            0         []                           |\n",
      "| )                                                                                              |\n",
      "|                                                                                                |\n",
      "| Dec_4_Conv2D_0 (Conv2D)    (None, 4, 8, 8)              80        []                           |\n",
      "|                                                                                                |\n",
      "| Dec_4_UpSampling2D (UpSam  (None, 8, 16, 8)             0         []                           |\n",
      "| pling2D)                                                                                       |\n",
      "|                                                                                                |\n",
      "| Dec_3_Conv2D_0 (Conv2D)    (None, 8, 16, 16)            1168      []                           |\n",
      "|                                                                                                |\n",
      "| Dec_3_UpSampling2D (UpSam  (None, 16, 32, 16)           0         []                           |\n",
      "| pling2D)                                                                                       |\n",
      "|                                                                                                |\n",
      "| Dec_3_Conv2D_1 (Conv2D)    (None, 16, 32, 16)           2320      []                           |\n",
      "|                                                                                                |\n",
      "| Dec_2_Conv2D_0 (Conv2D)    (None, 16, 32, 32)           4640      []                           |\n",
      "|                                                                                                |\n",
      "| Dec_2_UpSampling2D (UpSam  (None, 32, 64, 32)           0         []                           |\n",
      "| pling2D)                                                                                       |\n",
      "|                                                                                                |\n",
      "| Dec_2_Conv2D_1 (Conv2D)    (None, 32, 64, 32)           9248      []                           |\n",
      "|                                                                                                |\n",
      "| Dec_2_Conv2D_2 (Conv2D)    (None, 32, 64, 32)           9248      []                           |\n",
      "|                                                                                                |\n",
      "| Dec_1_Conv2D_0 (Conv2D)    (None, 32, 64, 2)            578       []                           |\n",
      "¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯\n",
      " Latent (Flatten)            (None, 32)                   0         ['Encoder[0][0]']             \n",
      "                                                                                                  \n",
      " dist_layer_1 (DistLayer)    (None, 2)                    0         ['Input[0][0]',               \n",
      "                                                                     'Shuffled_Input[0][0]',      \n",
      "                                                                     'Encoder[0][0]',             \n",
      "                                                                     'Encoder[1][0]']             \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 54563 (213.14 KB)\n",
      "Trainable params: 54563 (213.14 KB)\n",
      "Non-trainable params: 0 (0.00 Byte)\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "diresa = build_diresa(input_shape=(32, 64, 2), stack=(3, 2, 1), stack_filters=(32, 16, 8), latent_filters=1)\n",
    "diresa.summary(expand_nested=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3d06b52-991b-478f-a599-ec0a947bd6c8",
   "metadata": {},
   "source": [
    "### 9. Build DIRESA with custom encoder and decoder\n",
    "We can also build DIRESA models with custom encoder and decoder models. We define those two here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3fb3e14d-75d7-4aab-9ef3-794c9387e7a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras import layers, Input\n",
    "from keras.models import Model\n",
    "def encoder_model(input_shape=(3,), output_shape=2, units=40):\n",
    "    x = Input(shape=input_shape)\n",
    "    y = layers.Dense(units=units, activation=\"relu\")(x)\n",
    "    y = layers.Dense(units=units // 2, activation=\"relu\")(y)\n",
    "    y = layers.Dense(output_shape, activation=\"linear\")(y)\n",
    "    model = Model(x, y, name=\"Encoder\")\n",
    "    return model\n",
    "def decoder_model(input_shape=2, output_shape=3, units=40):\n",
    "    x = Input(shape=input_shape)\n",
    "    y = layers.Dense(units=units // 2, activation=\"relu\")(x)\n",
    "    y = layers.Dense(units=units, activation=\"relu\")(y)\n",
    "    y = layers.Dense(output_shape, activation=\"linear\")(y)\n",
    "    model = Model(x, y, name=\"Decoder\")\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7d0bbda-de40-4b84-9d8d-1585e000f163",
   "metadata": {},
   "source": [
    "Based on the custom encoder and decoder model, we now build the DIRESA model with the `diresa_model` function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c0875727-7e9f-402e-b14e-a25bed19601c",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_2\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                Output Shape                 Param #   Connected to                  \n",
      "==================================================================================================\n",
      " input_1 (InputLayer)        [(None, 3)]                  0         []                            \n",
      "                                                                                                  \n",
      " input_2 (InputLayer)        [(None, 3)]                  0         []                            \n",
      "                                                                                                  \n",
      " Encoder (Functional)        (None, 2)                    1022      ['input_1[0][0]',             \n",
      "                                                                     'input_2[0][0]']             \n",
      "|¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|\n",
      "| input_3 (InputLayer)       [(None, 3)]                  0         []                           |\n",
      "|                                                                                                |\n",
      "| dense (Dense)              (None, 40)                   160       []                           |\n",
      "|                                                                                                |\n",
      "| dense_1 (Dense)            (None, 20)                   820       []                           |\n",
      "|                                                                                                |\n",
      "| dense_2 (Dense)            (None, 2)                    42        []                           |\n",
      "¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯\n",
      " Decoder (Functional)        (None, 3)                    1023      ['Encoder[0][0]']             \n",
      "|¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|\n",
      "| input_4 (InputLayer)       [(None, 2)]                  0         []                           |\n",
      "|                                                                                                |\n",
      "| dense_3 (Dense)            (None, 20)                   60        []                           |\n",
      "|                                                                                                |\n",
      "| dense_4 (Dense)            (None, 40)                   840       []                           |\n",
      "|                                                                                                |\n",
      "| dense_5 (Dense)            (None, 3)                    123       []                           |\n",
      "¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯\n",
      " Latent (Flatten)            (None, 2)                    0         ['Encoder[0][0]']             \n",
      "                                                                                                  \n",
      " dist_layer_2 (DistLayer)    (None, 2)                    0         ['input_1[0][0]',             \n",
      "                                                                     'input_2[0][0]',             \n",
      "                                                                     'Encoder[0][0]',             \n",
      "                                                                     'Encoder[1][0]']             \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 2045 (7.99 KB)\n",
      "Trainable params: 2045 (7.99 KB)\n",
      "Non-trainable params: 0 (0.00 Byte)\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "from diresa.models import diresa_model\n",
    "from diresa.loss import mse_dist_loss, LatentCovLoss\n",
    "\n",
    "diresa = diresa_model(x=Input(shape=3), x_twin=Input(shape=3), encoder=encoder_model(), decoder=decoder_model())\n",
    "\n",
    "diresa.compile(loss=['MSE', LatentCovLoss(1.), mse_dist_loss], loss_weights=[1., 3., 1.])\n",
    "diresa.summary(expand_nested=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5995f8fc-d658-489b-8f0c-179910981fb6",
   "metadata": {},
   "source": [
    "### 9. Reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8e369e09-4daf-4ccf-bbb5-e4f5c5975938",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on module diresa.models in diresa:\n",
      "\n",
      "NAME\n",
      "    diresa.models\n",
      "\n",
      "DESCRIPTION\n",
      "    Creates DIRESA and (V)AE models out of an encoder and decoder model\n",
      "    Creates DIRESA and AE models from hyperparameters\n",
      "    :Author:  Geert De Paepe\n",
      "    :Email:   geert.de.paepe@vub.be\n",
      "    :License: Apache 2.0\n",
      "    \n",
      "    1. Creating (V)AE and Diresa models out of an encoder and decoder model:\n",
      "     - autoencoder_model(x, encoder, decoder, mask_layer=False)\n",
      "     - vae_model(x, encoder, decoder)\n",
      "     - cov_reg_ae_model(x, encoder, decoder)\n",
      "     - siamese_twin_model(x, x_twin, encoder, decoder)\n",
      "     - diresa_model(x, x_twin, encoder, decoder)\n",
      "    \n",
      "    2. Creating AE and Diresa models from hyperparameters\n",
      "       - build_ae(input_shape, stack, stack_filters, latent_filters, kernel_size=(3, 3),\n",
      "                  conv_transpose=False, up_first=False, residual=False, batchnorm=False,\n",
      "                  dense_units=(), mask_layer=False,\n",
      "                  activation='relu', encoder_activation='linear', decoder_activation='linear')\n",
      "       - build_diresa(input_shape, stack, stack_filters, latent_filters, kernel_size=(3, 3),\n",
      "                      conv_transpose=False, up_first=False, residual=False, batchnorm=False,\n",
      "                      dense_units=(),\n",
      "                      activation='relu', encoder_activation='linear', decoder_activation='linear')\n",
      "    \n",
      "       Encoder:\n",
      "        - 0 or more [blocks] with C (Conv2D) or residual units and a P (MaxPooling layer)\n",
      "        - 0 or 1 [block] of D (Dense layers)\n",
      "        - optionally a mask layer (only for AE)\n",
      "       Decoder:\n",
      "        - 0 or 1 [block] with D (Dense layers)\n",
      "        - 0 or more [blocks] with C (Conv2D) or residual units and an U (UpSampling layer)\n",
      "    \n",
      "       stack     dense_units  Encoder                Decoder (up_first=True)    Decoder (up_first=False)\n",
      "       [1]       ()           [C-P]-Cout             [U-C]-Cout                 [C-U]-Cout\n",
      "       [3]       ()           [C-C-C-P]-Cout         [U-C-C-C]-Cout             [C-U-C-C]-Cout\n",
      "       [1,1]     ()           [C-P]-[C-P]-Cout       [U-C]-[U-C]-Cout           [C-U]-[C-U]-Cout\n",
      "       ()        [20,10]      [D-Dout]               [D-Dout]                   [D-Dout]\n",
      "       [2]       [20,10]      [C-C-P]-[D-Dout]       [D-D]-[U-C]-Cout           [D-D]-[C-U]-Cout\n",
      "       [1,1]     [20,10]      [C-P]-[C-P]-[D-Dout]   [D-D]-[U]-[U-C]-Cout       [D-D]-[U]-[C-U]-Cout\n",
      "    \n",
      "       If conv_transpose=True, C is a ConvTranspose layer, only possible for up_first=True\n",
      "       If residual=True, C is a residual unit with a skip connection, only possible for up_first=True\n",
      "       Input shape should be 3 if Conv2D blocks, first 2 dimensions of input_shape should be a multiple of 2^len(stack)\n",
      "       Input shape should be 1 if only a Dense block\n",
      "\n",
      "FUNCTIONS\n",
      "    autoencoder_model(x, encoder, decoder, mask_layer=False)\n",
      "        Creates autoencoder model out of an encoder and a decoder model\n",
      "        :param x: keras input tensor (keras.Input())\n",
      "        :param encoder: encoder functional Keras model\n",
      "        :param decoder: decoder functional Keras model\n",
      "        :param mask_layer: if True, adds a MaskLayer to the encoder\n",
      "        :return: autoencoder model\n",
      "    \n",
      "    build_ae(input_shape=(), stack=(), stack_filters=(), latent_filters=1, kernel_size=(3, 3), conv_transpose=False, up_first=False, residual=False, batchnorm=False, dense_units=(), mask_layer=False, activation='relu', encoder_activation='linear', decoder_activation='linear')\n",
      "        Creates an AE model out of hyperparameters\n",
      "        :param input_shape: 3-dimensional with Conv2D layers, first 2 dimensions should be a multiple of 2^len(stack)\n",
      "                            1-dimensional if only Dense layers\n",
      "        :param stack: elements are nbr of Conv2D or residual units in a block\n",
      "        :param stack_filters: elements are nbr of filters in a block\n",
      "        :param latent_filters: nbr of filters in convolutional output (only used if no dense units)\n",
      "        :param kernel_size: kernel size for convolution\n",
      "        :param conv_transpose: if True ConvTranspose is used in decoder, only possible for up_first=True\n",
      "        :param up_first: if True UpSampling is first in decoder block, if False UpSampling is second\n",
      "        :param residual: if True, elements in blocks are residual units, if False elements are Conv2D layers\n",
      "        :param batchnorm: if True, each Conv2D is followed by a BatchNormalization layer, if False no BN is used\n",
      "        :param dense_units: elements are nbr of nodes of a Dense layer in the dense block\n",
      "        :param mask_layer: if True a MaskLayer is added between encoder and decoder\n",
      "        :param activation: activation function used (except for output of encoder/decoder)\n",
      "        :param encoder_activation: activation function used for output of encoder\n",
      "        :param decoder_activation: activation function used for output of decoder\n",
      "        :return: AE functional Keras model\n",
      "    \n",
      "    build_diresa(input_shape=(), stack=(), stack_filters=(), latent_filters=1, kernel_size=(3, 3), conv_transpose=False, up_first=False, residual=False, batchnorm=False, dense_units=(), activation='relu', encoder_activation='linear', decoder_activation='linear')\n",
      "        Creates a Diresa model out of hyperparameters\n",
      "        :param input_shape: 3-dimensional with Conv2D layers, first 2 dimensions should be a multiple of 2^len(stack)\n",
      "                            1-dimensional if only Dense layers\n",
      "        :param stack: elements are nbr of Conv2D or residual units in a block\n",
      "        :param stack_filters: elements are nbr of filters in a block\n",
      "        :param latent_filters: nbr of filters in convolutional output (only used if no dense units)\n",
      "        :param kernel_size: kernel size for convolution\n",
      "        :param conv_transpose: if True ConvTranspose is used in decoder, only possible for up_first=True\n",
      "        :param up_first: if True UpSampling is first in decoder block, if False UpSampling is second\n",
      "        :param residual: if True, elements in blocks are residual units, if False elements are Conv2D layers\n",
      "        :param batchnorm: if True, each Conv2D is followed by a BatchNormalization layer, if False no BN is used\n",
      "        :param dense_units: elements are nbr of nodes of a Dense layer in the dense block\n",
      "        :param activation: activation function used (except for output of encoder/decoder)\n",
      "        :param encoder_activation: activation function used for output of encoder\n",
      "        :param decoder_activation: activation function used for output of decoder\n",
      "        :return: Diresa functional Keras model\n",
      "    \n",
      "    cov_reg_ae_model(x, encoder, decoder)\n",
      "        Creates a covariance regulated autoencoder model out of an encoder and a decoder model\n",
      "        :param x: keras input tensor (keras.Input())\n",
      "        :param encoder: encoder functional Keras model\n",
      "        :param decoder: decoder functional Keras model\n",
      "        :return: covariance regulated autoencoder model\n",
      "    \n",
      "    diresa_model(x, x_twin, encoder, decoder)\n",
      "        Creates a Diresa model out of an encoder and a decoder model\n",
      "        :param x: keras input tensor (keras.Input())\n",
      "        :param x_twin: keras input tensor for shuffled input\n",
      "        :param encoder: encoder functional Keras model\n",
      "        :param decoder: decoder functional Keras model\n",
      "        :return: Diresa model\n",
      "    \n",
      "    siamese_twin_model(x, x_twin, encoder, decoder)\n",
      "        Creates a Diresa model (without cov regularization on latent space) out of an encoder and a decoder model\n",
      "        :param x: keras input tensor (keras.Input())\n",
      "        :param x_twin: keras input tensor for shuffled input\n",
      "        :param encoder: encoder functional Keras model\n",
      "        :param decoder: decoder functional Keras model\n",
      "        :return: Diresa model without cov regularization on latent space\n",
      "    \n",
      "    vae_model(x, encoder, decoder)\n",
      "        Creates a variational autoencoder model out of an encoder and a decoder model\n",
      "        :param x: keras input tensor (keras.Input())\n",
      "        :param encoder: variational encoder functional Keras model\n",
      "        :param decoder: decoder functional Keras model\n",
      "        :return: variational autoencoder model\n",
      "\n",
      "FILE\n",
      "    c:\\users\\depaepeg\\python\\venv\\lib\\site-packages\\diresa\\models.py\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from diresa import *\n",
    "help(models)\n",
    "#help(loss)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
