Fix upgrade feature bugs (#30)

Co-authored-by: mikivee <mikivee>
rewiringamerica · Jan 13, 2025 · 58cc254 · 58cc254
1 parent d953977
commit 58cc254
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/src/datagen.py b/src/datagen.py
@@ -131,6 +131,8 @@ class DataGenerator(tf.keras.utils.Sequence):
         "neighbor_distance_ft",
         "n_occupants",
         "vintage",
+        "has_heat_pump_dryer",
+        "has_induction_range",
         # fuel indicators -- these must be present for post-processing to work!!
         "has_methane_gas_appliance",
         "has_fuel_oil_appliance",

diff --git a/src/feature_utils.py b/src/feature_utils.py
@@ -1037,16 +1037,20 @@ def apply_upgrades(baseline_building_features: DataFrame, upgrade_id: int) -> Da
             .withColumn("heating_setpoint_offset_magnitude_degrees_f", F.lit(0.0))
         )
     if upgrade_id in [6, 9]:
-        upgrade_building_features = upgrade_building_features.withColumn(
-            "water_heater_efficiency",
-            F.when(  # electric tankless don't get upgraded due to likely size constraints
-                F.col("water_heater_efficiency") == "Electric Tankless",
-                F.col("water_heater_efficiency"),
+        upgrade_building_features = (
+            upgrade_building_features.withColumn("water_heater_fuel", F.lit("Electricity"))
+            .withColumn(
+                "water_heater_efficiency",
+                F.when(  # electric tankless don't get upgraded due to likely size constraints
+                    F.col("water_heater_efficiency") == "Electric Tankless",
+                    F.col("water_heater_efficiency"),
+                )
+                .when(F.col("n_bedrooms") <= 3, F.lit("Electric Heat Pump, 50 gal, 3.45 UEF"))
+                .when(F.col("n_bedrooms") == 4, F.lit("Electric Heat Pump, 66 gal, 3.35 UEF"))
+                .otherwise(F.lit("Electric Heat Pump, 80 gal, 3.45 UEF")),
             )
-            .when(F.col("n_bedrooms") <= 3, F.lit("Electric Heat Pump, 50 gal, 3.45 UEF"))
-            .when(F.col("n_bedrooms") == 4, F.lit("Electric Heat Pump, 66 gal, 3.35 UEF"))
-            .otherwise(F.lit("Electric Heat Pump, 80 gal, 3.45 UEF")),
-        ).transform(add_water_heater_features)
+            .transform(add_water_heater_features)
+        )
 
     if upgrade_id in [8.1, 9]:
         upgrade_building_features = upgrade_building_features.withColumn(