diff --git a/notebooks/capacity.ipynb b/notebooks/capacity.ipynb
index 463398d..3e3a412 100644
--- a/notebooks/capacity.ipynb
+++ b/notebooks/capacity.ipynb
@@ -67,16 +67,30 @@
    "source": [
     "# The assets.csv file contains info about which assets were invested in and when\n",
     "assets = pd.read_csv(OUTPUT_DIR / \"assets.csv\")\n",
+    "# The asset_capacities.csv file contains info about the capacities for each asset\n",
+    "# along the simulation\n",
+    "asset_capacities = pd.read_csv(OUTPUT_DIR / \"asset_capacities.csv\")\n",
     "\n",
-    "# Assets with no decommission_year are effectively decommissioned after time horizon\n",
-    "assets[\"decommission_year\"] = assets[\"decommission_year\"].fillna(years[-1] + 1)\n",
+    "\n",
+    "# We define a helper function to bring some useful information from 'assets' into\n",
+    "#'assets_capacity'.\n",
+    "def get_agent_and_process(x: pd.Series) -> pd.Series:\n",
+    "    \"\"\"Collects \"agent_id\", \"process_id\", \"commission_year\" from assets.\"\"\"\n",
+    "    col, val = (\n",
+    "        (\"asset_id\", x.asset_id) if x.asset_id is not None else (\"group_id\", x.group_id)\n",
+    "    )\n",
+    "    row = assets[assets[col] == val].iloc[0]\n",
+    "    return row[[\"agent_id\", \"process_id\", \"commission_year\"]]\n",
+    "\n",
+    "\n",
+    "asset_capacities = pd.concat(\n",
+    "    [asset_capacities, asset_capacities.apply(get_agent_and_process, axis=1)], axis=1\n",
+    ")\n",
     "\n",
     "# Calculate capacity for each type of process for each agent\n",
     "capacity = pd.DataFrame()\n",
     "for year in years:\n",
-    "    active = assets[\n",
-    "        (year >= assets[\"commission_year\"]) & (year < assets[\"decommission_year\"])\n",
-    "    ]\n",
+    "    active = asset_capacities[year >= asset_capacities[\"commission_year\"]]\n",
     "\n",
     "    # This only works because each agent is responsible for one and only one commodity\n",
     "    cap_sum = active.groupby([\"agent_id\", \"process_id\"])[\"capacity\"].sum().reset_index()\n",
@@ -112,7 +126,7 @@
     "import matplotlib.pyplot as plt\n",
     "\n",
     "agents = capacity[\"agent_id\"].unique()\n",
-    "_, axes = plt.subplots(1, len(agents))\n",
+    "_, axes = plt.subplots(1, len(agents), figsize=(4 * len(agents), 4))\n",
     "for ax, agent in zip(axes, agents):\n",
     "    capacity[capacity[\"agent_id\"] == agent].pivot(\n",
     "        index=\"year\", columns=\"process_id\", values=\"capacity\"\n",
@@ -120,13 +134,33 @@
     "    ax.set_title(agent)\n",
     "    ax.set_xlabel(\"Year\")\n",
     "    ax.set_ylabel(\"Capacity\")\n",
-    "    ax.legend(title=\"Process\")"
+    "    ax.legend(\n",
+    "        title=\"Process\", bbox_to_anchor=(1.05, 1), loc=\"upper left\", borderaxespad=0.0\n",
+    "    )\n",
+    "\n",
+    "plt.tight_layout()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "muse2-data-analysis",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -140,7 +174,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.14.2"
+   "version": "3.14.6"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/prices.ipynb b/notebooks/prices.ipynb
index 12deee7..94c58f3 100644
--- a/notebooks/prices.ipynb
+++ b/notebooks/prices.ipynb
@@ -50,13 +50,23 @@
     "\n",
     "ax.set_xlabel(\"Milestone year\")\n",
     "ax.set_ylabel(\"Price\")\n",
-    "ax.legend(title=\"Time slice\");"
+    "ax.legend(\n",
+    "    title=\"Time slice\", bbox_to_anchor=(1.05, 1), loc=\"upper left\", borderaxespad=0.0\n",
+    ");"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "muse2-data-analysis",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -70,7 +80,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.14.2"
+   "version": "3.14.6"
   }
  },
  "nbformat": 4,