Merge pull request #98 from washingtonpost/bug/fix_bad_normalized_margin_math

Handle division where `normalized_margin` numerator and denominator are 0

Author: dmnapolitano

src/elexmodel/handlers/data/Estimandizer.py

@@ -114,5 +114,7 @@ def margin(data_df, col_prefix):
     generated_normalized_margin_column_name = f"{col_prefix}normalized_margin"
     data_df[generated_weights_column_name] = data_df[f"{col_prefix}dem"] + data_df[f"{col_prefix}gop"]
     data_df[generated_margin_column_name] = data_df[f"{col_prefix}dem"] - data_df[f"{col_prefix}gop"]
-    data_df[generated_normalized_margin_column_name] = data_df[f"{col_prefix}margin"] / data_df[f"{col_prefix}weights"]
+    data_df[generated_normalized_margin_column_name] = np.nan_to_num(
+        data_df[f"{col_prefix}margin"] / data_df[f"{col_prefix}weights"], nan=0, posinf=0, neginf=0
+    )
     return data_df, [generated_weights_column_name, generated_normalized_margin_column_name]

tests/handlers/test_estimandizer.py

@@ -69,3 +69,18 @@ def test_add_turnout_factor(va_governor_county_data):
 
     assert "turnout_factor" in output_df.columns
     assert 0 == pytest.approx(output_df.loc[0, "turnout_factor"])
+
+
+def test_add_margin_estimand_zero_normalized_margin(va_governor_county_data):
+    estimand_baselines = {"margin": None}
+    estimandizer = Estimandizer()
+
+    # test that we're handling zeros ok
+    test_df = va_governor_county_data.copy()
+    test_df.loc[1, "baseline_dem"] = 0
+    test_df.loc[1, "baseline_gop"] = 0
+
+    output_df = estimandizer.add_estimand_baselines(test_df, estimand_baselines, False, include_results_estimand=False)
+
+    assert "baseline_normalized_margin" in output_df.columns
+    assert test_df.loc[1, "baseline_normalized_margin"] == 0