From 88f8ccfeadce5aa83132c8e8857b78d895ba2657 Mon Sep 17 00:00:00 2001
From: Xavier Weiss <xavierw@kth.se>
Date: Mon, 4 Nov 2024 19:12:57 +0100
Subject: [PATCH] Test: Add Entry for Allow_Shedding

Signed-off-by: Xavier Weiss <xavierw@kth.se>
---
 getting_started/13_Shedding.ipynb | 28 ++++------------------------
 grid2op/Backend/backend.py        |  1 +
 grid2op/Space/GridObjects.py      |  6 +++---
 grid2op/tests/test_Observation.py |  1 +
 4 files changed, 9 insertions(+), 27 deletions(-)

diff --git a/getting_started/13_Shedding.ipynb b/getting_started/13_Shedding.ipynb
index e5a11cbe..78605dcf 100644
--- a/getting_started/13_Shedding.ipynb
+++ b/getting_started/13_Shedding.ipynb
@@ -7,19 +7,7 @@
     "### Shedding\n",
     "In emergency conditions, it may be possible / necessary for a grid operator to disconnect certain loads, generators, or other components in order to prevent a larger blackout. This notebook explores how this can be achieved in Grid2OP. \n",
     "\n",
-    "By default shedding is disabled in all environments, to enable it provide set the environment parameter ALLOW_SHEDDING to True.\n",
-    "\n",
-    "Shed load:\n",
-    "* **Vector** :: $\\text{n\\_load}$\n",
-    "* **List** :: $[(\\text{load\\_id}, \\text{status})]$\n",
-    "* **Dictionary** :: $\\text{load\\_name}: \\text{status}$\n",
-    "\n",
-    "Shed generator:\n",
-    "* **Vector** :: $\\text{n\\_gen}$\n",
-    "* **List** :: $[(\\text{gen\\_id}, \\text{status})]$\n",
-    "* **Dictionary** :: $\\text{gen\\_name}: \\text{status}$\n",
-    "\n",
-    "Where $\\text{status}$ is a boolean (True/False)."
+    "By default shedding is disabled in all environments, to provide the keyword argument allow_shedding when initializing the environment."
    ]
   },
   {
@@ -38,7 +26,7 @@
     "from pathlib import Path\n",
     "\n",
     "data_path = Path.cwd() / \"grid2op\" / \"data\"\n",
-    "env = grid2op.make(data_path / \"rte_case5_example\", allow_shedding=True)\n",
+    "env = grid2op.make(data_path / \"rte_case5_example\", backend=LightSimBackend(), allow_shedding=True)\n",
     "plotter = PlotMatplot(env.observation_space, load_name=True, gen_name=True, dpi=150)\n",
     "print(f\"Loads: {env.n_load}, Generators: {env.n_gen}, Storage: {env.n_storage}, Allow Shedding: {env.allow_shedding}\")"
    ]
@@ -51,7 +39,9 @@
    "source": [
     "# Disconnect the load at substation 4\n",
     "load_lookup = {name:i for i,name in enumerate(env.name_load)}\n",
+    "gen_lookup = {name:i for i,name in enumerate(env.name_gen)}\n",
     "act = env.action_space({\"set_bus\":[(env.load_pos_topo_vect[load_lookup[\"load_4_2\"]], -1)]})\n",
+    "act = env.action_space({\"set_bus\":[(env.gen_pos_topo_vect[gen_lookup[\"gen_0_0\"]], -1)]})\n",
     "print(act)\n",
     "env.set_id(\"00\")\n",
     "init_obs = env.reset()\n",
@@ -60,16 +50,6 @@
     "plt.show()"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "load_shed_mask = (init_obs + act).topo_vect[env.load_pos_topo_vect] == -1\n",
-    "gen_shed_mask = (init_obs + act).topo_vect[env.gen_pos_topo_vect] == -1"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
diff --git a/grid2op/Backend/backend.py b/grid2op/Backend/backend.py
index 7e535af3..378945d3 100644
--- a/grid2op/Backend/backend.py
+++ b/grid2op/Backend/backend.py
@@ -247,6 +247,7 @@ def set_shedding(self, allow_shedding:bool=False):
         """
         Override if the Backend supports shedding.
         """
+        
         if allow_shedding:
             raise BackendError("Backend does not support shedding")
         else:
diff --git a/grid2op/Space/GridObjects.py b/grid2op/Space/GridObjects.py
index d9ce4cf1..5d854318 100644
--- a/grid2op/Space/GridObjects.py
+++ b/grid2op/Space/GridObjects.py
@@ -3930,6 +3930,9 @@ def _make_cls_dict(cls, res, as_list=True, copy_=True, _topo_vect_only=False):
         else:
             res["name_shunt"] = None
             res["shunt_to_subid"] = None
+
+        # Shedding
+        save_to_dict(res, cls, "allow_shedding", bool, copy_)
             
         if not _topo_vect_only:
             # all the attributes bellow are not needed for the "first call"
@@ -4072,9 +4075,6 @@ def _make_cls_dict(cls, res, as_list=True, copy_=True, _topo_vect_only=False):
             save_to_dict(
                 res, cls, "alertable_line_ids", (lambda li: [int(el) for el in li])  if as_list else None, copy_
             )
-            
-            # Shedding
-            save_to_dict(res, cls, "allow_shedding", str, copy_)
 
             # avoid further computation and save it
             if not as_list:
diff --git a/grid2op/tests/test_Observation.py b/grid2op/tests/test_Observation.py
index dff0b205..ed8feb0a 100644
--- a/grid2op/tests/test_Observation.py
+++ b/grid2op/tests/test_Observation.py
@@ -299,6 +299,7 @@ def setUp(self):
             "alertable_line_ids": [],
             "assistant_warning_type": None,
             "_PATH_GRID_CLASSES": None,
+            "allow_shedding": False,
         }
 
         self.json_ref = {