Topology-Only Graph of Grid State

[DEUCE1957]

Is your feature request related to a problem? Please describe.

In the documentation "A grid, a graph: grid2op representation of the powergrid", there are currently 5 ways of graphs/matrix representations of an observation in Grid2Op:

Type of graph	Same no. of Nodes	Encodes all of the Observation	Has Flow Info	Format
“energy graph”	no	almost	yes	Networkx DiGraph
“elements graph”	yes	yes	yes	Networkx DiGraph
“connectivity graph”	yes	no	no	Scipy Sparse Matrix
“bus connectivity graph”	no	no	no	Scipy Sparse Matrix
“flow bus graph”	no	no	yes	Scipy Sparse Matrix

While this seems to be complete, in certain cases it would be convenient to have a Networkx Graph that does NOT encode power flows, but only topology. I will refer to this as a topology graph. The advantage of this is that it should be quicker/easier to compute than the energy/elements graphs, but is still in a networkx format (so we can more conveniently do graph operations / algorithms on it).

Describe the solution you'd like

The following code recreates the energy graph without any power flow / extra attributes, only topological ones (connected / disconnected). This is a standalone class, but I would NOT recommend using this implementation directly. Instead there could be an optional argument when generating the energy graph that skips all power flow-related steps? Note the add_act method is identical to the one suggested in Issue 681

class TopoGraphConverter():
    
    def __init__(self, env:BaseEnv):
        self.n_busbar_per_sub = env.n_busbar_per_sub
        self.name_sub = env.name_sub
        self.n_sub = env.n_sub
        self._topo_vect_to_sub = env._topo_vect_to_sub
        self.local_bus_to_global = env.local_bus_to_global
        self.detachment_is_allowed = env.detachment_is_allowed
        self.name_load = env.name_load
        self.n_load = env.n_load
        self.load_to_subid = env.load_to_subid
        self.name_gen = env.name_gen
        self.n_gen = env.n_gen
        self.gen_to_subid = env.gen_to_subid
        self.name_line = env.name_line
        self.n_line = env.n_line
        self.line_or_to_subid = env.line_or_to_subid
        self.line_ex_to_subid = env.line_ex_to_subid
        self.name_storage = env.name_storage
        self.n_storage = env.n_storage
        self.storage_to_subid = env.storage_to_subid
        self.name_shunt = env.name_shunt
        self.n_shunt = env.n_shunt
        self.shunt_to_subid = env.shunt_to_subid
        self.NB_TIMESTEP_OVERFLOW_ALLOWED = env.parameters.NB_TIMESTEP_OVERFLOW_ALLOWED
        
    def add_act(self, obs:BaseObservation, act:BaseAction, issue_warn=True):
        cls = type(obs)
        cls_act = type(act)
        
        act = copy.deepcopy(act)
        res = cls()
        res.set_game_over(env=None)
        res.topo_vect[:] = obs.topo_vect
        res.line_status[:] = obs.line_status
        res.timestep_overflow[:] = obs.timestep_overflow
        res.timestep_overflow[obs.rho > 1.0] += 1
        overflow_mask = res.timestep_overflow > self.NB_TIMESTEP_OVERFLOW_ALLOWED
        res.line_status[overflow_mask] = -1
        line_or_topo_vect = res.topo_vect[obs.line_or_pos_topo_vect]
        line_or_topo_vect[overflow_mask] = -1
        line_ex_topo_vect = res.topo_vect[obs.line_ex_pos_topo_vect]
        line_ex_topo_vect[overflow_mask] = -1
        res.topo_vect[obs.line_or_pos_topo_vect] = line_or_topo_vect
        res.topo_vect[obs.line_ex_pos_topo_vect] = line_ex_topo_vect
        
        # If a powerline has been reconnected without specific bus, issue a warning
        if "set_line_status" in cls_act.authorized_keys:
            obs._aux_add_act_set_line_status(cls, cls_act, act, res, issue_warn)
            
        # topo vect
        if "set_bus" in cls_act.authorized_keys:
            res.topo_vect[act.set_bus != 0] = act.set_bus[act.set_bus != 0]

        if "change_bus" in cls_act.authorized_keys:
            do_change_bus_on = act.change_bus & (
                res.topo_vect > 0
            )  # change bus of elements that were on
            res.topo_vect[do_change_bus_on] = 3 - res.topo_vect[do_change_bus_on]

        # topo vect: reco of powerline that should be
        res.line_status = (res.topo_vect[cls.line_or_pos_topo_vect] >= 1) & (
            res.topo_vect[cls.line_ex_pos_topo_vect] >= 1
        )

        # powerline status
        if "set_line_status" in cls_act.authorized_keys:
            obs._aux_add_act_set_line_status2(cls, cls_act, act, res, issue_warn)
            
        if "change_line_status" in cls_act.authorized_keys:
            obs._aux_add_act_change_line_status2(cls, cls_act, act, res, issue_warn)
        return res
    
    def _aux_add_edges(self, el_ids, el_global_bus,
                       nb_el,el_connected,el_name,
                       edges_prop, graph):
        edges_el = [(el_ids[el_id], self.n_sub + el_global_bus[el_id]) if el_connected[el_id] else None
                    for el_id in range(nb_el) 
                   ]
        li_el_edges = [(*edges_el[el_id],
                        {"id": el_id,
                         "type": f"{el_name}_to_bus"})
                      for el_id in range(nb_el)
                      if el_connected[el_id]]
        if edges_prop is not None:
            ed_num = 0  # edge number
            for el_id in range(nb_el):
                if not el_connected[el_id]:
                    continue
                for prop_nm, prop_vect in edges_prop:
                    li_el_edges[ed_num][-1][prop_nm] = prop_vect[el_id]
                ed_num += 1        
        graph.add_edges_from(li_el_edges)
        return li_el_edges

    def _aux_add_el_to_comp_graph(self, graph,
                                  first_id,el_names_vect,
                                  el_name, nb_el,
                                  el_bus=None, el_to_sub_id=None,
                                  nodes_prop=None, edges_prop=None):  
        if el_bus is None and el_to_sub_id is not None:
            raise Grid2OpException("el_bus is None and el_to_sub_id is not None")
        
        if el_bus is not None and el_to_sub_id is None:
            raise Grid2OpException("el_bus is not None and el_to_sub_id is None")
        
        # add the nodes for the elements of this types
        el_ids = first_id + np.arange(nb_el)
        
        # add the properties for these nodes
        li_el_node = [(el_ids[el_id],
                        {"id": el_id,
                         "type": f"{el_name}",
                         "name": el_names_vect[el_id]
                        })
                      for el_id in range(nb_el)]
        if el_bus is not None:
            el_global_bus = self.local_bus_to_global(el_bus,
                                                    el_to_sub_id)
            el_connected = np.array(el_global_bus) >= 0
            for el_id in range(nb_el):
                li_el_node[el_id][-1]["connected"] = el_connected[el_id]

        if nodes_prop is not None:
            for el_id in range(nb_el):
                for prop_nm, prop_vect in nodes_prop:
                    li_el_node[el_id][-1][prop_nm] = prop_vect[el_id]
        
        if el_bus is None and el_to_sub_id is None:
            graph.add_nodes_from(li_el_node)
            return el_ids

        # Add the edges
        self._aux_add_edges(
            el_ids, el_global_bus, nb_el,
            el_connected, el_name, edges_prop,graph)
            
        graph.add_nodes_from(li_el_node)
        return el_ids
        
    def _aux_get_connected_buses(self, obs:BaseObservation):
        res = np.full(self.n_busbar_per_sub * self.n_sub, fill_value=False)
        global_bus = self.local_bus_to_global(obs.topo_vect,
                                                  self._topo_vect_to_sub)
        res[global_bus[global_bus != -1]] = True
        return res

    def _aux_add_buses(self, obs:BaseObservation, graph, first_id):
        bus_ids = first_id + np.arange(self.n_busbar_per_sub * self.n_sub)
        conn_bus = self._aux_get_connected_buses(obs)
        bus_li = [
            (bus_ids[bus_id],
                {"id": bus_id, 
                 "connected": conn_bus[bus_id]})
            for bus_id in range(self.n_busbar_per_sub * self.n_sub)
        ]
        graph.add_nodes_from(bus_li)
        edge_bus_li = [(bus_id,
                        bus_id % self.n_sub,
                        {"type": "bus_to_substation"})
                        for bus_id in bus_ids]
        graph.add_edges_from(edge_bus_li)
        return bus_ids

    def _aux_add_loads(self, obs:BaseObservation, graph, first_id):
        load_ids = self._aux_add_el_to_comp_graph(
            graph, first_id, self.name_load,
            "load", self.n_load, obs.load_bus, self.load_to_subid)
        return load_ids

    def _aux_add_gens(self, obs:BaseObservation, graph, first_id):
        gen_ids = self._aux_add_el_to_comp_graph(
            graph, first_id,  self.name_gen,
            "gen", self.n_gen, obs.gen_bus, self.gen_to_subid)
        return gen_ids
        
    def _aux_add_edge_line_side(self, graph, bus,
                                sub_id, line_node_ids,
                                side):
        global_bus = self.local_bus_to_global(bus, sub_id)
        conn_ = np.array(global_bus) >= 0
        res = self._aux_add_edges(
            line_node_ids, global_bus, self.n_line,
            conn_, "line", edges_prop=None, graph=graph)
        return res

    def _aux_add_lines(self, obs:BaseObservation, graph, first_id):        
        nodes_prop = [("connected", obs.line_status),
                      ("timestep_overflow", obs.timestep_overflow)]
        
        lin_ids = self._aux_add_el_to_comp_graph(
            graph, first_id, self.name_line,
            "line", self.n_line, el_bus=None,
            el_to_sub_id=None, nodes_prop=nodes_prop)
        
        # Add "or" edges
        self._aux_add_edge_line_side(graph, obs.line_or_bus,
                                     self.line_or_to_subid, lin_ids, "or")
        
        # Add "ex" edges
        self._aux_add_edge_line_side(graph, obs.line_ex_bus,
                                     self.line_ex_to_subid, lin_ids, "ex")
        return lin_ids

    def _aux_add_storages(self, obs:BaseObservation, graph, first_id):
        sto_ids = self._aux_add_el_to_comp_graph(
            graph, first_id, self.name_storage,
            "storage", self.n_storage, obs.storage_bus,
            el_to_sub_id=self.storage_to_subid)
        return sto_ids

    def _aux_add_shunts(self, obs:BaseObservation, graph, first_id): 
        sto_ids = self._aux_add_el_to_comp_graph(
            graph, first_id, self.name_shunt,
            "shunt", self.n_shunt, obs._shunt_bus,
            el_to_sub_id=self.shunt_to_subid)
        return sto_ids
    
    def convert(self, obs:BaseObservation):
        # Initialize the graph with "grid level" attributes
        graph = networkx.DiGraph()

        # Add the substations
        sub_li = [(sub_id,
                    {"id": sub_id,
                    "type": "substation",
                    "name": self.name_sub[sub_id]})
                for sub_id in range(self.n_sub)]
        graph.add_nodes_from(sub_li)

        # Handle Buses
        bus_ids = self._aux_add_buses(obs, graph, env.n_sub)

        # Handle Loads
        load_ids = self._aux_add_loads(obs, graph, bus_ids[-1] + 1)

        # Handle Gens
        gen_ids = self._aux_add_gens(obs, graph, load_ids[-1] + 1)

        # Handle Lines
        line_ids = self._aux_add_lines(obs, graph, gen_ids[-1] + 1)

        # Handle Energy Storage
        sto_ids = None
        if (has_storage := self.n_storage > 0):
            sto_ids = self._aux_add_storages(obs, graph, line_ids[-1] + 1)
        
        # Handle shunts
        if type(obs).shunts_data_available:
            _ = self._aux_add_shunts(obs, graph, 
                                     sto_ids[-1] + 1 if has_storage else line_ids[-1] + 1)

        # And now we use the data above to put the right properties to the nodes for the buses
        bus_v_theta = {}
        for bus_id in bus_ids:
            li_pred = list(graph.predecessors(n=bus_id))
            if li_pred:
                bus_v_theta[bus_id] = {"connected": True}
            else:
                bus_v_theta[bus_id] = {"connected": False}
        networkx.set_node_attributes(graph, bus_v_theta)

        # Extra layer of security: prevent accidental modification of this graph
        networkx.freeze(graph)
        return graph