diff --git a/src/elexmodel/handlers/data/Estimandizer.py b/src/elexmodel/handlers/data/Estimandizer.py index 6dd78192..ed61d449 100644 --- a/src/elexmodel/handlers/data/Estimandizer.py +++ b/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] diff --git a/tests/handlers/test_estimandizer.py b/tests/handlers/test_estimandizer.py index 69f1fd1e..cd502fae 100644 --- a/tests/handlers/test_estimandizer.py +++ b/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