some minor cleanup and test

.github/workflows/test.yaml

@@ -31,25 +31,4 @@ jobs:
 
     - name: Run tests
       run: |
-        pytest tests
-
-  test-deploy:
-    runs-on: ubuntu-latest
-    needs: test
-    steps:
-    - name: Checkout repository
-      uses: actions/checkout@v2
-
-    - name: Set up Python 3.11
-      uses: actions/setup-python@v2
-      with:
-        python-version: "3.11"
-
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install .[app]
-    
-    - name: Run deploy tests
-      run: |
-        streamlit run serve/app.py
+        pytest tests

mlip_arena/tasks/combustion/water.ipynb

@@ -143,6 +143,17 @@
       "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.18:39213 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
       "\n"
      ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2024-09-23 01:13:02,497 - distributed.scheduler - ERROR - Task md-7-afd01641c2bb41f1a26704e2fedf9ab4-1 marked as failed because 4 workers died while trying to run it\n",
+      "2024-09-23 03:36:38,684 - distributed.scheduler - ERROR - Task md-9-598f71df60d946639d37372b2b5d6dac-1 marked as failed because 4 workers died while trying to run it\n",
+      "2024-09-23 03:37:55,799 - distributed.scheduler - ERROR - Task md-6-17b6c0960c2c488682ec4cc080e23bef-1 marked as failed because 4 workers died while trying to run it\n",
+      "2024-09-23 03:57:30,367 - distributed.scheduler - ERROR - Task md-8-bce1a611143e49fc95b28e9c825fc257-1 marked as failed because 4 workers died while trying to run it\n",
+      "2024-09-23 03:57:59,183 - distributed.scheduler - ERROR - Task md-2-8e508991e36a4969ae355d1b0ed3d34d-1 marked as failed because 4 workers died while trying to run it\n"
+     ]
     }
    ],
    "source": [
@@ -222,7 +233,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {
     "tags": []
    },
@@ -538,6 +549,214 @@
      },
      "metadata": {},
      "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:57:59.188 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'prompt-newt'</span> - Encountered exception during execution:\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py\", line 867, in orchestrate_flow_run\n",
+       "    result = await flow_call.aresult()\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 327, in aresult\n",
+       "    return await asyncio.wrap_future(self.future)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 352, in _run_sync\n",
+       "    result = self.fn(*self.args, **self.kwargs)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/tmp/ipykernel_1275810/1920202683.py\", line 26, in combustion\n",
+       "    return [future.result() for future in futures]\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/tmp/ipykernel_1275810/1920202683.py\", line 26, in &lt;listcomp&gt;\n",
+       "    return [future.result() for future in futures]\n",
+       "            ^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py\", line 228, in result\n",
+       "    return from_sync.call_soon_in_loop_thread(result).result()\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 318, in result\n",
+       "    return self.future.result(timeout=timeout)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 188, in result\n",
+       "    return self.__get_result()\n",
+       "           ^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py\", line 401, in __get_result\n",
+       "    raise self._exception\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 389, in _run_async\n",
+       "    result = await coro\n",
+       "             ^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py\", line 237, in _result\n",
+       "    return await final_state.result(raise_on_failure=raise_on_failure, fetch=True)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py\", line 91, in _get_state_result\n",
+       "    raise await get_state_exception(state)\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 311, in wait\n",
+       "    return await future.result(timeout=timeout)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 336, in _result\n",
+       "    raise exc.with_traceback(tb)\n",
+       "distributed.scheduler.KilledWorker: Attempted to run task 'md-2-8e508991e36a4969ae355d1b0ed3d34d-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.67.15:38671. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "03:57:59.188 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'prompt-newt'\u001b[0m - Encountered exception during execution:\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py\", line 867, in orchestrate_flow_run\n",
+       "    result = await flow_call.aresult()\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 327, in aresult\n",
+       "    return await asyncio.wrap_future(self.future)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 352, in _run_sync\n",
+       "    result = self.fn(*self.args, **self.kwargs)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/tmp/ipykernel_1275810/1920202683.py\", line 26, in combustion\n",
+       "    return [future.result() for future in futures]\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/tmp/ipykernel_1275810/1920202683.py\", line 26, in <listcomp>\n",
+       "    return [future.result() for future in futures]\n",
+       "            ^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py\", line 228, in result\n",
+       "    return from_sync.call_soon_in_loop_thread(result).result()\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 318, in result\n",
+       "    return self.future.result(timeout=timeout)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 188, in result\n",
+       "    return self.__get_result()\n",
+       "           ^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py\", line 401, in __get_result\n",
+       "    raise self._exception\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 389, in _run_async\n",
+       "    result = await coro\n",
+       "             ^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py\", line 237, in _result\n",
+       "    return await final_state.result(raise_on_failure=raise_on_failure, fetch=True)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py\", line 91, in _get_state_result\n",
+       "    raise await get_state_exception(state)\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 311, in wait\n",
+       "    return await future.result(timeout=timeout)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 336, in _result\n",
+       "    raise exc.with_traceback(tb)\n",
+       "distributed.scheduler.KilledWorker: Attempted to run task 'md-2-8e508991e36a4969ae355d1b0ed3d34d-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.67.15:38671. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:58:27.163 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-2' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-2-8e508991e36a4969ae355d1b0ed3d34d-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.67.15:38671. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "03:58:27.163 | \u001b[36mINFO\u001b[0m    | Task run 'md-2' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-2-8e508991e36a4969ae355d1b0ed3d34d-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.67.15:38671. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:58:27.447 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-6' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-6-17b6c0960c2c488682ec4cc080e23bef-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.82.169:45277. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "03:58:27.447 | \u001b[36mINFO\u001b[0m    | Task run 'md-6' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-6-17b6c0960c2c488682ec4cc080e23bef-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.82.169:45277. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:58:27.714 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-7' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-7-afd01641c2bb41f1a26704e2fedf9ab4-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.68.221:42379. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "03:58:27.714 | \u001b[36mINFO\u001b[0m    | Task run 'md-7' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-7-afd01641c2bb41f1a26704e2fedf9ab4-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.68.221:42379. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:58:28.024 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-8' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-8-bce1a611143e49fc95b28e9c825fc257-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.175:34615. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "03:58:28.024 | \u001b[36mINFO\u001b[0m    | Task run 'md-8' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-8-bce1a611143e49fc95b28e9c825fc257-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.175:34615. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:58:28.301 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Task run 'md-9' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-9-598f71df60d946639d37372b2b5d6dac-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.50:40931. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "03:58:28.301 | \u001b[36mINFO\u001b[0m    | Task run 'md-9' - Crash detected! Execution was interrupted by an unexpected exception: KilledWorker: Attempted to run task 'md-9-598f71df60d946639d37372b2b5d6dac-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.66.50:40931. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:58:28.627 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'prompt-newt'</span> - Finished in state <span style=\"color: #d70000; text-decoration-color: #d70000\">Failed</span>(\"Flow run encountered an exception. KilledWorker: Attempted to run task 'md-2-8e508991e36a4969ae355d1b0ed3d34d-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.67.15:38671. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\")\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "03:58:28.627 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'prompt-newt'\u001b[0m - Finished in state \u001b[38;5;160mFailed\u001b[0m(\"Flow run encountered an exception. KilledWorker: Attempted to run task 'md-2-8e508991e36a4969ae355d1b0ed3d34d-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.67.15:38671. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\")\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "ename": "KilledWorker",
+     "evalue": "Attempted to run task 'md-2-8e508991e36a4969ae355d1b0ed3d34d-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.67.15:38671. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see https://distributed.dask.org/en/stable/killed.html.",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKilledWorker\u001b[0m                              Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[43mcombustion\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flows.py:1224\u001b[0m, in \u001b[0;36mFlow.__call__\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1219\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m task_viz_tracker:\n\u001b[1;32m   1220\u001b[0m     \u001b[38;5;66;03m# this is a subflow, for now return a single task and do not go further\u001b[39;00m\n\u001b[1;32m   1221\u001b[0m     \u001b[38;5;66;03m# we can add support for exploring subflows for tasks in the future.\u001b[39;00m\n\u001b[1;32m   1222\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m track_viz_task(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39misasync, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mname, parameters)\n\u001b[0;32m-> 1224\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43menter_flow_run_engine_from_flow_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1225\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1226\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1227\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_for\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1228\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1229\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py:297\u001b[0m, in \u001b[0;36menter_flow_run_engine_from_flow_call\u001b[0;34m(flow, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    290\u001b[0m     retval \u001b[38;5;241m=\u001b[39m from_async\u001b[38;5;241m.\u001b[39mwait_for_call_in_loop_thread(\n\u001b[1;32m    291\u001b[0m         begin_run,\n\u001b[1;32m    292\u001b[0m         done_callbacks\u001b[38;5;241m=\u001b[39mdone_callbacks,\n\u001b[1;32m    293\u001b[0m         contexts\u001b[38;5;241m=\u001b[39mcontexts,\n\u001b[1;32m    294\u001b[0m     )\n\u001b[1;32m    296\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 297\u001b[0m     retval \u001b[38;5;241m=\u001b[39m \u001b[43mfrom_sync\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwait_for_call_in_loop_thread\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    298\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbegin_run\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    299\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdone_callbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdone_callbacks\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    300\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcontexts\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcontexts\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    301\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    303\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m retval\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/api.py:243\u001b[0m, in \u001b[0;36mfrom_sync.wait_for_call_in_loop_thread\u001b[0;34m(_from_sync__call, timeout, done_callbacks, contexts)\u001b[0m\n\u001b[1;32m    241\u001b[0m     stack\u001b[38;5;241m.\u001b[39menter_context(context)\n\u001b[1;32m    242\u001b[0m waiter\u001b[38;5;241m.\u001b[39mwait()\n\u001b[0;32m--> 243\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:318\u001b[0m, in \u001b[0;36mCall.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mresult\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout: Optional[\u001b[38;5;28mfloat\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m    313\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    314\u001b[0m \u001b[38;5;124;03m    Wait for the result of the call.\u001b[39;00m\n\u001b[1;32m    315\u001b[0m \n\u001b[1;32m    316\u001b[0m \u001b[38;5;124;03m    Not safe for use from asynchronous contexts.\u001b[39;00m\n\u001b[1;32m    317\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 318\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:179\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    177\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m CancelledError()\n\u001b[1;32m    178\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;241m==\u001b[39m FINISHED:\n\u001b[0;32m--> 179\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__get_result\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    181\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_condition\u001b[38;5;241m.\u001b[39mwait(timeout)\n\u001b[1;32m    183\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;129;01min\u001b[39;00m [CANCELLED, CANCELLED_AND_NOTIFIED]:\n\u001b[1;32m    184\u001b[0m     \u001b[38;5;66;03m# Raise Prefect cancelled error instead of\u001b[39;00m\n\u001b[1;32m    185\u001b[0m     \u001b[38;5;66;03m# `concurrent.futures._base.CancelledError`\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py:401\u001b[0m, in \u001b[0;36mFuture.__get_result\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    399\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception:\n\u001b[1;32m    400\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 401\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    403\u001b[0m         \u001b[38;5;66;03m# Break a reference cycle with the exception in self._exception\u001b[39;00m\n\u001b[1;32m    404\u001b[0m         \u001b[38;5;28mself\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:389\u001b[0m, in \u001b[0;36mCall._run_async\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    387\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39menforce_async_deadline() \u001b[38;5;28;01mas\u001b[39;00m cancel_scope:\n\u001b[1;32m    388\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 389\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m coro\n\u001b[1;32m    390\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    391\u001b[0m         \u001b[38;5;66;03m# Forget this call's arguments in order to free up any memory\u001b[39;00m\n\u001b[1;32m    392\u001b[0m         \u001b[38;5;66;03m# that may be referenced by them; after a call has happened,\u001b[39;00m\n\u001b[1;32m    393\u001b[0m         \u001b[38;5;66;03m# there's no need to keep a reference to them\u001b[39;00m\n\u001b[1;32m    394\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/client/utilities.py:78\u001b[0m, in \u001b[0;36minject_client.<locals>.with_injected_client\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     76\u001b[0m \u001b[38;5;28;01masync\u001b[39;00m \u001b[38;5;28;01mwith\u001b[39;00m context \u001b[38;5;28;01mas\u001b[39;00m new_client:\n\u001b[1;32m     77\u001b[0m     kwargs\u001b[38;5;241m.\u001b[39msetdefault(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mclient\u001b[39m\u001b[38;5;124m\"\u001b[39m, new_client \u001b[38;5;129;01mor\u001b[39;00m client)\n\u001b[0;32m---> 78\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m fn(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py:400\u001b[0m, in \u001b[0;36mcreate_then_begin_flow_run\u001b[0;34m(flow, parameters, wait_for, return_type, client, user_thread)\u001b[0m\n\u001b[1;32m    398\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m state\n\u001b[1;32m    399\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mresult\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m--> 400\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m state\u001b[38;5;241m.\u001b[39mresult(fetch\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    401\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInvalid return type for flow engine \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mreturn_type\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py:91\u001b[0m, in \u001b[0;36m_get_state_result\u001b[0;34m(state, raise_on_failure)\u001b[0m\n\u001b[1;32m     84\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m UnfinishedRun(\n\u001b[1;32m     85\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRun is in \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstate\u001b[38;5;241m.\u001b[39mtype\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m state, its result is not available.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     86\u001b[0m     )\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure \u001b[38;5;129;01mand\u001b[39;00m (\n\u001b[1;32m     89\u001b[0m     state\u001b[38;5;241m.\u001b[39mis_crashed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_failed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_cancelled()\n\u001b[1;32m     90\u001b[0m ):\n\u001b[0;32m---> 91\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m get_state_exception(state)\n\u001b[1;32m     93\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(state\u001b[38;5;241m.\u001b[39mdata, DataDocument):\n\u001b[1;32m     94\u001b[0m     result \u001b[38;5;241m=\u001b[39m result_from_state_with_data_document(\n\u001b[1;32m     95\u001b[0m         state, raise_on_failure\u001b[38;5;241m=\u001b[39mraise_on_failure\n\u001b[1;32m     96\u001b[0m     )\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py:867\u001b[0m, in \u001b[0;36morchestrate_flow_run\u001b[0;34m(flow, flow_run, parameters, wait_for, interruptible, client, partial_flow_run_context, user_thread)\u001b[0m\n\u001b[1;32m    862\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    863\u001b[0m             from_async\u001b[38;5;241m.\u001b[39mcall_soon_in_new_thread(\n\u001b[1;32m    864\u001b[0m                 flow_call, timeout\u001b[38;5;241m=\u001b[39mflow\u001b[38;5;241m.\u001b[39mtimeout_seconds\n\u001b[1;32m    865\u001b[0m             )\n\u001b[0;32m--> 867\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m flow_call\u001b[38;5;241m.\u001b[39maresult()\n\u001b[1;32m    869\u001b[0m         waited_for_task_runs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m wait_for_task_runs_and_report_crashes(\n\u001b[1;32m    870\u001b[0m             flow_run_context\u001b[38;5;241m.\u001b[39mtask_run_futures, client\u001b[38;5;241m=\u001b[39mclient\n\u001b[1;32m    871\u001b[0m         )\n\u001b[1;32m    872\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m PausedRun \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    873\u001b[0m     \u001b[38;5;66;03m# could get raised either via utility or by returning Paused from a task run\u001b[39;00m\n\u001b[1;32m    874\u001b[0m     \u001b[38;5;66;03m# if a task run pauses, we set its state as the flow's state\u001b[39;00m\n\u001b[1;32m    875\u001b[0m     \u001b[38;5;66;03m# to preserve reschedule and timeout behavior\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:327\u001b[0m, in \u001b[0;36mCall.aresult\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    321\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    322\u001b[0m \u001b[38;5;124;03mWait for the result of the call.\u001b[39;00m\n\u001b[1;32m    323\u001b[0m \n\u001b[1;32m    324\u001b[0m \u001b[38;5;124;03mFor use from asynchronous contexts.\u001b[39;00m\n\u001b[1;32m    325\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    326\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 327\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m asyncio\u001b[38;5;241m.\u001b[39mwrap_future(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture)\n\u001b[1;32m    328\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m asyncio\u001b[38;5;241m.\u001b[39mCancelledError \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    329\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m CancelledError() \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mexc\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:352\u001b[0m, in \u001b[0;36mCall._run_sync\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    350\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39menforce_sync_deadline() \u001b[38;5;28;01mas\u001b[39;00m cancel_scope:\n\u001b[1;32m    351\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 352\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    353\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    354\u001b[0m         \u001b[38;5;66;03m# Forget this call's arguments in order to free up any memory\u001b[39;00m\n\u001b[1;32m    355\u001b[0m         \u001b[38;5;66;03m# that may be referenced by them; after a call has happened,\u001b[39;00m\n\u001b[1;32m    356\u001b[0m         \u001b[38;5;66;03m# there's no need to keep a reference to them\u001b[39;00m\n\u001b[1;32m    357\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "Cell \u001b[0;32mIn[4], line 26\u001b[0m, in \u001b[0;36mcombustion\u001b[0;34m(atoms)\u001b[0m\n\u001b[1;32m      6\u001b[0m     future \u001b[38;5;241m=\u001b[39m MD\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m      7\u001b[0m         atoms\u001b[38;5;241m=\u001b[39matoms,\n\u001b[1;32m      8\u001b[0m         calculator_name\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m         restart\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m[\u001b[49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfuture\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfutures\u001b[49m\u001b[43m]\u001b[49m\n",
+      "Cell \u001b[0;32mIn[4], line 26\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      6\u001b[0m     future \u001b[38;5;241m=\u001b[39m MD\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m      7\u001b[0m         atoms\u001b[38;5;241m=\u001b[39matoms,\n\u001b[1;32m      8\u001b[0m         calculator_name\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m         restart\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m [\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m future \u001b[38;5;129;01min\u001b[39;00m futures]\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py:228\u001b[0m, in \u001b[0;36mPrefectFuture.result\u001b[0;34m(self, timeout, raise_on_failure)\u001b[0m\n\u001b[1;32m    226\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m from_async\u001b[38;5;241m.\u001b[39mcall_soon_in_loop_thread(result)\u001b[38;5;241m.\u001b[39maresult()\n\u001b[1;32m    227\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 228\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfrom_sync\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_soon_in_loop_thread\u001b[49m\u001b[43m(\u001b[49m\u001b[43mresult\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:318\u001b[0m, in \u001b[0;36mCall.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mresult\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout: Optional[\u001b[38;5;28mfloat\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m    313\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    314\u001b[0m \u001b[38;5;124;03m    Wait for the result of the call.\u001b[39;00m\n\u001b[1;32m    315\u001b[0m \n\u001b[1;32m    316\u001b[0m \u001b[38;5;124;03m    Not safe for use from asynchronous contexts.\u001b[39;00m\n\u001b[1;32m    317\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 318\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:188\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    186\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m CancelledError()\n\u001b[1;32m    187\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;241m==\u001b[39m FINISHED:\n\u001b[0;32m--> 188\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__get_result\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    189\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    190\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py:401\u001b[0m, in \u001b[0;36mFuture.__get_result\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    399\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception:\n\u001b[1;32m    400\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 401\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    403\u001b[0m         \u001b[38;5;66;03m# Break a reference cycle with the exception in self._exception\u001b[39;00m\n\u001b[1;32m    404\u001b[0m         \u001b[38;5;28mself\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:389\u001b[0m, in \u001b[0;36mCall._run_async\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    387\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39menforce_async_deadline() \u001b[38;5;28;01mas\u001b[39;00m cancel_scope:\n\u001b[1;32m    388\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 389\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m coro\n\u001b[1;32m    390\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    391\u001b[0m         \u001b[38;5;66;03m# Forget this call's arguments in order to free up any memory\u001b[39;00m\n\u001b[1;32m    392\u001b[0m         \u001b[38;5;66;03m# that may be referenced by them; after a call has happened,\u001b[39;00m\n\u001b[1;32m    393\u001b[0m         \u001b[38;5;66;03m# there's no need to keep a reference to them\u001b[39;00m\n\u001b[1;32m    394\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/futures.py:237\u001b[0m, in \u001b[0;36mPrefectFuture._result\u001b[0;34m(self, timeout, raise_on_failure)\u001b[0m\n\u001b[1;32m    235\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m final_state:\n\u001b[1;32m    236\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCall timed out before task finished.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 237\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m final_state\u001b[38;5;241m.\u001b[39mresult(raise_on_failure\u001b[38;5;241m=\u001b[39mraise_on_failure, fetch\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py:91\u001b[0m, in \u001b[0;36m_get_state_result\u001b[0;34m(state, raise_on_failure)\u001b[0m\n\u001b[1;32m     84\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m UnfinishedRun(\n\u001b[1;32m     85\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRun is in \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstate\u001b[38;5;241m.\u001b[39mtype\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m state, its result is not available.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     86\u001b[0m     )\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure \u001b[38;5;129;01mand\u001b[39;00m (\n\u001b[1;32m     89\u001b[0m     state\u001b[38;5;241m.\u001b[39mis_crashed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_failed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_cancelled()\n\u001b[1;32m     90\u001b[0m ):\n\u001b[0;32m---> 91\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m get_state_exception(state)\n\u001b[1;32m     93\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(state\u001b[38;5;241m.\u001b[39mdata, DataDocument):\n\u001b[1;32m     94\u001b[0m     result \u001b[38;5;241m=\u001b[39m result_from_state_with_data_document(\n\u001b[1;32m     95\u001b[0m         state, raise_on_failure\u001b[38;5;241m=\u001b[39mraise_on_failure\n\u001b[1;32m     96\u001b[0m     )\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py:311\u001b[0m, in \u001b[0;36mDaskTaskRunner.wait\u001b[0;34m(self, key, timeout)\u001b[0m\n\u001b[1;32m    309\u001b[0m future \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_dask_future(key)\n\u001b[1;32m    310\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 311\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m future\u001b[38;5;241m.\u001b[39mresult(timeout\u001b[38;5;241m=\u001b[39mtimeout)\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m distributed\u001b[38;5;241m.\u001b[39mTimeoutError:\n\u001b[1;32m    313\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py:336\u001b[0m, in \u001b[0;36mFuture._result\u001b[0;34m(self, raiseit)\u001b[0m\n\u001b[1;32m    334\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m raiseit:\n\u001b[1;32m    335\u001b[0m     typ, exc, tb \u001b[38;5;241m=\u001b[39m exc\n\u001b[0;32m--> 336\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mwith_traceback(tb)\n\u001b[1;32m    337\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    338\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m exc\n",
+      "\u001b[0;31mKilledWorker\u001b[0m: Attempted to run task 'md-2-8e508991e36a4969ae355d1b0ed3d34d-1' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.67.15:38671. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see https://distributed.dask.org/en/stable/killed.html."
+     ]
     }
    ],
    "source": [

mlip_arena/tasks/diatomics/run.ipynb

@@ -23,6 +23,8 @@
     "\n",
     "from mlip_arena.models.utils import REGISTRY, MLIPEnum\n",
     "\n",
+    "%matplotlib inline\n",
+    "\n",
     "# model_name = \"MACE-MP(M)\"\n",
     "\n",
     "# calc = MLIPEnum[model_name].value()"
@@ -2800,12 +2802,234 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "id": "a0ac2c09-370b-4fdd-bf74-ea5c4ade0215",
    "metadata": {
+    "scrolled": true,
     "tags": []
    },
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== MACE-MP(M) ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "55562f662e56481ba346246ccc109626",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
+      "/tmp/ipykernel_433641/3530606248.py:156: RuntimeWarning: invalid value encountered in scalar divide\n",
+      "  \"tortuosity\": etv / (abs(es[0] - es.min()) + (es[-1] - es.min())),\n",
+      "/tmp/ipykernel_433641/3530606248.py:158: ConstantInputWarning: An input array is constant; the correlation coefficient is not defined.\n",
+      "  \"spearman-descending-force\": stats.spearmanr(rs[iminf:], fs[iminf:]).statistic,\n",
+      "/tmp/ipykernel_433641/3530606248.py:160: ConstantInputWarning: An input array is constant; the correlation coefficient is not defined.\n",
+      "  \"spearman-repulsion-energy\": stats.spearmanr(rs[imine:], es[imine:]).statistic,\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== CHGNet ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3a2235b0309e4dacab34cd8348b5963a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== M3GNet ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "34c651d05a494eb2a0e7988adc9a9835",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
+      "/tmp/ipykernel_433641/3530606248.py:159: ConstantInputWarning: An input array is constant; the correlation coefficient is not defined.\n",
+      "  \"spearman-ascending-force\": stats.spearmanr(rs[:iminf], fs[:iminf]).statistic,\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== ORB ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7d0004a7880d4aeeb42532e4718d34e2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== SevenNet ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8e181eeb14f24447933b76dac47b9514",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
+      "/tmp/ipykernel_433641/3530606248.py:161: ConstantInputWarning: An input array is constant; the correlation coefficient is not defined.\n",
+      "  \"spearman-attraction-energy\": stats.spearmanr(rs[:imine], es[:imine]).statistic,\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== EquiformerV2(OC22) ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0fda48c6de7b47ae93b62fde68263493",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
+      "/tmp/ipykernel_433641/3530606248.py:156: RuntimeWarning: invalid value encountered in scalar divide\n",
+      "  \"tortuosity\": etv / (abs(es[0] - es.min()) + (es[-1] - es.min())),\n",
+      "/tmp/ipykernel_433641/3530606248.py:160: ConstantInputWarning: An input array is constant; the correlation coefficient is not defined.\n",
+      "  \"spearman-repulsion-energy\": stats.spearmanr(rs[imine:], es[imine:]).statistic,\n",
+      "/tmp/ipykernel_433641/3530606248.py:161: ConstantInputWarning: An input array is constant; the correlation coefficient is not defined.\n",
+      "  \"spearman-attraction-energy\": stats.spearmanr(rs[:imine], es[:imine]).statistic,\n",
+      "/tmp/ipykernel_433641/3530606248.py:156: RuntimeWarning: divide by zero encountered in scalar divide\n",
+      "  \"tortuosity\": etv / (abs(es[0] - es.min()) + (es[-1] - es.min())),\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== EquiformerV2(OC20) ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f49debc4e85e4e2aa5c24de17cfa337a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
@@ -2816,7 +3040,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4b2d8cd4d5684d739de0cf3534849fe0",
+       "model_id": "55c068cb90a3463484ffd851c34a6ab8",
        "version_major": 2,
        "version_minor": 0
       },
@@ -2831,9 +3055,69 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/tmp/ipykernel_337193/1177231743.py:116: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
       "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n"
      ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== MACE-OFF(M) ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f23884741dd74e2a9853298cc6c631b5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "========== ALIGNN ==========\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3da8dc9a994a4bd7b95d62413a73306f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/118 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_433641/3530606248.py:164: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
+      "/tmp/ipykernel_433641/3530606248.py:161: ConstantInputWarning: An input array is constant; the correlation coefficient is not defined.\n",
+      "  \"spearman-attraction-energy\": stats.spearmanr(rs[:imine], es[:imine]).statistic,\n"
+     ]
     }
    ],
    "source": [
@@ -2847,7 +3131,7 @@
     "    \n",
     "#     if json_fpath.exists():\n",
     "#         continue\n",
-    "    \n",
+    "\n",
     "    print(f\"========== {model_name} ==========\")\n",
     "    \n",
     "    df = pd.DataFrame(columns=[\n",
@@ -2856,6 +3140,7 @@
     "        \"R\", \"E\", \"F\", \"S^2\", \n",
     "        \"force-flip-times\",\n",
     "        \"force-total-variation\",\n",
+    "        \"force-jump\",\n",
     "        \"energy-diff-flip-times\",\n",
     "        \"energy-grad-norm-max\",\n",
     "        \"energy-jump\",\n",
@@ -2904,6 +3189,8 @@
     "        indices = np.argsort(rs)[::-1]\n",
     "        rs = rs[indices]\n",
     "        es = es[indices]\n",
+    "        eshift = es[0]\n",
+    "        es -= eshift\n",
     "        fs = fs[indices]\n",
     "\n",
     "        iminf = np.argmin(fs)\n",
@@ -2915,10 +3202,20 @@
     "        rounded_fs = np.copy(fs)\n",
     "        rounded_fs[np.abs(rounded_fs) < 1e-2] = 0\n",
     "        fs_sign = np.sign(rounded_fs)\n",
-    "        fs_sign = fs_sign[fs_sign != 0]\n",
+    "        mask = fs_sign != 0\n",
+    "        rounded_fs = rounded_fs[mask]\n",
+    "        fs_sign = fs_sign[mask]\n",
+    "        f_flip = np.diff(fs_sign) != 0\n",
+    "        \n",
+    "        fdiff = np.diff(fs)\n",
+    "        fdiff_sign = np.sign(fdiff)\n",
+    "        mask = fdiff_sign != 0\n",
+    "        fdiff = fdiff[mask]\n",
+    "        fdiff_sign = fdiff_sign[mask]\n",
+    "        fdiff_flip = np.diff(fdiff_sign) != 0\n",
+    "        fjump = np.abs(fdiff[:-1][fdiff_flip]).sum() + np.abs(fdiff[1:][fdiff_flip]).sum()\n",
+    "        \n",
     "\n",
-    "        # rounded_ediff = np.diff(es)\n",
-    "        # rounded_ediff[np.abs(rounded_ediff) < zero_threshold] = 0\n",
     "        ediff = np.diff(es)\n",
     "        ediff[np.abs(ediff) < 1e-3] = 0\n",
     "        ediff_sign = np.sign(ediff)\n",
@@ -2927,6 +3224,40 @@
     "        ediff_sign = ediff_sign[mask]\n",
     "        ediff_flip = np.diff(ediff_sign) != 0\n",
     "        ejump = np.abs(ediff[:-1][ediff_flip]).sum() + np.abs(ediff[1:][ediff_flip]).sum()\n",
+    "        \n",
+    "        \n",
+    "#         edged_es = np.convolve(es, [1, -2, 1], mode='valid')\n",
+    "#         # edged_es[np.abs(edged_es) < 0.1] = 0\n",
+    "#         prob = np.exp(-es[1:-1]) / np.sum(np.exp(-es[1:-1]))\n",
+    "#         edged_es *= prob\n",
+    "#         # edged_es /= np.abs(es[1:-1])\n",
+    "#         ejump = np.linalg.norm(edged_es)\n",
+    "#         ejump = np.abs(edged_es).sum() / 2.0\n",
+    "        \n",
+    "#         edged_fs = np.convolve(fs, [1, -2, 1], mode='valid')\n",
+    "#         # edged_fs[np.abs(edged_fs) < 0.1] = 0\n",
+    "#         edged_fs *= prob\n",
+    "#         fjump = np.linalg.norm(edged_fs)\n",
+    "        # fjump = np.abs(edged_fs).sum() / 2.0\n",
+    "        \n",
+    "#         fig, axes = plt.subplot_mosaic(\n",
+    "#             \"\"\"\n",
+    "#             ac\n",
+    "#             bd\n",
+    "#             \"\"\",\n",
+    "#             constrained_layout=True\n",
+    "#         )\n",
+    "        \n",
+    "\n",
+    "#         axes['a'].plot(rs, es)\n",
+    "#         axes['b'].plot(rs[1:-1], edged_es)\n",
+    "#         # axes['b'].plot(0.5*(rs[1:] + rs[:-1]), np.diff(es))\n",
+    "#         axes['b'].text(0.7, 0.7, f\"{ejump:.3e}\", transform=axes['b'].transAxes)\n",
+    "        \n",
+    "#         axes['c'].plot(rs, fs)\n",
+    "#         axes['d'].plot(rs[1:-1], edged_fs)\n",
+    "#         axes['d'].text(0.7, 0.7, f\"{fjump:.3e}\", transform=axes['d'].transAxes)\n",
+    "        \n",
     "\n",
     "        conservation_deviation = np.mean(np.abs(fs + de_dr))\n",
     "        \n",
@@ -2936,11 +3267,12 @@
     "            \"name\": da,\n",
     "            \"method\": model_name,\n",
     "            \"R\": rs,\n",
-    "            \"E\": es,\n",
+    "            \"E\": es + eshift,\n",
     "            \"F\": fs,\n",
     "            \"S^2\": S2s,\n",
-    "            \"force-flip-times\": np.sum(np.diff(fs_sign)!=0),\n",
+    "            \"force-flip-times\": np.sum(f_flip),\n",
     "            \"force-total-variation\": np.sum(np.abs(np.diff(fs))),\n",
+    "            \"force-jump\": fjump,\n",
     "            \"energy-diff-flip-times\": np.sum(ediff_flip),\n",
     "            \"energy-grad-norm-max\": np.max(np.abs(de_dr)),\n",
     "            \"energy-jump\": ejump,\n",
@@ -2956,7 +3288,6 @@
     "\n",
     "        df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
     "\n",
-    "\n",
     "    json_fpath = Path(REGISTRY[model_name][\"family\"]) / \"homonuclear-diatomics.json\"\n",
     "\n",
     "    if json_fpath.exists():\n",
@@ -2969,7 +3300,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "id": "e0dd4367-3dca-440f-a7a9-7fdd84183f2c",
    "metadata": {
     "tags": []
@@ -3014,6 +3345,7 @@
        "      <th>spearman-repulsion-energy</th>\n",
        "      <th>spearman-attraction-energy</th>\n",
        "      <th>tortuosity</th>\n",
+       "      <th>force-jump</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
@@ -3037,6 +3369,7 @@
        "      <td>-0.986572</td>\n",
        "      <td>0.844786</td>\n",
        "      <td>1.797762</td>\n",
+       "      <td>458.210954</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>119</th>\n",
@@ -3058,6 +3391,7 @@
        "      <td>-0.720218</td>\n",
        "      <td>0.609519</td>\n",
        "      <td>3.921846</td>\n",
+       "      <td>1582.685197</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>120</th>\n",
@@ -3079,6 +3413,7 @@
        "      <td>-0.991010</td>\n",
        "      <td>-0.654715</td>\n",
        "      <td>1.386696</td>\n",
+       "      <td>1721.766498</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>121</th>\n",
@@ -3100,6 +3435,7 @@
        "      <td>-0.989962</td>\n",
        "      <td>0.342476</td>\n",
        "      <td>1.715953</td>\n",
+       "      <td>1775.001526</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>122</th>\n",
@@ -3121,6 +3457,7 @@
        "      <td>-0.985024</td>\n",
        "      <td>0.757880</td>\n",
        "      <td>1.617372</td>\n",
+       "      <td>1799.183469</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>...</th>\n",
@@ -3142,6 +3479,7 @@
        "      <td>...</td>\n",
        "      <td>...</td>\n",
        "      <td>...</td>\n",
+       "      <td>...</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>231</th>\n",
@@ -3163,6 +3501,7 @@
        "      <td>-0.946770</td>\n",
        "      <td>0.159576</td>\n",
        "      <td>2.694147</td>\n",
+       "      <td>3926.252725</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>232</th>\n",
@@ -3184,6 +3523,7 @@
        "      <td>-0.946770</td>\n",
        "      <td>0.159576</td>\n",
        "      <td>2.694147</td>\n",
+       "      <td>3926.252725</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>233</th>\n",
@@ -3205,6 +3545,7 @@
        "      <td>-0.946770</td>\n",
        "      <td>0.159576</td>\n",
        "      <td>2.694147</td>\n",
+       "      <td>3926.252725</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>234</th>\n",
@@ -3226,6 +3567,7 @@
        "      <td>-0.946770</td>\n",
        "      <td>0.159576</td>\n",
        "      <td>2.694147</td>\n",
+       "      <td>3926.252725</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>235</th>\n",
@@ -3247,10 +3589,11 @@
        "      <td>-0.946770</td>\n",
        "      <td>0.159576</td>\n",
        "      <td>2.694147</td>\n",
+       "      <td>3926.252725</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
-       "<p>118 rows × 18 columns</p>\n",
+       "<p>118 rows × 19 columns</p>\n",
        "</div>"
       ],
       "text/plain": [
@@ -3332,23 +3675,36 @@
        "234                  -0.249255                 -0.316556   \n",
        "235                  -0.249255                 -0.316556   \n",
        "\n",
-       "     spearman-repulsion-energy  spearman-attraction-energy  tortuosity  \n",
-       "118                  -0.986572                    0.844786    1.797762  \n",
-       "119                  -0.720218                    0.609519    3.921846  \n",
-       "120                  -0.991010                   -0.654715    1.386696  \n",
-       "121                  -0.989962                    0.342476    1.715953  \n",
-       "122                  -0.985024                    0.757880    1.617372  \n",
-       "..                         ...                         ...         ...  \n",
-       "231                  -0.946770                    0.159576    2.694147  \n",
-       "232                  -0.946770                    0.159576    2.694147  \n",
-       "233                  -0.946770                    0.159576    2.694147  \n",
-       "234                  -0.946770                    0.159576    2.694147  \n",
-       "235                  -0.946770                    0.159576    2.694147  \n",
+       "     spearman-repulsion-energy  spearman-attraction-energy  tortuosity  \\\n",
+       "118                  -0.986572                    0.844786    1.797762   \n",
+       "119                  -0.720218                    0.609519    3.921846   \n",
+       "120                  -0.991010                   -0.654715    1.386696   \n",
+       "121                  -0.989962                    0.342476    1.715953   \n",
+       "122                  -0.985024                    0.757880    1.617372   \n",
+       "..                         ...                         ...         ...   \n",
+       "231                  -0.946770                    0.159576    2.694147   \n",
+       "232                  -0.946770                    0.159576    2.694147   \n",
+       "233                  -0.946770                    0.159576    2.694147   \n",
+       "234                  -0.946770                    0.159576    2.694147   \n",
+       "235                  -0.946770                    0.159576    2.694147   \n",
        "\n",
-       "[118 rows x 18 columns]"
+       "      force-jump  \n",
+       "118   458.210954  \n",
+       "119  1582.685197  \n",
+       "120  1721.766498  \n",
+       "121  1775.001526  \n",
+       "122  1799.183469  \n",
+       "..           ...  \n",
+       "231  3926.252725  \n",
+       "232  3926.252725  \n",
+       "233  3926.252725  \n",
+       "234  3926.252725  \n",
+       "235  3926.252725  \n",
+       "\n",
+       "[118 rows x 19 columns]"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -3386,7 +3742,978 @@
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {
-    "state": {},
+    "state": {
+     "0ab5d0f3c74d4c4cba3a62d642cb99b1": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_db860024663d4cc7859e4808a518cc09",
+       "style": "IPY_MODEL_b1afe346307644d18fda893e52ea93b8",
+       "value": " 118/118 [00:25&lt;00:00,  3.97it/s]"
+      }
+     },
+     "0b6d7e7c471540b7b4b3771488c9a210": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "0c589f4517e446f3b2091c5283072e3d": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "0cdbae6c34704fd98cabfbdd5404c6db": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "0cded1953fb641969da232bc7dda3207": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "0ce144ef22854dc6bee81352022b3074": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_5b05067d2381420e91790696a2cd3519",
+       "max": 118,
+       "style": "IPY_MODEL_5d6e7ddfc71240c5b5ac9274c31b78bd",
+       "value": 118
+      }
+     },
+     "0ed26ebe44e24e5fa07e946503735553": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_f5210ca8dd5c4040a7368211d308abb4",
+       "style": "IPY_MODEL_5ed200ebc48c4d6382a8bec1f8836783",
+       "value": "100%"
+      }
+     },
+     "0fda48c6de7b47ae93b62fde68263493": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_c44f65cce025403ca9d39ad50044ea46",
+        "IPY_MODEL_cfd3ce4dce01407499562d73e1b40540",
+        "IPY_MODEL_7e7ac4bb45b9478f8aa6ce76d4a43b5a"
+       ],
+       "layout": "IPY_MODEL_541074eff2ca4e2cb14a5ae0931d8844"
+      }
+     },
+     "12c376723c444ab1bc7d86aaeebaa299": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_d9119d0d18c04a7eab1531dca1a71f9c",
+       "style": "IPY_MODEL_9a28bb5f413145fc8af90abbc055b2b1",
+       "value": "100%"
+      }
+     },
+     "1403f930c2324c4fa08125ea34f19123": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_23116387d8f64fd38b2af86dc47dd54c",
+       "style": "IPY_MODEL_bb4189f8cd2b4cb38d47be1e46893cd4",
+       "value": "100%"
+      }
+     },
+     "142c7a24cd4b4d898e29f93f4d6195df": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_288343a2cd814c478a62d7cc0e4e92fa",
+       "style": "IPY_MODEL_e047b2428ade4f76a403019d9e559093",
+       "value": " 118/118 [00:26&lt;00:00,  3.37it/s]"
+      }
+     },
+     "18d993ad92214f71aa4607d4f6bb47c6": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "18feade13b3c475782f5541d835eb714": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "1903fb3dfd654bd9b8740e65c4e4ec69": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "1d6ea3ca0867417981c636de639abf9f": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_d343048da6a04360931989f9b92a1fd8",
+       "style": "IPY_MODEL_18feade13b3c475782f5541d835eb714",
+       "value": "100%"
+      }
+     },
+     "1f51b02c7350488892f789cfebce7b1e": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "206172e87b3642c6ac6582b0ebbf6f80": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "20b3b53459634971b8ede3e6e42bdb22": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "23116387d8f64fd38b2af86dc47dd54c": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "24371d8e9f3f49b394b3f9be9eb40295": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "2717f12178e54818864eca87f9374dec": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_ea030b8eacad4f1a9a00fe1060e7f7a6",
+       "style": "IPY_MODEL_8bbcd88274414f01ae1cb8fcc43b1c3e",
+       "value": "100%"
+      }
+     },
+     "288343a2cd814c478a62d7cc0e4e92fa": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "29d3ca80c25b4a8090aa0d931c0fa3f0": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_baa62618ad75474c874991a8ee7efa28",
+       "max": 118,
+       "style": "IPY_MODEL_8cadb0d7721f47daaa500deee5f696a8",
+       "value": 118
+      }
+     },
+     "32692d276650455c9617184478d4be36": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "33917568acc94ed788eb2e40567251f9": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "34c651d05a494eb2a0e7988adc9a9835": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_12c376723c444ab1bc7d86aaeebaa299",
+        "IPY_MODEL_894a816d99da47dfb9ff8aab4c4e8063",
+        "IPY_MODEL_ad102543dcd24813a5cb6778017f1cbb"
+       ],
+       "layout": "IPY_MODEL_c57c7e23c4ee4728824c3582992d1a1b"
+      }
+     },
+     "35552075d3894247810490a5c7e157fb": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "37ee4dbe42f6450181aeefc9d917780c": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "3a2235b0309e4dacab34cd8348b5963a": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_3b4d7acfb2fa4226ae4800873fa58bda",
+        "IPY_MODEL_710afbb6b30c4332ac6a46e7f04f47d4",
+        "IPY_MODEL_142c7a24cd4b4d898e29f93f4d6195df"
+       ],
+       "layout": "IPY_MODEL_70861822edf24a0e80bb52f42fdd80ef"
+      }
+     },
+     "3a53326aa86b4c5ea26f15752b38ae67": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "3b4d7acfb2fa4226ae4800873fa58bda": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_422aa5d6dec748afa06e7c47bc2473a0",
+       "style": "IPY_MODEL_f60c31e316b741b288c553ecc69443b9",
+       "value": "100%"
+      }
+     },
+     "3b57a14b679b41a1b16c9c864a09617b": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "3b626570b96d40b2bccf6874afe72ab2": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "3da8dc9a994a4bd7b95d62413a73306f": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_86e7388e00c44e4aa5eecfdba0eb71bd",
+        "IPY_MODEL_4cbd608c3f5b44838586ef09f19b9f76",
+        "IPY_MODEL_739bae7033024608b64e34a29df54c4c"
+       ],
+       "layout": "IPY_MODEL_6167b7914d194cba8d6860f6a24fd037"
+      }
+     },
+     "422aa5d6dec748afa06e7c47bc2473a0": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "4a693d0251c64c849c8e3c57c4b5bf66": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "4cbd608c3f5b44838586ef09f19b9f76": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_f5dc54554e48411f926a60dd14a4bf61",
+       "max": 118,
+       "style": "IPY_MODEL_76925037e828414387d0eed2623415ed",
+       "value": 118
+      }
+     },
+     "4d744acd8d5042d78257c01108c81c65": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "5088f5dc07ab4ebf8a06c121884171ca": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_a94a375fe34646b892eaee9c4cb81c40",
+       "style": "IPY_MODEL_dc176dfbd3e14fb588b48290061ad612",
+       "value": " 118/118 [00:22&lt;00:00,  4.53it/s]"
+      }
+     },
+     "53c90d01f6de45e899ace60dab91d5b7": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "541074eff2ca4e2cb14a5ae0931d8844": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "5506630107dc4e71a5fd313a844f3976": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "55562f662e56481ba346246ccc109626": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_1d6ea3ca0867417981c636de639abf9f",
+        "IPY_MODEL_dc0b1e9a493f42f293bc5673ac162da7",
+        "IPY_MODEL_0ab5d0f3c74d4c4cba3a62d642cb99b1"
+       ],
+       "layout": "IPY_MODEL_18d993ad92214f71aa4607d4f6bb47c6"
+      }
+     },
+     "55c068cb90a3463484ffd851c34a6ab8": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_0ed26ebe44e24e5fa07e946503735553",
+        "IPY_MODEL_805ebd0e10ed4bfb98728a78028146d3",
+        "IPY_MODEL_d0ef47d1fd78494aa63b4c3538b90a67"
+       ],
+       "layout": "IPY_MODEL_b64e15735fb6459f92f9ff249fc21d0e"
+      }
+     },
+     "5b05067d2381420e91790696a2cd3519": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "5d6e7ddfc71240c5b5ac9274c31b78bd": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "5ed200ebc48c4d6382a8bec1f8836783": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "61361bf4390448a9941dbd37a280e28e": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "6167b7914d194cba8d6860f6a24fd037": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "674cd8b0d6264d0d958d4292fcf31d75": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_0cded1953fb641969da232bc7dda3207",
+       "style": "IPY_MODEL_fcab8f6cc0b8437fb41e4cd6b6953722",
+       "value": "100%"
+      }
+     },
+     "6af74179f34b49a79620fb72411ee763": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "70073ff888d642189a9767a874470b5b": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "70861822edf24a0e80bb52f42fdd80ef": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "710afbb6b30c4332ac6a46e7f04f47d4": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_b1069386e89341d49d790c3315f95a7c",
+       "max": 118,
+       "style": "IPY_MODEL_714eeb184a984cf3bcd1c3e0a7cdbf76",
+       "value": 118
+      }
+     },
+     "714eeb184a984cf3bcd1c3e0a7cdbf76": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "734b0f1c05ec429cbe4887417b6b6b55": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_a89ff36c048747c29d998f519c2efcdb",
+       "style": "IPY_MODEL_35552075d3894247810490a5c7e157fb",
+       "value": " 118/118 [00:34&lt;00:00,  3.14it/s]"
+      }
+     },
+     "739bae7033024608b64e34a29df54c4c": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_37ee4dbe42f6450181aeefc9d917780c",
+       "style": "IPY_MODEL_95f0d8ad3a8e4c55a3db1629f2d537fd",
+       "value": " 118/118 [00:33&lt;00:00,  3.78it/s]"
+      }
+     },
+     "76925037e828414387d0eed2623415ed": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "7cefefaa5c37402faecf27c7973157a5": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_6af74179f34b49a79620fb72411ee763",
+       "max": 118,
+       "style": "IPY_MODEL_ccb72d0da17d4184ae89125db001a9aa",
+       "value": 118
+      }
+     },
+     "7d0004a7880d4aeeb42532e4718d34e2": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_c38d653a12fa46849e28ceabab397e24",
+        "IPY_MODEL_0ce144ef22854dc6bee81352022b3074",
+        "IPY_MODEL_734b0f1c05ec429cbe4887417b6b6b55"
+       ],
+       "layout": "IPY_MODEL_4a693d0251c64c849c8e3c57c4b5bf66"
+      }
+     },
+     "7d6dcb54900b4dc1a657e89159799a9e": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "7e7ac4bb45b9478f8aa6ce76d4a43b5a": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_bf0a54995604496291dcfaa17bbdd504",
+       "style": "IPY_MODEL_7d6dcb54900b4dc1a657e89159799a9e",
+       "value": " 118/118 [00:23&lt;00:00,  4.40it/s]"
+      }
+     },
+     "805ebd0e10ed4bfb98728a78028146d3": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_1f51b02c7350488892f789cfebce7b1e",
+       "max": 118,
+       "style": "IPY_MODEL_61361bf4390448a9941dbd37a280e28e",
+       "value": 118
+      }
+     },
+     "86e7388e00c44e4aa5eecfdba0eb71bd": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_b6d72a6ce2c84a38afe2bc4a2ecf8385",
+       "style": "IPY_MODEL_c9084458f72c48bc8350824fc8a2161b",
+       "value": "100%"
+      }
+     },
+     "894a816d99da47dfb9ff8aab4c4e8063": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_0cdbae6c34704fd98cabfbdd5404c6db",
+       "max": 118,
+       "style": "IPY_MODEL_d42052c447c14dd499f35a3e135af714",
+       "value": 118
+      }
+     },
+     "89c3bd8612c74f71be28aea79cc9237d": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_5506630107dc4e71a5fd313a844f3976",
+       "style": "IPY_MODEL_0c589f4517e446f3b2091c5283072e3d",
+       "value": " 118/118 [00:02&lt;00:00, 42.18it/s]"
+      }
+     },
+     "8b89d0a64778494491fd7d6dcea6b7d9": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "8bbcd88274414f01ae1cb8fcc43b1c3e": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "8cadb0d7721f47daaa500deee5f696a8": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "8e181eeb14f24447933b76dac47b9514": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_674cd8b0d6264d0d958d4292fcf31d75",
+        "IPY_MODEL_e849d7345f1043158c1d8b7c4c5fc9df",
+        "IPY_MODEL_92ffef207b554a799cf754b9f33387e0"
+       ],
+       "layout": "IPY_MODEL_f89ba8b9da494ec0972d6fc800aefa0a"
+      }
+     },
+     "92ffef207b554a799cf754b9f33387e0": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_33917568acc94ed788eb2e40567251f9",
+       "style": "IPY_MODEL_a036de6b54d74784a9bac32af0db9873",
+       "value": " 118/118 [00:25&lt;00:00,  3.88it/s]"
+      }
+     },
+     "95f0d8ad3a8e4c55a3db1629f2d537fd": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "961228b9019d4b62a27fc84e373db5b1": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "9a28bb5f413145fc8af90abbc055b2b1": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "a036de6b54d74784a9bac32af0db9873": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "a89ff36c048747c29d998f519c2efcdb": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "a94a375fe34646b892eaee9c4cb81c40": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "ad102543dcd24813a5cb6778017f1cbb": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_20b3b53459634971b8ede3e6e42bdb22",
+       "style": "IPY_MODEL_206172e87b3642c6ac6582b0ebbf6f80",
+       "value": " 118/118 [00:24&lt;00:00,  4.12it/s]"
+      }
+     },
+     "b1069386e89341d49d790c3315f95a7c": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "b1afe346307644d18fda893e52ea93b8": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "b64e15735fb6459f92f9ff249fc21d0e": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "b6d72a6ce2c84a38afe2bc4a2ecf8385": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "b753d3dd93744e0aaec7dabe187e1b7d": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "baa62618ad75474c874991a8ee7efa28": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "bb4189f8cd2b4cb38d47be1e46893cd4": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "bf0a54995604496291dcfaa17bbdd504": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "c38d653a12fa46849e28ceabab397e24": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_32692d276650455c9617184478d4be36",
+       "style": "IPY_MODEL_53c90d01f6de45e899ace60dab91d5b7",
+       "value": "100%"
+      }
+     },
+     "c44f65cce025403ca9d39ad50044ea46": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_f4fa13dca645493984620cc32469e31b",
+       "style": "IPY_MODEL_961228b9019d4b62a27fc84e373db5b1",
+       "value": "100%"
+      }
+     },
+     "c57c7e23c4ee4728824c3582992d1a1b": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "c9084458f72c48bc8350824fc8a2161b": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "ccb72d0da17d4184ae89125db001a9aa": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "cfd3ce4dce01407499562d73e1b40540": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_24371d8e9f3f49b394b3f9be9eb40295",
+       "max": 118,
+       "style": "IPY_MODEL_70073ff888d642189a9767a874470b5b",
+       "value": 118
+      }
+     },
+     "d0ef47d1fd78494aa63b4c3538b90a67": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_3b626570b96d40b2bccf6874afe72ab2",
+       "style": "IPY_MODEL_8b89d0a64778494491fd7d6dcea6b7d9",
+       "value": " 118/118 [00:22&lt;00:00,  4.55it/s]"
+      }
+     },
+     "d343048da6a04360931989f9b92a1fd8": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "d42052c447c14dd499f35a3e135af714": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "d9119d0d18c04a7eab1531dca1a71f9c": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "db860024663d4cc7859e4808a518cc09": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "dc0b1e9a493f42f293bc5673ac162da7": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_1903fb3dfd654bd9b8740e65c4e4ec69",
+       "max": 118,
+       "style": "IPY_MODEL_3b57a14b679b41a1b16c9c864a09617b",
+       "value": 118
+      }
+     },
+     "dc176dfbd3e14fb588b48290061ad612": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "e047b2428ade4f76a403019d9e559093": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "e849d7345f1043158c1d8b7c4c5fc9df": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_0b6d7e7c471540b7b4b3771488c9a210",
+       "max": 118,
+       "style": "IPY_MODEL_4d744acd8d5042d78257c01108c81c65",
+       "value": 118
+      }
+     },
+     "ea030b8eacad4f1a9a00fe1060e7f7a6": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "f23884741dd74e2a9853298cc6c631b5": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_2717f12178e54818864eca87f9374dec",
+        "IPY_MODEL_29d3ca80c25b4a8090aa0d931c0fa3f0",
+        "IPY_MODEL_89c3bd8612c74f71be28aea79cc9237d"
+       ],
+       "layout": "IPY_MODEL_b753d3dd93744e0aaec7dabe187e1b7d"
+      }
+     },
+     "f49debc4e85e4e2aa5c24de17cfa337a": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_1403f930c2324c4fa08125ea34f19123",
+        "IPY_MODEL_7cefefaa5c37402faecf27c7973157a5",
+        "IPY_MODEL_5088f5dc07ab4ebf8a06c121884171ca"
+       ],
+       "layout": "IPY_MODEL_3a53326aa86b4c5ea26f15752b38ae67"
+      }
+     },
+     "f4fa13dca645493984620cc32469e31b": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "f5210ca8dd5c4040a7368211d308abb4": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "f5dc54554e48411f926a60dd14a4bf61": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "f60c31e316b741b288c553ecc69443b9": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "f89ba8b9da494ec0972d6fc800aefa0a": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "fcab8f6cc0b8437fb41e4cd6b6953722": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     }
+    },
     "version_major": 2,
     "version_minor": 0
    }

serve/tasks/combustion.py

@@ -145,14 +145,44 @@ fig.add_trace(
 fig.update_layout(
     title="Hydrogen Combustion (2H2 + O2 -> 2H2O, 64 units)",
     xaxis_title="Timestep",
-    yaxis_title="Temperatures",
-    yaxis2=dict(
-        title="Product Percentage (%)",
-        overlaying="y",
-        side="right",
-        range=[0, 100],
-        tickmode="sync",
-    ),
+    yaxis_title="Temperature (K)",
+    # yaxis2=dict(
+    #     title="Product Percentage (%)",
+    #     overlaying="y",
+    #     side="right",
+    #     range=[0, 100],
+    #     tickmode="sync",
+    # ),
+    # template="plotly_dark",
+)
+
+st.plotly_chart(fig)
+
+# Energy
+
+fig = go.Figure()
+
+for method in df["method"].unique():
+    row = df[df["method"] == method].iloc[0]
+    fig.add_trace(
+        go.Scatter(
+            x=row["timestep"],
+            y=np.array(row["energies"]) - row["energies"][0],
+            mode="lines",
+            name=method,
+            line=dict(
+                color=method_color_mapping[method],
+                # width=1
+            ),
+            marker=dict(color=method_color_mapping[method], size=3),
+            showlegend=True,
+        ),
+    )
+
+fig.update_layout(
+    title="Hydrogen Combustion (2H2 + O2 -> 2H2O, 64 units)",
+    xaxis_title="Timestep",
+    yaxis_title="ΔE (eV)",
     # template="plotly_dark",
 )
 
@@ -213,7 +243,7 @@ st.plotly_chart(fig)
 st.markdown("""### Center of mass drift
 
 The center of mass (COM) drift is a measure of the stability of the simulation. A well-behaved simulation should have a COM drift close to zero. The COM drift is calculated as the displacement of the COM of the system from the initial position.
-            """)
+""")
 
 
 @st.cache_data
@@ -229,7 +259,9 @@ def get_com_drifts(df):
     df_flat = df_exploded.drop(columns=["com_drifts"])
 
     df_flat["total_com_drift"] = np.sqrt(
-        df_flat["com_drift_x"] ** 2 + df_flat["com_drift_y"] ** 2 + df_flat["com_drift_z"] ** 2
+        df_flat["com_drift_x"] ** 2
+        + df_flat["com_drift_y"] ** 2
+        + df_flat["com_drift_z"] ** 2
     )
 
     return df_flat
@@ -290,7 +322,11 @@ def draw_com_drifts_plot():
         x="com_drift_x",
         y="com_drift_y",
         z="com_drift_z",
-        labels={"com_drift_x": "Δx (Å)", "com_drift_y": "Δy (Å)", "com_drift_z": "Δz (Å)"},
+        labels={
+            "com_drift_x": "Δx (Å)",
+            "com_drift_y": "Δy (Å)",
+            "com_drift_z": "Δz (Å)",
+        },
         category_orders={"method": df_exploded["method"].unique()},
         color_discrete_sequence=[
             method_color_mapping[method] for method in df_exploded["method"].unique()
@@ -313,31 +349,33 @@ def draw_com_drifts_plot():
             bgcolor="rgba(0, 0, 0, 0)",
         ),
     )
-    fig.add_traces([
-        go.Scatter3d(
-            x=[0],
-            y=[0],
-            z=[0],
-            mode="markers",
-            marker=dict(size=2, color="white"),
-            name="Origin",
-        ),
-        # add last point of each method and annotate the total drift
-        go.Scatter3d(
-            # df_filtered.groupby("method")["com_drift_x"].last(),
-            x=df_filtered.groupby("method")["com_drift_x"].last(),
-            y=df_filtered.groupby("method")["com_drift_y"].last(),
-            z=df_filtered.groupby("method")["com_drift_z"].last(),
-            mode="markers+text",
-            marker=dict(size=2, color="rgba(0, 0, 0, 128)"),
-            text=df_filtered.groupby("method")["total_com_drift"].last().round(3),
-            # size=5,
-            name="total drifts",
-            textposition="top center",
-        ),
-    ])
+    fig.add_traces(
+        [
+            go.Scatter3d(
+                x=[0],
+                y=[0],
+                z=[0],
+                mode="markers",
+                marker=dict(size=3, color="white"),
+                name="origin",
+            ),
+            # add last point of each method and annotate the total drift
+            go.Scatter3d(
+                # df_filtered.groupby("method")["com_drift_x"].last(),
+                x=df_filtered.groupby("method")["com_drift_x"].last(),
+                y=df_filtered.groupby("method")["com_drift_y"].last(),
+                z=df_filtered.groupby("method")["com_drift_z"].last(),
+                mode="markers+text",
+                marker=dict(size=3, color="white", opacity=0.5),
+                text=df_filtered.groupby("method")["total_com_drift"].last().round(3),
+                # size=5,
+                name="total drifts",
+                textposition="top center",
+            ),
+        ]
+    )
 
     st.plotly_chart(fig)
 
 
-draw_com_drifts_plot()
+draw_com_drifts_plot()