Accept zero number_of_people_per_household

SEL-Columbia · Feb 24, 2017 · 59627af · 59627af
1 parent 8113038
commit 59627af
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 7 deletions.
diff --git a/infrastructure_planning/demography/__init__.py b/infrastructure_planning/demography/__init__.py
@@ -1,10 +1,10 @@
 from pandas import Series
 
-from ..macros import get_final_value
+from ..macros import get_final_value, make_zero_by_year
 
 
 def estimate_population_profile(population_by_year):
     return {
         'final_population': get_final_value(population_by_year),
-        'zero_by_year': Series(0, index=population_by_year.index),
+        'zero_by_year': make_zero_by_year(population_by_year),
     }
diff --git a/infrastructure_planning/electricity/consumption/linear.py b/infrastructure_planning/electricity/consumption/linear.py
@@ -2,8 +2,10 @@
 from pandas import DataFrame, Series, concat, isnull, merge
 
 from ...growth import get_default_slope, get_future_years
-from ...growth.interpolated import get_interpolated_spline_extrapolated_linear_function as get_estimate_electricity_consumption  # noqa
-from ...macros import get_final_value
+from ...growth.interpolated import (
+    get_interpolated_spline_extrapolated_linear_function as
+    get_estimate_electricity_consumption)
+from ...macros import get_final_value, make_zero_by_year
 
 
 def estimate_consumption_from_connection_type(
@@ -14,10 +16,11 @@ def estimate_consumption_from_connection_type(
     if there is a local override for household_count.
     """
     d = {}
-    connection_count_by_year = Series(0, index=population_by_year.index)
-    consumption_by_year = Series(0, index=population_by_year.index)
+    connection_count_by_year = make_zero_by_year(population_by_year)
+    consumption_by_year = make_zero_by_year(population_by_year)
     estimated_household_connection_count_by_year = divide_safely(
-        population_by_year, number_of_people_per_household, 0)
+        population_by_year, number_of_people_per_household,
+        make_zero_by_year(population_by_year))
     for row_index, row in connection_type_table.iterrows():
         connection_type = row['connection_type']
         count_by_year = _get_connection_count_by_year(

diff --git a/infrastructure_planning/macros.py b/infrastructure_planning/macros.py
@@ -236,6 +236,10 @@ def get_final_value(value_by_year):
     return value_by_year.ix[sorted(value_by_year.index)[-1]]
 
 
+def make_zero_by_year(value_by_year):
+    return Series(0, index=value_by_year.index)
+
+
 def get_graph_from_table(table):
     graph = InfrastructureGraph()
     for index, row in table.iterrows():