Seed Alaska's Permanent Fund Dividend take-up and take reported dividends out of miscellaneous_income #1551

"""PROTOTYPE of the proposed Microcosm stage. No Populace release contains its output.

Writes copies of the Alaska-only Populace subset with
takes_up_ak_permanent_fund_dividend_if_eligible set, so the PR's impact comes from
real microsimulations of what the data stage would do:

- Anyone in a unit where someone reported a dividend gets the flag (anchor).
- Other units get a seeded blake2b draw on their stable source identity, the
  key microcosm's take_up.py uses, so support-channel clones share a flag.
- It picks the draw cutoff so that eligible flagged people equal Alaska's 2024
  paid applications, 624,489, on the file's 2024 weights.

With --clean, it also takes reported dividends out of miscellaneous_income
(microcosm#1443). Microcosm puts ASEC-channel OI_VAL in miscellaneous_income,
apart from the alimony and strike benefit codes.

usage: seed_flag.py SRC OUT LEVEL SEED MATCH CLEAN
  LEVEL spm_unit|person; MATCH exact|tolerant; CLEAN clean|keep
"""

import hashlib
import json
import sys

import numpy as np
import pandas as pd

VAR = "takes_up_ak_permanent_fund_dividend_if_eligible"
TARGET_2024 = 624_489
ELIGIBLE = {"CITIZEN", "LEGAL_PERMANENT_RESIDENT", "REFUGEE", "ASYLEE"}
# Whole payments Alaska made, and the 2022 and 2024 components people may report alone.
PAID = {
    2021: [1114],
    2022: [3284, 2622, 662],
    2023: [1312],
    2024: [1702, 1403.83, 298.17],
}
TABLES = ["_time_period", "household", "person", "tax_unit", "spm_unit", "family", "marital_unit"]

src, out, level, seed, match, clean = sys.argv[1:7]
assert level in ("spm_unit", "person") and match in ("exact", "tolerant") and clean in ("clean", "keep")
tables = {t: pd.read_hdf(src, t) for t in TABLES}
p = tables["person"].copy()
hh = tables["household"].set_index("household_id")
w = p["person_household_id"].map(hh["household_weight"]).astype(float).to_numpy()
eligible = p["immigration_status_str"].isin(ELIGIBLE).to_numpy()


def is_dividend(value, year):
    """A reported amount equal to one to eight dividends Alaska paid in this or the prior year."""
    if not value > 0:
        return False
    tol = 0.5 if match == "exact" else 100
    for y in (year, year - 1):
        for d in PAID.get(y, []):
            for k in range(1, 9):
                if abs(value - k * d) <= tol:
                    return True
    return False


reporter = np.array([is_dividend(v, int(y)) for v, y in zip(p["OI_VAL"], p["source_year"])])
sid = list(zip(p["source_year"], p["source_household_id"], p["source_person_id"]))
if level == "person":
    key = pd.Series([f"{a}:{b}:{c}" for a, b, c in sid], index=p.index)
else:
    frame = pd.DataFrame(
        {
            "unit": p["person_spm_unit_id"],
            "year": p["source_year"],
            "hh": p["source_household_id"],
            "pid": p["source_person_id"],
        }
    )
    first = frame.groupby("unit").agg(year=("year", "first"), hh=("hh", "first"), pid=("pid", "min"))
    unit_key = first["year"].astype(str) + ":" + first["hh"].astype(str) + ":" + first["pid"].astype(str)
    key = p["person_spm_unit_id"].map(unit_key)

units = pd.DataFrame({"key": key.to_numpy(), "w_e": w * eligible, "rep": reporter}).groupby("key").agg(
    w_e=("w_e", "sum"), anchored=("rep", "any")
)
units["u"] = [
    int.from_bytes(hashlib.blake2b(f"{seed}:{VAR}:{k}".encode(), digest_size=8).digest(), "big") / 2**64
    for k in units.index
]
anchored_weight = float(units.loc[units["anchored"], "w_e"].sum())
free = units.loc[~units["anchored"]].sort_values("u")
cum = anchored_weight + free["w_e"].cumsum().to_numpy()
# Keep the cutoff whose total falls closest to the target.
gaps = np.abs(np.concatenate(([anchored_weight], cum)) - TARGET_2024)
n_free = int(gaps.argmin())
granted = set(free.index[:n_free])
units["flag"] = units["anchored"] | units.index.isin(granted)
flag = key.map(units["flag"]).to_numpy(dtype=bool)
p[VAR] = flag

removed = 0.0
if clean == "clean":
    asec = (p["person_support_channel"] == "asec").to_numpy()
    mask = reporter & asec
    removed = float((w[mask] * p.loc[mask, "miscellaneous_income"].to_numpy()).sum())
    p.loc[mask, "miscellaneous_income"] = 0.0

tables["person"] = p
with pd.HDFStore(out, "w") as store:
    store.put("_time_period", tables["_time_period"], format="table")
    for t in TABLES[1:]:
        store.put(t, tables[t], format="table")

meta = {
    "PROTOTYPE_NOT_RELEASED": True,
    "source": src,
    "level": level,
    "seed": seed,
    "match": match,
    "clean": clean,
    "target_2024": TARGET_2024,
    "alaska_people_2024": float(w.sum()),
    "eligible_people_2024": float((w * eligible).sum()),
    "reporter_people": float(w[reporter].sum()),
    "reporter_records": int(reporter.sum()),
    "anchored_eligible_people": anchored_weight,
    "units": int(len(units)),
    "anchored_units": int(units["anchored"].sum()),
    "granted_units": n_free,
    "eligible_flagged_people_2024": float((w * eligible * flag).sum()),
    "eligible_share_flagged": float((w * eligible * flag).sum() / (w * eligible).sum()),
    "free_share_granted": float(free["w_e"].iloc[:n_free].sum() / free["w_e"].sum()),
    "misc_income_removed": removed,
    "misc_share_after": float(w[p["miscellaneous_income"].to_numpy() > 0].sum() / w.sum()),
}
json.dump(meta, open(out + ".json", "w"), indent=1)
print(json.dumps(meta))

"""Apply the amount rule to every state in the full file, to see how often it matches outside Alaska.

A record matches when OI_VAL falls within $100 of one to eight dividends Alaska
paid in its source year or the prior year (the seed_flag.py 'tolerant' rule), unless it
carries OI_OFF code 12 or 20, which microcosm's split already sends to strike_benefits
and alimony_income. No Alaska record with OI_VAL above 0 carries either code, so
the skip changes nothing in seed_flag.py's Alaska runs. A match outside
Alaska counts as a false positive unless the person received a dividend and moved before the
survey, so this bounds the rule's error from above.
"""
import sys
import numpy as np
import pandas as pd

src = sys.argv[1]
PAID = {2021: [1114], 2022: [3284, 2622, 662], 2023: [1312], 2024: [1702, 1403.83, 298.17]}
cols = ["person_household_id", "OI_VAL", "OI_OFF", "source_year", "person_support_channel", "miscellaneous_income"]
p = pd.read_hdf(src, "person", columns=cols)
h = pd.read_hdf(src, "household", columns=["household_id", "state_fips", "household_weight"]).set_index("household_id")
p["fips"] = p["person_household_id"].map(h["state_fips"])
p["w"] = p["person_household_id"].map(h["household_weight"]).astype(float)
amounts = {y: np.array([k * d for yy in (y, y - 1) for d in PAID.get(yy, []) for k in range(1, 9)]) for y in p["source_year"].unique()}
SEPARATE_CODES = (12, 20)
def match(v, code, y, tol):
    if not v > 0 or code in SEPARATE_CODES:
        return False
    return bool((np.abs(amounts[y] - v) <= tol).any())
pos = p[p["OI_VAL"] > 0].copy()
pos["tol100"] = [match(v, c, y, 100) for v, c, y in zip(pos["OI_VAL"], pos["OI_OFF"], pos["source_year"])]
pos["tol0"] = [match(v, c, y, 0.5) for v, c, y in zip(pos["OI_VAL"], pos["OI_OFF"], pos["source_year"])]
skipped = pos["OI_OFF"].isin(SEPARATE_CODES) & np.array(
    [match(v, 0, y, 100) for v, y in zip(pos["OI_VAL"], pos["source_year"])]
)
print("records the code skip removes from the $100 rule:", int(skipped.sum()), "in", pos.loc[skipped, "fips"].nunique(), "states")
tot = p.groupby("fips")["w"].sum()
rows = pos.groupby("fips").apply(lambda d: pd.Series({
    "oi_pos_share": d["w"].sum(), "tol100_share": d.loc[d.tol100, "w"].sum(), "tol0_share": d.loc[d.tol0, "w"].sum()}), include_groups=False)
rows = rows.div(tot, axis=0)
us_ex_ak = rows.drop(index=2)
print("AK shares of people", rows.loc[2].round(4).to_dict())
print("median other state", us_ex_ak.median().round(4).to_dict())
print("max other state", us_ex_ak.max().round(4).to_dict())
print("WA", rows.loc[53].round(4).to_dict(), "OR", rows.loc[41].round(4).to_dict())
ak = p[p["fips"] == 2]
print("AK unweighted share of records with miscellaneous_income > 0:", round(float((ak["miscellaneous_income"] > 0).mean()), 4))
print("WA", round(float((p[p.fips == 53]["miscellaneous_income"] > 0).mean()), 4), "OR", round(float((p[p.fips == 41]["miscellaneous_income"] > 0).mean()), 4))

"""Measure the share of people with miscellaneous_income above 0, as microcosm#1443's first criterion counts it.

Weights each person record by its household's weight in the file. Reads Alaska
before the removal from the Alaska subset, Alaska after it from the prototype's
cleaned copy (seed_flag.py with MATCH=tolerant and CLEAN=clean, which zeroes the
reported dividends the $100 rule matches, whatever the seed), and the other
states from the full file.
Reads inputs only. PROTOTYPE: no Populace release contains the cleaned copy.

usage: misc_share_check.py FULL_H5 AK_H5 AK_CLEANED_H5 OUT_JSON
"""

import json
import sys

import pandas as pd


def shares(path, fips=None):
    p = pd.read_hdf(path, "person", columns=["person_household_id", "miscellaneous_income"])
    h = pd.read_hdf(path, "household", columns=["household_id", "state_fips", "household_weight"]).set_index(
        "household_id"
    )
    p["fips"] = p["person_household_id"].map(h["state_fips"]).astype(int)
    p["w"] = p["person_household_id"].map(h["household_weight"]).astype(float)
    out = {}
    for code in fips or sorted(p["fips"].unique()):
        s = p[p["fips"] == code]
        positive = s["miscellaneous_income"] > 0
        out[int(code)] = {
            "records": int(len(s)),
            "records_positive": int(positive.sum()),
            "unweighted_share": float(positive.mean()),
            "people": float(s["w"].sum()),
            "people_positive": float(s.loc[positive, "w"].sum()),
            "weighted_share": float(s.loc[positive, "w"].sum() / s["w"].sum()),
        }
    return out


full, ak, ak_clean, out = sys.argv[1:5]
result = {
    "inputs": {"full": full, "alaska": ak, "alaska_cleaned": ak_clean},
    "alaska_before": shares(ak, [2])[2],
    "alaska_after": shares(ak_clean, [2])[2],
    "full_file": shares(full),
}
others = pd.Series({k: v["weighted_share"] for k, v in result["full_file"].items() if k != 2})
result["other_states_weighted_share"] = {
    "median": float(others.median()),
    "max": float(others.max()),
    "washington": result["full_file"][53]["weighted_share"],
    "oregon": result["full_file"][41]["weighted_share"],
}
json.dump(result, open(out, "w"), indent=1)
for key in ("alaska_before", "alaska_after"):
    print(key, {k: round(v, 4) for k, v in result[key].items()})
print("full-file Alaska", {k: round(v, 4) for k, v in result["full_file"][2].items()})
print("other states weighted", {k: round(v, 4) for k, v in result["other_states_weighted_share"].items()})

"""What a zero amount would add to the $100 match rule on Alaska's records.

match_rule_check.py's PAID table holds only amounts Alaska paid, so it skips
the years when Alaska paid no energy relief. A rule that took "each amount" from both
parameters would also take those zeros, and every multiple of 0 equals 0, so
the $100 rule would match any OI_VAL from $1 to $100.

usage: zero_amount_check.py AK_H5 OUT_JSON
"""

import json
import sys

import numpy as np
import pandas as pd

PAID = {2021: [1114], 2022: [3284, 2622, 662], 2023: [1312], 2024: [1702, 1403.83, 298.17]}
# gov.states.ak.dor.energy_relief holds 0 for 2015 to 2021, 2023, 2025 and 2027 on.
ZERO_RELIEF_YEARS = {2021, 2023}
TOLERANCE = 100


def amounts(source_year, with_zero):
    values = [
        k * d for y in (source_year, source_year - 1) for d in PAID.get(y, []) for k in range(1, 9)
    ]
    if with_zero and ZERO_RELIEF_YEARS & {source_year, source_year - 1}:
        values.append(0.0)
    return np.array(values)


def matches(value, source_year, with_zero):
    return bool(np.any(np.abs(amounts(source_year, with_zero) - value) <= TOLERANCE))


def main():
    path, out = sys.argv[1:3]
    person = pd.read_hdf(path, "person")
    household = pd.read_hdf(path, "household").set_index("household_id")
    weight = person.person_household_id.map(household.household_weight).astype(float)
    positive = person.OI_VAL > 0
    paid_only = np.array(
        [p and matches(v, y, False) for p, v, y in zip(positive, person.OI_VAL, person.source_year)]
    )
    with_zero = np.array(
        [p and matches(v, y, True) for p, v, y in zip(positive, person.OI_VAL, person.source_year)]
    )
    added = with_zero & ~paid_only
    units_paid = set(person.person_spm_unit_id[paid_only])
    units_added = set(person.person_spm_unit_id[added]) - units_paid
    in_added_units = person.person_spm_unit_id.isin(units_added)
    result = {
        "records_matched_paid_amounts_only": int(paid_only.sum()),
        "people_matched_paid_amounts_only": float(weight[paid_only].sum()),
        "records_added_by_a_zero_amount": int(added.sum()),
        "people_on_added_records": float(weight[added].sum()),
        "weighted_oi_val_on_added_records": float((weight * person.OI_VAL)[added].sum()),
        "added_records": person.loc[added, ["person_source_id", "source_year", "OI_OFF", "OI_VAL", "age"]]
        .assign(weight=weight[added])
        .to_dict("records"),
        "source_people_on_added_records": int(person.person_source_id[added].nunique()),
        "spm_units_newly_anchored": len(units_added),
        "people_in_newly_anchored_spm_units": float(weight[in_added_units].sum()),
    }
    json.dump(result, open(out, "w"), indent=1, default=float)
    print(json.dumps(result, indent=1, default=float))


if __name__ == "__main__":
    main()

"""Which Alaska people aged 0 can qualify for the 2024 dividend, by dating.

Populace's source_year holds the ASEC income year (AsecSource.year, which
microcosm's asec_pool.py passes as income_year). Census surveys that year's income in
February to April of the next year and records age as of the survey date. A
child born during the dividend year can't qualify: the child fails AS
43.23.005(a)(3), and under (c) only children born in the two calendar years
before the dividend year qualify. A child born earlier can qualify under (c) or (a).

Assumes births spread evenly over the 365 days before the survey date. Three
datings of a record's survey:

- own_survey: each record keeps its own survey, in source_year + 1;
- survey_in_dividend_year: every survey falls in the dividend year;
- birth_year: the dividend year minus age, as policyengine-us's birth_year
  formula reads it, which puts every child aged 0 in the dividend year.

usage: age_dating_check.py AK_H5 OUT_JSON
"""

import json
import sys

import pandas as pd

DIVIDEND_YEAR = 2024
# Day of the year of each survey date, in a 365-day year.
SURVEY_DAYS = {"mid_february": 46, "mid_march": 74, "mid_april": 105}


def share_born_before(survey_year, survey_day, year):
    """Share of children under one at the survey date born before `year`."""
    # Days from January 1 of the survey year: births fall in
    # (survey_day - 365, survey_day]; January 1 of `year` sits at
    # (year - survey_year) * 365 + 1.
    cutoff = (year - survey_year) * 365
    share = (cutoff - (survey_day - 365)) / 365
    return min(max(share, 0.0), 1.0)


def qualifying_share(people_by_source_year, dating, survey_day):
    total = sum(people_by_source_year.values())
    can = 0.0
    for source_year, people in people_by_source_year.items():
        if dating == "own_survey":
            survey_year = source_year + 1
        elif dating == "survey_in_dividend_year":
            survey_year = DIVIDEND_YEAR
        else:
            return 0.0
        can += people * share_born_before(survey_year, survey_day, DIVIDEND_YEAR)
    return can / total


def main():
    path, out = sys.argv[1:3]
    person = pd.read_hdf(path, "person")
    household = pd.read_hdf(path, "household").set_index("household_id")
    weight = person.person_household_id.map(household.household_weight).astype(float)
    zero = person.age == 0

    age0 = {
        int(year): float(people)
        for year, people in weight[zero].groupby(person.source_year[zero]).sum().items()
    }
    age0_records = {
        int(year): int(count)
        for year, count in person.source_year[zero].value_counts().sort_index().items()
    }
    everyone = {
        int(year): float(people)
        for year, people in weight.groupby(person.source_year).sum().items()
    }
    total0 = sum(age0.values())
    # Spread the same age-0 total over source years like all Alaska weight.
    total_all = sum(everyone.values())
    spread = {year: total0 * people / total_all for year, people in everyone.items()}

    result = {
        "dividend_year": DIVIDEND_YEAR,
        "age_0_records": int(zero.sum()),
        "age_0_people": total0,
        "age_0_records_by_source_year": age0_records,
        "age_0_people_by_source_year": age0,
        "age_0_share_by_source_year": {y: v / total0 for y, v in age0.items()},
        "all_people_share_by_source_year": {
            y: v / total_all for y, v in everyone.items()
        },
        "datings": {},
    }
    for dating in ("own_survey", "survey_in_dividend_year", "birth_year"):
        rows = {}
        for name, day in SURVEY_DAYS.items():
            share = qualifying_share(age0, dating, day)
            rows[name] = {
                "share_can_qualify": share,
                "people_who_can_qualify": share * total0,
                "people_who_cannot": (1 - share) * total0,
                "share_if_spread_like_all_weight": qualifying_share(
                    spread, dating, day
                ),
            }
        result["datings"][dating] = rows
    # Under own_survey, the share born two calendar years before the dividend
    # year: records from the survey in the year before it, with births before
    # that year.
    result["own_survey_share_born_two_years_before"] = {
        name: sum(
            people * share_born_before(year + 1, day, DIVIDEND_YEAR - 1)
            for year, people in age0.items()
            if year + 1 == DIVIDEND_YEAR - 1
        )
        / total0
        for name, day in SURVEY_DAYS.items()
    }
    json.dump(result, open(out, "w"), indent=1)
    print(json.dumps(result, indent=1))


if __name__ == "__main__":
    main()

"""Emulate choosing the prototype's fill cutoff on stage-time weights.

Microcosm assigns take-up before calibrate_l0_refit, so a stage that picks its
cutoff at that point sees different weights from the released ones. This script
sets the flag as seed_flag.py does, but chooses the cutoff on stand-in
stage-time weights that it scales to the released total, then counts eligible flagged
people on the released weights. The release has no design-weight column, so the script
uses stand-ins and doesn't run the real refit.

PROTOTYPE: these flags appear in no Populace release.

usage: stage_time_emulation.py AK_H5 [SEEDS]
"""

import hashlib
import sys

import numpy as np
import pandas as pd

path = sys.argv[1]
seeds = int(sys.argv[2]) if len(sys.argv) > 2 else 500
TARGET = 624_489
ELIGIBLE = {"CITIZEN", "LEGAL_PERMANENT_RESIDENT", "REFUGEE", "ASYLEE"}
PAID = {2021: [1114], 2022: [3284, 2622, 662], 2023: [1312], 2024: [1702, 1403.83, 298.17]}
VAR = "takes_up_ak_permanent_fund_dividend_if_eligible"

p = pd.read_hdf(path, "person")
h = pd.read_hdf(path, "household").set_index("household_id")
w = p.person_household_id.map(h.household_weight).astype(float).to_numpy()
elig = p.immigration_status_str.isin(ELIGIBLE).to_numpy()
rep = np.array(
    [
        v > 0 and any(abs(v - k * d) <= 100 for yy in (int(y), int(y) - 1) for d in PAID.get(yy, []) for k in range(1, 9))
        for v, y in zip(p.OI_VAL, p.source_year)
    ]
)
frame = pd.DataFrame(dict(u=p.person_spm_unit_id, y=p.source_year, h=p.source_household_id, i=p.source_person_id))
first = frame.groupby("u").agg(y=("y", "first"), h=("h", "first"), i=("i", "min"))
key = p.person_spm_unit_id.map(first.y.astype(str) + ":" + first.h.astype(str) + ":" + first.i.astype(str)).to_numpy()
asec = (p.person_support_channel == "asec").to_numpy()

raw = p.A_FNLWGT.astype(float).to_numpy()
standins = {
    "A_FNLWGT, all rows": raw,
    "A_FNLWGT, ASEC rows only": raw * asec,
    "Equal weight per ASEC person": asec.astype(float),
    "Released weights (the prototype's own basis)": w,
}

units = pd.DataFrame(dict(k=key, we=w * elig, a=rep)).groupby("k").agg(we=("we", "sum"), a=("a", "any"))
draws = {}
for s in range(1, seeds + 1):
    draws[s] = np.array(
        [int.from_bytes(hashlib.blake2b(f"{s}:{VAR}:{k}".encode(), digest_size=8).digest(), "big") / 2**64 for k in units.index]
    )

kish = units.we.sum() ** 2 / (units.we**2).sum()
print(f"Units by source identity: {len(units)}; effective units, Kish, eligible people on released weights: {kish:.1f}")
free = ~units.a.to_numpy()
print(f"Non-anchored eligible people: {units.we[free].sum():,.0f}; anchored: {units.we[~free].sum():,.0f}")
for name, d in standins.items():
    d = d / d.sum() * w.sum()
    de = pd.DataFrame(dict(k=key, de=d * elig)).groupby("k").de.sum().reindex(units.index).to_numpy()
    a = units.a.to_numpy()
    we = units.we.to_numpy()
    out = []
    for s in range(1, seeds + 1):
        idx = np.where(~a)[0]
        order = idx[np.argsort(draws[s][idx])]
        cum = np.r_[de[a].sum(), de[a].sum() + de[order].cumsum()]
        n = int(np.abs(cum - TARGET).argmin())
        flag = a.copy()
        flag[order[:n]] = True
        out.append(we[flag].sum())
    out = np.array(out)
    lo, hi = np.quantile(out, [0.05, 0.95])
    within = np.mean(np.abs(out / TARGET - 1) <= 0.02)
    print(
        f"{name}: mean {out.mean():,.0f}, SD {out.std():,.0f}, 5th to 95th {lo:,.0f} to {hi:,.0f} "
        f"({lo / TARGET - 1:+.1%} to {hi / TARGET - 1:+.1%}), seeds within 2%: {within:.0%}"
    )