0_molecular-design-with-parsl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "04038f13",
   "metadata": {},
   "source": [
    "# Molecular design ML-in-the-loop workflow with Parsl\n",
    "\n",
    "This notebook demonstrates a simple molecular design application where we use machine learning to guide which computations we perform.\n",
    "The objective of this application is to identify which molecules have the largest ionization energies (IE, the amount of energy required to remove an electron). \n",
    "\n",
    "IE can be computed using various simulation packages (here we use [xTB](https://xtb-docs.readthedocs.io/en/latest/contents.html) ); however, execution of these simulations is expensive, and thus, given a finite compute budget, we must carefully select which molecules to explore. We use machine learning to predict high IE molecules based on previous computations (a process often called [active learning](https://pubs.acs.org/doi/abs/10.1021/acs.chemmater.0c00768)). We iteratively retrain the machine learning model to improve the accuracy of predictions. The resulting ML-in-the-loop workflow proceeds as follows. \n",
    "\n",
    "![workflow](./figures/workflow.svg)\n",
    "\n",
    "In this notebook, we use Parsl to execute functions (simulation, model training, and inference) in parallel. Parsl allows us to establish dependencies in the workflow and to execute the workflow on arbitrary computing infrastructure, from laptops to supercomputers. We show how Parsl's integration with Python's native concurrency library (i.e., [`concurrent.futures`](https://docs.python.org/3/library/concurrent.futures.html#module-concurrent.futures)) let you write applications that dynamically respond to the completion of asynchronous tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb5a99f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt\n",
    "from chemfunctions import compute_vertical\n",
    "from concurrent.futures import as_completed\n",
    "from tqdm.notebook import tqdm\n",
    "from parsl.executors import HighThroughputExecutor\n",
    "from parsl.app.python import PythonApp\n",
    "from parsl.app.app import python_app\n",
    "from parsl.config import Config\n",
    "from time import monotonic\n",
    "from random import sample\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import parsl\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d066d4dd-bcd5-4426-b223-62e80166a87c",
   "metadata": {},
   "source": [
    "Configuration for the computing environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccf302e6-3790-4ad9-8c1d-e24e9c9482ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_workers = min(4, os.cpu_count())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38f0d354",
   "metadata": {},
   "source": [
    "## Define problem\n",
    "\n",
    "We first define configuration parameters for the app, specifically the search space of molecules (selected randomly from the QM9 database) and parameters controlling the optimization algorithm (the number of initial simulations, total moleucles to be evaluated, and the number of molecules to be evaluated in a batch)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5cb8267",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_space = pd.read_csv('data/QM9-search.tsv', delim_whitespace=True)  # Our search space of molecules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f170e63",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_count: int = 8  # Number of calculations to run at first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad59f50f",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_count: int = 64   # Number of molecules to evaluate in total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3bed4f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size: int = 4  # Number of molecules to evaluate in each batch of simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "414f939d",
   "metadata": {},
   "source": [
    "## Set up Parsl\n",
    "\n",
    "We now configure Parsl to make use of available resources. In this case we configure Parsl to run on the local machine with two workers. One of the benefits of Parsl is that we can change this configuration to make use of different resources without modifying the following workflow. For example, we can configure Parsl to use more cores on the local machine or to use many nodes on a Supercomputer or Cloud. The [Parsl website](https://parsl.readthedocs.io/en/stable/userguide/configuring.html) describes how Parsl can be configured for different resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7222dc6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Config(\n",
    "    executors=[HighThroughputExecutor(\n",
    "        max_workers=4, # Allows a maximum of two workers\n",
    "        cpu_affinity='block' # Prevents workers from using the same cores\n",
    "    )]\n",
    ")\n",
    "parsl.load(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c08379a3",
   "metadata": {},
   "source": [
    "## Make an initial dataset\n",
    "\n",
    "We need data to train our ML models. We'll do that by selecting a set of molecules at random from our search space, performing some simulations on those molecules, and training on the results.\n",
    "\n",
    "In [`chemfunctions.py`](./chemfunctions.py), we have defined a function `compute_vertical` that computes the \"vertical ionization energy\" of a molecule (a measure of how much energy it takes to strip an electron off the molecule). `compute_vertical` takes a string representation of a molecule in [SMILES format](https://en.wikipedia.org/wiki/Simplified_molecular-input_line-entry_system) as input and returns the ionization energy as a float. Under the hood, it is running [xTB](https://xtb-docs.readthedocs.io/en/latest/contents.html) to perform a series of quantum chemistry computations.\n",
    "\n",
    "### Execute a first simulation\n",
    "We need to prepare this function to run with Parsl. All we need to do is wrap this function with Parsl's `python_app`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7b97999",
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_vertical_app = python_app(compute_vertical)\n",
    "compute_vertical_app"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36525082",
   "metadata": {},
   "source": [
    "This new object is a Parsl `PythonApp`. It can be invoked like the original function, but instead of immediately executing, the function may be run asynchronously by Parsl. Instead of the result, the call will immediately return a `Future` which we can use to retrieve the result or obtain the status of the running task.\n",
    "\n",
    "For example, invoking the `compute_verticle_app` with the SMILES for water, `O`, returns a Future and schedules `compute_verticle` for execution in the background."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74c39daa",
   "metadata": {},
   "outputs": [],
   "source": [
    "future = compute_vertical_app('O') #  Run water as a demonstration (O is the SMILES for water)\n",
    "future"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf85cfad",
   "metadata": {},
   "source": [
    "We can access the result of this computation by asking the future for the `result()`. If the computation isn't finished yet, then the call to `.result()` will block until the result is ready."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "921c7a0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ie = future.result()\n",
    "print(f\"The ionization energy of {future.task_def['args'][0]} is {ie:.2f} eV\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b085a991",
   "metadata": {},
   "source": [
    "### Scale the simulation\n",
    "\n",
    "It is trivial now to scale our simulation and run it for several different molecules and gather their results.\n",
    "\n",
    "We use a standard Python loop to submit a set of simulations for execution. As above, each invocation returns a `Future` immediately, so this code should finish within a few milliseconds.\n",
    "\n",
    "Because we never call `.result()`, this code does not wait for any results to be ready. Instead, Parsl is running the computations in the background. Parsl manages sending work to each worker process, collecting results, and feeding new work to workers as new tasks are submitted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c74a1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "smiles = search_space.sample(initial_count)['smiles']\n",
    "futures = [compute_vertical_app(s) for s in smiles]\n",
    "print(f'Submitted {len(futures)} calculations to start with')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dda3b23",
   "metadata": {},
   "source": [
    "The futures produced by Parsl are based on Python's [native \"Future\"](https://docs.python.org/3/library/concurrent.futures.html#future-objects) object,\n",
    "so we can use Python's utility functions to work with them.\n",
    "\n",
    "As an example, we can build a loop that submits new computations if previous ones fail. This happens not too infrequently with our simulation application.\n",
    "\n",
    "We use `as_completed` to take an iterable (in this case a list) of futures and to yeild as each future completes.  Thus, we progress and handle each simulation as it completes\n",
    "\n",
    "We also use, `Future.exception()` rather than the similar `Future.result()`. `Future.exception()` behaves similarly in that it will block until the relevant task is completed, but rather than return the result, it returns any exception that was raised during execution (or `None` if not). In this case, if the future returns an exception we simply pick a new molecule and re-execute the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ac26a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = []\n",
    "while len(futures) > 0: \n",
    "    # First, get the next completed computation from the list\n",
    "    future = next(as_completed(futures))\n",
    "    \n",
    "    # Remove it from the list of still-running tasks\n",
    "    futures.remove(future)\n",
    "    \n",
    "    # Get the input \n",
    "    smiles = future.task_def['args'][0]\n",
    "    \n",
    "    # Check if the run completed successfully\n",
    "    if future.exception() is not None:\n",
    "        # If it failed, pick a new SMILES string at random and submit it    \n",
    "        print(f'Computation for {smiles} failed, submitting a replacement computation')\n",
    "        smiles = search_space.sample(1).iloc[0]['smiles'] # pick one molecule\n",
    "        new_future = compute_vertical_app(smiles) # launch a simulation in Parsl\n",
    "        futures.append(new_future) # store the Future so we can keep track of it\n",
    "    else:\n",
    "        # If it succeeded, store the result\n",
    "        print(f'Computation for {smiles} succeeded')\n",
    "        train_data.append({\n",
    "            'smiles': smiles,\n",
    "            'ie': future.result(),\n",
    "            'batch': 0,\n",
    "            'time': monotonic()\n",
    "        })"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f373bbc",
   "metadata": {},
   "source": [
    "We now have an initial set of training data. We load this training data into a pandas `DataFrame` containing the randomly samples molecules alongside the simulated ionization energy (`ie`). In addition, the code above has stored some metadata (`batch` and `time`) which we will use later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f15684d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = pd.DataFrame(train_data)\n",
    "train_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c3b0bad",
   "metadata": {},
   "source": [
    "## Train a machine learning model to screen candidate molecules\n",
    "Our next step is to create a machine learning model to estimate the outcome of new computations (i.e., ionization energy) and use it to rapidly scan the search space.\n",
    "\n",
    "To start, let's make a function that uses our prior simulations to train a model. We are going to use RDKit and scikit-learn to train a nearest-neighbor model that uses Morgan fingerprints to define similarity (see [notes from a UChicago AI course](https://github.com/WardLT/applied-ai-for-materials/blob/main/molecular-property-prediction/chemoinformatics/2_ml-with-fingerprints.ipynb) for more detail). In short, the function trains a model that first populates a list of certain substructures (Morgan fingerprints, specifically) and then trains a model which predicts the IE of a new molecule by averaging those with the most similar substructures.\n",
    "\n",
    "We want to use Parsl here to scale the model and to later combine it into our ML-in-the-loop workflow. To do so, we define the function using `python_app`. This time, `python_app` is used as a decorator directly on the function definition (earlier we defined a regular function, and then applied `python_app` afterwards)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21fc424b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python_app\n",
    "def train_model(train_data):\n",
    "    \"\"\"Train a machine learning model using Morgan Fingerprints.\n",
    "    \n",
    "    Args:\n",
    "        train_data: Dataframe with a 'smiles' and 'ie' column\n",
    "            that contains molecule structure and property, respectfully.\n",
    "    Returns:\n",
    "        A trained model\n",
    "    \"\"\"\n",
    "    # Imports for python functions run remotely must be defined inside the function\n",
    "    from chemfunctions import MorganFingerprintTransformer\n",
    "    from sklearn.neighbors import KNeighborsRegressor\n",
    "    from sklearn.pipeline import Pipeline\n",
    "    \n",
    "    \n",
    "    model = Pipeline([\n",
    "        ('fingerprint', MorganFingerprintTransformer()),\n",
    "        ('knn', KNeighborsRegressor(n_neighbors=4, weights='distance', metric='jaccard', n_jobs=-1))  # n_jobs = -1 lets the model run all available processors\n",
    "    ])\n",
    "    \n",
    "    return model.fit(train_data['smiles'], train_data['ie'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef88918a",
   "metadata": {},
   "source": [
    "Now let's execute the function and run it asynchronously with Parsl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cb0ff0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_future = train_model(train_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3986e1b6",
   "metadata": {},
   "source": [
    "One of the unique features of Parsl is that it can create workflows on-the-fly directly from Python. Parsl workflows are chains of functions, connected by dynamic depencies (i.e., data passed between Parsl `apps`), that can run in parallel when possible.\n",
    "\n",
    "To establish the workflow, we pass the future created by executing one function an input to another Parsl function.\n",
    "\n",
    "As an example, let's create a function that uses the trained model to run inference on a large set of molecules and then another that takes many predictions and concatenates them into a single collection. The sequential workflow is implemented as follows.\n",
    "\n",
    "        train_model --> run_model --> combine_inferences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e947854",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python_app\n",
    "def run_model(model, smiles):\n",
    "    \"\"\"Run a model on a list of smiles strings\n",
    "    \n",
    "    Args:\n",
    "        model: Trained model that takes SMILES strings as inputs\n",
    "        smiles: List of molecules to evaluate\n",
    "    Returns:\n",
    "        A dataframe with the molecules and their predicted outputs\n",
    "    \"\"\"\n",
    "    import pandas as pd\n",
    "    pred_y = model.predict(smiles)\n",
    "    return pd.DataFrame({'smiles': smiles, 'ie': pred_y})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fa1b2b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python_app\n",
    "def combine_inferences(inputs=[]):\n",
    "    \"\"\"Concatenate a series of inferences into a single DataFrame\n",
    "    Args:\n",
    "        inputs: a list of the component DataFrames\n",
    "    Returns:\n",
    "        A single DataFrame containing the same inferences\n",
    "    \"\"\"\n",
    "    import pandas as pd\n",
    "    return pd.concat(inputs, ignore_index=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b350ea4",
   "metadata": {},
   "source": [
    "Now we've created our Parsl `apps`, we can chop up the search space into chunks, and invoke `run_model`  once for each chunk of the search space.\n",
    "\n",
    "Note: we pass `train_future` (the future created from the training function above) as input to `run_model`. Parsl will wait for the training to be complete (i.e., the future to be resolved) before executing `run_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f327013",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chunk the search space into smaller pieces, so that each can run in parallel\n",
    "chunks = np.array_split(search_space['smiles'], 64)\n",
    "inference_futures = [run_model(train_future, chunk) for chunk in chunks]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e008cc2d",
   "metadata": {},
   "source": [
    "While we are running inferences in parallel we can define the final part of the workflow to combine results into a single DataFrame using `combine_inferences`.\n",
    "\n",
    "We pass the `inference_futures` as inputs to `combine_inferences` such that Parsl knows to establish a dependency between these two functions. That is, Parsl will ensure that `train_future` must complete before any of the `run_model` tasks start; and all of the `run_model` tasks must be finished before `combine_inferences` starts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4135687",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We pass the inputs explicitly as a named argument \"inputs\" for Parsl to recognize this as a \"reduce\" step\n",
    "#  See: https://parsl.readthedocs.io/en/stable/userguide/workflow.html#mapreduce\n",
    "predictions = combine_inferences(inputs=inference_futures).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eb78992",
   "metadata": {},
   "source": [
    "#### Results\n",
    "\n",
    "After completing the inference process we now have predicted IE values for all molecules in our search space. We can print out the best five molecules, according to the trained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9801c685",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions.sort_values('ie', ascending=False).head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a13b4ab",
   "metadata": {},
   "source": [
    "We have now created a Parsl workflow that is able to train a model and use it to identify molecules that are likely to be good next choices for simulations. Time to build a model-in-the-loop workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ad296c",
   "metadata": {},
   "source": [
    "## Model-in-the-Loop Workflow\n",
    "We are going to build an application that uses a machine learning model to pick a batch of simulations, runs the simulations in parallel, and then uses the data to retrain the model before repeating the loop.\n",
    "\n",
    "Our application uses `train_model`, `run_model`, and `combine_inferences` as above, but after running an iteration it picks the predicted best molecules and runs the `compute_vertical_app` to run the xTB simulation.  The workflow then repeatedly retrains the model using these results until a fixed number of molecule simulations have been trained. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50007b1d-b094-4933-8140-97c321e2ea9f",
   "metadata": {},
   "source": [
    "Make the search space a list so that we can remove completed molecules more easily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9dc018e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with tqdm(total=search_count) as prog_bar: # setup a graphical progress bar\n",
    "    # Mark when we started\n",
    "    start_time = monotonic()\n",
    "    \n",
    "    # Submit with some random guesses\n",
    "    train_data = []\n",
    "    init_mols = search_space.sample(initial_count)['smiles']\n",
    "    sim_futures = [compute_vertical_app(mol) for mol in init_mols]\n",
    "    already_ran = set()\n",
    "    \n",
    "    # Loop until you finish populating the initial set\n",
    "    while len(sim_futures) > 0: \n",
    "        # First, get the next completed computation from the list\n",
    "        future = next(as_completed(sim_futures))\n",
    "\n",
    "        # Remove it from the list of still-running tasks\n",
    "        sim_futures.remove(future)\n",
    "\n",
    "        # Get the input \n",
    "        smiles = future.task_def['args'][0]\n",
    "        already_ran.add(smiles)\n",
    "\n",
    "        # Check if the run completed successfully\n",
    "        if future.exception() is not None:\n",
    "            # If it failed, pick a new SMILES string at random and submit it    \n",
    "            smiles = search_space.sample(1).iloc[0]['smiles'] # pick one molecule\n",
    "            new_future = compute_vertical_app(smiles) # launch a simulation in Parsl\n",
    "            sim_futures.append(new_future) # store the Future so we can keep track of it\n",
    "        else:\n",
    "            # If it succeeded, store the result\n",
    "            prog_bar.update(1)\n",
    "            train_data.append({\n",
    "                'smiles': smiles,\n",
    "                'ie': future.result(),\n",
    "                'batch': 0,\n",
    "                'time': monotonic() - start_time\n",
    "            })\n",
    "    \n",
    "    \n",
    "    \n",
    "    # Create the initial training set as a \n",
    "    train_data = pd.DataFrame(train_data)\n",
    "    \n",
    "    # Loop until complete\n",
    "    batch = 1\n",
    "    while len(train_data) < search_count:\n",
    "        \n",
    "        # Train and predict as show in the previous section.\n",
    "        train_future = train_model(train_data)\n",
    "        inference_futures = [run_model(train_future, chunk) for chunk in np.array_split(search_space['smiles'], 64)]\n",
    "        predictions = combine_inferences(inputs=inference_futures).result()\n",
    "\n",
    "        # Sort the predictions in descending order, and submit new molecules from them\n",
    "        predictions.sort_values('ie', ascending=False, inplace=True)\n",
    "        sim_futures = []\n",
    "        for smiles in predictions['smiles']:\n",
    "            if smiles not in already_ran:\n",
    "                sim_futures.append(compute_vertical_app(smiles))\n",
    "                already_ran.add(smiles)\n",
    "                if len(sim_futures) >= batch_size:\n",
    "                    break\n",
    "\n",
    "        # Wait for every task in the current batch to complete, and store successful results\n",
    "        new_results = []\n",
    "        for future in as_completed(sim_futures):\n",
    "            if future.exception() is None:\n",
    "                prog_bar.update(1)\n",
    "                new_results.append({\n",
    "                    'smiles': future.task_def['args'][0],\n",
    "                    'ie': future.result(),\n",
    "                    'batch': batch, \n",
    "                    'time': monotonic() - start_time\n",
    "                })\n",
    "                \n",
    "        # Update the training data and repeat\n",
    "        batch += 1\n",
    "        train_data = pd.concat((train_data, pd.DataFrame(new_results)), ignore_index=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b66fd64b",
   "metadata": {},
   "source": [
    "We can plot the training data against the time of simulation, showing that the model is finding better molecules over time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b774aa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4.5, 3.))\n",
    "\n",
    "ax.scatter(train_data['time'], train_data['ie'])\n",
    "ax.step(train_data['time'], train_data['ie'].cummax(), 'k--')\n",
    "\n",
    "ax.set_xlabel('Walltime (s)')\n",
    "ax.set_ylabel('Ion. Energy (Ha)')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3930b49-cf7e-405a-bdf9-dd332addbeb6",
   "metadata": {},
   "source": [
    "Save that data for comparison with another application later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c295545-8f4c-476a-9300-cdc132add9cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "Path('run-data').mkdir(exist_ok=True)\n",
    "train_data.to_csv('run-data/parsl-results.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdccafdc-343a-47b6-a4b8-b712386e41e0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.16"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "0cacceed9db4464f9f414274519b4949": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "252de14e22c54080bfcaff0574d898d4": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "270ad7c288c14ba2af5fe1cd642a09eb": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_6fe9c731cdc640ad858660f64c387dc4",
       "placeholder": "​",
       "style": "IPY_MODEL_fe7006dbf4824d898f4f010f87b1c6f2",
       "value": " 70/? [01:05&lt;00:00,  1.26s/it]"
      }
     },
     "474ddeec4b9e40f5a89f9c7b67ea31e6": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_0cacceed9db4464f9f414274519b4949",
       "placeholder": "​",
       "style": "IPY_MODEL_d4c9ca989df549f09c7ce94ee74c6a57",
       "value": ""
      }
     },
     "6fe9c731cdc640ad858660f64c387dc4": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "84f635e5b62a44c0b110315fb375c321": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_tooltip": null,
       "layout": "IPY_MODEL_ad75de48e1324ca3993fc8cf8b31a853",
       "max": 64,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_252de14e22c54080bfcaff0574d898d4",
       "value": 64
      }
     },
     "91420bed062b46e9b5c2a914377ed68f": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "ad75de48e1324ca3993fc8cf8b31a853": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "1.2.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "overflow_x": null,
       "overflow_y": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "bb5f4153b6a44eac8cec5ed6fc6923ca": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "1.5.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_474ddeec4b9e40f5a89f9c7b67ea31e6",
        "IPY_MODEL_84f635e5b62a44c0b110315fb375c321",
        "IPY_MODEL_270ad7c288c14ba2af5fe1cd642a09eb"
       ],
       "layout": "IPY_MODEL_91420bed062b46e9b5c2a914377ed68f"
      }
     },
     "d4c9ca989df549f09c7ce94ee74c6a57": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     },
     "fe7006dbf4824d898f4f010f87b1c6f2": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "1.5.0",
       "_model_name": "DescriptionStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "1.2.0",
       "_view_name": "StyleView",
       "description_width": ""
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}