tutorials/data_driven/LSTM/rnn_calibration_GL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5c8468b6-76f6-4caa-a0b5-9a4e884ffb95",
   "metadata": {},
   "source": [
    "# Tutorial: Use a trained RNN in grainLearning calibration process."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2014adcf-1349-4544-b78f-91b1d2a5faf0",
   "metadata": {},
   "source": [
    "## ⚙️ Install grainlearning package\n",
    "(Not necessary if you are running jupyter-lab on an environment where grainlearning and rnn dependencies are installed)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3662d03a-b14e-47a5-bfd7-faceac6ee3a2",
   "metadata": {},
   "source": [
    "```bash \n",
    "pip install grainlearning --extras \"rnn\"\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "886054c8-6db6-49f2-8783-cfa38411cc19",
   "metadata": {},
   "source": [
    "## Introduction to this tutorial\n",
    "We hope that now you are familiar with grainLearning and the RNN module. In this tutorial we are going to explore how can we use a pretrained neural network as a surrogate model in the grainLearning calibration process. In such context, the RNN plays the role of a `DynamicSystem`.\n",
    "\n",
    "We consider the case of *Drained Triaxial Compression*, and have two resources:\n",
    "\n",
    "1. A RNN trained model on several DEM simulations of *Drained Triaxial Compression* tests. \n",
    "2. The experimental measurements of a *Drained Triaxial Compression* test.\n",
    "\n",
    "The *objective* is to find which set of `contact_params` of the DEM simulation would give us an equivalent to our real-world material.\n",
    "\n",
    "This tutorial has three main parts:\n",
    "\n",
    "1. Prepare the pre-trained model and observation data.\n",
    "2. Create a callback function to link to `DynamicSystem`.\n",
    "3. GrainLearning calibration loop.\n",
    "\n",
    "Let's start importing some useful packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "db3ce29e-eb0f-43be-97bd-153cd48d4045",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from grainlearning import BayesianCalibration\n",
    "import grainlearning.dynamic_systems\n",
    "import grainlearning.rnn.predict as predict_rnn\n",
    "import grainlearning.rnn.preprocessing as preprocessing_rnn"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d6cb044c-8f94-4623-878e-12b951f1197d",
   "metadata": {},
   "source": [
    "## 1.a Prepare the pre-trained model 🗄️\n",
    "For the purpose of this tutorial we are going to take an example from the a dataset containing DEM simulation results of a triaxial compression. \n",
    "Our RNN model was trained in a similar database, and we are going to take one example to infer from. In practice, such example could be the results of real-world experiments for wich we do not know the DEM _contact parameters_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "322c993c-dacf-4be1-a8be-ccdb66ce71c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_trained_model = '/Users/luisaorozco/Documents/Projects/GrainLearning/grainLearning/grainlearning/rnn/trained_models/rnn_triaxial_drained'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b54a08c6-1c1c-40d0-8995-c231ca7defd7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metal device set to: Apple M1 Pro\n",
      "\n",
      "systemMemory: 32.00 GB\n",
      "maxCacheSize: 10.67 GB\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-23 16:31:23.241391: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.\n",
      "2023-03-23 16:31:23.241521: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)\n"
     ]
    }
   ],
   "source": [
    "# Load the pre-trained model\n",
    "model_rnn, train_stats, config_rnn = predict_rnn.get_pretrained_model(path_to_trained_model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "33b62d30",
   "metadata": {},
   "source": [
    "## 1.b Prepare the observation data 🥼\n",
    "We load the data that we have obtained in an experiment (drained triaxial comrpession) and for which we want to calibrate contact parameters to use in DEM.\n",
    "In the first column there is the time-sequence of the control parameter $\\varepsilon_{axial}$, and the 3 next columns are the observed parameters $e,\\ p,\\ q$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1e320071",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_data = 'experimental_test_drained_s3=0.2.dat'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f83234de",
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_data = np.loadtxt(path_to_data) \n",
    "inputs = experiment_data[:,0]\n",
    "outputs = experiment_data[:,1:]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "38be8478-f63e-47e4-8de1-17d5e3a767a0",
   "metadata": {},
   "source": [
    "`extra_contact_params` are not in `system.param_data` i.e. they are not parameters that need to be inferred. \n",
    "However, these are control parameters necessary to predict from RNN that are added at the end of  `contact_params`.\n",
    "Given the values of `add_e0`, `add_pressure` and `add_experiment_type` in `config`, we can determine how many parmeters are extra."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cb0be40b",
   "metadata": {},
   "outputs": [],
   "source": [
    "extra_contact_params = [0.2] # confining pressure"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe626ae9-80b7-48be-acfe-78587a4bfb13",
   "metadata": {},
   "source": [
    "If no padding was taking into account during training, the predicitions of the RNN are going to be one `window_size` shorter.\n",
    "We adapt the labels so that during the calibration we compare tensors with the same dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e8bdd6e8-cd92-4346-8e22-ea63534da6be",
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'pad_length' not in config_rnn or config_rnn['pad_length'] == 0: # no padding\n",
    "    outputs = outputs[config_rnn['window_size']:, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b4fafd0-c0ef-43aa-b4e5-bf7a4d079ae5",
   "metadata": {},
   "source": [
    "In Grainlearning, the temporal dimension is always at the end, we need to switch the access of the data to comply with this format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ad8307a1-0ded-46a4-8919-517419f8d0a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert arrays from [num_time_steps, num_features] -> [num_features, num_time_steps]\n",
    "outputs = np.moveaxis(outputs, 0, -1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3444e321-4a4a-4c60-bd67-d2ffb22d62f7",
   "metadata": {},
   "source": [
    "In this preparation steps we got few important arrays to work with:\n",
    "- `inputs`: control of the system ($\\varepsilon_{axial}$)\n",
    "- `outputs`: observation experimental data ($e,\\ p,\\ q$)\n",
    "- `extra_contact_params`: Additional control parameters that don't need to be inferred but need ot be passed along with `contact_params` to the RNN.\n",
    "- `config_rnn`: dictionary containing the configration of the pre-trained RNN model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7599a32a-5b87-495d-bee7-b075c2ec2056",
   "metadata": {},
   "source": [
    "## 2. Create a callback function to link to `DynamicSystem`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "043d40db-6825-4782-9936-2f4d0e86cdb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_with_RNN(system, **_):\n",
    "    window_size = config_rnn['window_size']\n",
    "    sequence_length = np.shape(system.ctrl_data)[0]\n",
    "    \n",
    "    # For compatibility with the RNN (first dimension is sample)\n",
    "    load_sequence = np.expand_dims(inputs, axis=0) # only needed if system.ctrl_data ahs a single dimension\n",
    "    load_sequence = np.moveaxis(load_sequence, 0, -1)\n",
    "    load_sequence = np.repeat(load_sequence[np.newaxis, :], system.num_samples, axis=0)\n",
    "    \n",
    "    # add extra_contact_params to the contact_params used to draw a prediction\n",
    "    contact_params = np.array([np.append(i, extra_contact_params) for i in system.param_data])\n",
    "    \n",
    "    data_inputs = ({'load_sequence': load_sequence, 'contact_parameters': contact_params}, load_sequence)\n",
    "    data_inputs = tf.data.Dataset.from_tensor_slices(data_inputs)\n",
    "    predictions = predict_rnn.predict_macroscopics(model_rnn, \n",
    "                                                   data_inputs, \n",
    "                                                   train_stats, \n",
    "                                                   config_rnn, \n",
    "                                                   batch_size=system.num_samples)\n",
    "    # Getting rid of extra predictions of the rnn. \n",
    "    # We are only interested on the first 3 that we can compare against the observation data.\n",
    "    predictions = predictions.numpy()[:, :, :3]\n",
    "\n",
    "    # converting the predictions to GL format (temporal dimension at the end)\n",
    "    predictions = np.moveaxis(predictions, 1, -1)\n",
    "    # update sim_data in system\n",
    "    system.set_sim_data(predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98e4f819-fcb4-4ec6-a465-e24853806895",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 3. GrainLearning calibration loop 🔁"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "64f6c4e9-8a4b-471a-831f-945e3241da5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def grainlearning_calibration(inputs: np.array, outputs: np.array):\n",
    "    \"\"\"\n",
    "    Main function defining and driving the grainLearning calibration.\n",
    "    1. Define the BayesianCalibration and all its elements: DynamicSystem, inferences and sampling parameters.\n",
    "    2. Run the calibration and finally return the inferred parameters.\n",
    "\n",
    "    :param inputs: numpy.array with the control data.\n",
    "      In this example (drained triaxial compression) a time sequence of the axial strain.\n",
    "    :param outputs: numpy.array with the observation data.\n",
    "      In this example a time sequence with macroscopic observables such as e, p, q.\n",
    "\n",
    "    :return: most_prob_params numpy.array containing the inferred values of the contact parameters.\n",
    "    \"\"\"\n",
    "    curr_iter = 0\n",
    "    calibration = grainlearning.BayesianCalibration.from_dict(\n",
    "        {\n",
    "            \"curr_iter\": curr_iter,\n",
    "            \"num_iter\": 20,\n",
    "            \"system\": {\n",
    "                \"system_type\": grainlearning.dynamic_systems.DynamicSystem, # because I'm not reading files, my data is generated by an RNN.\n",
    "                \"obs_data\": outputs,    # experimental data\n",
    "                \"obs_names\": ['e', 'p', 'q'],\n",
    "                \"ctrl_data\": inputs[config_rnn['window_size']:], # Synthetic data.\n",
    "                \"ctrl_name\": ['e_z'],  # Only one of the strains (axial)\n",
    "                \"num_samples\": 50,     # num of samples (gaussian mixture) generated per iteration\n",
    "                \"sim_name\": 'Triaxial compression with RNN, from experiment.',\n",
    "                # to get these labels I opened the hdf5 file and query: file.attrs['contact_params']\n",
    "                \"param_names\": ['E', 'v', 'kr', 'eta', 'mu'],\n",
    "                \"param_min\": [7.00195859, 1.6e-4, 8.3e-4, 1.5e-4, 4.5e-2], # using the range of contact params used to train the RNN.\n",
    "                \"param_max\": [9.99607979, 0.49952, 0.99875, 0.99878, 59.955],\n",
    "                \"callback\": predict_with_RNN\n",
    "            },\n",
    "            \"calibration\": {\n",
    "                \"inference\": {\"ess_target\": 0.3},\n",
    "                \"sampling\": {\n",
    "                \"max_num_components\": 5,\n",
    "                \"prior_weight\": 0.01,\n",
    "                },\n",
    "            },\n",
    "            \"save_fig\": -1, # Not generating plots, but look at the end of this notebook for more info.\n",
    "        }\n",
    "    )\n",
    "\n",
    "    print(\"Run the calibration\")\n",
    "    calibration.run()\n",
    "    print(\"Calibration finished\")\n",
    "    return calibration.get_most_prob_params()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c4f879b-eb1d-43c5-927f-a33b485f3d35",
   "metadata": {},
   "source": [
    "## 🎛️ Start the calibration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "711de580-06a7-4661-9ac9-2d9035aae257",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run the calibration\n",
      "Bayesian calibration iter No. 0\n",
      "Bayesian calibration iter No. 1\n",
      "Bayesian calibration iter No. 2\n",
      "Bayesian calibration iter No. 3\n",
      "Bayesian calibration iter No. 4\n",
      "Bayesian calibration iter No. 5\n",
      "Bayesian calibration iter No. 6\n",
      "Bayesian calibration iter No. 7\n",
      "Bayesian calibration iter No. 8\n",
      "Bayesian calibration iter No. 9\n",
      "Bayesian calibration iter No. 10\n",
      "Bayesian calibration iter No. 11\n",
      "Bayesian calibration iter No. 12\n",
      "Bayesian calibration iter No. 13\n",
      "Bayesian calibration iter No. 14\n",
      "Bayesian calibration iter No. 15\n",
      "Bayesian calibration iter No. 16\n",
      "Bayesian calibration iter No. 17\n",
      "Bayesian calibration iter No. 18\n",
      "Bayesian calibration iter No. 19\n",
      "Calibration finished\n"
     ]
    }
   ],
   "source": [
    "# Run the calibration\n",
    "most_prob_params = grainlearning_calibration(inputs, outputs)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "02ee8252-b330-4d86-aca1-1412e09cd0b1",
   "metadata": {},
   "source": [
    "## ⚖️ Compare the results\n",
    "For comparison purposes, we kne beforehand, what would be the best `contact_params` for our experimental data. We are now going to compare those to the inferred ones (via grainlearning calibration process)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "079be3a4-0ff8-4d58-a030-6a13d03330ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Contact params (ground truth): [7.57811981, 0.4448, 0.47813, 0.22497, 56.0784] \n",
      "Contact parameters calibrated via GL: [ 7.59705764  0.47406657  0.28572617  0.23129764 56.9195252 ]\n"
     ]
    }
   ],
   "source": [
    "contact_params_reference = [7.57811981, 0.4448, 0.47813, 0.22497, 56.0784]\n",
    "print(f\"Contact params (ground truth): {contact_params_reference} \\nContact parameters calibrated via GL: {most_prob_params}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d8b0d51b-1e15-4522-a57c-684185d7da58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "error = [0.00249901 0.06579715 0.40240904 0.02812659 0.01499909]\n"
     ]
    }
   ],
   "source": [
    "# get the percentage error\n",
    "print(f\"error = {np.abs(most_prob_params - contact_params_reference) / contact_params_reference}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "638b08c1-34fd-43f5-a0c2-62a08103c6f7",
   "metadata": {},
   "source": [
    "### 📉 Plotting\n",
    "- Try with `save_fig` equals to 1: the plots for each one of the calibration iterations will be saved in a folder.\n",
    "- or try with `save_fig` equals to 0: the plots will be printed in this jupyter notebook on the calibration run."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b623f48-a916-4302-a4b5-2e0eba1c638a",
   "metadata": {},
   "source": [
    "## ✏️ Final tips\n",
    "\n",
    "- Check always the dimensions and the order of the parameters in the tensors that you are using. \n",
    "- Keep in mind that in GrainLearning the last dimension is the time."
   ]
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}