diff --git a/ethnicolr/dict_models.py b/ethnicolr/dict_models.py index 84ae630..f2eec26 100644 --- a/ethnicolr/dict_models.py +++ b/ethnicolr/dict_models.py @@ -24,6 +24,7 @@ import logging import sys from statistics import NormalDist +from typing import cast import numpy as np import pandas as pd @@ -100,6 +101,8 @@ def census_marginal(cls) -> np.ndarray: ) marginal = weights.sum(axis=0) cls._census_marginal = marginal / marginal.sum() + # Set by the branch above on first call, cached after. + assert cls._census_marginal is not None return cls._census_marginal @classmethod @@ -107,7 +110,8 @@ def rosenman(cls, which: str) -> pd.DataFrame: if which not in cls._rosenman: path = ROSENMAN_FIRST if which == "first" else ROSENMAN_LAST df = pd.read_csv(path).dropna(subset=["name"]) - cls._rosenman[which] = df.set_index("name")[VOTER_CATS] + # A list key always yields a DataFrame; the stubs widen it. + cls._rosenman[which] = cast(pd.DataFrame, df.set_index("name")[VOTER_CATS]) return cls._rosenman[which] @classmethod @@ -164,7 +168,7 @@ def census_fn( df = EthnicolrModelClass.test_and_norm_df(df, fname_col) table = _Tables.census_first() - keys = _norm_names(df[fname_col]) + keys = _norm_names(cast(pd.Series, df[fname_col])) matched = table.reindex(keys) rdf = df.copy() @@ -181,7 +185,7 @@ def census_fn( rdf[f"{col}_lb"] = (lb * 100).round(2) rdf[f"{col}_ub"] = (ub * 100).round(2) - matched_n = int(matched[CENSUS_PCT_COLS[0]].notna().sum()) + matched_n = int(cast(pd.Series, matched[CENSUS_PCT_COLS[0]]).notna().sum()) logger.info(f"Matched {matched_n} of {len(rdf)} first names") return rdf @@ -235,8 +239,8 @@ def pred_census_name( raise ValueError("lname_col and fname_col must exist in the DataFrame") rdf = df.copy() - last_keys = _norm_names(rdf[lname_col]) - first_keys = _norm_names(rdf[fname_col]) + last_keys = _norm_names(cast(pd.Series, rdf[lname_col])) + first_keys = _norm_names(cast(pd.Series, rdf[fname_col])) last_table = _Tables.census_last(year) first_table = _Tables.census_first() @@ -354,8 +358,8 @@ def pred_voter_name( raise ValueError("lname_col and fname_col must exist in the DataFrame") rdf = df.copy() - last_keys = _norm_names(rdf[lname_col]) - first_keys = _norm_names(rdf[fname_col]) + last_keys = _norm_names(cast(pd.Series, rdf[lname_col])) + first_keys = _norm_names(cast(pd.Series, rdf[fname_col])) last_table = _Tables.rosenman("last") first_table = _Tables.rosenman("first") diff --git a/ethnicolr/pred_census_ln.py b/ethnicolr/pred_census_ln.py index a8ea8dd..8686d6f 100644 --- a/ethnicolr/pred_census_ln.py +++ b/ethnicolr/pred_census_ln.py @@ -213,11 +213,15 @@ def pred_census_ln( "this model has no calibration stats file; run " "scripts/model-training/calibrate_model.py" ) - if coverage is not None and f"{coverage:.2f}" not in stats["conformal_quantiles"]: - raise ValueError( - f"coverage must be one of {sorted(stats['conformal_quantiles'])}, " - f"got {coverage}" - ) + if coverage is not None: + # Guaranteed by the check above: stats is None only when neither prior + # nor coverage was given. + assert stats is not None + if f"{coverage:.2f}" not in stats["conformal_quantiles"]: + raise ValueError( + f"coverage must be one of {sorted(stats['conformal_quantiles'])}, " + f"got {coverage}" + ) logger.info(f"Predicting {len(df)} names using Census {year} PyTorch model") @@ -255,6 +259,7 @@ def pred_census_ln( logits = model(X_tensor) mean_probs = torch.softmax(logits / temperature, dim=1).cpu().numpy() if prior is not None: + assert stats is not None # guaranteed by the check above mean_probs = apply_prior( mean_probs, RACES, prior, stats["train_class_distribution"] ) @@ -291,6 +296,7 @@ def pred_census_ln( result["race"] = [RACES[i] for i in pred_indices] if coverage is not None: + assert stats is not None # guaranteed by the check above qhat = stats["conformal_quantiles"][f"{coverage:.2f}"] result["race_set"] = conformal_sets(mean_probs, RACES, qhat) diff --git a/ethnicolr/pred_wiki_origin.py b/ethnicolr/pred_wiki_origin.py index f9468ba..1c0e503 100644 --- a/ethnicolr/pred_wiki_origin.py +++ b/ethnicolr/pred_wiki_origin.py @@ -73,6 +73,12 @@ def pred_wiki_origin( + working[fname_col].fillna("").astype(str).str.strip() ).str.strip() + # The base declares VOCABFN/RACEFN as `str | None` because non-LSTM + # models have neither. This is an LSTM model and sets both; narrowing + # them in the subclass does not work, since a mutable ClassVar is + # invariant and `str` is not assignable to `str | None`. + assert cls.VOCABFN is not None and cls.RACEFN is not None + rdf = cls.transform_and_pred( df=working, newnamecol="__name",