policyengine-us/policyengine_us/variables/input/geography.py at ff7bab5868659fe860012f90b06adc4b085f4222 · PolicyEngine/policyengine-us · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
from policyengine_us.model_api import *
from policyengine_core.parameters import homogenize_parameter_structures
from policyengine_core.simulations import Simulation
from policyengine_us.variables.household.demographic.geographic.state_name import (
    StateName,
)
from policyengine_us.parameters.gov.hhs.medicaid.geography import (
    medicaid_rating_areas,
    second_lowest_silver_plan_cost,
)
import warnings

warnings.filterwarnings("ignore")

label = "Geography"


class state_name(Variable):
    value_type = Enum
    possible_values = StateName
    default_value = StateName.CA
    entity = Household
    label = "State"
    definition_period = YEAR

    def formula(household, period, parameters):
        fips = household("state_fips", period)
        return (
            pd.Series(fips)
            .map(
                {
                    1: StateName.AL,
                    2: StateName.AK,
                    4: StateName.AZ,
                    5: StateName.AR,
                    6: StateName.CA,
                    8: StateName.CO,
                    9: StateName.CT,
                    10: StateName.DE,
                    11: StateName.DC,
                    12: StateName.FL,
                    13: StateName.GA,
                    15: StateName.HI,
                    16: StateName.ID,
                    17: StateName.IL,
                    18: StateName.IN,
                    19: StateName.IA,
                    20: StateName.KS,
                    21: StateName.KY,
                    22: StateName.LA,
                    23: StateName.ME,
                    24: StateName.MD,
                    25: StateName.MA,
                    26: StateName.MI,
                    27: StateName.MN,
                    28: StateName.MS,
                    29: StateName.MO,
                    30: StateName.MT,
                    31: StateName.NE,
                    32: StateName.NV,
                    33: StateName.NH,
                    34: StateName.NJ,
                    35: StateName.NM,
                    36: StateName.NY,
                    37: StateName.NC,
                    38: StateName.ND,
                    39: StateName.OH,
                    40: StateName.OK,
                    41: StateName.OR,
                    42: StateName.PA,
                    44: StateName.RI,
                    45: StateName.SC,
                    46: StateName.SD,
                    47: StateName.TN,
                    48: StateName.TX,
                    49: StateName.UT,
                    50: StateName.VT,
                    51: StateName.VA,
                    53: StateName.WA,
                    54: StateName.WV,
                    55: StateName.WI,
                    56: StateName.WY,
                    72: StateName.PR,
                }
            )
            .values
        )


class medicaid_rating_area(Variable):
    value_type = int
    entity = Household
    label = "Medicaid rating area"
    definition_period = YEAR
    hidden_input = True

    def formula(household, period, parameters):
        three_digit_zip_code = household("three_digit_zip_code", period)
        county = household("county_str", period)
        # Try a lookup on zip code, fill missing values with a lookup on county, fill missing with zero.
        df = pd.DataFrame(
            {
                "location": county,
            }
        )
        df_matched = pd.merge(
            df,
            medicaid_rating_areas,
            how="left",
            left_on="location",
            right_on="location",
        )
        county_lookup_failed = df_matched.rating_area.isna()
        df_matched.location[county_lookup_failed] = three_digit_zip_code[
            county_lookup_failed
        ]
        df_matched = pd.merge(
            df_matched,
            medicaid_rating_areas,
            how="left",
            left_on="location",
            right_on="location",
        )
        df_matched["rating_area"] = df_matched["rating_area_x"].fillna(
            df_matched["rating_area_y"]
        )
        return df_matched["rating_area"].fillna(1)


class county_fips(Variable):
    value_type = str
    label = "County FIPS code"
    entity = Household
    definition_period = YEAR
    documentation = "County FIPS code"


class state_fips(Variable):
    value_type = int
    label = "State FIPS code"
    entity = Household
    definition_period = YEAR
    documentation = "State FIPS code"
    default_value = 6