Add Phase 4: tidy output, MOE scaling, geometry for flows, and survey metadata

pypums/__init__.py

@@ -16,4 +16,5 @@
 from .moe import moe_sum as moe_sum
 from .moe import significance as significance
 from .pums import get_pums as get_pums
+from .survey import get_survey_metadata as get_survey_metadata
 from .variables import load_variables as load_variables

pypums/estimates.py

@@ -9,6 +9,9 @@
 from pypums.api.key import census_api_key
 from pypums.cache import CensusCache
 
+# Valid output formats.
+_VALID_OUTPUTS = frozenset({"tidy", "wide"})
+
 # Product-to-dataset mapping for PEP.
 _PRODUCT_DATASETS: dict[str, str] = {
     "population": "pep/population",
@@ -78,7 +81,6 @@ def get_estimates(
         compatibility but currently has no effect.
     output
         ``"tidy"`` (default) or ``"wide"``.
-        Not yet implemented — currently returns wide format only.
     geometry
         If True, return a GeoDataFrame with shapes.
     cache_table
@@ -93,6 +95,9 @@ def get_estimates(
     pd.DataFrame
         Population estimates data.
     """
+    if output not in _VALID_OUTPUTS:
+        raise ValueError(f"output must be 'tidy' or 'wide', got {output!r}")
+
     api_key = census_api_key(key) if key else census_api_key()
     for_clause, in_clause = build_geography_query(geography, state=state, county=county)
 
@@ -138,7 +143,7 @@ def get_estimates(
     breakdown_str = ",".join(breakdown) if breakdown is not None else ""
     cache_key = (
         f"est_{vintage}_{resolved_product}_{geography}_{state}_{county}"
-        f"_{year}_{vars_str}_{breakdown_str}"
+        f"_{year}_{vars_str}_{breakdown_str}_{output}"
     )
 
     # Check cache before calling API.
@@ -171,6 +176,23 @@ def get_estimates(
     for col in numeric_cols:
         df[col] = pd.to_numeric(df[col], errors="coerce")
 
+    # Format output.
+    if output == "tidy":
+        id_cols = ["GEOID", "NAME"] if "GEOID" in df.columns else ["NAME"]
+        # Include any breakdown columns in id_cols.
+        excluded = geo_set | set(id_cols) | set(numeric_cols)
+        breakdown_cols = [c for c in df.columns if c not in excluded]
+        id_cols = id_cols + breakdown_cols
+
+        value_cols = [c for c in numeric_cols if c in df.columns]
+        if value_cols:
+            df = df.melt(
+                id_vars=id_cols,
+                value_vars=value_cols,
+                var_name="variable",
+                value_name="value",
+            )
+
     if geometry:
         from pypums.spatial import attach_geometry
 

pypums/flows.py

@@ -9,15 +9,12 @@
 from pypums.api.key import census_api_key
 from pypums.cache import CensusCache
 
-# Core flow columns that should be numeric.
-_FLOW_NUMERIC_COLS = [
-    "MOVEDIN",
-    "MOVEDIN_M",
-    "MOVEDOUT",
-    "MOVEDOUT_M",
-    "MOVEDNET",
-    "MOVEDNET_M",
-]
+# Core flow estimate columns and their MOE counterparts.
+_FLOW_ESTIMATE_COLS = ["MOVEDIN", "MOVEDOUT", "MOVEDNET"]
+_FLOW_MOE_COLS = ["MOVEDIN_M", "MOVEDOUT_M", "MOVEDNET_M"]
+
+# All numeric flow columns.
+_FLOW_NUMERIC_COLS = _FLOW_ESTIMATE_COLS + _FLOW_MOE_COLS
 
 # Valid geography levels for migration flows.
 _VALID_GEOGRAPHIES = frozenset(
@@ -27,6 +24,13 @@
     }
 )
 
+# Z-scores for MOE confidence levels (same as acs.py).
+_Z_SCORES: dict[int, float] = {
+    90: 1.645,
+    95: 1.960,
+    99: 2.576,
+}
+
 _DEFAULT_CACHE_DIR = Path.home() / ".pypums" / "cache" / "api"
 
 
@@ -69,19 +73,16 @@ def get_flows(
         Data year (default 2019).
     output
         ``"tidy"`` (default) or ``"wide"``.
-        Not yet implemented — currently returns wide format only.
     state
         State FIPS code or abbreviation.
     county
         County FIPS code.
     msa
         Metropolitan Statistical Area code.
     geometry
-        If True, return a GeoDataFrame with shapes.
-        Not yet implemented for flows.
+        If True, return a GeoDataFrame with shapes for the origin geography.
     moe_level
         Confidence level for MOE: 90, 95, or 99 (default 90).
-        Not yet implemented — MOE columns are returned as-is at 90%.
     cache_table
         If True, cache the API response locally to avoid redundant calls.
     show_call
@@ -98,6 +99,12 @@ def get_flows(
         raise ValueError(
             f"geography must be one of {sorted(_VALID_GEOGRAPHIES)}, got {geography!r}"
         )
+    if output not in ("tidy", "wide"):
+        raise ValueError(f"output must be 'tidy' or 'wide', got {output!r}")
+    if moe_level not in _Z_SCORES:
+        raise ValueError(
+            f"moe_level must be one of {sorted(_Z_SCORES)}, got {moe_level!r}"
+        )
 
     api_key = census_api_key(key) if key else census_api_key()
 
@@ -156,7 +163,9 @@ def get_flows(
     vars_str = ",".join(variables) if variables is not None else ""
     breakdown_str = ",".join(breakdown) if breakdown is not None else ""
     cache_key = (
-        f"flows_{year}_{geography}_{state}_{county}_{msa}_{vars_str}_{breakdown_str}"
+        f"flows_{year}_{geography}_{state}_{county}"
+        f"_{msa}_{vars_str}_{breakdown_str}"
+        f"_{output}_{moe_level}"
     )
 
     # Check cache before calling API.
@@ -181,6 +190,60 @@ def get_flows(
         if col in df.columns:
             df[col] = pd.to_numeric(df[col], errors="coerce")
 
+    # Scale MOE if needed (Census API returns MOE at 90% confidence).
+    if moe_level != 90:
+        scale_factor = _Z_SCORES[moe_level] / _Z_SCORES[90]
+        moe_cols_present = [c for c in _FLOW_MOE_COLS if c in df.columns]
+        df[moe_cols_present] = df[moe_cols_present] * scale_factor
+
+    # Build GEOID for origin geography (state1 + county1).
+    geo_cols = [c for c in ("state1", "county1") if c in df.columns]
+    if geo_cols:
+        df["GEOID"] = df[geo_cols].apply(lambda row: "".join(row), axis=1)
+
+    # Format output.
+    if output == "tidy":
+        # Identify id columns (non-numeric, non-geo FIPS columns).
+        fips_cols = {"state1", "county1", "state2", "county2"}
+        id_cols = [
+            c for c in df.columns if c not in _FLOW_NUMERIC_COLS and c not in fips_cols
+        ]
+
+        est_cols = [c for c in _FLOW_ESTIMATE_COLS if c in df.columns]
+        moe_cols = [c for c in _FLOW_MOE_COLS if c in df.columns]
+
+        if est_cols:
+            est_long = df.melt(
+                id_vars=id_cols,
+                value_vars=est_cols,
+                var_name="variable",
+                value_name="estimate",
+            )
+            moe_long = df.melt(
+                id_vars=id_cols,
+                value_vars=moe_cols,
+                var_name="_moe_var",
+                value_name="moe",
+            )
+            # Map MOE variable back to estimate variable name.
+            moe_long["variable"] = moe_long["_moe_var"].str.replace(
+                "_M$", "", regex=True
+            )
+
+            df = est_long.merge(
+                moe_long[id_cols + ["variable", "moe"]],
+                on=id_cols + ["variable"],
+            )
+
+    if geometry:
+        from pypums.spatial import attach_geometry
+
+        # Map flows geography names to spatial module names.
+        geo_name = geography
+        if geography == "metropolitan statistical area":
+            geo_name = "cbsa"
+        df = attach_geometry(df, geo_name, state=state, year=year)
+
     if disk_cache is not None:
         disk_cache.set(cache_key, df, ttl_seconds=86400)
 

pypums/survey.py

@@ -0,0 +1,73 @@
+"""Survey metadata retrieval from the Census API discovery endpoint."""
+
+import pandas as pd
+
+from pypums.api.client import CENSUS_API_BASE, fetch_json
+
+
+def _fetch_data_json() -> dict:
+    """Thin wrapper so tests can mock ``pypums.survey._fetch_data_json``."""
+    return fetch_json(f"{CENSUS_API_BASE}.json")
+
+
+def get_survey_metadata(
+    year: int | None = None,
+) -> pd.DataFrame:
+    """Fetch available Census Bureau dataset metadata.
+
+    Queries the Census API discovery endpoint
+    (``https://api.census.gov/data.json``) and returns a DataFrame of
+    available datasets, optionally filtered by year.
+
+    Parameters
+    ----------
+    year
+        If provided, filter to datasets for that vintage/year.
+
+    Returns
+    -------
+    pd.DataFrame
+        Columns: ``title``, ``description``, ``vintage``, ``dataset_name``,
+        ``distribution_url``.
+    """
+    catalog = _fetch_data_json()
+
+    datasets = catalog.get("dataset", [])
+
+    rows = []
+    for ds in datasets:
+        vintage = ds.get("c_vintage")
+        title = ds.get("title", "")
+        description = ds.get("description", "")
+        dataset_name = "/".join(ds.get("c_dataset", []))
+
+        # Build the API distribution URL.
+        dist_url = ""
+        for dist in ds.get("distribution", []):
+            if dist.get("accessURL"):
+                dist_url = dist["accessURL"]
+                break
+
+        rows.append(
+            {
+                "title": title,
+                "description": description,
+                "vintage": int(vintage) if vintage is not None else None,
+                "dataset_name": dataset_name,
+                "distribution_url": dist_url,
+            }
+        )
+
+    _COLUMNS = [
+        "title",
+        "description",
+        "vintage",
+        "dataset_name",
+        "distribution_url",
+    ]
+    df = pd.DataFrame(rows, columns=_COLUMNS)
+
+    if year is not None:
+        df = df[df["vintage"] == year].reset_index(drop=True)
+
+    return df

tests/test_get_estimates.py

@@ -7,7 +7,9 @@
 * Returns a DataFrame.
 * ``product`` parameter selects population/components/housing/characteristics.
 * ``breakdown`` splits characteristics by AGEGROUP, RACE, SEX, HISP.
-* ``time_series=True`` returns multiple years of data.
+* ``time_series=True`` returns data across multiple years.
+* ``output="tidy"`` melts value columns into long format.
+* ``output="wide"`` keeps one row per geography (default pre-Phase 4 behavior).
 """
 
 from unittest.mock import patch
@@ -99,3 +101,67 @@ def test_time_series(self, estimates_api_response, fake_api_key):
                 key=fake_api_key,
             )
         assert isinstance(df, pd.DataFrame)
+
+
+class TestGetEstimatesOutput:
+    """Test output format parameter."""
+
+    def test_tidy_output_has_variable_and_value(
+        self, estimates_api_response, fake_api_key
+    ):
+        """Tidy output should have 'variable' and 'value' columns."""
+        with _mock_get_estimates(estimates_api_response):
+            df = get_estimates(
+                geography="state",
+                product="population",
+                variables=["POP_2023", "DENSITY_2023"],
+                output="tidy",
+                key=fake_api_key,
+            )
+        assert "variable" in df.columns
+        assert "value" in df.columns
+
+    def test_tidy_output_has_more_rows_than_wide(
+        self, estimates_api_response, fake_api_key
+    ):
+        """Tidy format melts variables so there are more rows."""
+        with _mock_get_estimates(estimates_api_response):
+            df_tidy = get_estimates(
+                geography="state",
+                product="population",
+                variables=["POP_2023", "DENSITY_2023"],
+                output="tidy",
+                key=fake_api_key,
+            )
+        with _mock_get_estimates(estimates_api_response):
+            df_wide = get_estimates(
+                geography="state",
+                product="population",
+                variables=["POP_2023", "DENSITY_2023"],
+                output="wide",
+                key=fake_api_key,
+            )
+        assert len(df_tidy) > len(df_wide)
+
+    def test_wide_output_keeps_columns(self, estimates_api_response, fake_api_key):
+        """Wide output keeps variable columns intact."""
+        with _mock_get_estimates(estimates_api_response):
+            df = get_estimates(
+                geography="state",
+                product="population",
+                variables=["POP_2023", "DENSITY_2023"],
+                output="wide",
+                key=fake_api_key,
+            )
+        assert "POP_2023" in df.columns
+        assert "DENSITY_2023" in df.columns
+
+    def test_invalid_output_raises(self, fake_api_key):
+        """Invalid output value should raise ValueError."""
+        with pytest.raises(ValueError, match="output must be"):
+            get_estimates(
+                geography="state",
+                product="population",
+                output="invalid",
+                key=fake_api_key,
+            )