diff --git a/gtep/model_library/components.py b/gtep/model_library/components.py
index a1dc680..a3b1a38 100644
--- a/gtep/model_library/components.py
+++ b/gtep/model_library/components.py
@@ -66,6 +66,18 @@ def add_model_sets(m, stages, rep_per=["a", "b"], com_per=2, dis_per=2):
         doc="All generators",
     )
 
+    m.generatorsByBus = pyo.Set(
+        m.buses,
+        initialize={
+            bus: [
+                gen
+                for gen in m.generators
+                if m.md.data["elements"]["generator"][gen]["bus"] == bus
+            ]
+            for bus in m.buses
+        },
+    )
+
     m.thermalGenerators = pyo.Set(
         within=m.generators,
         initialize=(
@@ -97,6 +109,22 @@ def add_model_sets(m, stages, rep_per=["a", "b"], com_per=2, dis_per=2):
             doc="Potential storage units",
         )
 
+    m.storageByBus = pyo.Set(
+        m.buses,
+        initialize={
+            bus: (
+                [
+                    batt
+                    for batt in m.storage
+                    if m.md.data["elements"]["storage"][batt]["bus"] == bus
+                ]
+                if m.md.data["elements"].get("storage")
+                else []
+            )
+            for bus in m.buses
+        },
+    )
+
 
 def add_model_parameters(m, num_commit, num_dispatch, duration_dispatch):
     """Creates and labels all the parameters in the GTEP model. This
diff --git a/gtep/model_library/dispatch.py b/gtep/model_library/dispatch.py
index 166334c..1013d78 100644
--- a/gtep/model_library/dispatch.py
+++ b/gtep/model_library/dispatch.py
@@ -229,16 +229,17 @@ def CP_flow_balance(b):
             loads = [l for l in b.loads]
             gens = [gen for gen in m.generators]
             batts = [bat for bat in m.storage]
+            # Add generation to constraint
             balance += sum(
                 b.thermalGeneration[g] for g in gens if g in m.thermalGenerators
             )
             balance += sum(
                 b.renewableGeneration[g] for g in gens if g in m.renewableGenerators
             )
-            """ Battery Storage added to flow balance constraint """
+            # Add battery storage to constraint
             balance += sum(b.storageDischarged[bt] for bt in batts)
             balance -= sum(b.storageCharged[bt] for bt in batts)
-
+            # Add the loads as a parameter (already includes units).
             balance -= sum(c_p.loads[l] for l in loads)
             balance += sum(b.loadShed[bus] for bus in buses)
 
@@ -249,32 +250,23 @@ def CP_flow_balance(b):
         @b.Constraint(m.buses, doc="Energy balance constraint")
         def flow_balance(b, bus):
             balance = 0
+            # Add power flow to constraint
             end_points = [line for line in m.lines if m.from_bus[line] == bus]
             start_points = [line for line in m.lines if m.to_bus[line] == bus]
-            gens = [
-                gen
-                for gen in m.generators
-                if m.md.data["elements"]["generator"][gen]["bus"] == bus
-            ]
-            batts = []
-            if m.config["storage"]:
-                batts = [
-                    bat
-                    for bat in m.storage
-                    if m.md.data["elements"]["storage"][bat]["bus"] == bus
-                ]
             balance -= sum(b.powerFlow[i] for i in end_points)
             balance += sum(b.powerFlow[i] for i in start_points)
+            # Add generation to constraint
             balance += sum(
-                b.thermalGeneration[g] for g in gens if g in m.thermalGenerators
+                b.thermalGeneration[g]
+                for g in m.thermalGenerators & m.generatorsByBus[bus]
             )
             balance += sum(
-                b.renewableGeneration[g] for g in gens if g in m.renewableGenerators
+                b.renewableGeneration[g]
+                for g in m.renewableGenerators & m.generatorsByBus[bus]
             )
-            """ Battery Storage added to flow balance constraint """
-            balance += sum(b.storageDischarged[bt] for bt in batts)
-            balance -= sum(b.storageCharged[bt] for bt in batts)
-
+            # Add battery storage to constraint
+            balance += sum(b.storageDischarged[bt] for bt in m.storageByBus[bus])
+            balance -= sum(b.storageCharged[bt] for bt in m.storageByBus[bus])
             # Add the loads as a parameter (already includes units).
             balance -= c_p.loads[bus]
             balance += b.loadShed[bus]