From f6cef2032d04745769a0d8695b94bbfca5a16f81 Mon Sep 17 00:00:00 2001
From: Daniel Olsen <daniel.olsen@breakthroughenergy.org>
Date: Fri, 4 Dec 2020 12:20:06 -0800
Subject: [PATCH] refactor: add cache for bus df

---
 powersimdata/input/change_table.py | 41 ++++++++++++++++++++----------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/powersimdata/input/change_table.py b/powersimdata/input/change_table.py
index d93b05dc9..0b052d591 100644
--- a/powersimdata/input/change_table.py
+++ b/powersimdata/input/change_table.py
@@ -119,6 +119,7 @@ def __init__(self, grid):
         """
         self.grid = grid
         self.ct = {}
+        self.new_bus_cache = {}
 
     @staticmethod
     def _check_resource(resource):
@@ -456,12 +457,12 @@ def add_storage_capacity(self, bus_id):
         :param dict bus_id: key(s) for the id of bus(es), value(s) is (are)
             capacity of the energy storage system in MW.
         """
-        new_grid = TransformGrid(self.grid, self.ct).get_grid()
+        anticipated_bus = self._get_new_bus()
 
         if "storage" not in self.ct:
             self.ct["storage"] = {}
 
-        diff = set(bus_id.keys()).difference(set(new_grid.bus.index))
+        diff = set(bus_id.keys()).difference(set(anticipated_bus.index))
         if len(diff) != 0:
             print("No bus with the following id:")
             for i in list(diff):
@@ -505,11 +506,11 @@ def _add_line(self, key, info):
             *'new_dcline'*
         :param list info: parameters of the line.
         """
+        anticipated_bus = self._get_new_bus()
         if key not in self.ct:
             self.ct[key] = []
 
         required_info = ["capacity", "from_bus_id", "to_bus_id"]
-        new_grid = TransformGrid(self.grid, self.ct).get_grid()
         for i, line in enumerate(info):
             if not isinstance(line, dict):
                 print("Each entry must be a dictionary")
@@ -522,13 +523,13 @@ def _add_line(self, key, info):
             else:
                 start = line["from_bus_id"]
                 end = line["to_bus_id"]
-                if start not in new_grid.bus.index:
+                if start not in anticipated_bus.index:
                     print(
                         "No bus with the following id for line #%d: %d" % (i + 1, start)
                     )
                     self.ct.pop(key)
                     return
-                elif end not in new_grid.bus.index:
+                elif end not in anticipated_bus.index:
                     print(
                         "No bus with the following id for line #%d: %d" % (i + 1, end)
                     )
@@ -544,16 +545,16 @@ def _add_line(self, key, info):
                     return
                 elif (
                     key == "new_branch"
-                    and new_grid.bus.interconnect[start]
-                    != new_grid.bus.interconnect[end]
+                    and anticipated_bus.interconnect[start]
+                    != anticipated_bus.interconnect[end]
                 ):
                     print("Buses of line #%d must be in same interconnect" % (i + 1))
                     self.ct.pop(key)
                     return
 
                 elif (
-                    new_grid.bus.lat[start] == new_grid.bus.lat[end]
-                    and new_grid.bus.lon[start] == new_grid.bus.lon[end]
+                    anticipated_bus.lat[start] == anticipated_bus.lat[end]
+                    and anticipated_bus.lon[start] == anticipated_bus.lon[end]
                 ):
                     print("Distance between buses of line #%d is 0" % (i + 1))
                     self.ct.pop(key)
@@ -571,10 +572,10 @@ def add_plant(self, info):
             print("Argument enclosing new plant(s) must be a list")
             return
 
+        anticipated_bus = self._get_new_bus()
         if "new_plant" not in self.ct:
             self.ct["new_plant"] = []
 
-        new_grid = TransformGrid(self.grid, self.ct).get_grid()
         for i, plant in enumerate(info):
             if not isinstance(plant, dict):
                 print("Each entry must be a dictionary")
@@ -594,7 +595,7 @@ def add_plant(self, info):
                 print("Missing key bus_id for plant #%d" % (i + 1))
                 self.ct.pop("new_plant")
                 return
-            elif plant["bus_id"] not in new_grid.bus.index:
+            elif plant["bus_id"] not in anticipated_bus.index:
                 print("No bus id %d available for plant #%d" % (plant["bus_id"], i + 1))
                 self.ct.pop("new_plant")
                 return
@@ -613,8 +614,8 @@ def add_plant(self, info):
                 self.ct.pop("new_plant")
                 return
             if plant["type"] in _renewable_resource:
-                lon = new_grid.bus.loc[plant["bus_id"]].lon
-                lat = new_grid.bus.loc[plant["bus_id"]].lat
+                lon = anticipated_bus.loc[plant["bus_id"]].lon
+                lat = anticipated_bus.loc[plant["bus_id"]].lat
                 plant_same_type = self.grid.plant.groupby("type").get_group(
                     plant["type"]
                 )
@@ -693,6 +694,20 @@ def add_bus(self, info):
             self.ct.pop("new_bus")
             raise
 
+    def _get_new_bus(self):
+        if "new_bus" not in self.ct:
+            return self.grid.bus
+        new_bus_tuple = tuple(
+            tuple((k, new_bus_dict[k]) for k in sorted(new_bus_dict.keys()))
+            for new_bus_dict in self.ct["new_bus"]
+        )
+        if new_bus_tuple in self.new_bus_cache:
+            return self.new_bus_cache[new_bus_tuple]
+        else:
+            bus = TransformGrid(self.grid, self.ct).get_grid().bus
+            self.new_bus_cache[new_bus_tuple] = bus
+            return bus
+
     def write(self, scenario_id):
         """Saves change table to disk.