molecular-design.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "04038f13",
   "metadata": {},
   "source": [
    "# Molecular design ML-in-the-loop workflow\n",
    "\n",
    "This notebook demonstrates an increasingly commmon ML-in-the-loop molecular design application. We use ML to guide the choice of simulations to 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 Globus Compute to execute functions (simulation, model training, and inference) in parallel on remote computers. We show how Globus Compute's use of (i.e., [`concurrent.futures`](https://docs.python.org/3/library/concurrent.futures.html#module-concurrent.futures)) allows applications to be easily written that dynamically respond to the completion of asynchronous tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90308d00-55ea-4f96-8336-119707ab17d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set this ID to your Globus Compute endpoint\n",
    "compute_endpoint = '' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1ba021c-1318-4a8a-a6c2-6010e3649594",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from concurrent.futures import as_completed\n",
    "from matplotlib import pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "from time import monotonic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c08379a3",
   "metadata": {},
   "source": [
    "## Run a quantum chemistry simulation to calcuate IE\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",
    "Below, we define a function: \n",
    "`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",
    "In preparation for running remotely with Globus Compute, we include all import statements within the function.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7b97999",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_vertical(mol_string: str) -> float:\n",
    "    \"\"\"Run the ionization potential computation\n",
    "\n",
    "    Args:\n",
    "        xyz: XYZ coordinates for the molecule to evaluate\n",
    "    Returns:\n",
    "        Ionization energy in Ha\n",
    "    \"\"\"\n",
    "    from rdkit import Chem, DataStructs\n",
    "    from rdkit.Chem import AllChem\n",
    "\n",
    "    from io import StringIO\n",
    "    from ase.optimize import LBFGSLineSearch\n",
    "    from xtb.ase.calculator import XTB\n",
    "    from ase.io import read\n",
    "    import numpy as np\n",
    "   \n",
    "    # Generate 3D coordinates for the molecule\n",
    "    mol = Chem.MolFromSmiles(mol_string)\n",
    "    if mol is None:\n",
    "        raise ValueError(f'Parse failure for {mol_string}')\n",
    "    mol = Chem.AddHs(mol)\n",
    "    AllChem.EmbedMolecule(mol, randomSeed=1)\n",
    "    AllChem.MMFFOptimizeMolecule(mol)\n",
    "\n",
    "    # Save geometry as 3D coordinates\n",
    "    xyz = f\"{mol.GetNumAtoms()}\\n\"\n",
    "    xyz += mol_string + \"\\n\"\n",
    "    conf = mol.GetConformer()\n",
    "    for i, a in enumerate(mol.GetAtoms()):\n",
    "        s = a.GetSymbol()\n",
    "        c = conf.GetAtomPosition(i)\n",
    "        xyz += f\"{s} {c[0]} {c[1]} {c[2]}\\n\"\n",
    "        \n",
    "    # Make the XTB calculator\n",
    "    calc = XTB(accuracy=0.05)\n",
    "    \n",
    "    # Parse the molecule\n",
    "    atoms = read(StringIO(xyz), format='xyz')\n",
    "\n",
    "    # Compute the neutral geometry\n",
    "    # Uses QCEngine (https://github.com/MolSSI/QCEngine) to handle interfaces to XTB\n",
    "    atoms.calc = calc\n",
    "    dyn = LBFGSLineSearch(atoms, logfile=None)\n",
    "    dyn.run(fmax=0.02, steps=250)\n",
    "    \n",
    "    neutral_energy = atoms.get_potential_energy()\n",
    "\n",
    "    # Compute the energy of the relaxed geometry in charged form\n",
    "    charges = np.ones((len(atoms),)) * (1 / len(atoms))\n",
    "    atoms.set_initial_charges(charges)\n",
    "    charged_energy = atoms.get_potential_energy()\n",
    "    \n",
    "    return charged_energy - neutral_energy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36525082",
   "metadata": {},
   "source": [
    "First, we can run these functions locally to compute the ionization potential ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74c39daa",
   "metadata": {},
   "outputs": [],
   "source": [
    "#  Run water as a demonstration (O is the SMILES for water)\n",
    "ie = compute_vertical('O') \n",
    "print(f\"The ionization energy of O is {ie} eV\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b085a991",
   "metadata": {},
   "source": [
    "### Run the simulation on a remote computer using Globus Compute\n",
    "\n",
    "We can now easily run the simulation on a remote computer using Globus Compute. First we need to choose a specific endpoint for execution and instantiate a Globus Compute Executor. \n",
    "\n",
    "Note: the Globus Compute endpoint must be configured with an environment with all dependencies. It can be configured using the same environment.yml file or with the Dockerfile in this repository. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eab5e9e-14e4-4121-9307-40856175191a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from globus_compute_sdk import Executor\n",
    "\n",
    "gce = Executor(endpoint_id=compute_endpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fdc65ea-575a-49aa-b3dc-ad33ec76d611",
   "metadata": {},
   "source": [
    "Now we use Globus Compute to submit the ``compute_vertical'' function for execution. \n",
    "\n",
    "Note: this code does not wait for any results to be ready. Instead, Globus Compute returns a future to track the process while it runs the computation on the remote endpoint. Globus Compute manages sending work to the endpoint, distributing to each worker process, collecting results, storing results in the cloud, and returning results to the notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2775127-e0f1-48eb-ad57-cdb874a75216",
   "metadata": {},
   "outputs": [],
   "source": [
    "ie_future = gce.submit(compute_vertical, 'O')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10ade0ba-bab4-4a2d-b57c-e3c41c03e486",
   "metadata": {},
   "source": [
    "Now we can wait for the result by calling ``.result()`` on the future."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24f0b9eb-c503-4563-bec6-a41626e1b31d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ie_future.result()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6fe5709-694d-40df-8beb-6239947e2c48",
   "metadata": {},
   "source": [
    "# Running many simulations\n",
    "\n",
    "We now want to run many simulations. We use the QM9 database to populate a search space of molecules. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a7f3e8b-8d4a-47b8-960d-989e3b6b23aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_space = pd.read_csv('data/QM9-search.tsv', sep='\\s+')  # Our search space of molecules"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b03c602a-0664-4599-881b-71bfda05fbd2",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c74a1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# smiles = search_space.sample(initial_count)['smiles']\n",
    "smiles = ['C', 'O', 'CO', 'OC']\n",
    "ie_futures = [gce.submit(compute_vertical, s) for s in smiles]\n",
    "smiles_futures = dict(zip(ie_futures,smiles)) # Mapping from future to smiles\n",
    "print(f'Submitted {len(ie_futures)} calculations')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c051b3-33bd-47a8-b754-700fd491f49e",
   "metadata": {},
   "source": [
    "The futures produced by Globus Compute 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",
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dda3b23",
   "metadata": {},
   "source": [
    "You may see that some functions fail (this is common with stochastic simulations). To overcome these errors, we can write a loop that runs a new computation if previous ones fail. \n",
    "\n",
    "We 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(ie_futures) > 0: \n",
    "    # First, get the next completed computation from the list\n",
    "    future = next(as_completed(ie_futures))\n",
    "    \n",
    "    # Remove it from the list of still-running tasks\n",
    "    ie_futures.remove(future)\n",
    "    \n",
    "    # Get the input \n",
    "    smiles = smiles_futures[future]  #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 = gce.submit(compute_vertical, smiles) # launch a new simulation\n",
    "        smiles_futures[future] = smiles\n",
    "        ie_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_df = pd.DataFrame(train_data)\n",
    "train_df"
   ]
  },
  {
   "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, we define a function that uses our prior simulations to train a simple k-nearest neighbors (KNN) model. We are going to use RDKit and scikit-learn to train the 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21fc424b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(train_data, n_neighbors=4):\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=n_neighbors, 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 to train the model using Globus Compute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cb0ff0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_future = gce.submit(train_model, train_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42afc89b-10ca-4c34-a813-2d449c2420f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = train_future.result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e947854",
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "ea8f2a91-09b7-41bd-9741-555644dddf0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference = gce.submit(run_model, model, 'O')\n",
    "inference.result()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b350ea4",
   "metadata": {},
   "source": [
    "Now we can chop up the search space into chunks, and invoke `run_model`  once for each chunk of the search space."
   ]
  },
  {
   "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 = [search_space['smiles'][i:i+2036].copy() for i in range(0, search_space.shape[0], 2036)]\n",
    "\n",
    "inference_futures = [gce.submit(run_model, model, chunk) for chunk in chunks]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e008cc2d",
   "metadata": {},
   "source": [
    "Finally, we need to wait for all the inferences to complete and combine the results into a single dataframe. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4135687",
   "metadata": {},
   "outputs": [],
   "source": [
    "def combine_predictions(inference_futures):\n",
    "    dfs = []\n",
    "    for r in inference_futures:\n",
    "        dfs.append(r.result())\n",
    "    predictions = pd.concat(dfs, ignore_index=True)\n",
    "    return predictions\n",
    "\n",
    "predictions = combine_predictions(inference_futures)"
   ]
  },
  {
   "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 workflow that is able to train a model and use it to identify molecules that are likely to be good next choices for simulations. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ad296c",
   "metadata": {},
   "source": [
    "## Model-in-the-Loop Workflow\n",
    "In our final step we will combine the steps from above into a single 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` and `run_model` as above, but after running an iteration it picks the predicted best molecules and runs `compute_vertical` to run new xTB simulations.  The workflow then repeatedly retrains the model using these results until a fixed number of molecule simulations have been trained. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9dc018e",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_count: int = 16   # Number of molecules to evaluate in total\n",
    "batch_size: int = 4  # Number of molecules to evaluate in each batch of simulations\n",
    "\n",
    "with tqdm(total=search_count) as prog_bar: # setup a graphical progress bar\n",
    "    # Mark when we started\n",
    "    start_time = monotonic()\n",
    "    result_data = []\n",
    "    \n",
    "    # Submit with some random guesses\n",
    "    smiles = search_space.sample(batch_size)['smiles']\n",
    "    ie_futures = [gce.submit(compute_vertical, s) for s in smiles]\n",
    "    smiles_futures = dict(zip(ie_futures, smiles)) # Mapping from future to smiles\n",
    "    print(f'Submitted {len(ie_futures)} calculations \\n({smiles})')\n",
    "    already_ran = set()\n",
    "   \n",
    "    # Loop until you finish populating the initial set\n",
    "    while len(ie_futures) > 0: \n",
    "        # First, get the next completed computation from the list\n",
    "        future = next(as_completed(ie_futures))\n",
    "\n",
    "        # Remove it from the list of still-running tasks\n",
    "        ie_futures.remove(future)\n",
    "\n",
    "        # Get the input \n",
    "        smiles = smiles_futures[future]  #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 = gce.submit(compute_vertical, smiles) # launch a new simulation \n",
    "            smiles_futures[new_future] = smiles # Mapping from future to smiles\n",
    "            ie_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",
    "            result_data.append({\n",
    "                'smiles': smiles,\n",
    "                'ie': future.result(),\n",
    "                'batch': 0,\n",
    "                'time': monotonic() - start_time\n",
    "            })\n",
    "            \n",
    "    \n",
    "    # Create the initial training set as a \n",
    "    train_df = pd.DataFrame(result_data)\n",
    "    \n",
    "    # Loop until complete the complete set of molecules \n",
    "    # Selecting molecules guided by the machine learning model\n",
    "    batch = 1\n",
    "    while len(train_df) < search_count:\n",
    "        \n",
    "        # Train and predict as shown in the previous section.\n",
    "        model = gce.submit(train_model, train_df).result()\n",
    "        chunks = [search_space['smiles'][i:i+2036].copy() for i in range(0, search_space.shape[0], 2036)]\n",
    "        inference_futures = [gce.submit(run_model, model, chunk) for chunk in chunks]\n",
    "        predictions = combine_predictions(inference_futures)\n",
    "\n",
    "        # Sort the predictions in descending order, and submit new molecules from them\n",
    "        predictions.sort_values('ie', ascending=False, inplace=True)\n",
    "        ie_futures = []\n",
    "        for smiles in predictions['smiles']:\n",
    "            if smiles not in already_ran:\n",
    "                new_future = gce.submit(compute_vertical, smiles)\n",
    "                ie_futures.append(new_future)\n",
    "                smiles_futures[new_future] = smiles # Mapping from future to smiles\n",
    "                already_ran.add(smiles)\n",
    "                print(f'Submitted new calculation: {smiles}')\n",
    "\n",
    "                if len(ie_futures) >= batch_size:\n",
    "                    print(f'Submitted {len(ie_futures)} calculations')\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(ie_futures):\n",
    "            if future.exception() is None:\n",
    "                prog_bar.update(1)\n",
    "                new_results.append({\n",
    "                    'smiles': smiles_futures[future],\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_df = pd.concat((train_df, 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_df['time'], train_df['ie'])\n",
    "ax.step(train_df['time'], train_df['ie'].cummax(), 'k--')\n",
    "\n",
    "ax.set_xlabel('Walltime (s)')\n",
    "ax.set_ylabel('Ion. Energy (Ha)')\n",
    "\n",
    "fig.tight_layout()"
   ]
  }
 ],
 "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.9"
  },
  "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
}