adding soc function map and fixing teste

evrouting/charge/factories.py

@@ -59,6 +59,16 @@ class ChargingFunctionMap:
         return cf
 
 
+class SoCFunctionMap:
+    """Maps Nodes to their charging functions."""
+
+    def __init__(self, cf: ChargingFunctionMap):
+        self.cf: ChargingFunctionMap = cf
+
+    def __getitem__(self, label: Label) -> SoCFunction:
+        return SoCFunction(label, self.cf[label.last_cs])
+
+
 class LabelsFactory:
 
     def __init__(self,

evrouting/charge/routing.py

@@ -5,11 +5,14 @@ import networkx as nx
 from evrouting.T import Node, SoC
 from evrouting.utils import PriorityQueue
 from evrouting.charge.factories import (
-    LabelsFactory, ChargingFunctionMap, soc_profile_factory
+    LabelsFactory,
+    ChargingFunctionMap,
+    SoCFunctionMap,
+    soc_profile_factory
 )
 
 from ..graph_tools import distance
-from .T import SoCFunction, SoCProfile, Label
+from .T import SoCProfile, Label
 from .utils import LabelPriorityQueue
 
 __all__ = ['shortest_path']
@@ -20,17 +23,19 @@ def shortest_path(G: nx.Graph, charging_stations: set, s: Node, t: Node,
     """
     Calculates shortest path using the CHarge algorithm.
 
-    :param G: Input Graph
-    :param s: Start Node identifier
-    :param t: End Node identifier
-    :param beta_s: Start SoC
-    :param beta_t: End SoC
-    :param U: Capacity
+    :param G:
+    :param charging_stations:
+    :param s:
+    :param t:
+    :param initial_soc:
+    :param final_soc:
+    :param capacity:
     :return:
     """
     t = _apply_final_constraints(G, t, final_soc)
 
     cf = ChargingFunctionMap(G=G, capacity=capacity, initial_soc=initial_soc)
+    f_soc = SoCFunctionMap(cf)
     label_factory = LabelsFactory(G, capacity, cf, initial_soc)
 
     # Init maps to manage labels
@@ -58,9 +63,7 @@ def shortest_path(G: nx.Graph, charging_stations: set, s: Node, t: Node,
         l_set[minimum_node].add(label_minimum_node)
 
         if minimum_node == t:
-            return SoCFunction(label_minimum_node,
-                               cf[label_minimum_node.last_cs]
-                               ).minimum
+            return f_soc[label_minimum_node].minimum
 
         # handle charging stations
         if minimum_node in charging_stations and \
@@ -72,7 +75,7 @@ def shortest_path(G: nx.Graph, charging_stations: set, s: Node, t: Node,
 
         # Update priority queue. This node might have gotten a new
         # minimum label spawned is th previous step.
-        _update_priority_queue(cf, prio_queue, l_uns, minimum_node)
+        _update_priority_queue(f_soc, prio_queue, l_uns, minimum_node)
 
         # scan outgoing arcs
         for n in G.neighbors(minimum_node):
@@ -101,15 +104,12 @@ def shortest_path(G: nx.Graph, charging_stations: set, s: Node, t: Node,
                     pass
                 else:
                     if l_new == l_uns[n].peak_min():
-                        key, count = _key(l_new, cf[l_new.last_cs])
+                        key, count = _key(l_new, f_soc)
                         prio_queue.insert(n, priority=key, count=count)
 
 
-def _key(label, cf):
-    soc_function = SoCFunction(
-        label,
-        cf
-    )
+def _key(label, f_soc):
+    soc_function = f_soc[label]
 
     t_min = soc_function.minimum
     soc_min = soc_function(t_min)
@@ -156,7 +156,7 @@ def _is_feasible_path(soc_profile: SoCProfile, capacity: SoC) -> bool:
 
 
 def _update_priority_queue(
-        cf: ChargingFunctionMap,
+        f_soc: SoCFunctionMap,
         prio_queue: PriorityQueue,
         l_uns: Dict[int, LabelPriorityQueue],
         node: Node):
@@ -170,7 +170,7 @@ def _update_priority_queue(
         # l_uns[v] empty
         prio_queue.delete_min()
     else:
-        key, count = _key(minimum_label, cf[minimum_label.last_cs])
+        key, count = _key(minimum_label, f_soc)
         prio_queue.insert(node, priority=key, count=count)
 
 

tests/charge/test_utils.py

@@ -6,8 +6,9 @@ from evrouting.charge.T import Label
 @pytest.fixture
 def q(label, ch_function):
     _, _, cf = ch_function
-    q = LabelPriorityQueue()
-    q.insert(label, cf)
+    dummy_cf = {label.last_cs: cf}
+    q = LabelPriorityQueue(dummy_cf)
+    q.insert(label)
 
     # create min
     label = Label(
@@ -17,7 +18,8 @@ def q(label, ch_function):
         last_cs=1
     )
 
-    q.insert(label, cf)
+    dummy_cf[label.last_cs] = cf
+    q.insert(label)
 
     yield q
     del q
@@ -39,4 +41,4 @@ class TestProrityQueue:
         _, _, cf = ch_function
         label = q.peak_min()
         q.remove_item(label)
-        q.insert(label, cf)
+        q.insert(label)