diff --git a/.env.example b/.env.example index 5b4cad9..15a63ec 100644 --- a/.env.example +++ b/.env.example @@ -1,13 +1,8 @@ -# GeoPops Environment Variables +# GeoPops environment variables # Copy this file to .env and fill in your values: # cp .env.example .env +# +# .env is gitignored; never commit real credentials. # Your Census API key (get one at https://api.census.gov/data/key_signup.html) CENSUS_API_KEY= - -# Path to your Julia environment -# Examples: -# Linux: /home/username/.julia/environments/v1.9/ -# macOS: /Users/username/.julia/environments/v1.9/ -# Windows: C:/Users/username/.julia/environments/v1.9/Project.toml -JULIA_ENV_PATH= diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..d867ab0 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,35 @@ +name: release + +on: + release: + types: [published] + workflow_dispatch: + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - run: pip install build twine + - run: python -m build + - run: twine check dist/* + - uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/ + + publish: + needs: build + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write # PyPI Trusted Publishing; no API token needed + steps: + - uses: actions/download-artifact@v4 + with: + name: dist + path: dist/ + - uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..429fb25 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,39 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - name: Install + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + - name: Import check + run: python -c "import geopops; print(geopops.__all__)" + - name: Fast tests + run: pytest tests/ -m "not slow" -v --durations=10 + + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - run: pip install ruff + - run: ruff check src/ tests/ diff --git a/.gitignore b/.gitignore index 6cc9cae..1ce0d16 100644 --- a/.gitignore +++ b/.gitignore @@ -5,13 +5,16 @@ build/ dist/ wheels/ *.egg-info +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ # Virtual environments .venv - .DS_Store +# Secrets: real credentials live in .env (see .env.example) .env # Ignore local/machine-specific config overrides (keep packaged config.json tracked) @@ -20,7 +23,8 @@ src/geopops/config.local.json uv.lock -tutorials/ -tests/ +# Test *data* is large and downloaded, but test *code* is tracked +tests/data/ +tests/**/data/ -*.pptx \ No newline at end of file +*.pptx diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..6c05733 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,99 @@ +# Changelog + +All notable changes to GeoPops are documented here. This project follows [semantic versioning](https://semver.org/). + +## [0.1.8] — unreleased + +This release is an engineering pass over the whole package: licensing and packaging, a working test suite and CI, a single functional API, and substantial performance work. **It contains breaking API changes** (see below); the science and the output file formats are unchanged except where noted. + +### Licensing + +- **Added `LICENSE` (AGPL-3.0-or-later) and `NOTICE`.** GeoPops was previously distributed with no license despite `src/geopops/process_data.py` being derived from [GREASYPOP-CO](https://github.com/CDDEP-DC/GREASYPOP-CO) (Copyright 2023 Alexander Tulchinsky), which is AGPL-3.0-or-later. `NOTICE` records the GREASYPOP-CO provenance and restores the MIT license attribution for the vendored `ipfn` module. Relicensing GeoPops permissively would require permission from the GREASYPOP-CO copyright holders. + +### Breaking changes + +- **Removed the "class as verb" API.** `WriteConfig`, `RunAll`, `ForStarsim`, and `RunJulia` are gone, along with `geopops/julia.py`. Constructors no longer run pipelines as a side effect. + + | Before | Now | + |---|---| + | `geopops.WriteConfig(**pars)` | `geopops.make_config(**pars)` | + | `geopops.RunAll(pars=pars)` | `geopops.run(cfg)` | + | `geopops.DownloadData(config=cfg)` | `geopops.download_data(cfg)` | + | `geopops.ProcessData(config_dict=cfg)` | `geopops.process_data(cfg)` | + | `geopops.GeneratePop(config_dict=cfg, auto_run=True)` | `geopops.generate_pop(cfg, seed=...)` | + | `geopops.ForStarsim.People()` | `geopops.to_starsim_people(pop.pop_export_dir)` | + | `geopops.ForStarsim.GPNetwork(name=...)` | `geopops.starsim_network(name, pop.pop_export_dir)` | + | `geopops.ForStarsim.SubgroupTracking(...)` | `geopops.SubgroupTracking(...)` | + | `geopops.RunJulia()` | removed; use `geopops.generate_pop` | + + `DownloadData`, `ProcessData`, and `GeneratePop` remain as objects for running or inspecting individual stages. + +- **Config is no longer written into the installed package.** `make_config()` returns a dict; `save=True` writes `config.json` into the run's own output directory. Previously `WriteConfig` wrote into `site-packages`, which broke read-only installs, was wiped by `pip install --upgrade`, and made concurrent runs race. The packaged `config.json` is now a read-only template. +- **`PersonData` traits are config-driven.** Traits named in `config["additional_traits"]` are carried in a `TraitSchema`-positioned tuple and reached by attribute (`person.hispanic`) as before. Adding a trait no longer requires editing the dataclass, which previously raised `TypeError: unexpected keyword argument`. +- **`people.csv` column order changed.** Trait columns now follow `commuter_workplace_category` and are whatever the config asked for, instead of a hardcoded list. Read columns by name, not position. +- **Downloads raise instead of exiting.** `try_download`, `try_curl_cffi`, and `try_download_text` are replaced by a single `download(src, dst, backend=..., mode=...)` that raises `DownloadError`. The old functions called `exit(1)`, which killed the host process — including Jupyter kernels. +- **Workplace assignment results differ for a given seed.** `drawCounts` now uses `Generator.multivariate_hypergeometric`, which draws from the same distribution but consumes the RNG stream differently. Population, household, school, group-quarters, and household-network outputs are bit-identical to 0.1.7; the workplace, school-worker, GQ-worker, and non-household network outputs are statistically equivalent but not identical. + +### Added + +- **Test suite and CI.** 78 unit tests covering config, utilities, CO, and networks, plus slow end-to-end and golden-output regression tests. The previous test files referenced `geopops.RunPython`, a class that no longer existed, so nothing ran. GitHub Actions workflows run tests on Python 3.11/3.12/3.13 and lint with ruff. +- **Exception hierarchy**: `GeoPopsError` and its subclasses `ConfigError`, `DownloadError`, `DataError`, `PipelineStateError`. +- **`validate_config()`**, run automatically by `make_config()`. Unknown override keys are now an error rather than silently ignored, an unparseable `main_year` raises instead of defaulting to a 2010 vintage, and an unseeded run warns. +- **`starsim_networks()`** builds all four layers, reading each matrix file once. +- **`GeneratePop.pop_export_dir`**, so downstream steps need not reconstruct the path. +- **`allow_insecure_downloads`** config flag (default `False`). +- **Project metadata**: authors, license, keywords, classifiers, and URLs in `pyproject.toml`; dependency lower bounds; ruff and pytest configuration. + +### Fixed + +- **Starsim networks were cached in class-level state and never invalidated.** Generating a second population in the same session silently reused the first population's edges. The new functions hold no cross-call state. +- **`ForStarsim.GPNetwork` read a different config from `ForStarsim.People`**, hardcoding the packaged `config.json` while `People()` honoured `config_dict`/`base_dir`. +- **Workplace size distributions were not reproducible**: `generate_work_sizes` used numpy's unseeded global RNG. It now takes `random_seed`, threaded from `config["random_seed"]`. +- **`QualityCheck.results` printed to stdout** as a side effect of attribute access. +- **`compute_decennial_year` silently returned 2010** for any unparseable `main_year`. +- **`save_config` crashed** when given a bare filename with no directory component. +- Loop-invariant constants were rebound inside a loop body in `pull_census_data`; exceptions raised inside `except` blocks were not chained. + +### Performance + +Measured on the Spartanburg County, SC fixture (195 CBGs, ~300k people), `CO_maxgens=20000`: + +| Stage | Before | After | Speedup | +|---|---|---|---| +| CO (`process_counties`) | 2192.9 s | 41.1 s | **53×** (identical output) | +| SynthPop | 200.0 s | 18.3 s | **11×** | +| Export | 4.0 s | 4.1 s | — | +| **Total** | **~40 min** | **~64 s** | **~37×** | + +The changes behind this: + +- `anneal()` updates its running column-sum incrementally instead of re-gathering and re-summing every selected household each generation. Because the sample matrix is integer, this is exact — verified bit-identical with an unchanged RNG call sequence. +- `sqrt(target + 1)` is hoisted out of the annealing loop. +- Candidate sample sub-pools are cached by lookup key, so CBGs sharing a PUMA no longer each re-extract the same sub-matrix. +- `read_workers_by_cat` sums a NumPy view column-wise instead of `.iloc[i][col]` per household per industry (~385× on the access pattern alone). +- `find_closest` uses `argpartition` instead of `iterrows` with per-cell access (~10–64×, verified identical output). +- `drawCounts` uses one `multivariate_hypergeometric` call instead of a loop of `rng.choice(p=...)` (~16×). +- `generate_people` precomputes person attributes column-wise instead of `.iloc`/`.apply(axis=1)` per person. +- `pull_inst_workers` maintains column sums incrementally rather than recomputing the full reduction per institution. +- `generate_commute_matrices` accumulates COO triplets instead of assigning into a `lil_matrix`, and **returns** its matrices — they were previously gzipped to `od_*.csv.gz` and read straight back four lines later. +- `PersonData` uses `__slots__` with a shared trait schema: ~240 bytes per person, down from ~1.5 kB (~600 MB saved on a 500k population). + +### Removed + +- ~450 lines of dead code: `generate_location_matrices` (never called), the `generate_test_targets`/`gen_samp_test_cols`/`test_cols` scaffolding shipped in `process_data.py`, `DownloadData.pipeline()` (a hardcoded prose description of the other methods), and `geopops/julia.py`. +- Two of the three near-identical download functions. +- `urllib3.disable_warnings()` at import time, which disabled TLS warnings process-wide for every library in the host process. + +### Known issues + +- `process_data.py` and `download_data.py` still set module-level globals (`config`, `OUTPUT_DIR`, `PROCESSED_DIR`) from their step objects, so those two stages are not reentrant or safely parallelizable. Threading an explicit path/config context through them is the next structural change. +- `process_data.py` writes seven `*_test.csv` debug files on every run. +- CO is not yet parallelized across CBGs. (Much less pressing now that it runs in 41 s rather than 37 minutes, and it depends on the module-global refactor above.) +- There are no type hints, and logging still goes through `print`. +- The `tests/data` fixture is stale and gitignored, so CI cannot run the end-to-end or golden-output tests. Its `processed/` files predate both the issue #2 trait change and the `st_puma` zero-padding normalization. A small committed fixture (one or two CBGs) would fix this. + +## [0.1.7] — 2026-08 + +- Replaced the Julia implementation with pure Python. +- Reworked race/ethnicity traits: removed the non-mutually-exclusive `race_ethnicity` variable in favour of the eight PUMS categories. See [issue #2](https://github.com/GeoPopsHub/geopops/issues/2). Note that these traits are carried through from PUMS but are *not* targeted by the CO step. +- Sanitized config handling so API keys are not written into the packaged template. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..8760168 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 by the GeoPopsHub + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/README.md b/README.md index 001cd77..b7af2d6 100644 --- a/README.md +++ b/README.md @@ -19,8 +19,54 @@ There are many packages for generating agents and households from Census data, b ## Get involved GeoPops is in development and we welcome feedback! Get in touch if you've tried making a population of your own or want to become a member. You can upload your own example as a respository in the [GeoPopsHub](https://github.com/GeoPopsHub). +## Installation + +```bash +pip install geopops +``` + ## How to use -[`1_run_geopops.ipynb`](https://github.com/GeoPopsHub/sc_spartanburg_measles/blob/main/1_run_geopops.ipynb) has instructions on how to build a GeoPops population. See the repo [sc_spartanburg_measles](https://github.com/GeoPopsHub/sc_spartanburg_measles) for a detailed example of how to build a population, simulate a disease, test out interventions, and track outcomes by subgroup. + +```python +import geopops as gp + +# 1. Describe the population you want +cfg = gp.make_config( + path="data", # where downloads and results are written + geos=["45083"], # state or county FIPS (Spartanburg County, SC) + main_year=2019, + commute_states=["45", "37"], # states whose commute data to download (SC, NC) + use_pums=["45", "37"], # states whose PUMS samples to draw from + random_seed=42, # set this for reproducible runs +) + +# 2. Fetch and prepare the input data (slow; both steps cache to `path`) +gp.download_data(cfg) +gp.process_data(cfg) + +# 3. Generate the population +pop = gp.generate_pop(cfg, seed=42) + +# 4. Hand it to Starsim +ppl = gp.to_starsim_people(pop.pop_export_dir) +networks = gp.starsim_networks(pop.pop_export_dir) +``` + +Or run the whole thing at once: + +```python +pop = gp.run(cfg, seed=42) +``` + +A `CENSUS_API_KEY` is required for the download step. Put it in a `.env` file (see [`.env.example`](.env.example)) or set it in the environment; it is read automatically and is never written into the package. + +Results land in `/pop_export/`: `people.csv` and `hh.csv` describe the agents and their households, and the `adj_upper_triang_*.mtx` files hold the household, school, workplace, and group-quarters contact networks. + +[`1_run_geopops.ipynb`](https://github.com/GeoPopsHub/sc_spartanburg_measles/blob/main/1_run_geopops.ipynb) walks through building a population. See [sc_spartanburg_measles](https://github.com/GeoPopsHub/sc_spartanburg_measles) for a full example that builds a population, simulates a disease, tests interventions, and tracks outcomes by subgroup. + +## License + +GeoPops is licensed under the [GNU Affero General Public License v3.0 or later](LICENSE). It builds on [GREASYPOP-CO](https://github.com/CDDEP-DC/GREASYPOP-CO) (Copyright 2023 Alexander Tulchinsky), which is AGPL-3.0-or-later; see [NOTICE](NOTICE) for full attribution. ## Support GeoPops is a collaboration between the following institutions: diff --git a/geopops-recommendations.md b/geopops-recommendations.md new file mode 100644 index 0000000..94b8091 --- /dev/null +++ b/geopops-recommendations.md @@ -0,0 +1,880 @@ +# GeoPops engineering review + +**Scope:** `src/geopops/` (7,207 lines, 17 modules), `tests/`, packaging. Reviewed 2026-08-28 against the goal of moving from research code to a production library. + +**Verdict:** The science and the pipeline decomposition are sound — the Julia→Python port is faithful, the module boundaries (`co` → `households` → `schools` → `workplaces` → `networks` → `export`) are the right ones, and the leaf modules are already mostly plain functions. What's missing is the production layer: there is no license despite AGPL-derived code, the test suite doesn't run, there's no CI, mutable module-level globals make the pipeline non-reentrant, and several hot loops are 10–400× slower than they need to be. None of these are deep problems; most are a day or two of work each. + +--- + +--- + +> **Status: implemented.** This review was carried out and then acted on in the same session. See `CHANGELOG.md` for the resulting 0.1.8 entry, and the [Implementation status](#implementation-status) section at the end for what was done, what was deliberately left, and the measured results. Sections A–G below are the original findings, kept as written. + +## Priority summary + +| # | Item | Severity | Effort | Section | +|---|---|---|---|---| +| 1 | No `LICENSE`; `process_data.py` is AGPL-3.0-derived | **Blocker** | 1 hr | [A1](#a1-licensing) | +| 2 | Test suite references a class that no longer exists — nothing runs | **Blocker** | 2 hr | [A2](#a2-the-test-suite-is-dead) | +| 3 | No CI | **Blocker** | 2 hr | [A3](#a3-no-ci) | +| 4 | `exit(1)` inside library code | **High** | 30 min | [A4](#a4-exit1-in-library-code) | +| 5 | Mutable module globals (`OUTPUT_DIR`, `PROCESSED_DIR`, `config`) | **High** | 1 day | [A5](#a5-mutable-module-level-globals) | +| 6 | Library writes `config.json` into its own installed package directory | **High** | 4 hr | [D1](#d1-config-is-written-into-site-packages) | +| 7 | `ForStarsim` class-level network cache is never invalidated | **High** | 1 hr | [B3](#b3-forstarsim-a-namespace-pretending-to-be-a-class) | +| 8 | Constructors that run 10-minute pipelines (`auto_run=True`) | **High** | 4 hr | [B1](#b1-classes-that-should-be-functions) | +| 9 | `anneal()` recomputes the full summary every generation | **High** | 1 hr | [C1](#c1-copy-fixes-co-annealing-10) | +| 10 | `read_workers_by_cat()` uses `.iloc[i][col]` in a triple loop | **High** | 30 min | [C2](#c2-copy-read_workers_by_cat-385) | +| 11 | Global SSL verification disabled at import time | **Medium** | 2 hr | [A6](#a6-tls-verification-disabled-globally) | +| 12 | `drawCounts()`, `find_closest()`, `households.iloc` hot loops | **Medium** | 3 hr | [C3](#c3-drawcounts-16)–[C5](#c5-per-person-iloc-in-households) | +| 13 | Default runs are silently non-reproducible (no `random_seed`) | **Medium** | 2 hr | [A7](#a7-reproducibility) | +| 14 | ~450 lines of dead code | **Medium** | 1 hr | [E1](#e1-dead-code-450-lines) | +| 15 | Zero type hints, zero logging, 175 `print()` calls | **Medium** | 2 days | [E5](#e5-printing-instead-of-logging), [E6](#e6-no-type-hints) | +| 16 | Redundant disk round-trips inside a single pipeline run | **Medium** | 1 day | [D2](#d2-round-trips-that-serve-no-purpose) | + +--- + +## A. Blockers for a production library + +### A1. Licensing + +`src/geopops/process_data.py` opens with: + +> Copyright 2023 Alexander Tulchinsky … Greasypop is free software: you can redistribute it and/or modify it under the terms of the **GNU Affero General Public License** … version 3 … + +The repository has **no `LICENSE` file**, and `pyproject.toml` declares no `license` field. GeoPops v0.1.7 is being published to PyPI in this state. AGPL-3.0 is strongly copyleft: a derivative work that links this module must itself be AGPL-3.0-or-later. Because `process_data.py` is imported by `__init__.py`, that covers the whole distributed package. + +Three options, in order of preference: + +1. **Adopt AGPL-3.0-or-later for GeoPops.** Add `LICENSE`, set `license = "AGPL-3.0-or-later"` and the matching classifier in `pyproject.toml`, and note the Greasypop provenance in the README. Simplest and unambiguously correct. Note this will constrain downstream users (including Starsim integrations, which are MIT). +2. **Get relicensing permission** from Alexander Tulchinsky / One Health Trust for the derived portions, then license GeoPops under MIT/BSD to match Starsim. +3. **Rewrite `process_data.py` clean-room** — expensive and probably not worth it. + +Also: `src/geopops/ipfn.py` is a vendored copy of the third-party [`ipfn`](https://pypi.org/project/ipfn/) package with its license header stripped. Either add `ipfn` as a dependency (preferred — it's actively maintained and it's 300 lines you now own) or restore the upstream copyright header and record it in a `NOTICE` file. + +**Do this first.** Everything else is engineering; this is legal exposure on an already-published artifact. + +### A2. The test suite is dead + +``` +tests/test_python_workflow.py:41: r = geopops.RunPython() +``` + +`RunPython` does not exist. It was renamed to `GeneratePop` and the tests were never updated: + +``` +$ python -c "import geopops; print('RunPython' in dir(geopops))" +False +``` + +Every test in `test_python_workflow.py` fails at fixture setup. `test_julia_workflow.py` is equally stale — it calls `GPNetwork(name='homenet', beta_value=1.0)`, but the parameter has been `edge_weight` since the rename. + +Three further problems with the test setup: + +- **Import-time side effects.** Both files call `geopops.WriteConfig(**pars_geopops)` at module scope (`test_python_workflow.py:19`, `test_julia_workflow.py:19`). Merely *collecting* the tests writes files to the installed package directory. Move this into a fixture. +- **`tests/` is in `.gitignore`.** Three files are tracked because they predate the rule, but any new test file is silently ignored. Remove that line — it is actively hostile to growing the suite. +- **No unit tests.** The only tests are end-to-end smoke tests over a full county. There is nothing covering `lrRound`, `ranges`, `drawCounts`, `FTdist`, `urbanization_lookup`, `split_lognormal`, or the config merge logic — all pure functions that are trivial to test and easy to break. + +Suggested structure: + +``` +tests/ + test_utils.py # pure functions, <1s, no data + test_config.py # merge/override/sanitize logic + test_co.py # anneal on synthetic targets, seeded determinism + test_networks.py # SBM/small-world/complete on toy inputs + test_regression.py # golden-output check on tests/data (marked slow) + test_workflow.py # the current end-to-end smoke test (marked slow) +``` + +Add a **golden-output regression test**: with a fixed `random_seed`, run the pipeline on the checked-in Spartanburg fixture and assert a hash of `people.csv` and the network `.mtx` files. That single test is what lets you refactor the hot loops in section C with confidence. + +### A3. No CI + +There is no `.github/` directory. For a package on PyPI that means nothing verifies that a commit imports cleanly, let alone passes tests. Minimum viable: + +```yaml +# .github/workflows/test.yml +on: [push, pull_request] +jobs: + test: + strategy: + matrix: {python-version: ["3.11", "3.12", "3.13"]} + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: {python-version: "${{ matrix.python-version }}"} + - run: pip install -e ".[dev]" + - run: pytest tests/ -m "not slow" -v +``` + +Add a `release.yml` that builds and publishes on tag via PyPI Trusted Publishing, and a scheduled weekly run of the slow tests (the Census/LODES endpoints change, and you want to know before a user does). + +`requires-python = ">=3.11"` is currently untested against any interpreter. + +### A4. `exit(1)` in library code + +``` +download_data.py:123: exit(1) +download_data.py:181: exit(1) +download_data.py:244: exit(1) +``` + +A failed download terminates the host process. In a Jupyter notebook — the primary documented workflow — this kills the kernel and discards everything in memory. `exit` is also the `site` builtin, not `sys.exit`; it isn't guaranteed to exist under `python -S` or in frozen environments. + +Replace with a real exception: + +```python +class GeoPopsDownloadError(RuntimeError): + """Raised when a required data file could not be downloaded.""" + +# ... +raise GeoPopsDownloadError(f"Download failed after {retries} attempts: {src}") +``` + +While you're there, define a small exception hierarchy (`GeoPopsError` → `ConfigError`, `DownloadError`, `DataError`) so callers can catch GeoPops failures specifically. Right now everything is `Exception`, `KeyError`, or `RuntimeError`, and there are 11 broad `except Exception` handlers that swallow context. + +### A5. Mutable module-level globals + +This is the single biggest structural problem in the codebase. + +```python +# process_data.py:37-40 +config = None +OUTPUT_DIR = BASE_DIR +PROCESSED_DIR = os.path.join(OUTPUT_DIR, "processed") + +# process_data.py:1676-1680, inside ProcessData.__init__ +global config, OUTPUT_DIR, PROCESSED_DIR +config = self.config +OUTPUT_DIR = self.config.get("path", self.base_dir) +PROCESSED_DIR = os.path.join(OUTPUT_DIR, "processed") +``` + +`download_data.py:1236-1237` does the same with `OUTPUT_DIR`. `PROCESSED_DIR` alone is referenced 38 times across `process_data.py`; the ~1,600 lines of module-level functions there are not actually functions of their arguments — they're functions of hidden global state that a constructor happens to set. + +Consequences: + +- **Non-reentrant.** Two `ProcessData` instances cannot coexist. The second silently repoints the first's file paths. +- **Cannot be parallelized.** This forecloses the `sc.parallelize()` work in issue #6 for anything that touches these modules. +- **Cannot be unit-tested** without monkeypatching module globals. +- **Import-order dependent.** `read_acs()` called before any `ProcessData` exists reads from the *package* directory. + +The fix is mechanical, if tedious: thread a small immutable context through the call chain. + +```python +@dataclass(frozen=True) +class Paths: + root: Path + @property + def processed(self) -> Path: return self.root / "processed" + @property + def census(self) -> Path: return self.root / "census" + @property + def pums(self) -> Path: return self.root / "pums" + # ... + +def read_acs(table, paths: Paths, geos=None): ... +``` + +Every function that currently reads a global takes `paths` (and `config` where needed) as an explicit parameter. Nothing about the pipeline logic changes. Do this behind the golden-output regression test from A2 and it's a safe refactor. + +### A6. TLS verification disabled globally + +```python +# download_data.py:18 — at import time +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) +``` + +Importing `geopops` silently disables InsecureRequestWarning **for the entire host process**, including unrelated libraries. Combined with `verify=False` at three call sites (`download_data.py:104, 157, 222` — one of them unconditional in `try_curl_cffi`), downloaded Census/LODES data is not authenticated. + +- Move `disable_warnings` out of module scope into the narrowest possible `warnings.catch_warnings()` block. +- Make the insecure fallback opt-in: `config["allow_insecure_downloads"]`, defaulting to `False`, and emit a loud `warnings.warn()` when it engages. +- `try_curl_cffi` passes `verify=False` unconditionally — that one should default to verifying and only fall back if configured to. + +The browser-impersonation headers are pragmatic given how the federal data portals behave; keep them, but they don't require disabling verification. + +### A7. Reproducibility + +`GeneratePop._make_stage_random_seeds` (`generate_pop.py:72-88`) is a well-built piece of work — `SeedSequence.spawn(5)` gives each stage an independent, deterministic stream. But: + +- **`random_seed` is absent from the shipped `config.json`.** `GeneratePop` reads `config.get("random_seed")` → `None` → every stage gets an OS-entropy RNG. So the default path is non-reproducible, and `RunAll` never exposes a seed parameter at all. +- **`process_data.py:1569` uses the unseeded legacy global RNG:** + ```python + sim_dist = [np.concatenate([np.random.randint(l,h,np.int64(s)) for ... + ``` + Workplace size distributions are therefore non-reproducible regardless of what seed the caller passes. +- **`geopops_starsim.py:218`** defaults the edge-flip seed to `0` when `random_seed` is missing, so the network endpoint shuffling is deterministic while everything upstream of it isn't. Inconsistent. + +Fixes: add `"random_seed": null` to the shipped config with a documented meaning; add `random_seed=` to `RunAll`; convert line 1569 to a passed-in `Generator`; and have the pipeline record the *effective* seed (including an auto-generated one when the user passes `None`) into the run's output directory so any run can be replayed. + +--- + +## B. Classes vs functions + +This addresses GeoPopsHub issue #4 directly. The codebase has 11 classes. Four earn their keep; seven are namespaces, and two of those actively cause bugs. + +### The three patterns in play + +**Pattern 1 — the constructor that runs the pipeline.** `WriteConfig`, `RunAll`, `DownloadData`, `ProcessData`, `QualityCheck` all end `__init__` with `if auto_run: self.run_all()`. `WriteConfig` doesn't even offer the escape hatch — constructing it always writes to disk. + +This is the core problem. `DownloadData(config=cfg)` performs ten minutes of network I/O as a side effect of object construction. Constructors should construct. Beyond the principle, it produces concrete friction: + +- Defaults are inconsistent (`DownloadData`/`ProcessData`/`RunAll`/`QualityCheck` default `auto_run=True`; `GeneratePop` defaults `False`; `WriteConfig` has no flag). +- `GeneratePop` carries *both* `auto_run` and a legacy `run_all` parameter that shadows the `run_all` **method** name (`generate_pop.py:28, 63-66`). +- The returned object is discarded at every call site. In `run_all.py:73-87`, `WriteConfig(...)`, `DownloadData(...)`, and `ProcessData(...)` are constructed purely for effect — three objects created and dropped. +- You can't inspect what a step *would* do before running it. + +**Pattern 2 — the class as namespace.** `ForStarsim` is the clearest case (see B3). `WriteConfig` and `RunAll` are the same thing with extra steps. + +**Pattern 3 — the class as counter.** `Indexer` (`utils.py:54-63`) and `DummyGenerator` (`workplaces.py:70-78`) are each a class wrapping a single integer. + +### B1. Classes that should be functions + +| Class | Verdict | Replacement | +|---|---|---| +| `WriteConfig` (`config.py:94`) | **Function.** `__init__` builds an overrides dict and calls `run_all()`. `get_pars()` just prints. | `write_config(path=None, **overrides) -> dict` | +| `RunAll` (`run_all.py:31`) | **Function.** Holds three attributes, all consumed once by `run_all()`. | `run_all(pars=None, config=None, seed=None, verbose=1) -> Population` | +| `QualityCheck` (`process_data.py:2078`) | **Function.** Returns a dict of diagnostics; the class adds nothing. Its `results` **property prints to stdout** (line 2148-2153) — a property with I/O side effects is a trap. | `quality_check(paths) -> dict` + `print_quality_check(results)` | +| `ForStarsim` (`geopops_starsim.py:257`) | **Module of functions.** See B3. | `to_starsim_people()`, `starsim_network()`, `SubgroupTracking` | +| `Indexer` (`utils.py:54`) | **Delete.** `d.setdefault(k, len(d) + 1)` is the whole implementation. | inline | +| `DummyGenerator` (`workplaces.py:70`) | **Closure or `itertools.count`.** | `count = itertools.count(1)` | +| `ipfn` (`ipfn.py:9`) | **Vendored third party.** | Depend on upstream `ipfn` | + +`ProcessData` and `DownloadData` are the interesting middle cases. Both are *currently* namespaces around globals (A5), but both have a genuine reason to exist once the globals go: they expose per-step methods (`pull_pums_data()`, `generate_targets()`, …) so users can re-run one stage. Keep them, but: + +- Make them hold real state (resolved `Paths`, validated config, an RNG) rather than mirroring it into module globals. +- Drop `auto_run`; make `run_all()` an explicit call. +- Provide thin function wrappers for the common case: `download_data(config)` / `process_data(config)`. + +`GeneratePop` is the one class that unambiguously earns its keep — it holds 18 genuine pipeline intermediates between `CO()`, `SynthPop()`, and `Export()`, and users legitimately want to inspect them. Even so, I'd restructure it so the stages are free functions and the class is the *result* container: + +```python +@dataclass +class Population: + """Result of a GeoPops run; holds all pipeline intermediates.""" + config: dict + co_results: dict | None = None + people: dict | None = None + # ... + def export(self, path): ... + def to_starsim(self): ... + +def generate_pop(config, *, seed=None, verbose=1) -> Population: ... +``` + +That gives you `pop = geopops.generate_pop(cfg)` for the 95% case and keeps the object for inspection. `_ForStarsimGPNetwork` and `_ForStarsimSubgroupTracking` should stay classes — they subclass `ss.Network` and `ss.Analyzer` and are genuinely polymorphic. + +### B2. Suggested public API + +The current API mixes CamelCase classes-as-verbs with an inconsistent parameter vocabulary: `DownloadData(config=)`, `ProcessData(config_dict=)`, `GeneratePop(config_dict=)`, `ForStarsim(config_dict=)`, `RunAll(config_dict=, pars=)`. Pick one name (`config`) and use it everywhere. + +```python +import geopops as gp + +cfg = gp.make_config(geos=["45083"], main_year=2019, path="data") + +gp.download_data(cfg) # verbs, not nouns +gp.process_data(cfg) +pop = gp.generate_pop(cfg, seed=42) # returns a Population +pop.export() +ppl = pop.to_starsim() + +# or, all at once: +pop = gp.run(cfg, seed=42) +``` + +Keep the existing class names as thin deprecated aliases for one minor version so you don't break the GeoPopsHub example repos. + +### B3. `ForStarsim`: a namespace pretending to be a class + +`ForStarsim` (`geopops_starsim.py:257-407`) illustrates every cost of the pattern: + +**Its `__init__` is entirely vestigial.** Lines 270-291 set `self.base_dir`, `self.config`, `self.path` — and no method reads any of them. `People` is a `@classmethod` that calls `cls._load_config()`; `GPNetwork` and `SubgroupTracking` are `@staticmethod`s. `main()` (line 409) constructs an instance and returns it unused. + +**Class-level mutable state is a live bug.** Lines 264-268 and 204-207: + +```python +class ForStarsim: + _net_h = None; _net_s = None; _net_w = None; _net_g = None + +def _ensure_networks_created(self): + if ForStarsim._net_h is None: + self._create_networks() +``` + +This cache is never invalidated. Generate one population, call `GPNetwork('homenet')`, then generate a *second* population in the same session and call `GPNetwork('homenet')` again — you silently get the first population's edges. In a notebook comparing two counties this produces wrong results with no error. As a module-level function taking an explicit path, the bug cannot occur. + +**The config source is inconsistent.** `People()` honours `config_dict` and `base_dir`, but `_create_networks` hardcodes the package directory (line 212): + +```python +cfg_path = os.path.join(BASE_DIR, "config.json") +``` + +So in `run_all.py:96-100`, `ForStarsim.People(config_dict=effective_config, base_dir=self.base_dir)` uses the caller's config while the four subsequent `GPNetwork(...)` calls read a different one from site-packages. If `base_dir` was overridden, these disagree. + +Replace the whole class with module functions taking explicit paths: + +```python +def to_starsim_people(pop_export_dir, *, save=True) -> ss.People: ... +def starsim_network(name, pop_export_dir, *, edge_weight=1.0, seed=None) -> ss.Network: ... +``` + +--- + +## C. Performance + +All figures below are measured on this machine, not estimated. Benchmark scripts are reproducible from the snippets given. + +### C1. `anneal()`: recomputes the full summary every generation — **9.6× (bit-identical)** + +This is the highest-value single change in the codebase and it directly answers issue #6. + +```python +# co.py:35-48 — the inner loop +while True: + gen += 1 + cidx = rng.integers(len(c0)) + orig = c0[cidx] + c0[cidx] = rng.integers(n_samples) + summary = samples[c0, :].sum(axis=0, keepdims=True) # <-- O(n_hh x n_cols), every generation + E1 = FTdist(summary, targ) +``` + +Exactly **one** of `n_hh` selected households changes per generation, yet the code re-gathers all of them (allocating an `n_hh × n_cols` array) and re-sums. With `n_hh ≈ 800` and `n_cols ≈ 120` that's ~96,000 element operations per generation to reflect a change to 120 of them. + +Because `samples` is `int64`, the incremental update is *exactly* equivalent — no floating-point drift: + +```python +summary += samples[new] +summary -= samples[orig] +# on reject: +summary += samples[orig] +summary -= samples[new] +``` + +Also hoist `np.sqrt(targ + 1.0)` out of the loop — it's constant, and `FTdist` recomputes it every generation. + +Measured, with `maxgens` forced to 5,000 and an identical RNG stream: + +``` +bit-identical indices: True | same gens: True | E0 0.2967676919127237 == 0.2967676919127237 +time 0.686s -> 0.071s (9.6x, RNG stream preserved) +``` + +**The output is byte-for-byte identical and the RNG call sequence is unchanged**, so this is a drop-in replacement requiring no revalidation. The speedup grows with `n_hh`, so it's largest exactly where runs are slowest — dense urban CBGs. + +Batching the RNG draws (`rng.integers(..., size=4096)` outside the loop) gets a further ~1.4× to **13.8×** total, but it *does* change the RNG stream, so results shift (statistically equivalent, not identical). Take it as a second, separately-validated step. + +Do this before reaching for Numba. Numba on top of the incremental version is worth maybe another 3–5× (the loop becomes ~10 µs/gen of mostly-NumPy overhead on ~120 elements), but it adds a compiled dependency. Get the algorithmic 10× free first, then measure. + +**Parallelization (issue #6):** parallelize over **counties**, at `co.py:169`. Each county iteration is fully independent — it reads shared arrays and writes to its own `all_co_results[c]`. `sc.parallelize` over counties is clean. Parallelizing inside `optimize()` over CBGs would be finer-grained and better for single-county runs (the common case!), and is also safe — each `anneal()` call is independent — but you'd need to spawn per-CBG child RNG streams from a `SeedSequence` to keep results reproducible. Given a typical run is one county, **the per-CBG level is the more useful axis**; do that one. + +### C2. `read_workers_by_cat()`: `.iloc[i][col]` in a triple loop — **~385×** + +```python +# workplaces.py:96-100 +total = sum( + hh_samps.iloc[hh_idx[x] - 1][cat_col] + for x in hhvec + if x in hh_idx and pd.notna(hh_samps.iloc[hh_idx[x] - 1][cat_col]) +) +``` + +`hh_samps.iloc[i]` constructs a fresh pandas `Series` for the whole row, then `[cat_col]` pulls one scalar out of it. The `pd.notna` guard does it **a second time**. This sits inside `for county → for cbg → for cat_code`, so for a 500-CBG county with 15 industry codes and ~1,500 households per CBG, that's roughly **22 million Series constructions**. + +Measured cost of the access pattern in isolation: + +``` +.iloc[i][col] x2000: 28.9 ms vs numpy 75 us -> 385x +``` + +Restructure to hoist the array conversion and do all categories at once: + +```python +arr = hh_samps[cat_cols].to_numpy(dtype=float) # once +row_of = {s: i for i, s in enumerate(hh_samps['SERIALNO'])} +for ori, hhvec in cbg_dict.items(): + rows = [row_of[x] for x in hhvec if x in row_of] + totals = np.nansum(arr[rows], axis=0) # all 15 categories at once + for cat_code, t in zip(ind_codes, totals): + workers_by_cat[cat_code][ori] = int(t) +``` + +This also removes the 15× redundancy of re-walking `hhvec` per category. Expect this loop to go from minutes to well under a second. + +### C3. `drawCounts()`: **16×** + +```python +# utils.py:104-117 +for _ in range(n): + probs = v.astype(float) / v.sum() + idx = rng.choice(len(v), p=probs) + v[idx] -= 1 +``` + +This is sampling *n* items without replacement from a multiset of counts — which is exactly `Generator.multivariate_hypergeometric`, a single C call. The Python version instead does `n` calls to `rng.choice(p=...)`, each of which internally normalizes and cumsums the full probability vector. + +```python +def drawCounts(v, n=1, rng=None): + rng = rng or np.random.default_rng() + n = min(int(n), int(v.sum())) + if n <= 0: + return [] + drawn = rng.multivariate_hypergeometric(v, n) + v -= drawn + return np.repeat(np.arange(len(v)), drawn).tolist() +``` + +Measured (n=200 draws over 1,500 origin bins): **6.19 ms → 0.38 ms, 16×**. This is called once per workplace and once per institution, so it runs tens of thousands of times per county. + +The distribution is identical; the *order* of returned indices differs (the current version returns draw order, the replacement returns bin order). Check whether `pull_inst_workers` and `generate_workplaces` depend on that ordering — they appear to consume it as a bag, but confirm against the golden-output test. If order matters, `rng.permutation()` the result. + +Two related fixes in the same call sites: + +- `pull_inst_workers` (`workplaces.py:218`) recomputes `colsums = count_matrix.sum(axis=0)` — a full O(rows×cols) reduction — **once per institution**. Compute it once and decrement it as columns are drawn down. +- `generate_workplaces` (`workplaces.py:265-267`) does `count_matrix[:, col].copy()` and writes back per workplace. Slice once per destination column, mutate in place, write back once. + +### C4. `find_closest()`: **64× (verified identical output)** + +```python +# schools.py:40-45 +for _, row in distmat.iterrows(): + dists = [(s, row[s]) for s in valid_cols if pd.notna(row[s])] + dists.sort(key=lambda x: x[1]) + top = dists[:n] +``` + +`iterrows()` with per-cell `row[s]` scalar access, wrapped in an outer loop over **14 grade levels** — so the full CBG × school distance matrix is walked 14 times, cell by cell, to find 4 nearest schools per row. + +```python +sub = distmat[valid_cols].to_numpy(dtype=float) +sub = np.where(np.isnan(sub), np.inf, sub) +order = np.argpartition(sub, n, axis=1)[:, :n] +rows = np.arange(len(sub))[:, None] +order = order[rows, np.argsort(sub[rows, order], axis=1)] # sort just the top-n +``` + +Measured on a 600 CBG × 400 school matrix, and **verified to produce identical output**: + +``` +find_closest per grade: iterrows 489 ms vs vectorized 7.6 ms -> 64x +identical: True +x14 grades: 6.8s -> 0.11s +``` + +### C5. Per-person `.iloc` in `households.py` + +```python +# households.py:203, in the innermost per-person loop +row = p_samps.iloc[r - 1] +``` + +Same 385× pattern as C2, executed once per **person** in the synthetic population — hundreds of thousands of times per county. Convert `p_samps` to a dict of NumPy arrays (or `itertuples()`) once, before the loop. + +The two preceding lines are also row-wise `apply(axis=1)` over the whole sample frame: + +```python +# households.py:175-183 +p_samps['ind_code'] = p_samps[ind_colnames].apply(lambda row: first_true(row.values), axis=1) +p_samps['com_cat'] = p_samps.apply(lambda row: (row['ind_code'] + 1) if ... , axis=1) +``` + +`first_true` over a boolean row is `argmax`: + +```python +ind = p_samps[ind_colnames].to_numpy(dtype=bool) +has_any = ind.any(axis=1) +p_samps['ind_code'] = np.where(has_any, ind.argmax(axis=1), None) +p_samps['com_cat'] = np.where(p_samps['commuter'].to_numpy(bool) & has_any, + ind.argmax(axis=1) + 1, None) +``` + +### C6. Memory: `sample_lookup` and `PersonData` + +**`sample_lookup`** (`co.py:108-115`) returns *one full-length boolean mask per CBG*. For a 500-CBG county against 100,000 PUMS samples that's 500 × 100,000 = 50 MB of masks, most of them duplicates — many CBGs share a PUMA. Then `anneal` (`co.py:20`) does `all_samples[mask, :]`, materializing a fresh sub-matrix copy **per CBG**, even when the mask is identical to the previous one. + +Group once, index many: + +```python +groups = samp_geo.groupby(col).indices # dict[value -> index array] +# then per CBG: idx = groups.get(puma, EMPTY) +``` + +Cache the extracted `samples[idx]` sub-matrix keyed by group value. For a single county most CBGs fall in a handful of PUMAs, so this collapses ~500 gathers into ~5. + +**`PersonData`** (`utils.py:19-39`) is a plain dataclass with 18 fields, instantiated once per person. Without `slots`, each instance carries a `__dict__` — roughly 600+ bytes per person, so ~300 MB for a 500k population, plus a dict of 500k tuple keys on top. `@dataclass(slots=True)` is a one-line change that cuts the per-instance overhead by more than half. The larger win would be a columnar representation (a DataFrame or a dict of arrays) instead of half a million small objects, but that's a bigger refactor — do `slots=True` now, consider columnar later. + +### C7. Other hot spots + +- **`generate_commute_matrices`** (`workplaces.py:346, 385`) accumulates into `sparse.lil_matrix` inside a per-origin loop. `lil` assignment is slow; collect `(row, col, val)` triplets in Python lists and build one `coo_matrix` at the end. +- **`connect_SBM`** (`networks.py:27`) builds group membership with a nested comprehension that is O(n_groups × n_keys); use a `defaultdict(list)` single pass. The zero-degree fix-up loop (lines 66-69) calls `g.degree(v)` per node — fine for small workplaces, but it's inside the per-employer loop that runs tens of thousands of times. +- **`process_data.py` writes 7 `*_test.csv` debug files** unconditionally on every run (lines 678, 737, 1223, 1286-1291). Pure I/O waste in production; gate behind a `debug` flag or delete. + +--- + +## D. Files vs in-memory + +The pipeline currently persists **everything** at every stage boundary. Some of that is right; a lot of it isn't. + +### The current data flow + +``` +DownloadData --> {path}/census, /pums, /geo, /work, /school [~GB, network] +ProcessData --> {path}/processed/*.csv (~20 files, incl. 7 *_test.csv) +GeneratePop.CO <-- reads processed/*.csv + .SynthPop --> processed/od_*.csv.gz then immediately reads them back + .Export --> {path}/pop_export/*.csv, *.mtx +ForStarsim.People <-- re-reads pop_export/*.csv, --> people_all.csv, ppl.pkl +ForStarsim.GPNetwork<-- re-reads *.mtx, --> starsim/net_*.csv +``` + +**What should stay on disk:** + +- **Downloaded raw data.** Expensive, remote, rate-limited, and the whole point is to cache it. Correct as is. +- **`processed/`.** This is a genuine checkpoint — it takes ~5 minutes, it's the boundary where a user might swap in their own inputs, and CO is re-run against it many times during tuning. Correct as is. +- **`pop_export/`.** The deliverable. Correct as is. + +**What shouldn't:** + +### D1. Config is written into site-packages + +`WriteConfig.run_all()` (`config.py:150-154`) writes `config.json` **into the installed package directory** (`BASE_DIR = os.path.dirname(__file__)`). `load_config()`, `ForStarsim._load_config()`, `julia.load_config()`, `ProcessData.__init__`, and `DownloadData.__init__` all read it back from there by default. + +For a library on PyPI this is the wrong model: + +- Breaks on read-only installs, containers, and system-managed site-packages. +- `pip install --upgrade geopops` silently wipes the user's settings. +- Two users of the same shared install clobber each other; two concurrent runs race. +- The config in site-packages is invisible to the user's version control, so runs aren't reproducible from the repo alone. +- It required inventing the `sanitize=True` machinery (`config.py:38-52`) specifically to keep API keys from being written into the package — a workaround for a problem created by the design. + +Config should live in the **user's working directory or output directory**, defaulting to `./geopops.json` or `{path}/config.json`, with the packaged file treated as a read-only template. `make_config()` should return a dict and let the caller decide whether to persist it. Note that the `.env` / `CENSUS_API_KEY` handling (`config.py:71-74`) is already the right pattern — keep that. + +### D2. Round-trips that serve no purpose + +**`od_*.csv.gz`.** `generate_jobs_and_workers` calls `generate_commute_matrices(data_dir)` (`workplaces.py:417`), which computes per-industry OD matrices and gzips them to `processed/od_*.csv.gz`. Four lines later, `calc_od_counts` reads all 15 of them straight back. The matrices never leave the function's own call stack. This is gzip compression + CSV serialization + parsing of a multi-megabyte sparse matrix, entirely for nothing. + +Have `generate_commute_matrices` **return** the matrices, and make writing them optional (`save=True` for the checkpointing/debugging value): + +```python +od_matrices = generate_commute_matrices(paths, save=save_intermediates) +origin_labels, dest_labels, od_counts = calc_od_counts(..., od_matrices=od_matrices) +``` + +**`GeneratePop` → `Export` → `ForStarsim`.** After `SynthPop()`, the full population is in memory (`self.people`, `self.households`, `self.adj_hh`, …). `Export()` writes it to CSV/MTX. `ForStarsim.People()` then re-reads those CSVs, does five merges, and reconstructs the same information. In `run_all.py:89-100` this happens within a single function call — the in-memory objects are alive the entire time and are simply ignored. + +`ForStarsim.People` should accept either a `Population` object *or* a path: + +```python +def to_starsim_people(pop_or_path): ... +``` + +Keep the from-disk path — loading a previously-generated population is a real use case — but don't force a serialization round-trip when the data is already in hand. Same for `_create_networks`, which re-reads the `.mtx` files that `export_networks` just wrote from `adj_hh`/`adj_sch`/`adj_wp`/`adj_gq` still in memory. + +### D3. Format choices + +- **CSV for adjacency data.** `adj_mat_keys.csv` and `people.csv` have one row per agent; for a large county these are hundreds of MB of text that then get re-parsed by `ForStarsim`. `pyarrow` is already a dependency — Parquet would be several times smaller and much faster to read, with dtypes preserved (which would also eliminate the `low_memory=False` and `.astype(str).replace({'na': ...})` string-coercion dance at `geopops_starsim.py:323-326`). Keep CSV as an export option for interoperability; make Parquet the internal format. +- **`.mtx` for networks.** Fine and portable, but `scipy.sparse.save_npz` is faster and smaller if the consumer is always Python. +- **`ppl.pkl`** (`geopops_starsim.py:389`) — pickle is version-fragile across Starsim releases. Fine as a convenience cache; make sure nothing in the pipeline *depends* on being able to read it back. + +### D4. Intermediate results as return values + +More broadly: the leaf modules (`co`, `households`, `schools`, `workplaces`, `networks`) are already written as functions that take data and return data — that's good design and it's why the pipeline is testable in principle. The friction is that several of them reach out to the filesystem for config in the middle of their work: + +```python +workplaces.py:406: config = tryJSON(os.path.join(data_dir, 'config.json')) +schools.py:62: config = tryJSON(os.path.join(data_dir, 'config.json')) +co.py:155: config = tryJSON(os.path.join(data_dir, 'config.json')) +``` + +`households.generate_people` already does this correctly — it takes `config=None` and only falls back to disk (`_resolve_config`, lines 154-158). Extend that pattern to the other three. Note that `tryJSON` **silently returns `{}` on any failure**, so a typo'd path means every tuning parameter silently reverts to its hardcoded default and the run completes with wrong numbers and no warning. That's the worst kind of failure mode for a scientific tool. + +--- + +## E. Conciseness and clarity + +### E1. Dead code (~450 lines) + +| Location | Lines | What | +|---|---|---| +| `process_data.py:2198-2370` | ~172 | `generate_test_targets`, `gen_samp_test_cols`, `test_cols` — test scaffolding shipped in the package, called by nothing | +| `download_data.py:1438-1585` | ~147 | `DownloadData.pipeline()` — prints a hardcoded description of what the other methods do; guaranteed to drift out of sync. This is documentation, not code | +| `networks.py:247-335` | 89 | `generate_location_matrices` — never called | +| `julia.py` | 123 | Legacy Julia path. `main()` (line 117) has its only real statement commented out | +| `process_data.py`, various | — | ~25 commented-out `# print(...)` and `# ....to_csv(...)` lines (223, 232, 238, 246, 490, 496, 507, 591, 632, 707, 735, …) | + +`julia.py` deserves a decision rather than deletion by default: it's the reference implementation the Python port was validated against. Either keep it deliberately (document it as "reference only", exclude from `__all__`) or drop it and rely on the git history plus a golden-output test. Shipping it in `__all__` as a peer of `GeneratePop` implies it's supported, and `RunJulia.__init__` raises `ValueError` for any user who hasn't configured a Julia environment. + +### E2. Duplicated helpers + +- **`tryJSON`** is defined twice: `utils.py:11` (silently returns `{}`) and `process_data.py:25` (prints a warning). Different behaviour, same name. +- **`lrRound`** is defined twice with *different implementations*: `utils.py:66` (NumPy) and `process_data.py:44` (pandas Series, mutating in place via label indexing). Two largest-remainder rounders that can disagree on ties is a latent scientific bug. +- **Three near-identical download functions** — `try_download` (line 77), `try_curl_cffi` (127), `try_download_text` (185) — ~180 lines implementing the same retry/SSL-fallback loop three times. One function with `backend=` and `mode=` parameters replaces all three. + +### E3. Repeated blocks that want a loop + +**`co.py:176-215`** — the four optimization passes (PUMA → county → CBSA → urbanization) are the same ten lines copy-pasted four times, differing only in the column name and lookup dict. ~55 lines becomes ~15: + +```python +LEVELS = [("PUMA", "st_puma", cbg_puma, params), + ("county", "county", cbg_county, params), + ("CBSA", "cbsa", cbg_cbsa, params), + ("urbanization", "U", cbg_urban, params_slow)] + +for label, col, lookup, p in LEVELS: + rerun = [i for i, r in enumerate(x) if r[2] > c_val] + if not rerun: + continue + masks = sample_lookup(samp_geo, col, [lookup[geos[i]] for i in rerun]) + reoptimize(x, rerun, samples, masks, targs, n_hhs, p, rng) + _report(label, x, c_val) +``` + +Note this also fixes a latent bug: the current code prints `"Optimizing 0 CBG(s) at county level"` and *then* checks `if rerun:` — so at the CBSA and urbanization stages it announces work it doesn't do. + +**`geopops_starsim.py:327-336`** — ten near-identical `.loc[]` assignments to build age groups: + +```python +ppl_df['agegroup'] = np.clip(ppl_df['age'] // 10, 0, 9) +``` + +**`export.py`** — nine blocks of `sorted([...]) → pd.DataFrame(..., columns=[...]) → to_csv → _log_export`. A small helper collapses each to one line: + +```python +def _write(rows, columns, name, key=None): + pd.DataFrame(sorted(rows, key=key) if key else rows, columns=columns).to_csv(export_dir / name, index=False) + _log_export(verbose, f"-- {rel}/{name}") +``` + +**`run_all.py:96-100`** — four `ForStarsim.GPNetwork(name=..., edge_weight=1.0)` calls; `for name in ('homenet', 'schoolnet', 'worknet', 'gqnet')`. + +### E4. Convoluted logic + +**`co.py:130-132`:** + +```python +enough = [samp_masks[i].sum() > (n_hhs[rerun[j]] // 2) for j, i in enumerate(range(len(samp_masks)))] +valid = [rerun[j] for j, ok in enumerate(enough) if ok] +valid_mask_idx = [j for j, ok in enumerate(enough) if ok] +``` + +`for j, i in enumerate(range(len(samp_masks)))` means `j == i` unconditionally, and `valid_mask_idx` is just the surviving `j` values. The whole thing is: + +```python +for j, ri in enumerate(rerun): + if samp_masks[j].sum() <= n_hhs[ri] // 2: + continue + r = anneal(samples, samp_masks[j], targs[ri:ri+1], n_hhs[ri], params, rng) + if r[2] < x[ri][2]: + x[ri] = r +``` + +**`co.py:121`:** `zip(samp_masks, range(len(targs)), n_hhs)` → `zip(samp_masks, targs, n_hhs)`. + +**`co.py:170-171`:** `cmask = [co == c for co in county_of]` then `idxs = [i for i, m in enumerate(cmask) if m]` → `idxs = [i for i, co in enumerate(county_of) if co == c]`. (The loop variable `co` also shadows the imported `co` module.) + +**`generate_pop.py:100-176`** — five `_log_*_summary` methods, ~75 lines, are pure logging and make up 30% of the class. Move them to a `summary.py` module, or better, have the stages return small summary dicts and log them uniformly. + +### E5. Printing instead of logging + +175 `print()` calls across the package (80 in `download_data.py`, 48 in `process_data.py`). The `verbose` flag is threaded by hand through every function that needs it, with inconsistent semantics — `DownloadData` documents `verbose` as `1`/`0`, `export_synthpop` defaults it to `True`, `GeneratePop` uses truthiness. + +Use `logging` with a package logger: + +```python +logger = logging.getLogger("geopops") +``` + +Then `verbose=` on the public entry points just sets a level, users can redirect or silence output through standard mechanisms, and library code stops writing to stdout unconditionally. `ipfn.py` prints 7 times from inside a numerical inner loop. + +`process_data.py:1319` has a live example of the cost: `x_c = test_c.sub(ref_c).apply(abs).apply(lambda s: (s > 5.0).any(), axis=1)` computes a diagnostic whether or not anyone will look at it. + +### E6. No type hints + +Zero annotated function signatures in 7,207 lines. For a library whose core data structures are undocumented nested tuples — `people` is `dict[tuple[int,int,int], PersonData]`, `company_workers` is `dict[tuple[int,int,str], list[tuple[int,int,int,int]]]`, worker tuples are variously sliced `[:3]` and indexed `[3]` — this is the difference between an API a new contributor can use and one they have to reverse-engineer. + +Start with the public entry points and the inter-module boundaries (the return signature of `generate_networks` is an 8-tuple; `generate_jobs_and_workers` returns a 5-tuple). Named tuples or small dataclasses for the key types would help as much as the annotations: + +```python +class PersonKey(NamedTuple): + p_id: int + hh_id: int + cbg_id: int +``` + +Then `w[:3]` becomes `w.person` and `k[2]` becomes `k.cbg_id` throughout. Add `mypy` (or `ty`) to CI in non-strict mode and tighten over time. + +### E7. Fragile config→code coupling + +`households.py:204-207` passes `additional_traits` from config into `PersonData(**trait_kwargs)`, but `PersonData` (`utils.py:19-39`) has a **hardcoded** field list. The shipped config's 11 `additional_traits` happen to match exactly. Add a twelfth trait — the natural thing to do after reading issue #2's discussion of race/ethnicity categories — and the pipeline dies with an opaque `TypeError` deep in the person loop. + +Either derive the traits dynamically, or validate up front with a clear message: + +```python +known = {f.name for f in dataclasses.fields(PersonData)} +unknown = set(additional_traits) - known +if unknown: + raise ConfigError(f"additional_traits not supported by PersonData: {sorted(unknown)}") +``` + +Given issue #2's caveat — that these traits are carried through from PUMS but are *not* used in the CO step, so their distributions aren't matched to ACS — this is worth surfacing in code as well as in the issue: a warning when a trait is requested that CO doesn't target. + +### E8. Small correctness items + +- **`generate_pop.py:96-98`** — `_county_from_cbg_idx` reads `self._cbg_by_idx`, which is only created inside `SynthPop()` (line 199) and isn't declared in `__init__` alongside the other 18 state attributes. Calling `Export()` on a partially-run instance gives `AttributeError` rather than the clear `RuntimeError` the other stages raise. +- **`geopops_starsim.py:386-387`** — `sim = ss.Sim(people=ppl).init()` then `_ = sim # keep side-effect parity`. A comment admitting the code depends on an unnamed side effect. Work out what `init()` is actually doing to `ppl` and call that directly, or document it. +- **`config.py:55-59`** — `compute_decennial_year` catches bare `Exception` and returns `2010` for any unparseable input, so `main_year="twenty-nineteen"` yields a plausible-looking config that produces wrong data. +- **`utils.py:66-75`** — `lrRound` handles `vrem > 0` but not `vrem < 0`, which can occur with negative inputs. Probably unreachable given the call sites, but it should assert rather than silently under-round. +- **`config.py:50`** — `os.makedirs(os.path.dirname(cfg_path), exist_ok=True)` raises `FileNotFoundError` when `cfg_path` is a bare filename (`dirname` returns `""`). + +--- + +## F. Packaging and project metadata + +This covers issue #3. + +`pyproject.toml` is missing: `authors`, `maintainers`, `license`, `keywords`, `classifiers`, `[project.urls]` (Homepage / Documentation / Repository / Issues / Changelog). PyPI currently shows a package with no author, no license, and no links. Compare `starsim`'s for a template. + +Also: + +- **No `LICENSE`** (see A1). +- **No `CHANGELOG.md`.** Version is at 0.1.7 with no record of what changed. Given that the class rename (`RunPython` → `GeneratePop`) and the parameter rename (`beta_value` → `edge_weight`) both broke callers — the tests among them — this matters now. +- **No `CONTRIBUTING.md`** despite the README's "Get involved" section actively soliciting contributions. +- **Unpinned dependencies.** `geopandas`, `networkx`, `numpy`, `sciris`, `scipy`, `shapely` have no lower bounds. `numpy` in particular matters: `Generator.multivariate_hypergeometric` (recommended in C3) needs ≥1.18, and `SeedSequence.spawn` needs ≥1.17. +- **`build/` is committed** and contains a stale copy of the package (`build/lib/geopops/pyjulia/`, `census.py`) that no longer matches `src/`. It's gitignored but present in the working tree; it will confuse grep-based navigation and any tooling that walks the tree. Delete it. +- **`README.md.orig`** is an untracked leftover — remove or restore it. +- **`.gitignore` contains `tests/`** (see A2) and `tutorials/`, which means the tutorials the README points to can't live in this repo. +- **No `py.typed`** marker (add once E6 is underway). +- **No linter config.** Add `ruff` with a modest rule set and run it in CI; it will catch the shadowed `co` variable, the unused imports, and the bare excepts automatically. + +--- + +## G. Suggested sequencing + +**Week 1 — stop the bleeding** + +1. Add `LICENSE` + `pyproject.toml` license metadata (A1). Decide on the `ipfn` vendoring. +2. Fix the tests so they run at all: `RunPython` → `GeneratePop`, `beta_value` → `edge_weight`, move the `WriteConfig` call into a fixture, un-ignore `tests/` (A2). +3. Add the golden-output regression test on the Spartanburg fixture with a fixed seed. **This is the enabler for everything downstream.** +4. Add CI (A3). +5. Replace `exit(1)` with exceptions (A4). + +**Week 2 — free performance** + +6. Incremental `anneal()` summary (C1). Bit-identical, so the golden test proves it immediately. ~10× on the dominant cost. +7. Vectorize `read_workers_by_cat` (C2), `find_closest` (C4), and the `households` per-person `.iloc` (C5). +8. `drawCounts` → `multivariate_hypergeometric` (C3), plus the `colsums` hoist in `pull_inst_workers`. +9. Delete the dead code (E1) and the `*_test.csv` writes. + +At this point the runtime should be dramatically better and you'll know exactly where the remaining time goes — profile again before considering Numba or parallelism. + +**Weeks 3–4 — structure** + +10. Kill the module globals; introduce `Paths` and thread config explicitly (A5). Biggest single refactor, safest with the golden test in place. +11. Move config out of site-packages (D1). +12. Convert the namespace classes to functions (B1, B3), keeping deprecated aliases. Fix the `ForStarsim` cache bug on the way. +13. Remove the `od_*.csv.gz` round-trip; let `ForStarsim` accept in-memory objects (D2). +14. Add `random_seed` to the shipped config and to `RunAll`; seed `process_data.py:1569` (A7). + +**Ongoing** + +15. Type hints from the public API inward; `NamedTuple` for `PersonKey` and the worker tuples (E6). +16. `logging` instead of `print` (E5). +17. Consolidate the duplicated helpers and the three download functions (E2). +18. Parallelize CO over CBGs with reproducible per-CBG seeds (C1). +19. Docs (issue #5) and project metadata (issue #3 / section F). + +--- + +## Appendix: measurements + +All benchmarks run on this machine, 2026-08-28. Sizes chosen to approximate a mid-sized county. + +| Change | Before | After | Speedup | Output | +|---|---|---|---|---| +| `anneal` incremental summary (5k gens, n_hh=800, 120 cols) | 0.686 s | 0.071 s | **9.6×** | bit-identical, RNG stream preserved | +| `anneal` + batched RNG draws (20k gens) | 2.74 s | 0.20 s | **13.8×** | statistically equivalent | +| `.iloc[i][col]` → NumPy indexing (2,000 lookups) | 28.9 ms | 0.075 ms | **385×** | identical | +| `drawCounts` → `multivariate_hypergeometric` (n=200, 1,500 bins) | 6.19 ms | 0.38 ms | **16×** | same distribution, different order | +| `find_closest` per grade (600 CBG × 400 schools) | 489 ms | 7.6 ms | **64×** | verified identical | +| `find_closest` × 14 grades | 6.8 s | 0.11 s | **64×** | verified identical | + + +--- + +## Implementation status + +Everything below was done in this session and verified against the Spartanburg County, SC fixture (195 CBGs, ~297k people, ~357k agents including dummies). The work is uncommitted in the working tree. + +### Done + +| Item | What changed | +|---|---| +| [A1](#a1-licensing) | Added `LICENSE` (AGPL-3.0-or-later, verbatim FSF text) and `NOTICE` recording the GREASYPOP-CO provenance and restoring the `ipfn` MIT attribution. Declared `license` in `pyproject.toml`; both files ship in the wheel. | +| [A2](#a2-the-test-suite-is-dead) | Replaced the two dead test files with 96 tests across `test_utils`, `test_config`, `test_co`, `test_networks`, `test_sources`, `test_workflow`, `test_regression`. Added `conftest.py` fixtures; import-time side effects gone; `tests/` un-ignored (only `tests/data/` is ignored now). | +| [A3](#a3-no-ci) | `.github/workflows/test.yml` (pytest on 3.11/3.12/3.13 + ruff) and `release.yml` (build, `twine check`, PyPI Trusted Publishing). | +| [A4](#a4-exit1-in-library-code) | All three `exit(1)` calls replaced by `DownloadError`. Added a `GeoPopsError` hierarchy. | +| [A6](#a6-tls-verification-disabled-globally) | Removed the import-time `urllib3.disable_warnings()`. The unverified-TLS fallback is now opt-in via `allow_insecure_downloads`, warns loudly, and scopes warning suppression to the single request. | +| [A7](#a7-reproducibility) | `random_seed` added to the shipped config; `run()`/`generate_pop()` take `seed=`; `generate_work_sizes` no longer uses numpy's unseeded global RNG; unseeded runs warn. | +| [B1](#b1-classes-that-should-be-functions), [B2](#b2-suggested-public-api), [B3](#b3-forstarsim-a-namespace-pretending-to-be-a-class) | `WriteConfig`, `RunAll`, `ForStarsim`, `RunJulia` removed in favour of `make_config`, `run`, `to_starsim_people`/`starsim_network`/`starsim_networks`. Fixed the never-invalidated `ForStarsim` network cache and the config mismatch between `People()` and `GPNetwork()`. `DownloadData`/`ProcessData`/`GeneratePop` kept as step objects. | +| [C1](#c1-copy-fixes-co-annealing-10)–[C7](#c7-other-hot-spots) | All the hot-path work: incremental annealing summary, hoisted `sqrt(targ+1)`, cached sample sub-pools, vectorized `read_workers_by_cat` / `find_closest` / `generate_people`, hypergeometric `drawCounts`, incremental column sums in `pull_inst_workers`, COO instead of LIL in `generate_commute_matrices`, `__slots__` + shared trait schema on `PersonData`. | +| [D1](#d1-config-is-written-into-site-packages) | Config no longer written into site-packages. `make_config()` returns a dict; `save=True` writes into the run directory. The packaged `config.json` is a read-only template. | +| [D2](#d2-round-trips-that-serve-no-purpose) | `generate_commute_matrices` returns its matrices instead of gzipping them and reading them straight back. `to_starsim_people`/`starsim_networks` take an explicit `pop_export_dir`, and `GeneratePop.pop_export_dir` supplies it. | +| [E1](#e1-dead-code-450-lines) | Removed `generate_location_matrices`, the `generate_test_targets`/`gen_samp_test_cols`/`test_cols` scaffolding, `DownloadData.pipeline()`, `julia.py`, the 11 `julia/*.jl` files, and the stale empty `pyjulia/` directory. | +| [E2](#e2-duplicated-helpers) | Three near-identical download functions collapsed into one `download(src, dst, backend=, mode=)`. | +| [E3](#e3-repeated-blocks-that-want-a-loop), [E4](#e4-convoluted-logic) | CO's four copy-pasted passes became a loop over levels (also fixing the misleading "Optimizing 0 CBGs" log); `reoptimize`'s `enumerate(range(len(...)))` untangled; ten `agegroup` assignments became one `np.clip`. | +| [E7](#e7-fragile-configcode-coupling) | Traits are config-driven via `TraitSchema` and reached by attribute as before. **This was a live bug, not a latent one:** the checked-in fixture config lists `white_non_hispanic`, so the fixture could not run against 0.1.7 at all — `TypeError: unexpected keyword argument`. | +| [E8](#e8-small-correctness-items) | `compute_decennial_year` raises instead of defaulting to 2010; `save_config` handles a bare filename; `QualityCheck.results` no longer prints; `_cbg_by_idx` declared up front; exception chaining and loop-invariant hoisting fixed. | +| [F](#f-packaging-and-project-metadata) | Full `pyproject.toml` metadata (authors, license, keywords, classifiers, URLs), dependency lower bounds, ruff + pytest config, `CHANGELOG.md`, updated `README.md`, cleaned `.gitignore`, removed the Julia entry from `.env.example`, deleted the stale `build/` tree. `ruff check src/ tests/` passes clean. | + +One thing surfaced during implementation that wasn't in the original review: exporting `download_data`/`process_data`/`generate_pop` as functions shadowed the modules of the same name. The modules were renamed to `sources.py`, `census.py`, `population.py`, `starsim_bridge.py`, and `pipeline.py`. + +### Measured results + +Spartanburg County, SC (195 CBGs, 292k people); `CO_maxgens=20000`; seed 42. + +| Stage | Before | After | Speedup | Output | +|---|---|---|---|---| +| CO (`process_counties`) | 2192.9 s | 41.1 s | **53.3×** | **identical** | +| SynthPop | 200.0 s | 18.3 s | **10.9×** | see below | +| Export | 4.0 s | 4.1 s | — | — | +| **Total** | **~36.6 min** | **~63 s** | **~35×** | | + +Each rewritten component was checked against the original implementation on the real fixture, not just benchmarked: + +| Component | Before | After | Speedup | Equivalence check | +|---|---|---|---|---| +| `co.process_counties`, 195 CBGs | 2192.9 s | 41.1 s | 53.3× | **identical** — `co_results` and `co_scores` compared whole | +| `co.anneal`, 6 seeded trials incl. empty pool and early exit | 0.686 s | 0.071 s | 9.6× | **bit-identical**, RNG call sequence preserved | +| `households.generate_people`, 292,039 people | 46.9 s | 6.3 s | 7.4× | **identical** — every core field, every trait, all households, all GQs, and `gq_summary` | +| `schools.find_closest`, 14 grades | 0.36 s | 0.035 s | 10× | **identical** across 2,730 CBG/grade entries | +| `.iloc[i][col]` → NumPy (2,000 lookups) | 28.9 ms | 0.075 ms | 385× | identical | +| `drawCounts` (n=200, 1,500 bins) | 6.19 ms | 0.38 ms | 16× | same distribution, different draw order | +| `PersonData` memory | ~1,515 B/person | ~240 B/person | 6.3× smaller | — | + +The 53× on CO exceeds the 9.6× that `anneal` alone gives, because caching the candidate sub-pools also removes a per-CBG re-extraction of the same PUMA sub-matrix. + +`census.py` (the former `process_data.py`) was validated by running it end to end: it completes in 149 s and **24 of its 27 outputs are byte-identical** to the reference. Of the three that differ, `work_sizes.csv` is by design (it is now seeded rather than using numpy's unseeded global RNG), and `p_samples.csv`/`samp_geo.csv` differ only in `st_puma`/`PUMA` zero-padding — a **pre-existing** discrepancy, since `HEAD` already applies `zfill(5)` and the checked-in fixture predates it. See "Two things worth a decision" below. + +### Output equivalence + +Running the full pipeline before and after the workplace changes, same seed: + +- **Bit-identical**: `people.csv`, `hh.csv`, `cbg_idxs.csv`, `sch_students.csv`, `gqs.csv`, `gq_residents.csv`, `adj_mat_keys.csv`, `adj_dummy_keys.csv`, `adj_out_workers.csv`, `adj_upper_triang_hh.mtx`. +- **Changed, as designed**: `company_workers.csv`, `sch_workers.csv`, `gq_workers.csv`, and the school/workplace/GQ/non-household network matrices — everything downstream of `drawCounts`, which now consumes the RNG stream differently. The distribution is unchanged; the specific assignment is not. This is the one behavioural change in the release and it is called out in `CHANGELOG.md`. + +`tests/golden_hashes.json` pins the post-change output, and `test_same_seed_gives_the_same_population` asserts run-to-run reproducibility independently of that baseline. + +### Deliberately not done + +- **[A5](#a5-mutable-module-level-globals) — the module-global refactor in `census.py`/`sources.py`.** `PROCESSED_DIR` alone is referenced 38 times across 2,200 lines. It is the right next change and it is what unblocks parallelizing CO ([C1](#c1-copy-fixes-co-annealing-10)), but it is a large mechanical edit to the least-tested module, and the golden test that would make it safe only exists as of this session. Doing it now would have put the riskiest change on the least evidence. It is recorded under "Known issues" in the changelog. +- **Parallelizing CO across CBGs.** CO went from 37 minutes to 41 seconds without it, so the payoff is much smaller than it was, and it depends on A5. +- **[E5](#e5-printing-instead-of-logging) `logging` and [E6](#e6-no-type-hints) type hints.** Both are broad, low-risk, and mostly mechanical; neither blocks anything. Worth doing incrementally from the public API inward. +- **[D3](#d3-format-choices) Parquet internally.** A real win for large counties, but it changes the on-disk contract for downstream consumers and deserves its own release. +- **The seven `*_test.csv` debug writes** in `census.py`, which live in the middle of the module-global code A5 covers. +- **Documentation ([issue #5](https://github.com/GeoPopsHub/geopops/issues/5)).** The README and changelog are updated, but a built docs site is a separate piece of work. + +### Two things worth a decision + +1. **The license.** AGPL-3.0-or-later is the only option available without third-party permission, given that `census.py` is derived from GREASYPOP-CO. It is also viral, which will constrain downstream users — including Starsim integrations, which are MIT. If a permissive license matters, that is a conversation to have with the GREASYPOP-CO copyright holders, and it should happen before more releases go out. +2. **The fixture data is stale, in two independent ways.** `tests/data/processed/p_samples.csv` carries the pre-[issue #2](https://github.com/GeoPopsHub/geopops/issues/2) trait columns (`white_non_hispanic` rather than the eight PUMS categories) — which is why the fixture could not run against 0.1.7 at all. Separately, its `st_puma` values are unpadded (`45501`) while the current code emits `zfill(5)`-normalized ones (`4500101`); regenerating `processed/` changes those two files. Neither is caused by this session's work, but together they mean the fixture no longer represents what the pipeline produces, which limits what the golden test proves. It is also large and gitignored, so it is not reproducible from a fresh clone — **a small committed fixture (one or two CBGs) would let CI run the end-to-end and regression tests at all**, which is currently the biggest remaining gap in the test story. diff --git a/pyproject.toml b/pyproject.toml index 08d1e91..76e419c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,29 +4,69 @@ build-backend = "setuptools.build_meta" [project] name = "geopops" -version = "0.1.7" -description = "GeoPops" +version = "0.1.8" +description = "Geographically and demographically realistic synthetic populations for any US Census location" readme = "README.md" requires-python = ">=3.11" +license = "AGPL-3.0-or-later" +license-files = ["LICENSE", "NOTICE"] +authors = [ + { name = "GeoPops development team" }, +] +maintainers = [ + { name = "GeoPops development team" }, +] +keywords = [ + "synthetic population", + "agent-based modeling", + "census", + "demography", + "epidemiology", + "contact networks", + "starsim", +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "Natural Language :: English", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Bio-Informatics", + "Topic :: Scientific/Engineering :: Information Analysis", +] dependencies = [ "curl-cffi>=0.13.0", - "geopandas", - "networkx", - "numpy", + "geopandas>=0.14", + "networkx>=3.0", + # >=1.18 for Generator.multivariate_hypergeometric; >=1.17 for SeedSequence.spawn + "numpy>=1.18", + "openpyxl>=3.1", "pandas>=2.3.2", "pyarrow>=14.0.0", "python-dotenv>=1.0.0", "requests>=2.32.5", - "sciris", - "scipy", - "shapely", + "sciris>=3.0", + "scipy>=1.10", + "shapely>=2.0", "starsim>=3.0.3", "urllib3>=2.5.0", ] +[project.urls] +Homepage = "https://github.com/GeoPopsHub/geopops" +Repository = "https://github.com/GeoPopsHub/geopops" +Issues = "https://github.com/GeoPopsHub/geopops/issues" +Changelog = "https://github.com/GeoPopsHub/geopops/blob/main/CHANGELOG.md" +Examples = "https://github.com/GeoPopsHub" + [project.optional-dependencies] dev = [ "ipykernel", + "pytest>=7.0", + "ruff>=0.6", ] [tool.setuptools.packages.find] @@ -34,4 +74,26 @@ where = ["src"] # patterns are relative to the *package root* (src/geopops) [tool.setuptools.package-data] -geopops = ["geocorr/*", "julia/*", "config.json"] +geopops = ["geocorr/*", "config.json", "py.typed"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +markers = [ + "slow: end-to-end tests that need downloaded/processed data (deselect with '-m \"not slow\"')", + "network: tests that hit external data providers", +] + +[tool.ruff] +line-length = 140 +target-version = "py311" +extend-exclude = ["*.ipynb", "build"] + +[tool.ruff.lint] +select = ["E", "F", "W", "B", "UP"] +ignore = [ + "E501", # line length is handled by the formatter, not enforced + "E741", # ambiguous variable names (l, I) are pervasive in the ported numeric code +] + +[tool.ruff.lint.per-file-ignores] +"src/geopops/ipfn.py" = ["ALL"] # vendored third-party code, kept close to upstream diff --git a/src/geopops/__init__.py b/src/geopops/__init__.py index 47173fe..945af65 100644 --- a/src/geopops/__init__.py +++ b/src/geopops/__init__.py @@ -1,10 +1,45 @@ -from .config import WriteConfig -from .download_data import DownloadData -from .process_data import ProcessData, QualityCheck -from .julia import RunJulia -from .generate_pop import GeneratePop -from .run_all import RunAll -from .geopops_starsim import ForStarsim +"""GeoPops: geographically and demographically realistic synthetic populations. -__all__ = ["WriteConfig", "DownloadData", "ProcessData", "RunJulia", "GeneratePop", "RunAll", "ForStarsim", "QualityCheck"] +Typical use:: + import geopops as gp + + cfg = gp.make_config(path="data", geos=["45083"], main_year=2019, + commute_states=["45", "37"], use_pums=["45", "37"], + random_seed=42) + + gp.download_data(cfg) # fetch Census/PUMS/LODES/school data (slow, cached) + gp.process_data(cfg) # build the CO targets and sample pools + pop = gp.generate_pop(cfg, seed=42) + ppl = gp.to_starsim_people(pop.pop_export_dir) + +or, all at once:: + + pop = gp.run(cfg, seed=42) +""" +from .exceptions import (GeoPopsError, ConfigError, DownloadError, DataError, + PipelineStateError) +from .config import make_config, load_config, save_config, validate_config +from .sources import DownloadData, download_data, download +from .census import ProcessData, process_data, quality_check +from .population import GeneratePop, generate_pop +from .pipeline import run +from .starsim_bridge import (GPNetwork, SubgroupTracking, to_starsim_people, + starsim_network, starsim_networks, load_network_edges) + +__version__ = "0.1.8" + +__all__ = [ + "__version__", + # Errors + "GeoPopsError", "ConfigError", "DownloadError", "DataError", "PipelineStateError", + # Config + "make_config", "load_config", "save_config", "validate_config", + # Pipeline + "download_data", "process_data", "generate_pop", "quality_check", "run", + # Starsim bridge + "to_starsim_people", "starsim_network", "starsim_networks", "load_network_edges", + "GPNetwork", "SubgroupTracking", + # Step objects, for running or inspecting individual stages + "DownloadData", "ProcessData", "GeneratePop", "download", +] diff --git a/src/geopops/process_data.py b/src/geopops/census.py similarity index 89% rename from src/geopops/process_data.py rename to src/geopops/census.py index 58e5f4a..f49e9de 100644 --- a/src/geopops/process_data.py +++ b/src/geopops/census.py @@ -12,11 +12,8 @@ import pandas as pd import numpy as np -from numpy import array as arr import geopandas as gpd -from shapely.geometry.point import Point import os -import shutil from glob import glob import json @@ -24,7 +21,7 @@ def tryJSON(filename): try: - with open(filename, 'r') as f: + with open(filename) as f: d = json.load(f) except Exception as e: print("warning: ",e) @@ -47,7 +44,7 @@ def lrRound(v): vrem = np.int64(np.round(sum(v) - sum(vrnd))) vidxs = verr.index[np.flip(np.argsort(verr.values))] for i in range(vrem): - vrnd[vidxs[i]] += 1 + vrnd[vidxs[i]] += 1 return vrnd def read_census(file_list, geos, usecols, columns): @@ -67,16 +64,16 @@ def read_census(file_list, geos, usecols, columns): def read_acs(table,geos=None): filematch = ''.join(['*.', table, '-Data.*']) with os.scandir(os.path.join(OUTPUT_DIR, "census")) as d: - files = [glob(os.path.join(f,filematch))[0] for f in d + files = [glob(os.path.join(f,filematch))[0] for f in d if f.is_dir() and not f.name.startswith(".")] - return read_census(files, geos, + return read_census(files, geos, (lambda x: (x=='Geography') or (str.split(x,'!!')[0]=='Estimate')), (lambda x: ''.join([table,':',*str.split(x,'!!')[2:]]))) def read_decennial(table,geos=None): filematch = ''.join(['*.', table, '-Data.*']) with os.scandir(os.path.join(OUTPUT_DIR, "census")) as d: - files = [glob(os.path.join(f,filematch))[0] for f in d + files = [glob(os.path.join(f,filematch))[0] for f in d if f.is_dir() and not f.name.startswith(".")] return read_census(files, geos, (lambda x: (x=='Geography') or (str.split(x,'!!')[0]=='Total')), @@ -216,15 +213,13 @@ def read_psamp(LODES_cutoff, ind_codes, occ_codes): 'CIT':str,'FER':str,'LANX':str,'DIS':str,'RAC1P':str,'HISP':str,'PAP':"Int64"} psamp = read_pums("psam_p",psamp_dtype) - + #### # Convert all string columns from float format to proper strings # Some columns need leading zeros preserved (2-digit codes) # print('SCHG before',psamp['SCHG'].unique()) - columns_needing_zeros = ['SCHG', 'SCHL'] # 'NAICSP', 'SOCP', 'COW', 'POWPUMA', 'POWSP', 'JWTRNS', 'WKL', 'WKW', 'WRK', 'SFN', 'SFR', 'CIT', 'FER', 'LANX' - for col, dtype in psamp_dtype.items(): - if col in psamp.columns and dtype == str: + if col in psamp.columns and dtype is str: psamp[col] = psamp[col].astype(str).str.replace('.0', '') # Add leading zeros for specific columns that need them # if col in columns_needing_zeros: @@ -237,12 +232,12 @@ def read_psamp(LODES_cutoff, ind_codes, occ_codes): # psamp['ESR'] = psamp['ESR'].astype(str).str.replace('.0', '').replace('nan', '6').replace('', '6') # print('test',psamp['SCHG'].unique()) #### - + # Normalize PUMA to fixed-width code so keys align with crosswalk-derived st_puma (STATE + 5-digit PUMA). psamp['ST'] = psamp['ST'].astype(str).str.strip() psamp['PUMA'] = psamp['PUMA'].astype(str).str.strip().str.zfill(5) psamp['st_puma'] = psamp['ST'] + psamp['PUMA'] - psamp['sch_grade'] = psamp['SCHG'].map(dict(zip([str(x) for x in range(1,17)], ['p','k',*[str(x) for x in range(1,13)],'c','g']))) + psamp['sch_grade'] = psamp['SCHG'].map(dict(zip([str(x) for x in range(1,17)], ['p','k',*[str(x) for x in range(1,13)],'c','g'], strict=False))) # print('sch_grade',psamp['sch_grade'].unique()) ## using 2-digit industry and occupation codes psamp["industry"] = psamp["NAICSP"].fillna("").str.slice(0,2) @@ -371,9 +366,9 @@ def read_hsamp(ptotals, ADJINC, inc_cats, inc_cols): hsamp = read_pums("psam_h",hsamp_dtype) hsamp = hsamp[(hsamp['NP']>0) & (hsamp['TYPE'] == '1')] - + for col, dtype in hsamp_dtype.items(): - if col in hsamp.columns and dtype == str: + if col in hsamp.columns and dtype is str: hsamp[col] = hsamp[col].astype(str).str.replace('.0', '') ## pre-compute census columns so we can generate them quickly from a PUMS subsample: @@ -519,7 +514,7 @@ def generate_samples(sample_columns, ADJINC, inc_cats, inc_cols, LODES_cutoff, i "Expected a file matching '*puma_to_cbsa*.*'." ) file = matches[0] - + #### puma_to_cbsa #### # puma_to_cbsa = pd.read_csv(file,dtype=str,skiprows=[1],usecols=["state","puma12","cbsa","afact"]) candidate_cols = {"state", "puma12", "puma22", "cbsa", "cbsa20", "afact"} @@ -532,7 +527,7 @@ def generate_samples(sample_columns, ADJINC, inc_cats, inc_cols, LODES_cutoff, i raise ValueError(f"Missing required columns in puma_to_cbsa crosswalk {file}: {missing}. Found: {list(puma_to_cbsa.columns)}") puma_to_cbsa = puma_to_cbsa[list(required_found)].copy() puma_to_cbsa.rename(columns={col_puma: "puma12", col_cbsa: "cbsa"}, inplace=True) - # now use only "puma" and "cbsa" + # now use only "puma" and "cbsa" puma_to_cbsa['st_puma'] = puma_to_cbsa['state'] + puma_to_cbsa['puma12'].str.zfill(5) puma_to_cbsa.loc[puma_to_cbsa['cbsa']==" ", 'cbsa'] = "none" puma_to_cbsa = puma_to_cbsa.groupby('st_puma',group_keys=True).apply(lambda g: g[g['afact'] == g['afact'].max()], include_groups=False).reset_index() @@ -541,7 +536,7 @@ def generate_samples(sample_columns, ADJINC, inc_cats, inc_cols, LODES_cutoff, i samp_geo = samp_geo.merge(puma_to_cbsa, how='left', on='st_puma') ## some pumas span counties; associate each with its primary county ## -- we're only looking up samples by county when it contains multiple pumas - + #### puma_to_county #### matches = glob(os.path.join(OUTPUT_DIR, "geo", "*puma_to_county*.*")) if not matches: @@ -570,7 +565,7 @@ def generate_samples(sample_columns, ADJINC, inc_cats, inc_cols, LODES_cutoff, i "Expected a file matching '*puma_urban_rural*.*'." ) file = matches[0] - + #### puma_urban_rural #### # puma_ur = pd.read_csv(file,dtype=str,skiprows=[1],usecols=["state","puma12","ur","afact"]) candidate_cols = {"state", "puma12", "puma22", "ur", "afact"} @@ -670,14 +665,14 @@ def read_geo_xwalk(cbg_index): tpum['st_puma'] = tpum['STATEFP']+tpum['PUMA5CE'].str.zfill(5) tpum['tract'] = tpum['STATEFP']+tpum['COUNTYFP']+tpum['TRACTCE'] ## note, population 0 cbgs are not in geocorr - cbg_geo = pd.DataFrame({'Geo':cbg_index, + cbg_geo = pd.DataFrame({'Geo':cbg_index, 'tract':cbg_index.map(lambda x: x[0:-1]), 'county':cbg_index.map(lambda x: x[0:5])}) cbg_geo = cbg_geo.merge(tpum, how='left', on='tract').set_index('Geo', verify_integrity=True) - + cbg_geo.to_csv(os.path.join(PROCESSED_DIR,'cbg_geo_test.csv')) - - + + #### cbg_to_cbsa #### if main_year is not None and main_year >= 2020: cbg_cbsa_pattern = "*geocorr2022*cbg_to_cbsa*.*" @@ -705,7 +700,7 @@ def read_geo_xwalk(cbg_index): cbg_to_cbsa['bg'] cbg_to_cbsa.set_index('Geo', inplace=True, verify_integrity=True) # cbg_to_cbsa.to_csv(os.path.join(PROCESSED_DIR,'cbg_to_cbsa_test.csv')) - + #### cbg_urban_rural #### if main_year is not None and main_year >= 2020: cbg_ur_pattern = "*geocorr2022*cbg_urban_rural*.*" @@ -773,7 +768,7 @@ def generate_gq(geos, df_adults_in_hh, geo_xwalk, p_summary, ind_codes, occ_code df_gq['group quarters:65 and over'] = df_gq['adult_65o_in_group_quarters'] ## ignore cbgs with less than 20 in gq's - df_gq = df_gq.loc[df_gq['group quarters:'] > 19, + df_gq = df_gq.loc[df_gq['group quarters:'] > 19, ['group quarters:', 'group quarters:under 18', 'group quarters:18 to 64', 'group quarters:65 and over']].copy(deep=True) ## from previous decennial census @@ -856,9 +851,9 @@ def generate_gq(geos, df_adults_in_hh, geo_xwalk, p_summary, ind_codes, occ_code ## use employment status code (ESR): ## assume that everyone in military GQ is employed in armed forces ## and anyone in GQ not employed by armed forces is in civilian GQ - s_noninst_civ = p_summary.loc[(p_summary['RELSHIPP'] == '38') + s_noninst_civ = p_summary.loc[(p_summary['RELSHIPP'] == '38') & (p_summary['armed_forces'] != 1) - & (p_summary['AGEP'] > 17) + & (p_summary['AGEP'] > 17) & (p_summary['AGEP'] < 65)].copy(deep=True) ## note, employment stats generated here still need to be reconciled with census data in next step cols = ['commuter','work_from_home','com_LODES_low','com_LODES_high', @@ -867,16 +862,16 @@ def generate_gq(geos, df_adults_in_hh, geo_xwalk, p_summary, ind_codes, occ_code pcols = [c+'_p|ninst1864civ' for c in cols] ## PWGTP = individual "weight" according to the pums sample data - for c,pc in zip(cols,pcols): + for c,pc in zip(cols,pcols, strict=False): s_noninst_civ[pc] = s_noninst_civ[c] * s_noninst_civ['PWGTP'] - s_noninst_mil = p_summary.loc[(p_summary['RELSHIPP'] == '38') + s_noninst_mil = p_summary.loc[(p_summary['RELSHIPP'] == '38') & (p_summary['armed_forces'] == 1) - & (p_summary['AGEP'] > 17) + & (p_summary['AGEP'] > 17) & (p_summary['AGEP'] < 65)].copy(deep=True) cols_mil = [c+'_p|milGQ' for c in cols] - for c,pc in zip(cols,cols_mil): + for c,pc in zip(cols,cols_mil, strict=False): s_noninst_mil[pc] = s_noninst_mil[c] * s_noninst_mil['PWGTP'] ## not many GQ residents represented in PUMS samples @@ -900,11 +895,11 @@ def generate_gq(geos, df_adults_in_hh, geo_xwalk, p_summary, ind_codes, occ_code agg_mil = s_noninst_mil[['PWGTP',*cols_mil]].sum() p_agg_mil = agg_mil.apply(lambda x: x/agg_mil['PWGTP']).drop("PWGTP").fillna(0.0) - + df_gq[p_agg_civ.index] = p_agg_civ.values df_gq[p_agg_mil.index] = p_agg_mil.values - keep_cols = ['group quarters:', 'group quarters:under 18', 'group quarters:18 to 64', 'group quarters:65 and over', + keep_cols = ['group quarters:', 'group quarters:under 18', 'group quarters:18 to 64', 'group quarters:65 and over', 'p_u18_inst', 'p_18_64_inst', 'p_65o_inst', 'p_18_64_noninst_civil', 'p_18_64_noninst_mil', *pcols, *cols_mil] @@ -924,7 +919,7 @@ def generate_gq(geos, df_adults_in_hh, geo_xwalk, p_summary, ind_codes, occ_code def read_ind_df(geos, ind_codes): - + ## industries C24030 = read_acs("C24030",geos) ## civilian vs military workforce @@ -947,12 +942,12 @@ def read_ind_df(geos, ind_codes): ind_df = C24030[["C24030:"]].copy(deep=True) for k in ind_codes: ind_df = ind_df.join(C24030['C24030:All:'+ind_codes[k][1]].rename('C24030:'+k)) - + return ind_df def read_occ_df(geos, occ_codes): - + ## occupations C24010 = read_acs("C24010",geos) ## civilian vs military workforce @@ -978,7 +973,7 @@ def read_occ_df(geos, occ_codes): occ_df = C24010[["C24010:"]].copy(deep=True) for k in occ_codes: occ_df = occ_df.join(C24010['C24010:All:'+occ_codes[k][1]].rename('C24010:'+k)) - + return occ_df @@ -999,24 +994,24 @@ def read_occ_df(geos, occ_codes): # ## family / nonfamily households by size # B11016 = read_acs('B11016',geos) - + # ############### START HERE ################ # # B11012 -# # 0. ['', -# # 1. 'With children of the householder under 18 years', -# # 2. 'With no children of the householder under 18 years', -# # 3. '', -# # 4. 'With children of the householder under 18 years', -# # 5. 'With no children of the householder under 18 years', -# # 6. '', -# # 7. 'Living alone', -# # 8. 'With children of the householder under 18 years', -# # 9. 'With relatives, no children of the householder under 18 years', -# # 10. 'With only nonrelatives present', -# # 11. '', -# # 12. 'Living alone', -# # 13. 'With children of the householder under 18 years', -# # 14. 'With relatives, no children of the householder under 18 years', +# # 0. ['', +# # 1. 'With children of the householder under 18 years', +# # 2. 'With no children of the householder under 18 years', +# # 3. '', +# # 4. 'With children of the householder under 18 years', +# # 5. 'With no children of the householder under 18 years', +# # 6. '', +# # 7. 'Living alone', +# # 8. 'With children of the householder under 18 years', +# # 9. 'With relatives, no children of the householder under 18 years', +# # 10. 'With only nonrelatives present', +# # 11. '', +# # 12. 'Living alone', +# # 13. 'With children of the householder under 18 years', +# # 14. 'With relatives, no children of the householder under 18 years', # # 15. 'With only nonrelatives present'] def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_cols, ind_codes, occ_codes): @@ -1033,45 +1028,45 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co B11016 = read_acs('B11016',geos) ## household types married/cohab/single/alone/own_ch_u18/other_rels/only_nonrels B11012 = read_acs('B11012',geos) - + # Make a list of column names for B11012. Wording is different in 2019 and 2020. B11012_lst = [col.split(':')[2] for col in B11012.columns if len(col.split(':')) > 2] # print('***** B11012_lst ******', B11012_lst) #2019 - # 0 ['', - # 1 'With own children under 18 years', - # 2 'With no own children under 18 years', - # 3 '', - # 4 'With own children of the householder under 18 years', - # 5 'With no own children of the householder under 18 years', - # 6 '', - # 7 'Living alone', - # 8 'With own children under 18 years', - # 9 'With relatives, no own children under 18 years', - # 10 'With only nonrelatives present', - # 11 '', - # 12 'Living alone', - # 13 'With own children under 18 years', - # 14 'With relatives, no own children under 18 years', + # 0 ['', + # 1 'With own children under 18 years', + # 2 'With no own children under 18 years', + # 3 '', + # 4 'With own children of the householder under 18 years', + # 5 'With no own children of the householder under 18 years', + # 6 '', + # 7 'Living alone', + # 8 'With own children under 18 years', + # 9 'With relatives, no own children under 18 years', + # 10 'With only nonrelatives present', + # 11 '', + # 12 'Living alone', + # 13 'With own children under 18 years', + # 14 'With relatives, no own children under 18 years', # 15 'With only nonrelatives present'] #2020 - # 0 ['', - # 1 'With children of the householder under 18 years', - # 2 'With no children of the householder under 18 years', - # 3 '', - # 4 'With children of the householder under 18 years', - # 5 'With no children of the householder under 18 years', - # 6 '', - # 7 'Living alone', - # 8 'With children of the householder under 18 years', - # 9 'With relatives, no children of the householder under 18 years', - # 10 'With only nonrelatives present', - # 11 '', - # 12 'Living alone', - # 13 'With children of the householder under 18 years', - # 14 'With relatives, no children of the householder under 18 years', + # 0 ['', + # 1 'With children of the householder under 18 years', + # 2 'With no children of the householder under 18 years', + # 3 '', + # 4 'With children of the householder under 18 years', + # 5 'With no children of the householder under 18 years', + # 6 '', + # 7 'Living alone', + # 8 'With children of the householder under 18 years', + # 9 'With relatives, no children of the householder under 18 years', + # 10 'With only nonrelatives present', + # 11 '', + # 12 'Living alone', + # 13 'With children of the householder under 18 years', + # 14 'With relatives, no children of the householder under 18 years', # 15 'With only nonrelatives present'] - + ## family households, by # workers _in family_ (not other workers in hh), presence of own_ch_u18, and marriage status B23009 = read_acs('B23009',geos) ## family households, by marriage status and presence of _related_ children in age groups @@ -1117,12 +1112,12 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co B11004[':'.join(['B11004','Other family:Female householder, no spouse present:With related children of the householder under 18 years',x])] # 2019 and 2020 have slightly different wording so pull from B11012_lst - # Ex 2019: for x in ['Living alone','With own children under 18 years','With relatives, no own children under 18 years','With only nonrelatives present']: + # Ex 2019: for x in ['Living alone','With own children under 18 years','With relatives, no own children under 18 years','With only nonrelatives present']: for x in [B11012_lst[12],B11012_lst[13],B11012_lst[14],B11012_lst[15]]: B11012[':'.join(['B11012','Single householder',x])] = \ B11012[':'.join(['B11012','Female householder, no spouse or partner present',x])] + \ B11012[':'.join(['B11012','Male householder, no spouse or partner present',x])] - + ## combine married and cohab households B11012['B11012:Two-partner household:With children of the householder under 18 years'] = \ B11012[f'B11012:Married-couple household:{B11012_lst[1]}'] + \ @@ -1131,7 +1126,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co B11012['B11012:Two-partner household:With no children of the householder under 18 years'] = \ B11012[f'B11012:Married-couple household:{B11012_lst[2]}'] + \ B11012[f'B11012:Cohabiting couple household:{B11012_lst[5]}'] # With no own children of the householder under 18 years - + ## combine married and unmarried partners B09021['B09021:Householder living with partner or partner of householder'] = \ B09021['B09021:Householder living with spouse or spouse of householder'] + B09021['B09021:Householder living with unmarried partner or unmarried partner of householder'] @@ -1152,7 +1147,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co .apply(lambda s: s * gq_stats['p_18_64_noninst_civil'] * gq_stats['group quarters:18 to 64']) \ .rename(columns={ 'ind_'+k+'_p|ninst1864civ':'C24030:'+k for k in ind_codes.keys()}) - + ind_mil_gq = gq_stats[['ind_'+k+'_p|milGQ' for k in ind_codes.keys()]].copy(deep=True) \ .apply(lambda s: s * gq_stats['p_18_64_noninst_mil'] * gq_stats['group quarters:18 to 64']) \ .rename(columns={ @@ -1190,7 +1185,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co ).rename("min_hh") df_bounds = pd.DataFrame({'total_workers':ind_df["C24030:"]}) - df_bounds = df_bounds.join(min_hh_B23009) + df_bounds = df_bounds.join(min_hh_B23009) ## in some places minimum # "workers" per table B23009 = more than total people in labor force?? df_bounds["min_hh"] = df_bounds[["total_workers","min_hh"]].min(axis=1) ## using total people in hh as max (no other reliable max; B23009 only counts families) @@ -1229,7 +1224,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co occ_hh = occ_df[['C24010:'+k for k in occ_codes.keys()]].copy(deep=True) \ .apply(lambda s: s * occ_df["_p_hh"]) - + ## IPF for each location that has gq workers idxs = gq_stats.index.intersection(has_puma_idx).intersection(df_bounds.index[df_bounds["_gq_workers"]>0]) df_bounds = df_bounds.loc[idxs,:].copy(deep=True) @@ -1238,7 +1233,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co (df_bounds["hh_workers"] <= df_bounds["max_hh"]).all(), (df_bounds["hh_workers"] >= df_bounds["min_hh"]).all()]): print("warning: could not reconcile household and GQ employment sums for some locations") - + ## industries for i in idxs: ## rows are: hh, civ gq, mil gq @@ -1248,7 +1243,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co colsums = ind_col_totals.loc[i,:].values m = np.stack([ind_hh.loc[i,:].astype(float).values, - ind_civ_gq.loc[i,:].astype(float).values, + ind_civ_gq.loc[i,:].astype(float).values, ind_mil_gq.loc[i,:].astype(float).values]) ## ipf handles 0's poorly; replace with 0.5 @@ -1269,7 +1264,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co colsums = occ_col_totals.loc[i,:].values m = np.stack([occ_hh.loc[i,:].astype(float).values, - occ_civ_gq.loc[i,:].astype(float).values, + occ_civ_gq.loc[i,:].astype(float).values, occ_mil_gq.loc[i,:].astype(float).values]) ## ipf handles 0's poorly; replace with 0.5 @@ -1280,8 +1275,8 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co occ_hh.loc[i,:] = new_m[0,:] occ_civ_gq.loc[i,:] = new_m[1,:] occ_mil_gq.loc[i,:] = new_m[2,:] - - + + ## check for NAN ind_hh.to_csv(os.path.join(PROCESSED_DIR,'ind_hh_test.csv')) ind_civ_gq.to_csv(os.path.join(PROCESSED_DIR,'ind_civ_gq_test.csv')) @@ -1289,9 +1284,9 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co occ_hh.to_csv(os.path.join(PROCESSED_DIR,'occ_hh_test.csv')) occ_civ_gq.to_csv(os.path.join(PROCESSED_DIR,'occ_civ_gq_test.csv')) occ_mil_gq.to_csv(os.path.join(PROCESSED_DIR,'occ_mil_gq_test.csv')) - - - + + + ## round to integer (preserve row sums) ind_hh = ind_hh.apply(lrRound,axis=1) ind_civ_gq = ind_civ_gq.apply(lrRound,axis=1) @@ -1299,8 +1294,8 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co occ_hh = occ_hh.apply(lrRound,axis=1) occ_civ_gq = occ_civ_gq.apply(lrRound,axis=1) occ_mil_gq = occ_mil_gq.apply(lrRound,axis=1) - - + + ## save gq employment counts for synth pop construction ind_civ_gq.join(occ_civ_gq).to_csv(os.path.join(PROCESSED_DIR,'gq_civilian_workers.csv')) @@ -1327,7 +1322,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co acs_tables['county'] = acs_tables.index.map(lambda x: x[0:5]) ## income data - for k,v in zip(inc_cats,inc_cols): + for k,v in zip(inc_cats,inc_cols, strict=False): acs_tables['B19001:'+k] = acs_tables[['B19001:'+x for x in v]].sum(axis=1) # print("Writing Census targets") @@ -1358,7 +1353,7 @@ def generate_targets(target_columns, geos, geo_xwalk, gq_stats, inc_cats, inc_co ## use JT01, "primary" jobs (because JT00 counts 2+ jobs for the same individual) ## "main" = work and live in state ## "aux" = work in state, live outside state -## +## ## need: ## main file for every state in the synth area, named *od_main_JT01* ## aux file for every state in the synth area, named *od_aux_JT01* @@ -1388,9 +1383,9 @@ def read_origin_destination(geos, commute_states): work_in_area = np.any([od['w_geocode'].str.startswith(g) for g in geos],axis=0) ## live in the synth area, work anywhere in the state(s) with od files provided live_in_area = np.any([od['h_geocode'].str.startswith(g) for g in geos],axis=0) - else: + else: ## if no synth area geos specified, assume work area is synth area - main_states = od.loc[od["main"]==True, "w_state"].unique() + main_states = od.loc[od["main"].astype(bool), "w_state"].unique() work_in_area = od["w_state"].isin(main_states) live_in_area = od["h_state"].isin(main_states) @@ -1530,7 +1525,7 @@ def calc_commute_marginals(geos,ind_codes,ind_keys,commute_states): wac_by_dest.loc["outside",ind_keys] = n_commuting_out * p_by_ind ## save for next step - + od_matrix.to_csv(os.path.join(PROCESSED_DIR,"work_od_prop.csv"),float_format="%.8g") ind_df.round(2).to_csv(os.path.join(PROCESSED_DIR,"work_io_sums.csv")) wac_by_dest.round(2).to_csv(os.path.join(PROCESSED_DIR,"work_id_est_sums.csv")) @@ -1555,7 +1550,15 @@ def calc_commute_marginals(geos,ind_codes,ind_keys,commute_states): # 2,500-4,999 Employees #N1000_4 N Number of Establishments: Employment Size Class: # 5,000 or More Employees -def generate_work_sizes(): +def generate_work_sizes(random_seed=None): + """Infer a lognormal employer-size distribution per county from CBP size classes. + + Args: + random_seed: seed for the size simulation. Passing one makes the inferred + distributions reproducible; the previous implementation used numpy's + unseeded global RNG, so work sizes varied between otherwise identical runs. + """ + rng = np.random.default_rng(random_seed) work_counties = pd.read_csv(os.path.join(PROCESSED_DIR,'work_sizes.csv'),dtype=str)["county"].values cbp = read_cbp(work_counties) ## @@ -1566,11 +1569,10 @@ def generate_work_sizes(): b = [5,10,20,50,100,250,500,1000,1500,2500,5000,30000] ## don't believe 0 cells? adj_z = 0.5 - sim_dist = [np.concatenate([np.random.randint(l,h,np.int64(s)) for (l,h,s) in zip(a,b, 1000*(adj_z+cbp.iloc[r,2:]))]) for r in range(cbp.shape[0])] + sim_dist = [np.concatenate([rng.integers(l,h,np.int64(s)) for (l,h,s) in zip(a,b, 1000*(adj_z+cbp.iloc[r,2:]), strict=False)]) for r in range(cbp.shape[0])] mu_l = [np.mean(np.log(s)) for s in sim_dist] - var_l = [np.var(np.log(s)) for s in sim_dist] mu_sz = [np.mean(s) for s in sim_dist] - imp_v = [2.0*(np.log(a) - b) for (a,b) in zip(mu_sz, mu_l)] + imp_v = [2.0*(np.log(a) - b) for (a,b) in zip(mu_sz, mu_l, strict=False)] cbp['mu_ln'] = mu_l cbp['sigma_ln'] = np.sqrt(imp_v) @@ -1668,7 +1670,7 @@ def __init__(self, config_dict=None, base_dir=None, verbose=1, auto_run=True): cfg_path = os.path.join(self.base_dir, "config.json") if not os.path.exists(cfg_path): raise FileNotFoundError(f"config.json file not found at {cfg_path}. Please create this file with the required configuration.") - with open(cfg_path, "r") as f: + with open(cfg_path) as f: self.config = json.load(f) # Set module-level variables for use by other functions @@ -1846,13 +1848,13 @@ def _init_parameters(self): # Generated sample columns that match each of target_columns, in the same order # (generated in read_psamp and read_hsamp) - self.sample_columns = ['fam_hh_2', 'fam_hh_3', 'fam_hh_4', 'fam_hh_5', 'fam_hh_6', 'fam_hh_7o', - 'non_fam_hh_1', 'non_fam_hh_2', 'non_fam_hh_3', 'non_fam_hh_4', 'non_fam_hh_5', 'non_fam_hh_6', 'non_fam_hh_7o', - 'w_own_ch_u18_married_fam_work0', + self.sample_columns = ['fam_hh_2', 'fam_hh_3', 'fam_hh_4', 'fam_hh_5', 'fam_hh_6', 'fam_hh_7o', + 'non_fam_hh_1', 'non_fam_hh_2', 'non_fam_hh_3', 'non_fam_hh_4', 'non_fam_hh_5', 'non_fam_hh_6', 'non_fam_hh_7o', + 'w_own_ch_u18_married_fam_work0', 'w_own_ch_u18_married_fam_work1', 'w_own_ch_u18_married_fam_work2', 'w_own_ch_u18_married_fam_work3o', - 'w_own_ch_u18_unmar_fam_work0', + 'w_own_ch_u18_unmar_fam_work0', 'w_own_ch_u18_unmar_fam_work1', 'w_own_ch_u18_unmar_fam_work2', 'w_own_ch_u18_unmar_fam_work3o', @@ -1865,33 +1867,33 @@ def _init_parameters(self): 'no_own_ch_u18_unmar_fam_work2', 'no_own_ch_u18_unmar_fam_work3o', 'fam_married_w_rel_ch_u6_only', - 'fam_married_w_rel_ch_u6_and_6_17', + 'fam_married_w_rel_ch_u6_and_6_17', 'fam_married_w_rel_ch_6_17_only', 'fam_married_no_rel_ch_u18', - 'fam_unmar_w_rel_ch_u6_only', - 'fam_unmar_w_rel_ch_u6_and_6_17', + 'fam_unmar_w_rel_ch_u6_only', + 'fam_unmar_w_rel_ch_u6_and_6_17', 'fam_unmar_w_rel_ch_6_17_only', 'fam_unmar_no_rel_ch_u18', - 'partner_hh_ch_u18', - 'partner_hh_no_ch_u18', + 'partner_hh_ch_u18', + 'partner_hh_no_ch_u18', 'hh_alone', - 'hh_single_ch_u18', - 'hh_single_other_rel', + 'hh_single_ch_u18', + 'hh_single_other_rel', 'hh_nonrel_only', 'ch_u18_in_hh', 'grandch_u18', 'age_18_34_alone', - 'age_18_34_partner', + 'age_18_34_partner', 'age_18_34_child_of_hh', - 'age_18_34_other_rel', - 'age_18_34_non_rel', - 'age_35_64_alone', - 'age_35_64_partner', + 'age_18_34_other_rel', + 'age_18_34_non_rel', + 'age_35_64_alone', + 'age_35_64_partner', 'age_35_64_child_of_hh', - 'age_35_64_other_rel', - 'age_35_64_non_rel', + 'age_35_64_other_rel', + 'age_35_64_non_rel', 'age_65o_alone', - 'age_65o_partner', + 'age_65o_partner', 'age_65o_other_rel', 'age_65o_non_rel', *self.inc_cats, @@ -2004,7 +2006,7 @@ def generate_work_sizes(self): """Run the workplace size generation step.""" self._log("\n*** Running ProcessData.generate_work_sizes() ***") self._log("-- Generating employer sizes") - generate_work_sizes() + generate_work_sizes(random_seed=self.config.get("random_seed")) self._log("-- processed/work_sizes.csv") def generate_schools(self): @@ -2031,7 +2033,7 @@ def run_all(self): self.generate_schools() self._log("") self._log("All ProcessData() steps complete") - + def quality_check(self, *, auto_print=True): """Run processed geography alignment checks and optionally print a summary. @@ -2053,11 +2055,10 @@ def quality_check(self, *, auto_print=True): after `auto_run=True` as well. """ self._log("\n*** Running ProcessData.quality_check() ***") - qc = QualityCheck(config_dict=self.config, base_dir=self.base_dir, auto_run=False) - qc._results = qc.run_all() # ensure we return the same object we print + qc = QualityCheck(config_dict=self.config, base_dir=self.base_dir, auto_run=True) if auto_print: qc.print_results() - return qc._results + return qc.results def _missingness_summary(df, cols): total = len(df) @@ -2084,7 +2085,7 @@ def __init__(self, config_dict=None, config_path=None, base_dir=None, path_overr self.config = config_dict else: cfg_path = config_path if config_path is not None else os.path.join(self.base_dir, "config.json") - with open(cfg_path, "r") as f: + with open(cfg_path) as f: self.config = json.load(f) self.data_dir = path_override if path_override is not None else self.config.get("path", self.base_dir) @@ -2146,9 +2147,9 @@ def run_all(self): @property def results(self): + """The diagnostics dict, computing it on first access.""" if self._results is None: self.run_all() - self.print_results() return self._results def print_results(self): @@ -2185,186 +2186,31 @@ def print_results(self): for note in summary["notes"]: print(f" - {note}") -def main(): - runner = ProcessData(auto_run=False) - runner.run_all() - - - - - - - -def generate_test_targets(geos,geo_xwalk): - - ## B09002 own ch u 18 by age and family type - B09002 = read_acs('B09002',geos) - ## B19123 Families; by size and public assistance/snap - B19123 = read_acs('B19123',geos) - ## B22010 households; person with disability - B22010 = read_acs('B22010',geos) - ## B23008 Own children under 18 years in families and subfamilies - B23008 = read_acs('B23008',geos) - ## households internet access - B28002 = read_acs('B28002',geos) - ## B28006 Household population 25 years and over; educational achievement - B28006 = read_acs('B28006',geos) - - for x in ["Under 6 years","6 to 17 years"]: - B23008[":".join(["B23008",x,"Living with two parents:One parent in labor force"])] = \ - B23008[":".join(["B23008",x,"Living with two parents:Father only in labor force"])] + \ - B23008[":".join(["B23008",x,"Living with two parents:Mother only in labor force"])] - for y in ["In labor force","Not in labor force"]: - B23008[":".join(["B23008",x,"Living with one parent",y])] = \ - B23008[":".join(["B23008",x,"Living with one parent:Living with father",y])] + \ - B23008[":".join(["B23008",x,"Living with one parent:Living with mother",y])] - - for x in ["Under 3 years","3 and 4 years","5 years","6 to 11 years","12 to 17 years"]: - B09002[":".join(["B09002:In other families:Single householder, no spouse present",x])] = \ - B09002[":".join(["B09002:In other families:Male householder, no spouse present",x])] + \ - B09002[":".join(["B09002:In other families:Female householder, no spouse present",x])] - - for x in ["In married-couple families","In other families:Single householder, no spouse present"]: - B09002[":".join(["B09002",x,"3 to 5 years"])] = \ - B09002[":".join(["B09002",x,"3 and 4 years"])] + \ - B09002[":".join(["B09002",x,"5 years"])] - - for x in ["Households with 1 or more persons with a disability","Households with no persons with a disability"]: - B22010[":".join(["B22010",x])] = \ - B22010[":".join(["B22010:Household received Food Stamps/SNAP in the past 12 months",x])] + \ - B22010[":".join(["B22010:Household did not receive Food Stamps/SNAP in the past 12 months",x])] - - ## join all census tables together - acs_tables = B09002.join([B19123,B22010,B23008,B28002,B28006]) - acs_tables['state'] = acs_tables.index.map(lambda x: x[0:2]) - acs_tables['county'] = acs_tables.index.map(lambda x: x[0:5]) - - ## which census columns to match: - target_columns = [ - 'B09002:In married-couple families:Under 3 years', - 'B09002:In married-couple families:3 to 5 years', - 'B09002:In married-couple families:6 to 11 years', - 'B09002:In married-couple families:12 to 17 years', - 'B09002:In other families:Single householder, no spouse present:Under 3 years', - 'B09002:In other families:Single householder, no spouse present:3 to 5 years', - 'B09002:In other families:Single householder, no spouse present:6 to 11 years', - 'B09002:In other families:Single householder, no spouse present:12 to 17 years', - 'B19123:2-person families:With cash public assistance income or households receiving Food Stamps/SNAP benefits in the past 12 months', - 'B19123:3-person families:With cash public assistance income or households receiving Food Stamps/SNAP benefits in the past 12 months', - 'B19123:4-person families:With cash public assistance income or households receiving Food Stamps/SNAP benefits in the past 12 months', - 'B19123:5-person families:With cash public assistance income or households receiving Food Stamps/SNAP benefits in the past 12 months', - 'B19123:6-person families:With cash public assistance income or households receiving Food Stamps/SNAP benefits in the past 12 months', - 'B19123:7-or-more-person families:With cash public assistance income or households receiving Food Stamps/SNAP benefits in the past 12 months', - 'B22010:Households with 1 or more persons with a disability', - 'B23008:Under 6 years:Living with two parents:Both parents in labor force', - 'B23008:Under 6 years:Living with two parents:One parent in labor force', - 'B23008:Under 6 years:Living with two parents:Neither parent in labor force', - 'B23008:Under 6 years:Living with one parent:In labor force', - 'B23008:Under 6 years:Living with one parent:Not in labor force', - 'B23008:6 to 17 years:Living with two parents:Both parents in labor force', - 'B23008:6 to 17 years:Living with two parents:One parent in labor force', - 'B23008:6 to 17 years:Living with two parents:Neither parent in labor force', - 'B23008:6 to 17 years:Living with one parent:In labor force', - 'B23008:6 to 17 years:Living with one parent:Not in labor force', - "B28002:With an Internet subscription", - "B28002:Internet access without a subscription", - "B28002:No Internet access", - "B28006:Less than high school graduate or equivalency:", - "B28006:High school graduate (includes equivalency) , some college or associate's degree :", - "B28006:Bachelor's degree or higher:"] - - missing_puma = geo_xwalk.index[geo_xwalk['st_puma'].isna()] - acs_tables = acs_tables[~acs_tables.index.isin(missing_puma)].copy(deep=True) - - ## drop cbgs with less than 20 hh - acs20 = acs_tables[acs_tables['B22010:'] > 19] - acs20[target_columns].to_csv(os.path.join(PROCESSED_DIR,'test_targets.csv')) - - return None - - -def gen_samp_test_cols(ADJINC, inc_cats, inc_cols, LODES_cutoff): - - psamp, ptotals = read_psamp(LODES_cutoff, [], []) - hsamp = read_hsamp(ptotals, ADJINC, inc_cats, inc_cols) - ## generated sample columns that match each of target_columns, in the same order: - sample_columns = ['own_ch_u3_in_fam_married', - 'own_ch_3_5_in_fam_married', - 'own_ch_6_11_in_fam_married', - 'own_ch_12_17_in_fam_married', - 'own_ch_u3_in_fam_unmar', - 'own_ch_3_5_in_fam_unmar', - 'own_ch_6_11_in_fam_unmar', - 'own_ch_12_17_in_fam_unmar', - 'fam_hh_2_snap_pap', - 'fam_hh_3_snap_pap', - 'fam_hh_4_snap_pap', - 'fam_hh_5_snap_pap', - 'fam_hh_6_snap_pap', - 'fam_hh_7o_snap_pap', - 'h_1_or_more_with_disab', - 'esp_2p_2w_age_u6', - 'esp_2p_1w_age_u6', - 'esp_2p_nw_age_u6', - 'esp_1p_1w_age_u6', - 'esp_1p_nw_age_u6', - 'esp_2p_2w_age_6_17', - 'esp_2p_1w_age_6_17', - 'esp_2p_nw_age_6_17', - 'esp_1p_1w_age_6_17', - 'esp_1p_nw_age_6_17', - 'h_internet_sub', - 'h_internet_nosub', - 'h_no_internet', - 'edu_not_hsgrad_age_25o', - 'edu_hs_or_somecoll_age_25o', - 'edu_bach_or_higher_age_25o'] - - ## csv of household samples - hsamp[sample_columns].to_csv(os.path.join(PROCESSED_DIR,'test_samples.csv')) - - return None - - -def test_cols(): - ## default income categories - inc_cats_def = ['q1_1', 'q1_2', 'q1_3', 'q2', 'q3', 'q4', 'q5'] - inc_cols_def = [['Less than $10,000'], - ['$10,000 to $14,999', '$15,000 to $19,999', '$20,000 to $24,999'], - ['$25,000 to $29,999', '$30,000 to $34,999', '$35,000 to $39,999'], - ['$40,000 to $44,999', '$45,000 to $49,999', '$50,000 to $59,999', '$60,000 to $74,999'], - ['$75,000 to $99,999', '$100,000 to $124,999'], - ['$125,000 to $149,999', '$150,000 to $199,999'], - ['$200,000 or more']] - - ## parameters from json files - d = tryJSON("config.json") - - ## read states/counties to include - geos = d.get("geos", d.get("geos",None)) - - ## read ADJINC - ADJINC = d.get("inc_adj",1.010145) - ## read income categories - inc_cats = d.get("inc_cats",inc_cats_def) - inc_cols = d.get("inc_cols",inc_cols_def) - LODES_cutoff = d.get("LODES_annual_income_boundary",40000) - adults_hh_cbg = read_acs('B09021',geos)[['B09021:']] - cbg_index = adults_hh_cbg.index - geo_xwalk = read_geo_xwalk(cbg_index) - - generate_test_targets(geos,geo_xwalk) - gen_samp_test_cols(ADJINC, inc_cats, inc_cols, LODES_cutoff) - - - - +def process_data(config, *, base_dir=None, verbose=1): + """Build the processed CO targets and sample pools for a run. + Args: + config: a config dict (see :func:`geopops.make_config`). + verbose: 0 for quiet, 1 for progress logging. + Returns: + ProcessData: the completed step, for inspection. + """ + return ProcessData(config_dict=config, base_dir=base_dir, verbose=verbose, auto_run=True) +def quality_check(config, *, base_dir=None, print_results=True): + """Check that the processed geographies line up before running CO. + Returns: + dict: PUMA coverage and missingness diagnostics. + """ + qc = QualityCheck(config_dict=config, base_dir=base_dir, auto_run=True) + if print_results: + qc.print_results() + return qc.results -if __name__ == "__main__": - main() +def main(): + runner = ProcessData(auto_run=False) + runner.run_all() diff --git a/src/geopops/co.py b/src/geopops/co.py index e11568c..cc7a0da 100644 --- a/src/geopops/co.py +++ b/src/geopops/co.py @@ -1,6 +1,14 @@ """ Combinatorial optimization via simulated annealing. -Translated from julia/CO.jl — parallelism removed, uses numpy. +Translated from julia/CO.jl — uses numpy. + +Each Census Block Group (CBG) is fitted independently: a set of ``n`` PUMS +household samples is chosen so that its aggregate demographic profile matches +the ACS target profile for that CBG, minimizing the Freeman-Tukey distance. + +Fitting proceeds in four passes over progressively broader sample pools --- PUMA, +county, CBSA, then urbanization level --- with each pass re-running only the CBGs +that did not yet meet the stopping criterion. """ import numpy as np import pandas as pd @@ -13,11 +21,21 @@ def FTdist(v1, v2): return float(np.sum((np.sqrt(v1 + 1.0) - np.sqrt(v2 + 1.0)) ** 2)) -def anneal(all_samples, mask, targ, n, params, rng): - """Simulated annealing on a subset of samples defined by mask. - Returns (global_indices, generations, score, temperature). +def anneal(samples, global_idxs, targ, n, params, rng): + """Simulated annealing over a pool of candidate household samples. + + Args: + samples: ``(n_pool, n_cols)`` int array of candidate sample profiles. + global_idxs: ``(n_pool,)`` int array mapping pool rows back to the full + sample table. + targ: ``(1, n_cols)`` target profile for this CBG. + n: number of households to select. + params: dict with ``maxgens``, ``critval``, ``cooldown``. + rng: numpy Generator. + + Returns: + tuple: ``(global_indices, generations, score, temperature)``. """ - samples = all_samples[mask, :] n_samples = samples.shape[0] if n_samples == 0: return (np.array([], dtype=int), 0, float('inf'), 0.0) @@ -26,9 +44,16 @@ def anneal(all_samples, mask, targ, n, params, rng): critval = params['critval'] cooldown = params['cooldown'] + # sqrt(targ + 1) is constant across the whole run; hoist it out of the loop. + sqrt_targ = np.sqrt(np.asarray(targ, dtype=float).ravel() + 1.0) + c0 = rng.integers(0, n_samples, size=n) - summary = samples[c0, :].sum(axis=0, keepdims=True) - E0 = FTdist(summary, targ) + # `summary` is the column-wise sum of the currently selected samples. Only one + # selection changes per generation, so it is updated incrementally rather than + # recomputed --- the inner loop becomes O(n_cols) instead of O(n * n_cols). + # `samples` is integer, so the incremental update is exact. + summary = samples[c0, :].sum(axis=0) + E0 = float(np.sum((np.sqrt(summary + 1.0) - sqrt_targ) ** 2)) T = 0.5 * E0 gen = 0 @@ -36,9 +61,11 @@ def anneal(all_samples, mask, targ, n, params, rng): gen += 1 cidx = rng.integers(len(c0)) orig = c0[cidx] - c0[cidx] = rng.integers(n_samples) - summary = samples[c0, :].sum(axis=0, keepdims=True) - E1 = FTdist(summary, targ) + new = rng.integers(n_samples) + c0[cidx] = new + summary += samples[new] + summary -= samples[orig] + E1 = float(np.sum((np.sqrt(summary + 1.0) - sqrt_targ) ** 2)) neg_dE = E0 - E1 if neg_dE >= 0 or rng.random() < np.exp(neg_dE / max(T, 1e-30)): @@ -46,13 +73,13 @@ def anneal(all_samples, mask, targ, n, params, rng): E0 = E1 else: c0[cidx] = orig + summary += samples[orig] + summary -= samples[new] if E0 < critval or gen > maxgens: break - global_indices = np.where(mask)[0] - res = global_indices[c0] - return (res, gen, E0, T) + return (global_idxs[c0], gen, E0, T) def read_targets(data_dir): @@ -65,7 +92,7 @@ def read_targets(data_dir): def read_hh_counts(data_dir): df = pd.read_csv(os.path.join(data_dir, 'processed', 'hh_counts.csv'), dtype={'Geo': str}) - return dict(zip(df.iloc[:, 0], df.iloc[:, 1])) + return dict(zip(df.iloc[:, 0], df.iloc[:, 1], strict=False)) def read_samples(data_dir): @@ -78,10 +105,10 @@ def read_samples(data_dir): def read_targ_geo(data_dir): cols = ['Geo', 'st_puma', 'cbsa', 'county', 'R', 'U'] df = pd.read_csv(os.path.join(data_dir, 'processed', 'cbg_geo.csv'), usecols=cols, dtype={'Geo': str, 'st_puma': str, 'cbsa': str, 'county': str}) - cbg_puma = dict(zip(df['Geo'], df['st_puma'])) - cbg_county = dict(zip(df['Geo'], df['county'])) - cbg_cbsa = dict(zip(df['Geo'], df['cbsa'])) - cbg_urban = dict(zip(df['Geo'], df['U'])) + cbg_puma = dict(zip(df['Geo'], df['st_puma'], strict=False)) + cbg_county = dict(zip(df['Geo'], df['county'], strict=False)) + cbg_cbsa = dict(zip(df['Geo'], df['cbsa'], strict=False)) + cbg_urban = dict(zip(df['Geo'], df['U'], strict=False)) return cbg_puma, cbg_county, cbg_cbsa, cbg_urban @@ -106,43 +133,81 @@ def urbanization_lookup(U_values, x): def sample_lookup(samp_geo, col, target_vals): - """For each target value, return a boolean mask over samples.""" + """Map each target value to the sample rows that match it. + + Returns a list of ``(key, index_array)`` pairs, one per target value. Targets + that share a key share the *same* index array object, so the caller can cache + the extracted sub-matrix by key --- CBGs in one county usually fall into only a + handful of PUMAs, so this collapses hundreds of gathers into a few. + """ if col == 'U': - return [urbanization_lookup(samp_geo['U'], x) for x in target_vals] + def mask_for(x): + return urbanization_lookup(samp_geo['U'], x) else: vals = samp_geo[col].values nan_mask = pd.isna(vals) - return [np.where(nan_mask, False, vals == x) for x in target_vals] + def mask_for(x): + return np.where(nan_mask, False, vals == x) + + cache = {} + out = [] + for x in target_vals: + idx = cache.get(x) + if idx is None: + idx = cache[x] = np.flatnonzero(mask_for(x)) + out.append((x, idx)) + return out -def optimize(samples, samp_masks, targs, n_hhs, params, rng): + +def _subpool(samples, key, idxs, cache): + """Extract (and memoize) the candidate sample sub-matrix for one lookup key.""" + sub = cache.get(key) + if sub is None: + sub = samples[idxs, :] + cache[key] = sub + return sub + + +def optimize(samples, samp_lookups, targs, n_hhs, params, rng): """Run annealing for each target. Returns list of (indices, gen, score, temp).""" + cache = {} results = [] - for mask, targ_row, n in zip(samp_masks, range(len(targs)), n_hhs): - targ = targs[targ_row:targ_row+1, :] - r = anneal(samples, mask, targ, n, params, rng) - results.append(r) + for (key, idxs), targ, n in zip(samp_lookups, targs, n_hhs, strict=False): + sub = _subpool(samples, key, idxs, cache) + results.append(anneal(sub, idxs, targ[np.newaxis, :], n, params, rng)) return results -def reoptimize(x, rerun, samples, samp_masks, targs, n_hhs, params, rng): - """Re-run optimization on targets that scored poorly; update x in-place.""" - enough = [samp_masks[i].sum() > (n_hhs[rerun[j]] // 2) for j, i in enumerate(range(len(samp_masks)))] - valid = [rerun[j] for j, ok in enumerate(enough) if ok] - valid_mask_idx = [j for j, ok in enumerate(enough) if ok] - if not valid: - return - for j, ri in enumerate(valid): - mi = valid_mask_idx[j] - targ = targs[ri:ri+1, :] - r = anneal(samples, samp_masks[mi], targ, n_hhs[ri], params, rng) +def reoptimize(x, rerun, samples, samp_lookups, targs, n_hhs, params, rng): + """Re-run optimization on targets that scored poorly; update `x` in place. + + Targets whose candidate pool is smaller than half the households they need are + skipped --- the pool is too thin for the fit to improve. + """ + cache = {} + for j, ri in enumerate(rerun): + key, idxs = samp_lookups[j] + if len(idxs) <= (n_hhs[ri] // 2): + continue + sub = _subpool(samples, key, idxs, cache) + r = anneal(sub, idxs, targs[ri:ri + 1, :], n_hhs[ri], params, rng) if r[2] < x[ri][2]: x[ri] = r -def process_counties(data_dir, counties=None, random_seed=None): +def _score_report(label, x, c_val, log=print): + """Print min/mean/max fit score and how many CBGs are still above threshold.""" + scores = [a[2] for a in x] + n_bad = sum(1 for s in scores if s > c_val) + log(f"-- After {label} pass: {n_bad} still above threshold; " + f"E0 min/mean/max: {min(scores):.2f} / {sum(scores) / len(scores):.2f} / {max(scores):.2f}") + + +def process_counties(data_dir, counties=None, random_seed=None, config=None, verbose=1): """Run CO for all counties. Returns (co_results, co_scores).""" rng = np.random.default_rng(random_seed) + log = print if verbose else (lambda *a, **k: None) samples, hh_ids = read_samples(data_dir) cbg_puma, cbg_county, cbg_cbsa, cbg_urban = read_targ_geo(data_dir) @@ -152,13 +217,16 @@ def process_counties(data_dir, counties=None, random_seed=None): n_hhs_all = [hh_counts[g] for g in geos_all] county_of = [g[:5] for g in geos_all] - config = tryJSON(os.path.join(data_dir, 'config.json')) + if config is None: + config = tryJSON(os.path.join(data_dir, 'config.json')) c_val = config.get('CO_crit_val', 10.0) CO_cooldown = config.get('CO_cooldown', 0.99) - CO_cooldown_slow = 0.5 + 0.5 * CO_cooldown CO_maxgens = config.get('CO_maxgens', 200000) params = dict(maxgens=CO_maxgens, critval=c_val, cooldown=CO_cooldown) + # The final (urbanization) pass cools more slowly, exploring longer, because by + # then the sample pool is broad and the fit is hard. + params_slow = dict(maxgens=CO_maxgens, critval=c_val, cooldown=0.5 + 0.5 * CO_cooldown) if counties is None: counties = sorted(set(county_of)) @@ -167,52 +235,30 @@ def process_counties(data_dir, counties=None, random_seed=None): all_co_scores = {} for c in counties: - cmask = [co == c for co in county_of] - idxs = [i for i, m in enumerate(cmask) if m] + idxs = [i for i, county in enumerate(county_of) if county == c] geos = [geos_all[i] for i in idxs] targs = targs_all[idxs, :] n_hhs = [n_hhs_all[i] for i in idxs] - print(f"\nCounty {c}: {len(geos)} CBGs") - print("") - print(f"Optimizing {len(geos)} CBGs at PUMA level") - samp_masks = sample_lookup(samp_geo, 'st_puma', [cbg_puma[g] for g in geos]) - x = optimize(samples, samp_masks, targs, n_hhs, params, rng) - scores = [a[2] for a in x] - n_bad = sum(1 for s in scores if s > c_val) - print(f"-- After PUMA pass; {n_bad} above threshold (will rerun); " - f"E0 min/mean/max: {min(scores):.2f} / {sum(scores)/len(scores):.2f} / {max(scores):.2f}") - - rerun = [i for i, r in enumerate(x) if r[2] > c_val] - print(f"Optimizing {len(rerun)} CBG(s) at county level") - if rerun: - re_masks = sample_lookup(samp_geo, 'county', [cbg_county[geos[i]] for i in rerun]) - reoptimize(x, rerun, samples, re_masks, targs, n_hhs, params, rng) - scores = [a[2] for a in x] - n_bad = sum(1 for s in scores if s > c_val) - print(f"-- After county pass: {n_bad} still above threshold; " - f"E0 min/mean/max: {min(scores):.2f} / {sum(scores)/len(scores):.2f} / {max(scores):.2f}") - - rerun = [i for i, r in enumerate(x) if r[2] > c_val] - print(f"Optimizing {len(rerun)} CBG(s) at CBSA level") - if rerun: - re_masks = sample_lookup(samp_geo, 'cbsa', [cbg_cbsa[geos[i]] for i in rerun]) - reoptimize(x, rerun, samples, re_masks, targs, n_hhs, params, rng) - scores = [a[2] for a in x] - n_bad = sum(1 for s in scores if s > c_val) - print(f"-- After CBSA pass: {n_bad} still above threshold; " - f"E0 min/mean/max: {min(scores):.2f} / {sum(scores)/len(scores):.2f} / {max(scores):.2f}") - - rerun = [i for i, r in enumerate(x) if r[2] > c_val] - print(f"Optimizing {len(rerun)} CBG(s) at urbanization level") - if rerun: - params_slow = dict(maxgens=CO_maxgens, critval=c_val, cooldown=CO_cooldown_slow) - re_masks = sample_lookup(samp_geo, 'U', [cbg_urban[geos[i]] for i in rerun]) - reoptimize(x, rerun, samples, re_masks, targs, n_hhs, params_slow, rng) - scores = [a[2] for a in x] - n_bad = sum(1 for s in scores if s > c_val) - print(f"-- After urbanization pass: {n_bad} still above threshold; " - f"E0 min/mean/max: {min(scores):.2f} / {sum(scores)/len(scores):.2f} / {max(scores):.2f}") + log(f"\nCounty {c}: {len(geos)} CBGs\n") + log(f"Optimizing {len(geos)} CBGs at PUMA level") + x = optimize(samples, sample_lookup(samp_geo, 'st_puma', [cbg_puma[g] for g in geos]), + targs, n_hhs, params, rng) + _score_report("PUMA", x, c_val, log) + + # Progressively broaden the candidate pool for CBGs that still fit poorly. + fallbacks = [("county", 'county', cbg_county, params), + ("CBSA", 'cbsa', cbg_cbsa, params), + ("urbanization", 'U', cbg_urban, params_slow)] + for label, col, lookup, level_params in fallbacks: + rerun = [i for i, r in enumerate(x) if r[2] > c_val] + if not rerun: + log(f"All CBGs met the criterion; skipping {label} pass") + continue + log(f"Optimizing {len(rerun)} CBG(s) at {label} level") + samp_lookups = sample_lookup(samp_geo, col, [lookup[geos[i]] for i in rerun]) + reoptimize(x, rerun, samples, samp_lookups, targs, n_hhs, level_params, rng) + _score_report(label, x, c_val, log) co_results_county = {} co_scores_county = {} @@ -225,7 +271,7 @@ def process_counties(data_dir, counties=None, random_seed=None): all_co_scores[c] = co_scores_county n_good = sum(1 for s in co_scores_county.values() if s <= c_val) - print("") - print(f"{n_good}/{len(geos)} CBGs met the stopping criterion (E0 <= {c_val}) for the Freeman-Tukey distance score.") + log(f"\n{n_good}/{len(geos)} CBGs met the stopping criterion " + f"(E0 <= {c_val}) for the Freeman-Tukey distance score.") return all_co_results, all_co_scores diff --git a/src/geopops/config.json b/src/geopops/config.json index 2c21cfc..653efc7 100644 --- a/src/geopops/config.json +++ b/src/geopops/config.json @@ -1,7 +1,7 @@ { "path": "data", "census_api_key": null, - "julia_env_path": null, + "random_seed": null, "main_year": 2019, "geos": [ "45083" @@ -110,4 +110,4 @@ "gq_K": 12, "netw_K": 8, "netw_B": 0.25 -} +} \ No newline at end of file diff --git a/src/geopops/config.py b/src/geopops/config.py index f598bcb..fe941b6 100644 --- a/src/geopops/config.py +++ b/src/geopops/config.py @@ -1,10 +1,40 @@ +"""Configuration loading, merging, and validation for GeoPops. + +The packaged ``config.json`` is a read-only *template*. A run's config is an +ordinary dict, and :func:`make_config` returns one; persisting it is the caller's +choice and, by default, writes into the run's own output directory rather than into +the installed package. + +Secrets (the Census API key) come from the environment or a ``.env`` file and are +never written into the packaged template. +""" import json import os +import warnings + from dotenv import load_dotenv, find_dotenv +from .exceptions import ConfigError + load_dotenv(find_dotenv()) BASE_DIR = os.path.dirname(os.path.abspath(__file__)) +TEMPLATE_PATH = os.path.join(BASE_DIR, "config.json") + +#: Config keys that must never be written into the packaged template +SENSITIVE_CONFIG_KEYS = ("census_api_key",) + +#: Config keys that may be supplied as user-facing overrides +OVERRIDE_KEYS = ("census_api_key", "main_year", "geos", "commute_states", + "use_pums", "path", "random_seed") + +#: ACS / decennial tables required by the pipeline, used to backfill minimal configs +DEFAULT_ACS_REQUIRED = [ + "B01001", "B09019", "B09020", "C24030", "B23025", "C24010", "B11016", + "B11012", "B23009", "B11004", "B19001", "B22010", "B09021", "B09018", + "B11001H", "B11001I", "B25006", +] +DEFAULT_DEC_REQUIRED = ["P43", "P18"] def _merge_dict(base, override): @@ -16,27 +46,40 @@ def _merge_dict(base, override): return base -def load_config(base_dir=None): - cfg_dir = base_dir if base_dir is not None else BASE_DIR - cfg_path = os.path.join(cfg_dir, "config.json") - with open(cfg_path, "r") as f: +def load_config(base_dir=None, path=None): + """Load a config file, applying any sibling ``config.local.json`` overrides. + + Args: + base_dir: directory holding ``config.json``. Defaults to the package + directory (the shipped template). + path: explicit path to a config file; overrides `base_dir`. + + Returns: + dict: the config. + """ + if path is not None: + cfg_path = path + cfg_dir = os.path.dirname(os.path.abspath(path)) + else: + cfg_dir = base_dir if base_dir is not None else BASE_DIR + cfg_path = os.path.join(cfg_dir, "config.json") + + if not os.path.exists(cfg_path): + raise ConfigError(f"Config file not found: {cfg_path}") + with open(cfg_path) as f: config = json.load(f) # Optional untracked local overrides for machine-specific values. local_cfg_path = os.path.join(cfg_dir, "config.local.json") if os.path.exists(local_cfg_path): - with open(local_cfg_path, "r") as f: - local_cfg = json.load(f) - config = _merge_dict(config, local_cfg) + with open(local_cfg_path) as f: + config = _merge_dict(config, json.load(f)) return config -SENSITIVE_CONFIG_KEYS = ("census_api_key", "julia_env_path") - - def _template_config(config): - """Return a copy safe to ship/write as the package template (no secrets).""" + """A copy safe to ship as the package template (no secrets).""" template = dict(config) for key in SENSITIVE_CONFIG_KEYS: template[key] = None @@ -44,126 +87,118 @@ def _template_config(config): def save_config(config, config_path=None, *, sanitize=False): - cfg_path = config_path if config_path is not None else os.path.join(BASE_DIR, "config.json") - payload = _template_config(config) if sanitize else config - # Ensure target directory exists when writing to a custom location - os.makedirs(os.path.dirname(cfg_path), exist_ok=True) - with open(cfg_path, "w") as f: - json.dump(payload, f, indent=4) + """Write `config` as JSON. + + Args: + config: the config dict. + config_path: destination path, or a directory to write ``config.json`` into. + Defaults to the run's own ``path`` directory. + sanitize: blank out secrets first (used only for the packaged template). + """ + if config_path is None: + config_path = os.path.join(config.get("path", "."), "config.json") + elif os.path.isdir(config_path): + config_path = os.path.join(config_path, "config.json") + + parent = os.path.dirname(os.path.abspath(config_path)) + os.makedirs(parent, exist_ok=True) + with open(config_path, "w") as f: + json.dump(_template_config(config) if sanitize else config, f, indent=4) + return config_path def compute_decennial_year(main_year): + """The decennial census vintage that applies to `main_year`.""" try: - return 2020 if int(main_year) >= 2020 else 2010 - except Exception: - return 2010 - - -def update_config_values(config, - census_api_key=None, - main_year=None, - geos=None, - commute_states=None, - use_pums=None, - path=None, - julia_env_path=None): - # Fall back to environment variables for sensitive/user-specific values - if census_api_key is None: - census_api_key = os.environ.get("CENSUS_API_KEY") - if julia_env_path is None: - julia_env_path = os.environ.get("JULIA_ENV_PATH") - - if census_api_key is not None: - config["census_api_key"] = census_api_key - if main_year is not None: - config["main_year"] = main_year - config["decennial_year"] = compute_decennial_year(main_year) - if geos is not None: - config["geos"] = geos - if commute_states is not None: - config["commute_states"] = commute_states - if use_pums is not None: - config["use_pums"] = use_pums - if path is not None: - config["path"] = path - if julia_env_path is not None: - config["julia_env_path"] = julia_env_path + year = int(main_year) + except (TypeError, ValueError): + raise ConfigError( + f"main_year must be an integer year, got {main_year!r}." + ) from None + return 2020 if year >= 2020 else 2010 + + +def update_config_values(config, **overrides): + """Apply user overrides to `config` in place, falling back to the environment. + + Only keys in :data:`OVERRIDE_KEYS` are accepted; anything else is a typo and is + reported rather than silently ignored. ``None`` values mean "leave unchanged". + """ + unknown = set(overrides) - set(OVERRIDE_KEYS) + if unknown: + raise ConfigError( + f"Unknown config override(s): {sorted(unknown)}. " + f"Valid overrides: {list(OVERRIDE_KEYS)}" + ) + + # Sensitive/user-specific values fall back to the environment + if overrides.get("census_api_key") is None: + overrides["census_api_key"] = os.environ.get("CENSUS_API_KEY") + + for key, value in overrides.items(): + if value is not None: + config[key] = value + if overrides.get("main_year") is not None: + config["decennial_year"] = compute_decennial_year(overrides["main_year"]) + + config.setdefault("acs_required", list(DEFAULT_ACS_REQUIRED)) + config.setdefault("dec_required", list(DEFAULT_DEC_REQUIRED)) return config -class WriteConfig: - def __init__(self, - census_api_key=None, - main_year=None, - geos=None, - commute_states=None, - use_pums=None, - path=None, - julia_env_path=None, - pars=None, - config_dict=None, - base_dir=None): - pars = pars or {} - self.base_dir = base_dir if base_dir is not None else BASE_DIR - # Load base template from the package directory unless a dict is provided - self.template_config_path = os.path.join(self.base_dir, "config.json") - path = pars.get("path") if path is None else path - if path is not None: - # If path is a directory, append 'config.json' to it - if os.path.isdir(path): - self.path = os.path.join(path, "config.json") - else: - self.path = path - else: - self.path = self.template_config_path - if config_dict is None: - config_dict = pars.get("config_dict") - self.config = config_dict if config_dict is not None else load_config(self.base_dir) - self.overrides = { - "census_api_key": pars.get("census_api_key") if census_api_key is None else census_api_key, - "main_year": pars.get("main_year") if main_year is None else main_year, - "geos": pars.get("geos") if geos is None else geos, - "commute_states": pars.get("commute_states") if commute_states is None else commute_states, - "use_pums": pars.get("use_pums") if use_pums is None else use_pums, - "path": path, - "julia_env_path": pars.get("julia_env_path") if julia_env_path is None else julia_env_path, - } - - self.run_all() - - def run_all(self): - print("") - print("============================================================") - print("Running WriteConfig()") - print("============================================================") - update_config_values( - self.config, - census_api_key=self.overrides["census_api_key"], - main_year=self.overrides["main_year"], - geos=self.overrides["geos"], - commute_states=self.overrides["commute_states"], - use_pums=self.overrides["use_pums"], - path=self.overrides["path"], - julia_env_path=self.overrides["julia_env_path"], +def validate_config(config): + """Check a config for the mistakes that otherwise surface deep in the pipeline. + + Raises: + ConfigError: if a required key is missing or a value is unusable. + """ + for key in ("path", "main_year", "geos"): + if not config.get(key): + raise ConfigError(f"config is missing required key {key!r}.") + + if not isinstance(config["geos"], list | tuple) or not config["geos"]: + raise ConfigError("config['geos'] must be a non-empty list of state/county FIPS codes.") + + compute_decennial_year(config["main_year"]) # raises if unparseable + + # Traits are carried generically, but only PUMS-derived ones will have values, + # and CO does not target them --- worth saying once, up front. + traits = config.get("additional_traits") or [] + if not isinstance(traits, list | tuple): + raise ConfigError("config['additional_traits'] must be a list of column names.") + + if config.get("random_seed") is None: + warnings.warn( + "No random_seed set: this run will not be reproducible. " + "Pass seed=... or set config['random_seed'].", + stacklevel=3, ) - # User output path may include secrets for local runs; package template never should. - if os.path.abspath(self.path) != os.path.abspath(self.template_config_path): - save_config(self.config, self.path) - save_config(self.config, self.template_config_path, sanitize=True) - else: - save_config(self.config, self.template_config_path, sanitize=True) - print("-- Updated config.json with parameter dictionary") + return config - def get_pars(self): - with open(self.path, "r") as f: - cfg = json.load(f) - print(json.dumps(cfg, indent=2)) +def make_config(path=None, *, base_dir=None, template=None, save=False, **overrides): + """Build a run configuration from the packaged template plus overrides. -def main(): - runner = WriteConfig() - runner.run_all() + Args: + path: output directory for the run (also where results are written). + base_dir: directory to load the template from; defaults to the package. + template: a full config dict to use instead of loading a template. + save: if True, also write ``/config.json``. + **overrides: any of :data:`OVERRIDE_KEYS`. + Returns: + dict: the effective config. -if __name__ == "__main__": - main() \ No newline at end of file + Example:: + + cfg = geopops.make_config(path="data", geos=["45083"], main_year=2019, + commute_states=["45", "37"], use_pums=["45", "37"]) + """ + config = dict(template) if template is not None else load_config(base_dir) + if path is not None: + overrides["path"] = path + update_config_values(config, **overrides) + validate_config(config) + if save: + save_config(config) + return config diff --git a/src/geopops/export.py b/src/geopops/export.py index 473f305..3d0f8a3 100644 --- a/src/geopops/export.py +++ b/src/geopops/export.py @@ -15,6 +15,13 @@ def _log_export(verbose, msg=""): print(msg) +def _trait_names(people): + """The config-driven trait names carried by this population, in column order.""" + for person in people.values(): + return list(person.schema.names) + return [] + + def _mcon(val): if val is None: return "" @@ -46,22 +53,19 @@ def export_synthpop(data_dir, cbgs, households, people, sch_students, sch_worker os.path.join(export_dir, 'hh.csv'), index=False) _log_export(verbose, f"-- {rel}/hh.csv") + # Trait columns are config-driven, so read them off the population's schema + # rather than hardcoding a list that silently drifts out of sync. + trait_names = _trait_names(people) p_rows = sorted([ (int(k[0]), int(k[1]), int(k[2]), _mcon(v.sample), _mcon(v.age), - _mcon(v.female), _mcon(v.working), _mcon(v.commuter), - _mcon(v.com_inc), _mcon(v.com_cat), - _mcon(v.race_white_alone), _mcon(v.race_black_alone), _mcon(v.race_amerindian_or_alaskan), - _mcon(v.race_asian_alone), _mcon(v.race_pacific_alone), _mcon(v.race_other_alone), - _mcon(v.race_two_or_more), _mcon(v.hispanic), - _mcon(v.sch_grade)) + _mcon(v.working), _mcon(v.commuter), _mcon(v.com_inc), _mcon(v.com_cat), + *(_mcon(t) for t in v.trait_values), _mcon(v.sch_grade)) for k, v in people.items() ], key=lambda x: (x[2], x[1], x[0])) pd.DataFrame(p_rows, columns=[ - 'p_id', 'hh_id', 'cbg_id', 'sample_index', 'age', 'female', 'working', 'commuter', + 'p_id', 'hh_id', 'cbg_id', 'sample_index', 'age', 'working', 'commuter', 'commuter_income_category', 'commuter_workplace_category', - 'race_white_alone', 'race_black_alone', 'race_amerindian_or_alaskan', - 'race_asian_alone', 'race_pacific_alone', 'race_other_alone', - 'race_two_or_more', 'hispanic', 'sch_grade' + *trait_names, 'sch_grade' ]).to_csv(os.path.join(export_dir, 'people.csv'), index=False) _log_export(verbose, f"-- {rel}/people.csv") diff --git a/src/geopops/geopops_starsim.py b/src/geopops/geopops_starsim.py deleted file mode 100644 index 2ccb868..0000000 --- a/src/geopops/geopops_starsim.py +++ /dev/null @@ -1,412 +0,0 @@ -import pandas as pd -import numpy as np -import starsim as ss -from scipy.io import mmread -import os -import json - -BASE_DIR = os.path.dirname(os.path.abspath(__file__)) - - -def _load_age_by_matrix_index(pop_export_dir): - """Map matrix row/col index (``index_zero``) to age, same merge as ``ForStarsim.People``.""" - adj = pd.read_csv(os.path.join(pop_export_dir, "adj_mat_keys.csv"), low_memory=False) - people = pd.read_csv(os.path.join(pop_export_dir, "people.csv"), low_memory=False) - merged = adj.merge(people, on=["p_id", "hh_id", "cbg_id"], how="left") - merged = merged.drop_duplicates(subset=["index_zero"], keep="first") - return merged.set_index("index_zero")["age"] - - -def _canonicalize_undirected_edges_df(net_df, age_by_idx): - """Reorder ``p1``, ``p2`` so ``age(p1) <= age(p2)`` when both ages exist; else smaller index is ``p1``.""" - if net_df.empty: - return net_df.copy() - out = net_df.copy() - p1 = out["p1"].to_numpy(dtype=np.int64, copy=True) - p2 = out["p2"].to_numpy(dtype=np.int64, copy=True) - a1 = age_by_idx.reindex(p1).to_numpy() - a2 = age_by_idx.reindex(p2).to_numpy() - a1 = np.where(pd.isna(a1), np.nan, np.asarray(a1, dtype=float)) - a2 = np.where(pd.isna(a2), np.nan, np.asarray(a2, dtype=float)) - both = np.isfinite(a1) & np.isfinite(a2) - swap = np.zeros(len(out), dtype=bool) - swap[both] = (a1[both] > a2[both]) | ((a1[both] == a2[both]) & (p1[both] > p2[both])) - swap[~both] = p1[~both] > p2[~both] - out.loc[swap, ["p1", "p2"]] = np.column_stack([p2[swap], p1[swap]]) - return out - - -def _random_flip_undirected_edges_df(net_df, rng): - """ - Randomly swap (p1, p2) per edge with probability 0.5. - This only changes endpoint labeling and therefore affects plots that treat (p1_age, p2_age) as ordered. - """ - if net_df.empty: - return net_df.copy() - out = net_df.copy() - p1 = out["p1"].to_numpy(dtype=np.int64, copy=True) - p2 = out["p2"].to_numpy(dtype=np.int64, copy=True) - flip = rng.random(len(out)) < 0.5 - # swap endpoints for flipped edges - p1_new = p1.copy() - p2_new = p2.copy() - p1_new[flip] = p2[flip] - p2_new[flip] = p1[flip] - out["p1"] = p1_new - out["p2"] = p2_new - return out - - -class _ForStarsimSubgroupTracking(ss.Analyzer): - def __init__(self, subgroup, outcome, name=None, state_id=None, *args, **kwargs): - super().__init__(*args, **kwargs) - self.has_product = False - self.subgroup = subgroup - self.outcome = outcome - self.state_id = state_id - self.n_outcome = {} - if name: - self.name = name - - def step(self): - sim = self.sim - - if not self.n_outcome: - groups = np.unique(sim.people[self.subgroup]) - self.n_outcome = {group: [] for group in groups} - - disease_name = sim.diseases[0].name.lower() - disease_obj = getattr(sim.people, disease_name, None) - - for group in self.n_outcome.keys(): - if self.state_id is not None: - count = len( - ss.uids( - (sim.people[self.subgroup] == group) - & (disease_obj[self.outcome] == 1) - & (sim.people.state == self.state_id) - ) - ) - else: - count = len(ss.uids((sim.people[self.subgroup] == group) & (disease_obj[self.outcome] == 1))) - self.n_outcome[group].append(count) - - def get_subgroup_data(self): - """Return a DataFrame where rows are subgroups and columns are time steps.""" - df = pd.DataFrame.from_dict(self.n_outcome, orient='index') - df.columns = [f't_{i}' for i in range(len(df.columns))] - df.index.name = self.subgroup - df = df.reset_index() - return df - - -class _ForStarsimGPNetwork(ss.Network): - def __init__(self, name, edge_weight=1.0, csv_path=None, network_df=None, p1_col='p1', p2_col='p2', beta_col=None): - super().__init__() - self.name = name - self.edge_weight = edge_weight - self.csv_path = csv_path - self.network_df_input = network_df - self.p1_col = p1_col - self.p2_col = p2_col - self.beta_col = beta_col - - if self.csv_path is not None and self.network_df_input is not None: - raise ValueError("Provide only one of csv_path or network_df, not both.") - - if self.network_df_input is not None: - self.network_df = self._normalize_network_dataframe(self.network_df_input) - elif self.csv_path is not None: - self.network_df = self._load_custom_network_csv() - else: - self._ensure_networks_created() - - self.network_map = { - 'homenet': ForStarsim._net_h, - 'schoolnet': ForStarsim._net_s, - 'worknet': ForStarsim._net_w, - 'gqnet': ForStarsim._net_g, - } - - if name not in self.network_map: - raise ValueError( - f"Unknown network name '{name}'. Available built-in names: {list(self.network_map.keys())}. " - "To use a custom CSV network, provide csv_path=..." - ) - - self.network_df = self.network_map[name] - - self._populate_edges() - - def _load_custom_network_csv(self): - """Load and validate a custom edge-list CSV for network creation.""" - if not os.path.exists(self.csv_path): - raise FileNotFoundError(f"Custom network file not found: {self.csv_path}") - - df = pd.read_csv(self.csv_path) - return self._normalize_network_dataframe(df, source_desc=f"CSV '{self.csv_path}'") - - def _normalize_network_dataframe(self, df, source_desc="provided dataframe"): - """Validate and normalize custom network data into p1/p2/edge_weight columns.""" - if not isinstance(df, pd.DataFrame): - raise TypeError(f"network_df must be a pandas DataFrame, got {type(df)}") - - tmp = df.copy() - if 'Unnamed: 0' in tmp.columns: - tmp = tmp.drop(columns=['Unnamed: 0']) - - missing_cols = [c for c in [self.p1_col, self.p2_col] if c not in tmp.columns] - if missing_cols: - raise ValueError( - f"{source_desc} is missing required column(s): {missing_cols}. " - f"Available columns: {list(tmp.columns)}" - ) - - out = pd.DataFrame({ - 'p1': pd.to_numeric(tmp[self.p1_col], errors='coerce'), - 'p2': pd.to_numeric(tmp[self.p2_col], errors='coerce'), - }).dropna(subset=['p1', 'p2']) - - out['p1'] = out['p1'].astype(np.int64) - out['p2'] = out['p2'].astype(np.int64) - - if self.beta_col is not None: - if self.beta_col not in tmp.columns: - raise ValueError( - f"beta_col '{self.beta_col}' not found in {source_desc}. " - f"Available columns: {list(tmp.columns)}" - ) - beta_vals = pd.to_numeric(tmp[self.beta_col], errors='coerce') - beta_vals = beta_vals.loc[out.index].fillna(float(self.edge_weight)) - out['edge_weight'] = beta_vals.astype(float).values - else: - out['edge_weight'] = float(self.edge_weight) - - if out.empty: - raise ValueError(f"{source_desc} has no valid edges after parsing.") - - return out.reset_index(drop=True) - - def _populate_edges(self): - """Populate network edges from dataframe.""" - self.edges.p1 = self.network_df['p1'].values - self.edges.p2 = self.network_df['p2'].values - # Honor an explicitly provided scalar edge_weight (e.g., 2.0 for homenet) - # so users can rescale built-in networks without editing CSV/dataframes. - if float(self.edge_weight) != 1.0: - self.edges.beta = np.full(len(self.network_df), float(self.edge_weight)) - elif 'edge_weight' in self.network_df.columns: - self.edges.beta = self.network_df['edge_weight'].values.astype(float) - else: - self.edges.beta = np.full(len(self.network_df), self.edge_weight) - self.validate() - - def _ensure_networks_created(self): - """Create networks if they haven't been created yet.""" - if ForStarsim._net_h is None: - self._create_networks() - - def _create_networks(self): - """Create network dataframes from matrix files.""" - print("\n*** Running ForStarsim.GPNetwork._create_networks() ***") - cfg_path = os.path.join(BASE_DIR, "config.json") - if not os.path.exists(cfg_path): - raise FileNotFoundError(f"config.json file not found at {cfg_path}") - with open(cfg_path, "r") as f: - config = json.load(f) - path = config.get("path") - flip_seed = int(config.get("random_seed", 0)) - - def new_layer(file): - m = mmread(file) - mat = pd.DataFrame({ - "p1": np.asarray(m.col, dtype=np.int64), - "p2": np.asarray(m.row, dtype=np.int64), - }) - mat["edge_weight"] = np.int64(1) - return mat - - pop_export = os.path.join(path, "pop_export") - starsim_dir = os.path.join(pop_export, "starsim") - os.makedirs(starsim_dir, exist_ok=True) - - ForStarsim._net_h = new_layer(os.path.join(pop_export, "adj_upper_triang_hh.mtx")) - ForStarsim._net_s = new_layer(os.path.join(pop_export, "adj_upper_triang_sch.mtx")) - ForStarsim._net_w = new_layer(os.path.join(pop_export, "adj_upper_triang_wp.mtx")) - ForStarsim._net_g = new_layer(os.path.join(pop_export, "adj_upper_triang_gq.mtx")) - - # For plotting/visualization: treat undirected layers as having ~50/50 endpoint ordering. - # This avoids having plots that interpret (p1_age, p2_age) as ordered appear "triangular" or asymmetric. - flip_rng = np.random.default_rng(flip_seed) - ForStarsim._net_h = _random_flip_undirected_edges_df(ForStarsim._net_h, flip_rng) - ForStarsim._net_s = _random_flip_undirected_edges_df(ForStarsim._net_s, flip_rng) - ForStarsim._net_w = _random_flip_undirected_edges_df(ForStarsim._net_w, flip_rng) - ForStarsim._net_g = _random_flip_undirected_edges_df(ForStarsim._net_g, flip_rng) - - ForStarsim._net_h.to_csv(os.path.join(starsim_dir, "net_h.csv"), index=False) - ForStarsim._net_s.to_csv(os.path.join(starsim_dir, "net_s.csv"), index=False) - ForStarsim._net_w.to_csv(os.path.join(starsim_dir, "net_w.csv"), index=False) - ForStarsim._net_g.to_csv(os.path.join(starsim_dir, "net_g.csv"), index=False) - - print("Network csv files created and saved successfully") - - def step(self): - self.validate() - - -class ForStarsim: - """Creates and initializes Starsim People objects from GeoPops data. - - This class orchestrates the creation of Starsim People objects using processed - GeoPops data. It follows the same pattern as other GeoPops classes. - """ - - # Class-level variables to store network dataframes - _net_h = None - _net_s = None - _net_w = None - _net_g = None - - def __init__(self, config_dict=None, config_path=None, base_dir=None): - """Create a Starsim runner. - - Args: - config_dict: Optional dict with configuration. If provided, takes precedence over config_path. - config_path: Optional path to a JSON config file. Defaults to config.json in package directory. - base_dir: Optional base dir to use for relative paths. Defaults to package directory. - """ - # Use package directory as base, same as process_data.py - self.base_dir = base_dir if base_dir is not None else BASE_DIR - - if config_dict is not None: - self.config = config_dict - else: - cfg_path = config_path if config_path is not None else os.path.join(self.base_dir, "config.json") - if not os.path.exists(cfg_path): - raise FileNotFoundError(f"config.json file not found at {cfg_path}. Please create this file with the required configuration.") - with open(cfg_path, "r") as f: - self.config = json.load(f) - - # Set path from config (no hardcoded default) - self.path = self.config.get("path") - - @staticmethod - def _load_config(config_dict=None, config_path=None, base_dir=None): - base_dir = base_dir if base_dir is not None else BASE_DIR - if config_dict is not None: - return config_dict - cfg_path = config_path if config_path is not None else os.path.join(base_dir, "config.json") - if not os.path.exists(cfg_path): - raise FileNotFoundError( - f"config.json file not found at {cfg_path}. Please create this file with the required configuration." - ) - with open(cfg_path, "r") as f: - return json.load(f) - - @classmethod - def People(cls, config_dict=None, config_path=None, base_dir=None): - """Create and return a Starsim People object directly.""" - print("\n*** Running ForStarsim.People() ***") - config = cls._load_config(config_dict=config_dict, config_path=config_path, base_dir=base_dir) - path = config.get("path") - - adj_mat_keys = pd.read_csv(f'{path}/pop_export/adj_mat_keys.csv') - people = pd.read_csv(f'{path}/pop_export/people.csv') - ppl_df = adj_mat_keys.merge(people, on=['p_id', 'hh_id', 'cbg_id'], how='left') - schools = pd.read_csv(f'{path}/pop_export/sch_students.csv') - ppl_df = ppl_df.merge(schools, on=['p_id', 'hh_id', 'cbg_id'], how='left') - ppl_df.loc[ppl_df['sch_code'].isnull(), 'sch_code'] = 0 - ppl_df.insert(0, 'uid', ppl_df['index_zero'].values) - - cbg_idxs = pd.read_csv(f'{path}/pop_export/cbg_idxs.csv') - ppl_df = ppl_df.merge(cbg_idxs, on='cbg_id', how='left') - ppl_df['state'] = ppl_df['cbg_geocode'].astype(str).str[:2].replace({'na': '0.0', 'nan': '0.0'}).astype(float) - ppl_df['county'] = ppl_df['cbg_geocode'].astype(str).str[:5].replace({'na': '0.0', 'nan': '0.0'}).astype(float) - ppl_df['tract'] = ppl_df['cbg_geocode'].astype(str).str[:11].replace({'na': '0.0', 'nan': '0.0'}).astype(float) - ppl_df['cbg_geocode'] = ppl_df['cbg_geocode'].astype(str).str[:12].replace({'na': '0.0', 'nan': '0.0'}).astype(float) - ppl_df.loc[(ppl_df['age'] >= 0) & (~ppl_df['age'].isnull()), 'agegroup'] = 0.0 - ppl_df.loc[(ppl_df['age'] >= 10) & (~ppl_df['age'].isnull()), 'agegroup'] = 1.0 - ppl_df.loc[(ppl_df['age'] >= 20) & (~ppl_df['age'].isnull()), 'agegroup'] = 2.0 - ppl_df.loc[(ppl_df['age'] >= 30) & (~ppl_df['age'].isnull()), 'agegroup'] = 3.0 - ppl_df.loc[(ppl_df['age'] >= 40) & (~ppl_df['age'].isnull()), 'agegroup'] = 4.0 - ppl_df.loc[(ppl_df['age'] >= 50) & (~ppl_df['age'].isnull()), 'agegroup'] = 5.0 - ppl_df.loc[(ppl_df['age'] >= 60) & (~ppl_df['age'].isnull()), 'agegroup'] = 6.0 - ppl_df.loc[(ppl_df['age'] >= 70) & (~ppl_df['age'].isnull()), 'agegroup'] = 7.0 - ppl_df.loc[(ppl_df['age'] >= 80) & (~ppl_df['age'].isnull()), 'agegroup'] = 8.0 - ppl_df.loc[(ppl_df['age'] >= 90) & (~ppl_df['age'].isnull()), 'agegroup'] = 9.0 - - hh = pd.read_csv(f'{path}/pop_export/hh.csv') - hh['household'] = hh.index + 1 - hh.drop(columns=['sample_index'], inplace=True) - ppl_df = ppl_df.merge(hh, on=['cbg_id', 'hh_id'], how='left') - ppl_df.loc[ppl_df['household'].isnull(), 'household'] = 0 - race_trait_cols = [ - 'race_white_alone', 'race_black_alone', 'race_amerindian_or_alaskan', - 'race_asian_alone', 'race_pacific_alone', 'race_other_alone', - 'race_two_or_more', 'hispanic', - ] - ppl_df = ppl_df[['uid', 'p_id', 'hh_id', 'cbg_id', 'sample_index', 'state', 'county', 'tract', 'cbg_geocode', - 'household', 'age', 'agegroup', 'female', *race_trait_cols, 'working', 'commuter', - 'commuter_income_category', 'commuter_workplace_category', 'sch_grade', 'sch_code']] - ppl_df.to_csv(f'{path}/pop_export/people_all.csv', index=False) - - age = ss.FloatArr('age', default=ss.BaseArr(ppl_df['age'].values)) - agegroup = ss.FloatArr('agegroup', default=ss.BaseArr(ppl_df['agegroup'].values)) - female = ss.FloatArr('female', default=ss.BaseArr(ppl_df['female'].values)) - race_trait_states = [ - ss.FloatArr(col, default=ss.BaseArr(ppl_df[col].values)) for col in race_trait_cols - ] - state = ss.IntArr('state', default=ss.BaseArr(ppl_df['state'].values)) - county = ss.IntArr('county', default=ss.BaseArr(ppl_df['county'].values)) - tract = ss.IntArr('tract', default=ss.BaseArr(ppl_df['tract'].values)) - cbg_geocode = ss.IntArr('cbg_geocode', default=ss.BaseArr(ppl_df['cbg_geocode'].values)) - household = ss.IntArr('household', default=ss.BaseArr(ppl_df['household'].values)) - commuter = ss.FloatArr('commuter', default=ss.BaseArr(ppl_df['commuter'].values)) - commuter_income_category = ss.FloatArr( - 'commuter_income_category', default=ss.BaseArr(ppl_df['commuter_income_category'].values) - ) - commuter_workplace_category = ss.FloatArr( - 'commuter_workplace_category', default=ss.BaseArr(ppl_df['commuter_workplace_category'].values) - ) - sch_code = ss.IntArr('sch_code', default=ss.BaseArr(ppl_df['sch_code'].values)) - - ppl = ss.People(n_agents=len(ppl_df), extra_states=[ - agegroup, *race_trait_states, state, county, tract, cbg_geocode, - household, commuter, commuter_income_category, commuter_workplace_category, sch_code - ]) - - ppl.states.append(age, overwrite=True) - setattr(ppl, age.name, age) - age.link_people(ppl) - - ppl.states.append(female, overwrite=True) - setattr(ppl, female.name, female) - female.link_people(ppl) - - sim = ss.Sim(people=ppl).init() - _ = sim # keep side-effect parity with previous implementation - os.makedirs(f'{path}/pop_export/starsim', exist_ok=True) - ss.save(f'{path}/pop_export/starsim/ppl.pkl', ppl) - print("Starsim People object created and saved successfully") - return ppl - - @staticmethod - def GPNetwork(name, edge_weight=1.0, csv_path=None, network_df=None, p1_col='p1', p2_col='p2', beta_col=None): - return _ForStarsimGPNetwork( - name=name, - edge_weight=edge_weight, - csv_path=csv_path, - network_df=network_df, - p1_col=p1_col, - p2_col=p2_col, - beta_col=beta_col, - ) - - @staticmethod - def SubgroupTracking(subgroup, outcome, name=None, state_id=None, *args, **kwargs): - return _ForStarsimSubgroupTracking(subgroup, outcome, name=name, state_id=state_id, *args, **kwargs) - -def main(): - """Main function for command-line usage.""" - runner = ForStarsim() - return runner \ No newline at end of file diff --git a/src/geopops/households.py b/src/geopops/households.py index 963ae39..1ff82f6 100644 --- a/src/geopops/households.py +++ b/src/geopops/households.py @@ -5,8 +5,8 @@ import numpy as np import pandas as pd import os -from .utils import (PersonData, Household, GQres, Indexer, tryJSON, - thresh, ranges, first_true, lrRound) +from .utils import (PersonData, Household, GQres, Indexer, TraitSchema, + tryJSON, thresh, ranges) def read_counties(data_dir): @@ -18,7 +18,7 @@ def read_counties(data_dir): def read_hh_serials(data_dir): df = pd.read_csv(os.path.join(data_dir, 'processed', 'hh_samples.csv'), usecols=['SERIALNO'], dtype={'SERIALNO': str}) - return dict(zip(df['SERIALNO'], range(1, len(df) + 1))) + return dict(zip(df['SERIALNO'], range(1, len(df) + 1), strict=False)) def read_psamp_df(data_dir, ind_codes, additional_traits): @@ -46,9 +46,13 @@ def _row_gq_employment(n, jobtype, ind_codes, row): return [int(round(row.get(prefix + k, 0))) for k in ind_codes] -def generate_group_quarters(config, cbgs, cbg_indexer, ind_codes, data_dir, rng): +def generate_group_quarters(config, cbgs, cbg_indexer, ind_codes, data_dir, rng, schema=None): min_gq_residents = config.get('min_gq_residents', 20) - add_trait_cols = config.get('additional_traits', []) + if schema is None: + schema = TraitSchema(config.get('additional_traits', [])) + # Group-quarters residents come from ACS aggregates rather than PUMS person + # records, so none of the PUMS-derived traits are known for them. + gq_trait_values = (None,) * len(schema) gq_cols = ['Geo', 'group quarters:', 'group quarters:under 18', 'group quarters:18 to 64', 'group quarters:65 and over', 'p_u18_inst', 'p_18_64_inst', 'p_65o_inst', @@ -123,7 +127,7 @@ def generate_group_quarters(config, cbgs, cbg_indexer, ind_codes, data_dir, rng) working=has_job, commuter=is_commuter, com_cat=emp_cat_final, com_inc=inc_cat, sch_grade=None, - **{trait: None for trait in add_trait_cols} + schema=schema, trait_values=gq_trait_values, ) summary_rows = [] @@ -158,6 +162,51 @@ def _resolve_config(config, data_dir): return tryJSON(os.path.join(data_dir, 'config.json')) +def _person_columns(p_samps, ind_codes, additional_traits): + """Precompute the per-person columns needed to build PersonData, as arrays. + + Everything here used to be done row-by-row with ``.iloc``/``.apply(axis=1)`` + inside the person loop, which costs a fresh pandas Series per person. Doing it + once, column-wise, is several hundred times faster. + """ + ind_cols = ['ind_' + k for k in ind_codes] + ind = p_samps[ind_cols].to_numpy(dtype=bool) if ind_cols else np.zeros((len(p_samps), 0), bool) + has_ind = ind.any(axis=1) if ind.shape[1] else np.zeros(len(p_samps), bool) + # first_true(row) == argmax for a boolean row, but only where some value is True + first_ind = ind.argmax(axis=1) if ind.shape[1] else np.zeros(len(p_samps), int) + + commuter = p_samps['commuter'].fillna(False).to_numpy(dtype=bool) + com_cat = np.where(commuter & has_ind, first_ind + 1, -1) + + income = p_samps[['com_LODES_low', 'com_LODES_high']].to_numpy(dtype=bool) + has_income = income.any(axis=1) + com_inc = np.where(has_income, income.argmax(axis=1) + 1, -1) + + sch_grade = p_samps['sch_grade'].to_numpy(dtype=object) + sch_grade = np.where(pd.isna(sch_grade), None, sch_grade) + + # Traits become a tuple per person, positioned by a schema shared population-wide + schema = TraitSchema(additional_traits) + if additional_traits: + raw = p_samps[list(additional_traits)].to_numpy(dtype=object) + trait_vals = np.where(pd.isna(raw), None, raw.astype(bool, copy=False) + if raw.dtype != object else raw) + trait_rows = [tuple(None if v is None else bool(v) for v in row) for row in trait_vals] + else: + trait_rows = [()] * len(p_samps) + + return dict( + age=p_samps['AGEP'].to_numpy(dtype=np.int64), + working=p_samps['has_job'].fillna(False).to_numpy(dtype=bool), + commuter=commuter, + com_cat=com_cat, + com_inc=com_inc, + sch_grade=sch_grade, + schema=schema, + trait_rows=trait_rows, + ) + + def generate_people(co_results, data_dir, config=None, random_seed=None): rng = np.random.default_rng(random_seed) config = _resolve_config(config, data_dir) @@ -170,17 +219,10 @@ def generate_people(co_results, data_dir, config=None, random_seed=None): p_samps = read_psamp_df(data_dir, ind_codes, additional_traits) p_idx = people_by_serial(p_samps) - - ind_colnames = ['ind_' + k for k in ind_codes] - p_samps['ind_code'] = p_samps[ind_colnames].apply( - lambda row: first_true(row.values), axis=1) - p_samps['com_cat'] = p_samps.apply( - lambda row: (row['ind_code'] + 1) if (row['commuter'] and row['ind_code'] is not None) else None, axis=1) - - income_cols = ['com_LODES_low', 'com_LODES_high'] - p_samps['income_code'] = p_samps[income_cols].apply( - lambda row: first_true(row.values), axis=1) - p_samps['com_inc'] = p_samps['income_code'].apply(lambda x: (x + 1) if x is not None else None) + cols = _person_columns(p_samps, ind_codes, additional_traits) + age, working, commuter = cols['age'], cols['working'], cols['commuter'] + com_cat, com_inc, sch_grade = cols['com_cat'], cols['com_inc'], cols['sch_grade'] + schema, trait_rows = cols['schema'], cols['trait_rows'] cbgs = {} cbg_indexer = Indexer() @@ -199,23 +241,19 @@ def generate_people(co_results, data_dir, config=None, random_seed=None): hh_key = (hh_i, cbg_i) p_vec = p_idx.get(hh_serial, []) for p_i_0, r in enumerate(p_vec): - p_i = p_i_0 + 1 - row = p_samps.iloc[r - 1] - trait_kwargs = {} - for trait in additional_traits: - val = row.get(trait) - trait_kwargs[trait] = bool(val) if pd.notna(val) else None - sch_grade = row['sch_grade'] if pd.notna(row['sch_grade']) else None - people[(p_i, hh_i, cbg_i)] = PersonData( + j = r - 1 # p_idx is 1-based for Julia parity + grade = sch_grade[j] + people[(p_i_0 + 1, hh_i, cbg_i)] = PersonData( hh=hh_key, sample=r, - age=int(row['AGEP']), - working=bool(row['has_job']), - commuter=bool(row['commuter']), - com_cat=int(row['com_cat']) if row['com_cat'] is not None and pd.notna(row['com_cat']) else None, - com_inc=int(row['com_inc']) if row['com_inc'] is not None and pd.notna(row['com_inc']) else None, - sch_grade=str(sch_grade) if sch_grade is not None else None, - **trait_kwargs, + age=int(age[j]), + working=bool(working[j]), + commuter=bool(commuter[j]), + com_cat=int(com_cat[j]) if com_cat[j] > 0 else None, + com_inc=int(com_inc[j]) if com_inc[j] > 0 else None, + sch_grade=str(grade) if grade is not None else None, + schema=schema, + trait_values=trait_rows[j], ) households[hh_key] = Household( sample=hh_idx.get(hh_serial, 0), @@ -223,6 +261,6 @@ def generate_people(co_results, data_dir, config=None, random_seed=None): ) cbgs, gqs, gq_people, gq_summary = generate_group_quarters( - config, cbgs, cbg_indexer, ind_codes, data_dir, rng) + config, cbgs, cbg_indexer, ind_codes, data_dir, rng, schema=schema) people.update(gq_people) return cbgs, people, households, gqs, gq_summary diff --git a/src/geopops/ipfn.py b/src/geopops/ipfn.py index 55ae073..74b0aba 100644 --- a/src/geopops/ipfn.py +++ b/src/geopops/ipfn.py @@ -1,4 +1,15 @@ #!/usr/bin/env python +""" +Iterative proportional fitting (IPF / IPFN). + +Vendored from the `ipfn` package: https://github.com/Dirguis/ipfn +Copyright (c) Damien Forthomme and contributors. Licensed under the MIT License; +the full license text is reproduced in the NOTICE file at the root of this +repository. + +Kept close to upstream to ease future syncing. Prefer depending on the upstream +`ipfn` package rather than maintaining this copy. +""" from __future__ import print_function import numpy as np import pandas as pd diff --git a/src/geopops/julia.py b/src/geopops/julia.py deleted file mode 100644 index 2515ab9..0000000 --- a/src/geopops/julia.py +++ /dev/null @@ -1,123 +0,0 @@ -""" -Legacy Julia execution path for GeoPops. - -This module is kept for backward compatibility. The active implementation -has been migrated to pure Python modules and should be used via -`geopops.GeneratePop`. -""" - -import os -import subprocess -import json -from pathlib import Path -from shutil import which - -BASE_DIR = os.path.dirname(os.path.abspath(__file__)) - -def load_config(): - """Load config file from package directory.""" - config_path = os.path.join(BASE_DIR, "config.json") - if not os.path.exists(config_path): - raise FileNotFoundError(f"config.json file not found at {config_path}. Please create this file with the required configuration.") - with open(config_path, "r") as f: - return json.load(f) - -class RunJulia: - """Orchestrates Julia script execution using the existing functions in this module. - """ - - def __init__(self, output_dir=None, julia_env_path=None): - """Create a Julia runner. - - Args: - output_dir: Optional output directory. If None, uses output_dir from config.json. - julia_env_path: Optional Julia environment path. If None, uses julia_env_path from config.json. - """ - # Load config to get paths - config = load_config() - - # Julia scripts are always in the package directory - self.julia_scripts_dir = os.path.join(BASE_DIR, "julia") - - # Output directory from parameter or config - if output_dir is not None: - self.output_dir = output_dir - else: - self.output_dir = config.get("path", BASE_DIR) - - # Julia environment path from parameter or config - if julia_env_path is not None: - self.julia_env_path = julia_env_path - else: - self.julia_env_path = config.get("julia_env_path") - if not self.julia_env_path: - raise ValueError("julia_env_path not found in config.json") - - # Check if the environment exists - if not os.path.exists(self.julia_env_path): - raise RuntimeError(f"Julia environment not found at {self.julia_env_path}") - - # Use standard julia command - self.julia_cmd = ["julia"] - - print(f"Using Julia environment: {self.julia_env_path}") - print(f"Using output directory: {self.output_dir}") - - def CO(self): - """Run the CO.jl Julia script.""" - print("*** Running RunJulia.CO() ***") - subprocess.run([ - *self.julia_cmd, - f"--project={self.julia_env_path}", - os.path.join(self.julia_scripts_dir, "CO.jl") - ], check=True, text=True, cwd=self.output_dir) - - def SynthPop(self): - """Run the synthpop.jl Julia script.""" - print("\n*** Running RunJulia.SynthPop() ***") - try: - result = subprocess.run([ - *self.julia_cmd, - f"--project={self.julia_env_path}", - os.path.join(self.julia_scripts_dir, "synthpop.jl") - ], check=True, text=True, cwd=self.output_dir, capture_output=True) - print("synthpop.jl completed successfully") - return result - except subprocess.CalledProcessError as e: - print(f"synthpop.jl failed with exit code {e.returncode}") - if e.stdout: - print("STDOUT:", e.stdout) - if e.stderr: - print("STDERR:", e.stderr) - raise - - def Export(self): - """Run the export_synthpop.jl and the export_network.jl Julia scripts.""" - print("\n*** Running RunJulia.Export() ***") - subprocess.run([ - *self.julia_cmd, - f"--project={self.julia_env_path}", - os.path.join(self.julia_scripts_dir, "export_synthpop.jl") - ], check=True, text=True, cwd=self.output_dir) - subprocess.run([ - *self.julia_cmd, - f"--project={self.julia_env_path}", - os.path.join(self.julia_scripts_dir, "export_network.jl") - ], check=True, text=True, cwd=self.output_dir) - - def run_all(self): - """Run all Julia scripts in sequence.""" - print("============================================================") - print("Running RunJulia()") - print("============================================================") - self.CO() - self.SynthPop() - self.Export() - -def main(): - runner = RunJulia() - # runner.run_all() - -if __name__ == "__main__": - main() - diff --git a/src/geopops/julia/CO.jl b/src/geopops/julia/CO.jl deleted file mode 100644 index 5745276..0000000 --- a/src/geopops/julia/CO.jl +++ /dev/null @@ -1,333 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -using Distributed -include("fileutils.jl") - -## the definitions in this block are sent to all processors -@everywhere begin - - ## for sharing data between local processors without copying - using SharedArrays - - ## mutation function replaces one element in combs with an element from opts - ## avoid making an array copy, it only takes 1 assignment to revert - function mutate!(comb, opts) - cidx = rand(eachindex(comb)) - orig = comb[cidx] - comb[cidx] = rand(opts) ## mutate passed vector - return (cidx,orig) ## return info needed to revert - end - - function revert!(comb,(cidx,orig)) - comb[cidx] = orig - return nothing - end - - ## pick the rows given by idxs from the pop, return summary of columns - ## use view() to avoid making a copy - summarize(idxs::Vector{Int64}, pop::Matrix{Int64}) = sum(view(pop,idxs,:), dims=1) - ## use sum!() to save on array allocations - function summarize!(dest::Matrix{Int64}, idxs::Vector{Int64}, pop::Matrix{Int64}) - sum!(dest, view(pop,idxs,:)) - return nothing - end - - ## distance functions - ## note if coming from Numpy: Julia functions do not automatically broadcast over collections - ## --- must specify that intent using dot-syntax: .+ for operators, f.() for functions - ## (also, Julia docs said to avoid collections of abstract types, hence the polymorphic - ## type specification "Matrix{T} where T<:Real" instead of the simpler Matrix{Real}) - ## FTdist is 1/4 * Freeman-Tukey distance - ## -- multiply by 4 to get something with a chi-sq distribution with len(v)-1 d.f. - #FTdist(v1::Matrix{T},v2::Matrix{T}) where T<:Real = sum((sqrt.(v1) .- sqrt.(v2)) .^ 2) - ## zeros are common in target cells, and get penalized too harshly -- add 1 to relax this - FTdist(v1::Matrix{T},v2::Matrix{T}) where T<:Real = sum((sqrt.(v1 .+ one(T)) .- sqrt.(v2 .+ one(T))) .^ 2) - ## weighted - #FTdist(v1::Matrix{T},v2::Matrix{T},wts::Matrix{Float64}) where T<:Real = sum(wts .* ((sqrt.(v1 .+ one(T)) .- sqrt.(v2 .+ one(T))) .^ 2)) - ## or, good ol' euclidean-squared - #sqdist(v1::Matrix{T},v2::Matrix{T}) where T<:Real = sum((v1 .- v2) .^ 2) - - ## standard simulated annealing acceptance function - function accept(neg_dE::Float64, T::Float64) - if neg_dE >= 0 - return true - else - p = exp(neg_dE / T) - return rand() < p - end - end - - ## function for stepping down the temperature; adaptive would probably be better? - function sched(gen, temp, E0, E1, c) - return temp * c - end - - ## termination criteria - function term(gen, E0, temp, max_gens, crit_val, report=1000) - t = E0 < crit_val || gen > max_gens - if gen % report == 0 || t - println(E0, " ", gen, " ", temp) - end - return t - end - - ## optimize with simulated annealing - ## returns a vector of indices from samples - function anneal(glob_samp_ref::SharedMatrix{Int64}, mask::BitVector, targ::Matrix{Int64}, n::Int64, params::Dict{Symbol, R}) where R<:Real - ## make a local copy of needed samples (this isn't done that often, and should be the fastest approach) - samples = glob_samp_ref[mask, :] - ## indices to sample - ## note, these index the local subset, not the shared samples - idxs = axes(samples,1) - ## if no valid samples, return a bad score (this shouldn't happen) - if isempty(idxs) - return (Vector{Int64}(), 0, Inf, 0.0) - end - ## initialize with a random combination of samples - c0 = rand(idxs,n) - summary = summarize(c0, samples) - E0 = FTdist(summary, targ) - T = 0.5 * E0 ## set initial temperature based on starting distance, I guess? - gen = 0 - - done = false - while !done - gen += 1 - orig = mutate!(c0, idxs) - summarize!(summary, c0, samples) - E1 = FTdist(summary, targ) - if accept(E0-E1, T) - T = sched(gen, T, E0, E1, params[:cooldown]) - E0 = E1 ## only update the score if we accept, lol - else - revert!(c0, orig) - end - done = term(gen, E0, T, params[:maxgens], params[:critval], params[:report]) - end - - ## convert the answer back to global indices - ## I wasn't really going to forget to do this, was I? - res = findall(mask)[c0] - - ## force garbage collection, distributed GC isn't always smart - GC.gc() - - return (res, gen, E0, T) - end - - ## closes over params and a reference to shared matrix of samples - ## returns a fn that uses anneal with given params, and can see the shared samples - function annealer(shared_samples::SharedMatrix{Int64}, params::Dict{Symbol, R}) where R<:Real - ## call this function to do the work: - ## mask = which samples from shared_samples to use for a given target - function f(mask::BitVector, targ::Matrix{Int64}, n::Int64) - return anneal(shared_samples, mask, targ, n, params) - end - return f - end - -## end @everywhere -end - -## targets are census block group (cbg) summary statistics -function read_targets(targ_idxs=[]) - acs_targets = read_df("processed/acs_targets.csv"; types=Dict("Geo"=>String15)) - if isempty(targ_idxs) - targ_idxs = axes(acs_targets,1) - end - ## pull target stats from the dataframe and convert to matrix or array for faster math - targs = [Matrix{Int64}(acs_targets[[targ_id],2:end]) for targ_id in targ_idxs] - geos = [acs_targets[targ_id,1] for targ_id in targ_idxs] - return (targs, geos, names(acs_targets)[2:end]) -end - -## number of households in each cbg -function read_hh_counts() - hhcdf = read_df("processed/hh_counts.csv"; types=Dict("Geo"=>String15)) - ## convert to dicts for easier lookup - ## note the .=> syntax; this broadcasts associations between two arrays, similar to dict(zip()) in python - return Dict(hhcdf[:,1] .=> hhcdf[:,2]) -end - -## returns a sharedmatrix of samples -- can pass it to local parallel processes without copying data -function read_samples() - puma_samples_all = read_df("processed/census_samples.csv"; types=Dict("SERIALNO"=>String15)) - hh_ids = puma_samples_all[:,1] ## for looking up the household serial numbers later - shared_samples = convert(SharedMatrix, Matrix{Int64}(puma_samples_all[:,2:end])) - return (shared_samples, hh_ids) -end - -## geographic data for target cbg's -- will use these to determine which samples to use -function read_targ_geo() - cbg_geo_cols = Dict("Geo"=>String15,"st_puma"=>String7,"cbsa"=>String7,"county"=>String7,"R"=>Float64,"U"=>Float64) - cbg_geo = read_df("processed/cbg_geo.csv"; select=collect(keys(cbg_geo_cols)), types=cbg_geo_cols) - cbg_puma = Dict(cbg_geo[:,"Geo"] .=> cbg_geo[:,"st_puma"]) - cbg_county = Dict(cbg_geo[:,"Geo"] .=> cbg_geo[:,"county"]) - cbg_cbsa = Dict(cbg_geo[:,"Geo"] .=> cbg_geo[:,"cbsa"]) - cbg_urban = Dict(cbg_geo[:,"Geo"] .=> cbg_geo[:,"U"]) - return (cbg_puma, cbg_county, cbg_cbsa, cbg_urban) -end - -## for finding samples with similar urbanization to target urbanization x -function urbanization_lookup(df::DataFrame, x::Float64) - if x > 0.999 - ## if df has a missing cell, comparison will return "missing" - ## -- coalesce() makes it return "false" instead - return (coalesce.(df.U .> 0.999, false)) - elseif x < 0.334 - return (coalesce.(df.U .< 0.334, false)) - else - a = x - 0.1; b = min(x+0.1, 0.999) - return (coalesce.(a .< df.U .< b, false)) - end -end - -## looks up which samples in df to use, based on column name k and associated geo data of targets -## don't usually need to specify function return types but in this case, target_geo_codes might be empty and -## this way it still returns a vector{BitVector} -function sample_lookup(df::DataFrame, k::Symbol, target_geo_codes)::Vector{BitVector} - if k == :U - return [urbanization_lookup(df,x) for x in target_geo_codes] - else - ## note .== syntax for element-wise comparison - return [coalesce.(df[!,k] .== x, false) for x in target_geo_codes] - end -end - -## creates annealing function and executes it on available processors -## returns a vector of whatever anneal() returns -function optimize(samples::SharedMatrix{Int64}, samp_masks::Vector{BitVector}, targs::Vector{Matrix{Int64}}, n_hhs::Vector{Int64}, params) - a_fn = annealer(samples, params) - ## use pmap for this, not @distributed for - return pmap(a_fn, samp_masks, targs, n_hhs) -end - -## like optimize, but overwrites previous results in vector "x" at the indices given by "rerun" -function reoptimize!(x, rerun::Vector{Int64}, samples::SharedMatrix{Int64}, - samp_masks::Vector{BitVector}, targs::Vector{Matrix{Int64}}, n_hhs::Vector{Int64}, params_r) - - ## make sure we found enough suitable samples for each target - enough_samps = sum.(samp_masks) .> (n_hhs[rerun] .// 2) - rerun = rerun[enough_samps] - ## make sure there's actually something to rerun - if lastindex(rerun) > 0 - a_fn = annealer(samples, params_r) - x_r = pmap(a_fn, samp_masks, targs[rerun], n_hhs[rerun]) - ## only overwrite the ones whose score improved - improved = [a[3] for a in x_r] .< [a[3] for a in x[rerun]] - ## note .= syntax for broadcasting assignment to an array - x[rerun[improved]] .= x_r[improved] - end - return nothing -end - -## performs several optimization runs on the target cbg's in each county -## writes each county's results to a separate file -function process_counties(counties=[]) - - samples, hh_ids = read_samples() - cbg_puma, cbg_county, cbg_cbsa, cbg_urban = read_targ_geo() - samp_geo_cols = Dict("SERIALNO"=>String15,"st_puma"=>String7,"cbsa"=>String7,"county"=>String7,"R"=>Float64,"U"=>Float64) - samp_geo = read_df("processed/samp_geo.csv"; select=collect(keys(samp_geo_cols)), types=samp_geo_cols) - ## samp geo indices match those in census_samples - ## all(samp_geo.SERIALNO .== hh_ids) - targs_all, geos_all, targ_colnames = read_targets() - hh_counts = read_hh_counts() - ## number of households in each target - n_hhs_all = [hh_counts[g] for g in geos_all] - ## each target's county - county = [g[1:5] for g in geos_all] - - ## weight households more than individuals? (not currently implemented) - #targ_weights = [ones(1,43).*2.0 ones(1,length(targ_colnames)-43)] - dConfig = tryJSON("config.json") - c_val::Float64 = get(dConfig, "CO_crit_val", 10.0) - CO_cooldown::Float64 = get(dConfig, "CO_cooldown", 0.99) - CO_cooldown_slow = 0.5 + 0.5 * CO_cooldown - CO_maxgens::Int = get(dConfig, "CO_maxgens", 200000) - CO_report = CO_maxgens + 1 - - params = Dict(:maxgens => CO_maxgens, :critval => c_val, :cooldown => CO_cooldown, :report => CO_report) - - ## run all counties by default - if isempty(counties) - counties = unique(county) - end - - mkpath("jlse/CO") - for c in counties - targs = targs_all[county .== c] - geos = geos_all[county .== c] - n_hhs = n_hhs_all[county .== c] - - println("\n puma \n") - ## look up samples for each target based on target's puma code - ## puma is the most local sample level, has the fewest samples available - samp_masks = sample_lookup(samp_geo, :st_puma, [cbg_puma[g] for g in geos]) - x = optimize(samples, samp_masks, targs, n_hhs, params) - scores = [a[3] for a in x] - n_bad = sum(s > c_val for s in scores) - println(" county ", c, ": ", length(x), " targets; ", n_bad, " above threshold (will rerun); E0 min/mean/max: ", minimum(scores), " / ", sum(scores)/length(scores), " / ", maximum(scores)) - - println("\n county \n") - ## rerun poor matches using county-level - rerun = findall([a[3]>c_val for a in x]) - println(" county pass: ", length(rerun), " targets to rerun") - samp_masks = sample_lookup(samp_geo, :county, [cbg_county[g] for g in geos[rerun]]) - params_r = Dict(:maxgens => CO_maxgens, :critval => c_val, :cooldown => CO_cooldown, :report => CO_report) - reoptimize!(x, rerun, samples, samp_masks, targs, n_hhs, params_r) - scores = [a[3] for a in x] - n_bad = sum(s > c_val for s in scores) - println(" after county: ", n_bad, " still above threshold; E0 min/mean/max: ", minimum(scores), " / ", sum(scores)/length(scores), " / ", maximum(scores)) - - println("\n cbsa \n") - ## cbsa-level - rerun = findall([a[3]>c_val for a in x]) - println(" cbsa pass: ", length(rerun), " targets to rerun") - samp_masks = sample_lookup(samp_geo, :cbsa, [cbg_cbsa[g] for g in geos[rerun]]) - params_r = Dict(:maxgens => CO_maxgens, :critval => c_val, :cooldown => CO_cooldown, :report => CO_report) - reoptimize!(x, rerun, samples, samp_masks, targs, n_hhs, params_r) - scores = [a[3] for a in x] - n_bad = sum(s > c_val for s in scores) - println(" after cbsa: ", n_bad, " still above threshold; E0 min/mean/max: ", minimum(scores), " / ", sum(scores)/length(scores), " / ", maximum(scores)) - - println("\n urb \n") - ## urbanization level, has the most samples; more likely to match but longer to search - ## (also the least associated with the target's local geography) - rerun = findall([a[3]>c_val for a in x]) - println(" urb pass: ", length(rerun), " targets to rerun") - samp_masks = sample_lookup(samp_geo, :U, [cbg_urban[g] for g in geos[rerun]]) - params_r = Dict(:maxgens => CO_maxgens, :critval => c_val, :cooldown => CO_cooldown_slow, :report => CO_report) - reoptimize!(x, rerun, samples, samp_masks, targs, n_hhs, params_r) - scores = [a[3] for a in x] - n_bad = sum(s > c_val for s in scores) - println(" after urb: ", n_bad, " still above threshold; E0 min/mean/max: ", minimum(scores), " / ", sum(scores)/length(scores), " / ", maximum(scores)) - - ## store results as dict keyed by cbg code; look up actual hh id's of sample indices - scores = Dict(geos .=> [a[3] for a in x]) - households = Dict(geos .=> [hh_ids[a[1]] for a in x]) - - serialize(abspath("jlse/CO/hh"*c*".jlse"),households) - serialize(abspath("jlse/CO/hh"*c*"_scores.jlse"),scores) - - ## force garbage collection, distributed GC isn't always smart - GC.gc() - end - - return nothing -end - - -process_counties() - - diff --git a/src/geopops/julia/export_network.jl b/src/geopops/julia/export_network.jl deleted file mode 100644 index 1955d0e..0000000 --- a/src/geopops/julia/export_network.jl +++ /dev/null @@ -1,87 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -using SparseArrays -include("fileutils.jl") - -## Matrix Market native exchange format for sparse matrices -## https://math.nist.gov/MatrixMarket/formats.html#MMformat -using MatrixMarket - -mkpath("pop_export") - -println("household adjacency matrix") - -m = dser_path("jlse/adj_mat_hh.jlse") -MatrixMarket.mmwrite(abspath("pop_export/adj_upper_triang_hh.mtx"), m) - -d = nothing -GC.gc() - -println("non-household adjacency matrix") - -m = dser_path("jlse/adj_mat_non_hh.jlse") -MatrixMarket.mmwrite(abspath("pop_export/adj_upper_triang_non_hh.mtx"), m) - -d = nothing -GC.gc() - -println("workplace only adjacency matrix") - -m = dser_path("jlse/adj_mat_wp.jlse") -MatrixMarket.mmwrite(abspath("pop_export/adj_upper_triang_wp.mtx"), m) - -d = nothing -GC.gc() - -println("school only adjacency matrix") - -m = dser_path("jlse/adj_mat_sch.jlse") -MatrixMarket.mmwrite(abspath("pop_export/adj_upper_triang_sch.mtx"), m) - -d = nothing -GC.gc() - -println("groups quarters only adjacency matrix") - -m = dser_path("jlse/adj_mat_gq.jlse") -MatrixMarket.mmwrite(abspath("pop_export/adj_upper_triang_gq.mtx"), m) - -d = nothing -GC.gc() - -println("index keys") - -d = dser_path("jlse/adj_mat_keys.jlse") -df = DataFrame([(index_one=i, index_zero=i-1, p_id=Int(v[1]), hh_id=Int(v[2]), cbg_id=Int(v[3])) for (i,v) in enumerate(d)]) -write_df("pop_export/adj_mat_keys.csv",df) - -## indices of people commuting fron outside the synth pop area -## these have no household connections -d = dser_path("jlse/adj_dummy_keys.jlse") -df = DataFrame(sort([ - (index_one=Int64(k), index_zero=Int64(k)-1, p_id=Int(v[1]), hh_id=Int(v[2]), cbg_id=Int(v[3])) - for (k,v) in d - ], by=x->x.index_one)) -write_df("pop_export/adj_dummy_keys.csv",df) - -## indices of people working outside the synth pop area -## these have no workplace connections -d = dser_path("jlse/adj_out_workers.jlse") -df = DataFrame(sort([ - (index_one=Int64(k), index_zero=Int64(k)-1, p_id=Int(v[1]), hh_id=Int(v[2]), cbg_id=Int(v[3])) - for (k,v) in d - ], by=x->x.index_one)) -write_df("pop_export/adj_out_workers.csv",df) - -println("done") - diff --git a/src/geopops/julia/export_synthpop.jl b/src/geopops/julia/export_synthpop.jl deleted file mode 100644 index 29a981e..0000000 --- a/src/geopops/julia/export_synthpop.jl +++ /dev/null @@ -1,106 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -include("utils.jl") -include("fileutils.jl") - -mkpath("pop_export") - -println("households and people") - -d = dser_path("jlse/cbg_idxs.jlse") -df = DataFrame(sort([ - (cbg_id=Int64(k), cbg_geocode=String(v)) - for (k,v) in d - ], by=x->x.cbg_id)) -write_df("pop_export/cbg_idxs.csv",df) - -d = dser_path("jlse/hh.jlse") -df = DataFrame(sort([ - (hh_id=Int(k[1]), cbg_id=Int(k[2]), sample_index=Int(v.sample), n_people=length(v.people)) - for (k,v) in d - ], by=x->(x.cbg_id,x.hh_id))) -write_df("pop_export/hh.csv",df) - -d = dser_path("jlse/people.jlse") -df = DataFrame(sort([ - (p_id=Int(k[1]), hh_id=Int(k[2]), cbg_id=Int(k[3]), sample_index=mcon(Int,v.sample), - age=mcon(Int,v.age), female=mcon(Int,v.female), working=mcon(Int,v.working), commuter=mcon(Int,v.commuter), - commuter_income_category=mcon(Int,v.com_inc) , commuter_workplace_category=mcon(Int,v.com_cat), - race_black_alone=mcon(Int,v.race_black_alone), white_non_hispanic=mcon(Int,v.white_non_hispanic), hispanic=mcon(Int,v.hispanic), - sch_grade=mcon(String,v.sch_grade)) - for (k,v) in d - ], by=x->(x.cbg_id,x.hh_id,x.p_id))) -write_df("pop_export/people.csv",df) - -d = nothing -GC.gc() - -println("schools") - -d = dser_path("jlse/sch_students.jlse") -df = DataFrame(sort([ - (sch_code=String(k), p_id=Int(v[1]), hh_id=Int(v[2]), cbg_id=Int(v[3])) - for (k,v) in dflat(d) - ], by=x->x.sch_code)) -write_df("pop_export/sch_students.csv",df) - -d = dser_path("jlse/sch_workers.jlse") -df = DataFrame(sort([ - (sch_code=String(k), p_id=Int(v[1]), hh_id=Int(v[2]), cbg_id=Int(v[3])) - for (k,v) in dflat(d) - ], by=x->x.sch_code)) -write_df("pop_export/sch_workers.csv",df) - -println("group quarters") - -d = dser_path("jlse/gqs.jlse") -df = DataFrame(sort([ - (gq_id=Int(k[1]), cbg_id=Int(k[2]), gq_type=String(v.type), n_residents=length(v.residents)) - for (k,v) in d - ], by=x->(x.cbg_id,x.gq_id))) -write_df("pop_export/gqs.csv",df) - -d = Dict(k=>v.residents for (k,v) in d) -df = DataFrame(sort([ - (gq_id=Int(k[1]), cbg_id=Int(k[2]), p_id=Int(v[1]), hh_id=Int(v[2])) - for (k,v) in dflat(d) - ], by=x->(x.cbg_id,x.gq_id))) -write_df("pop_export/gq_residents.csv",df) - -println("workplaces") - -d = dser_path("jlse/gq_workers.jlse") -df = DataFrame(sort([ - (gq_id=Int(k[1]), gq_cbg_id=Int(k[2]), p_id=Int(v[1]), p_hh_id=Int(v[2]), p_cbg_id=Int(v[3])) - for (k,v) in dflat(d) - ], by=x->(x.gq_cbg_id,x.gq_id))) -write_df("pop_export/gq_workers.csv",df) - -d = dser_path("jlse/company_workers.jlse") -df = DataFrame(sort([ - (employer_geo_code=String(k[3]), employer_type=Int(k[2]), employer_num=Int(k[1]), p_id=Int(v[1]), p_hh_id=Int(v[2]), p_cbg_id=Int(v[3])) - for (k,v) in dflat(d) - ], by=x->(x.employer_geo_code,x.employer_type,x.employer_num))) -write_df("pop_export/company_workers.csv",df) - -d = dser_path("jlse/outside_workers.jlse") -df = DataFrame(sort([ - (p_id=Int(v[1]), p_hh_id=Int(v[2]), p_cbg_id=Int(v[3])) - for (k,v) in dflat(d) - ], by=x->(x.p_cbg_id,x.p_hh_id,x.p_id))) -write_df("pop_export/outside_workers.csv",df) - -d = nothing -GC.gc() - -println("done") diff --git a/src/geopops/julia/fileutils.jl b/src/geopops/julia/fileutils.jl deleted file mode 100644 index b9daaf3..0000000 --- a/src/geopops/julia/fileutils.jl +++ /dev/null @@ -1,51 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -using CSV -using DataFrames -using SparseArrays -using Serialization -using JSON - -function tryJSON(f::AbstractString)::Dict{String,Any} - try - return JSON.parsefile(abspath(f)) - catch e - return Dict{String,Any}() - end -end - -function dser_path(f::AbstractString) - #println("reading ", f) - return deserialize(abspath(f)) -end - -function ser_path(f::AbstractString,obj::Any) - #println("writing ", f) - serialize(abspath(f), obj) - return nothing -end - -function read_df(f::AbstractString; kwargs...) - return CSV.read(abspath(f), DataFrame; kwargs...) -end - -function write_df(f::AbstractString, df; kwargs...) - CSV.write(abspath(f), df; kwargs...) -end - -## creates a sparse dataframe from a sparse matrix -spDataFrame(m::SparseMatrixCSC, labels::Union{Vector,Symbol}=:auto) = DataFrame(collect(findnz(m)), labels) - -function write_df(f::AbstractString, m::SparseMatrixCSC, labels::Union{Vector,Symbol}=:auto; kwargs...) - CSV.write(abspath(f), spDataFrame(m, labels); kwargs...) -end diff --git a/src/geopops/julia/households.jl b/src/geopops/julia/households.jl deleted file mode 100644 index de32a78..0000000 --- a/src/geopops/julia/households.jl +++ /dev/null @@ -1,318 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -#using Statistics - -include("utils.jl") -include("fileutils.jl") - -function read_counties() - cbg_geo_cols = Dict("Geo"=>String15,"st_puma"=>String7,"cbsa"=>String7,"county"=>String7,"R"=>Float64,"U"=>Float64) - cbg_geo = read_df("processed/cbg_geo.csv"; select=collect(keys(cbg_geo_cols)), types=cbg_geo_cols) - return unique(cbg_geo.county) -end - -## for looking up row # by hh serial -function read_hh_serials() - hh_samps = read_df("processed/hh_samples.csv"; types=Dict("SERIALNO"=>String15)) - return Dict(hh_samps.SERIALNO .=> eachindex(hh_samps.SERIALNO)) -end - -## read person samples into a dataframe -function read_psamp_df(ind_codes::Vector{<:AbstractString}, additional_traits::Vector{<:AbstractString}) - nonbool_cols = ["SERIALNO","AGEP","sch_grade"] - bool_cols = ["commuter","has_job","com_LODES_low","com_LODES_high"] - ind_cols = ["ind_"*k for k in ind_codes] - typedict = Dict(["SERIALNO"=>String15; - [x=>Bool for x in bool_cols]; - ## additional traits are all bools for now - [x=>Bool for x in additional_traits]; - [x=>Bool for x in ind_cols] - ]) - df = read_df("processed/p_samples.csv"; select=[nonbool_cols; bool_cols; additional_traits; ind_cols], types=typedict) - return df[:, [nonbool_cols; bool_cols; additional_traits; ind_cols]] -end - -## for looking up all person samps by hh serial -## (this way is much faster than a search loop through a million serials) -function people_by_serial(p_samps::DataFrame) - ## make a column with the row #'s, group by serial - p_samps_by_hh = groupby(insertcols(p_samps, 1, :idx => 1:nrow(p_samps)), :SERIALNO) - ## Ref() is a trick of the combine function (found this in the dataframes.jl docs) - df = combine(p_samps_by_hh, :idx => Ref => :idxs) - return Dict(df.SERIALNO .=> collect.(df.idxs)) -end - -## returns a vector of number employed by industry index from dataframe row r -function row_gq_employment(n::Integer, jobtype::Symbol, ind_codes::Vector{<:AbstractString}, r::DataFrameRow)::Vector{Int} - if jobtype == :none || n < 1 - return [0 for x in ind_codes] - else - pref = Dict(:civ=>"civ_ind_",:mil=>"mil_ind_") - return int.([r[pref[jobtype]*k] for k in ind_codes]) - end -end - -## for evaluating results -function cbg_summary(cbg_code, cbg_dict, hh_samps, hh_idx) - hhvec = cbg_dict[cbg_code] - cbg_samps = empty(hh_samps) - for x in hhvec - push!(cbg_samps, hh_samps[hh_idx[x],:]) - end - return combine(cbg_samps, Not(:SERIALNO) .=> x->sum(skipmissing(x)); renamecols=false)[1,:] -end - - -function generate_group_quarters(dConfig::Dict, cbgs::Dict{<:AbstractString, I}, - cbg_indexer::Indexer{I}, ind_codes::Vector{<:AbstractString}) where {I<:Integer} - - min_gq_residents::Int = get(dConfig, "min_gq_residents", 20) - add_trait_cols::Vector{String} = get(dConfig, "additional_traits", String[]) - - ## loop through gq, create individuals - ## -- assume only 18-64 noninst have jobs - ## -- assume children in gq don't attend public schools - ## will need to add workers to gq, and generate a within-org network - df_gq_cols = ["Geo", "group quarters:", "group quarters:under 18", "group quarters:18 to 64", - "group quarters:65 and over", "p_u18_inst", "p_18_64_inst", "p_65o_inst", "p_18_64_noninst_civil", - "p_18_64_noninst_mil", "commuter_p|ninst1864civ", "work_from_home_p|ninst1864civ", - "com_LODES_low_p|ninst1864civ", "com_LODES_high_p|ninst1864civ", "commuter_p|milGQ", - "work_from_home_p|milGQ", "com_LODES_low_p|milGQ", "com_LODES_high_p|milGQ"] - - df_gq = read_df("processed/group_quarters.csv"; select=df_gq_cols, types=Dict("Geo"=>String15)) - - ## assume 5 gq's max per cbg: - ## -- u18 inst, 18-64 inst, 18-64 noninst civilian, 18-64 noninst military, 65+ inst - ## -- ignore noninst for u18 and 65+ - ## -- ignore gq with less than min_gq_size people - gq_types = [:instu18, :inst1864, :ninst1864civ, :milGQ, :inst65o] - assumed_ages = [15, 30, 30, 30, 75] ## ages unknown (could generate these if needed) - job_types = [:none, :none, :civ, :mil, :none] ## assume only 18-64 noninst have jobs - transform!(df_gq, ["group quarters:under 18","p_u18_inst"] => ((a,b)->thresh.(int.(a.*b),min_gq_residents)) => "pop_instu18") - transform!(df_gq, ["group quarters:18 to 64","p_18_64_inst"] => ((a,b)->thresh.(int.(a.*b),min_gq_residents)) => "pop_inst1864") - transform!(df_gq, ["group quarters:18 to 64","p_18_64_noninst_civil"] => ((a,b)->thresh.(int.(a.*b),min_gq_residents)) => "pop_ninst1864civ") - transform!(df_gq, ["group quarters:18 to 64","p_18_64_noninst_mil"] => ((a,b)->thresh.(int.(a.*b),min_gq_residents)) => "pop_milGQ") - transform!(df_gq, ["group quarters:65 and over","p_65o_inst"] => ((a,b)->thresh.(int.(a.*b),min_gq_residents)) => "pop_inst65o") - - ## read gq worker counts by industry, which were derived from census data - df_civil_emp = rename(x->replace(x,"C24030:"=>"civ_ind_","C24010:"=>"civ_occ_"), - read_df("processed/gq_civilian_workers.csv"; types=Dict("Geo"=>String15))) - df_mil_emp = rename(x->replace(x,"C24030:"=>"mil_ind_","C24010:"=>"mil_occ_"), - read_df("processed/gq_military_workers.csv"; types=Dict("Geo"=>String15))) - - ## not all of those are commuters -- assume p(work from home) is the same across industries - ## randomly assign some people as wfh; do same for income category - transform!(df_gq, ["commuter_p|ninst1864civ","work_from_home_p|ninst1864civ"] => ((a,b)->a./(a.+b)) => "commuter_p|civ_worker") - transform!(df_gq, ["commuter_p|milGQ","work_from_home_p|milGQ"] => ((a,b)->a./(a.+b)) => "commuter_p|mil_worker") - transform!(df_gq, ["com_LODES_high_p|ninst1864civ","com_LODES_low_p|ninst1864civ"] => ((a,b)->a./(a.+b)) => "LODES_high|civ_commuter") - transform!(df_gq, ["com_LODES_high_p|milGQ","com_LODES_low_p|milGQ"] => ((a,b)->a./(a.+b)) => "LODES_high|mil_commuter") - df_gq = innerjoin(df_gq, df_civil_emp, df_mil_emp; on=:Geo) - - gqs = Dict{GQkey, GQres}() - gq_people = Dict{Pkey, PersonData}() - - for r in eachrow(df_gq) - ## group quarters includes some cbgs without households, need to add them - cbg_index = cbg_indexer(cbgs, r.Geo) - ## pop of each gq type in the cbg (0 if below threshhold) - gq_pops = [r["pop_"*string(x)] for x in gq_types] - ## employment category counts; a vector for each gq type - emp_stats = map((a,b)->row_gq_employment(a,b,ind_codes,r), gq_pops, job_types) - ## for creating a unique key for each person - p_idxs = ranges(gq_pops) - - ## - ## assume that all residents of military quarters work at the quarters - ## only civilians in non-inst GQs commute to jobs - ## (this assumption is also in workplaces.jl) - ## - commuter_p = Dict(:instu18 => 0.0, :inst1864 => 0.0, - :ninst1864civ => r["commuter_p|civ_worker"], - :milGQ => 0.0, :inst65o => 0.0) - - LODES_high_p = Dict(:instu18 => 0.0, :inst1864 => 0.0, - :ninst1864civ => r["LODES_high|civ_commuter"], - :milGQ => r["LODES_high|mil_commuter"], :inst65o => 0.0) - - for (t_idx, t_code) in enumerate(gq_types) - if gq_pops[t_idx] > 0 - ## use 0 as hh index for gq people (so looking up household will return nothing) - pkeys = [(p_i,0,cbg_index) for p_i in p_idxs[t_idx]] - ## create group quarter entry and add to dict - gqs[(t_idx,cbg_index)] = GQres(t_code, pkeys) - ## for mapping person index to employment category - emp_idxs = cumsum(emp_stats[t_idx]) - ## create people and add to dict - for (i,k) in enumerate(pkeys) - ## indices mapped to emp cats based on category counts; missing = no job - emp_cat = something(findfirst(x -> x >= i, emp_idxs), missing) - has_job = !ismissing(emp_cat) - is_commuter = has_job ? (rand() < commuter_p[t_code]) : false - ## income category, for commuters only - inc_cat = is_commuter ? (rand() < LODES_high_p[t_code] ? 2 : 1) : missing - ## employment category is only for assigning commuters to workplaces - emp_cat = is_commuter ? emp_cat : missing - ## using 0 for household index and sample#; age,sex,race,etc. unknown (could generate these if needed) - gq_people[k] = PersonData((0,cbg_index), 0, assumed_ages[t_idx], - has_job, is_commuter, emp_cat, inc_cat, - missing, [missing for x in add_trait_cols]...) - end - end - end - end - - ## write summary statistics for all gq's; will be used to assign ppl in gq to jobs - df_gq_summary = DataFrame([ - "geo" => String15[]; - [string(x) => Int64[] for x in gq_types]; - #["civ_ind_"*i => Int64[] for i in ind_codes]; - #["mil_ind_"*i => Int64[] for i in ind_codes] - ["ind_"*i => Int64[] for i in ind_codes] - ]) - - for c in df_gq.Geo - row_summ = Dict{String,Any}("geo" => c) - merge!(row_summ, Dict(string.(gq_types) .=> 0)) - #merge!(row_summ, Dict(["civ_ind_"*i for i in ind_codes] .=> 0)) - #merge!(row_summ, Dict(["mil_ind_"*i for i in ind_codes] .=> 0)) - merge!(row_summ, Dict(["ind_"*i for i in ind_codes] .=> 0)) - cbg_index = cbgs[c] - for (t_idx, t_code) in enumerate(gq_types) - gq = get(gqs, (t_idx,cbg_index), missing) - if !ismissing(gq) - ppl = [gq_people[k] for k in gq.residents] - row_summ[string(gq.type)] = length(ppl) - if gq.type == :ninst1864civ - #merge!(row_summ, Dict(["civ_ind_"*i for i in ind_codes] .=> counts([coalesce(x.com_cat,0) for x in ppl], 1:length(ind_codes)))) - merge!(row_summ, Dict(["ind_"*i for i in ind_codes] .=> counts([coalesce(x.com_cat,0) for x in ppl], 1:length(ind_codes)))) - #elseif gq.type == :milGQ - # merge!(row_summ, Dict(["mil_ind_"*i for i in ind_codes] .=> counts([coalesce(x.com_cat,0) for x in ppl], 1:length(ind_codes)))) - end - end - end - push!(df_gq_summary, row_summ) - end - ser_path("jlse/df_gq_summary.jlse", df_gq_summary) - - return (cbgs, gqs, gq_people) -end - - -function generate_people() - - dConfig = tryJSON("config.json") - additional_traits::Vector{String} = get(dConfig, "additional_traits", String[]) - println("DEBUG: additional_traits from config: ", additional_traits) - println("DEBUG: additional_traits type: ", typeof(additional_traits)) - println("DEBUG: additional_traits length: ", length(additional_traits)) - wp_codes = tryJSON("processed/codes.json") - ind_codes::Vector{String} = get(wp_codes, "ind_codes", String[]) - - counties = read_counties() - hh_idx = read_hh_serials() ## for linking household to row # in hh samps - println("reading person samples") - p_samps = read_psamp_df(ind_codes, additional_traits) ## for looking up person traits from sample data - p_idx = people_by_serial(p_samps) ## for linking people in households to row #s in p_samps - - ind_colnames = Symbol.(["ind_"*k for k in ind_codes]) ## industry code columns in p_samps - ind_col_idxs = Dict(ind_colnames .=> eachindex(ind_colnames)) ## assign an integer index to each, in order - - # Debug: Print PersonData struct info - println("DEBUG: PersonData struct fields:") - for (i, field) in enumerate(fieldnames(PersonData)) - println(" $i: $field (", fieldtype(PersonData, field), ")") - end - println("DEBUG: Expected total fields: ", length(fieldnames(PersonData))) - - ## pre-compute some traits so we don't have to do it inside the loop - ## each person has only one industry; append its code and index to each record - transform!(p_samps, AsTable(ind_colnames) => ByRow(first_true) => "ind_code") - transform!(p_samps, "ind_code" => ByRow(k->get(ind_col_idxs,k,missing)) => "ind_cat_idx") - ## this determines job category for commuters - transform!(p_samps, ["commuter", "ind_cat_idx"] => - ByRow((a,b)-> a ? b : missing) => - "com_cat") - - ## same for income categories (but these already exclude non-commuters) - income_colnames = [:com_LODES_low, :com_LODES_high] - income_col_idxs = Dict(:com_LODES_low => 1, :com_LODES_high => 2) - transform!(p_samps, AsTable(income_colnames) => ByRow(first_true) => "income_code") - transform!(p_samps, "income_code" => ByRow(k->get(income_col_idxs,k,missing)) => "com_inc") - - add_trait_cols = Symbol.(additional_traits) - println("DEBUG: add_trait_cols after conversion: ", add_trait_cols) - println("DEBUG: add_trait_cols length: ", length(add_trait_cols)) - - cbgs = Dict{String15, CBGkey}() ## assign an index to each cbg processed - cbg_indexer = Indexer{CBGkey}() - households = Dict{Hkey, Household}() ## create a unique hh id for each hh - people = Dict{Pkey, PersonData}() ## create unique person id for each person - - println("generating people") - for c in counties - println("county ",c) - cbg_hhs = dser_path("jlse/CO/hh"*c*".jlse"); - - for (cbg_code, hh_vec) in cbg_hhs ## vector of households in each cbg - cbg_i = cbg_indexer(cbgs, cbg_code) - for (hh_i, hh_serial) in enumerate(hh_vec) - hh_key = (hh_i, cbg_i) - p_vec = p_idx[hh_serial] ## person sample indices in each household - for (p_i, r) in enumerate(p_vec) ## create each person from data in sample df - # Debug: Print information about the first person to see what's happening - if p_i == 1 && hh_i == 1 - println("DEBUG: First person data:") - println(" hh_key: ", hh_key) - println(" r (sample index): ", r) - println(" AGEP: ", p_samps[r,:AGEP], " (type: ", typeof(p_samps[r,:AGEP]), ")") - println(" has_job: ", p_samps[r,:has_job], " (type: ", typeof(p_samps[r,:has_job]), ")") - println(" commuter: ", p_samps[r,:commuter], " (type: ", typeof(p_samps[r,:commuter]), ")") - println(" com_cat: ", p_samps[r,:com_cat], " (type: ", typeof(p_samps[r,:com_cat]), ")") - println(" com_inc: ", p_samps[r,:com_inc], " (type: ", typeof(p_samps[r,:com_inc]), ")") - println(" sch_grade: ", p_samps[r,:sch_grade], " (type: ", typeof(p_samps[r,:sch_grade]), ")") - println(" add_trait_cols: ", add_trait_cols) - println(" additional traits values: ", [p_samps[r,x] for x in add_trait_cols]) - println(" additional traits types: ", [typeof(p_samps[r,x]) for x in add_trait_cols]) - println(" total arguments: ", 8 + length(add_trait_cols)) - end - - people[(p_i,hh_i,cbg_i)] = PersonData( - hh_key, - r, - p_samps[r,:AGEP], - p_samps[r,:has_job], - p_samps[r,:commuter], - p_samps[r,:com_cat], - p_samps[r,:com_inc], - p_samps[r,:sch_grade], - [p_samps[r,x] for x in add_trait_cols]...) - end - households[hh_key] = Household(hh_idx[hh_serial], [(i,hh_i,cbg_i) for i in eachindex(p_vec)]) - end - end - end - - println("generating group quarters") - cbgs, gqs, gq_people = generate_group_quarters(dConfig, cbgs, cbg_indexer, ind_codes) - people = merge(gq_people, people) - - ## write to files for next step - println("writing people to file") - ser_path("jlse/cbg_idxs.jlse", Dict([v=>k for (k,v) in cbgs])) - ser_path("jlse/hh.jlse", households) - ser_path("jlse/gqs.jlse", gqs) - ser_path("jlse/people.jlse", people) - - return nothing -end - diff --git a/src/geopops/julia/netw.jl b/src/geopops/julia/netw.jl deleted file mode 100644 index f07b447..0000000 --- a/src/geopops/julia/netw.jl +++ /dev/null @@ -1,399 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -using Graphs -#using GraphPlot -using SparseArrays -using LinearAlgebra -using Random -using Distributions - -include("utils.jl") -include("fileutils.jl") - -## neighbors for each vertex in graph g -function neighVec(g) - return [neighbors(g,x) for x in vertices(g)] -end - -function read_gq_ppl() - ## gq residents - gq_res = Dict(k => v.residents for (k,v) in dser_path("jlse/gqs.jlse")) - ## gq employees; remove extra worker tag - gq_workers = Dict(k => [x[1:3] for x in v] for (k,v) in dser_path("jlse/gq_workers.jlse")) - return vecmerge(gq_res, gq_workers) -end - -function assign_teachers_to_grades(school_key::String15,students_by_grade::Vector{S},sch_workers::Dict{String15,Vector{Pkey}}) where S<:Any - ## person ids of teachers in this school - teachers = sch_workers[school_key] - ## a school could have been created with teachers but no students - if isempty(students_by_grade) - return [(t...,String3("0")) for t in teachers] - else - ## proportionmap from statsbase: p students in each grade in this school - pstudents = proportionmap([x[4]::String3 for x in students_by_grade]) - ## n teachers by grade, assuming it's exactly proportional to students - n_tmp = Dict(k=>v*length(teachers) for (k::String3,v::Float64) in pstudents) - ## round to integers and expand into a vector having the same length as "teachers" - teacher_grades = reduce(vcat, fill(k,v) for (k::String3,v::Int) in Dict(keys(n_tmp) .=> lrRound(collect(values(n_tmp))))) - ## append grade to each teacher's person id (teacher order is already random) - return [(t...,g) for (t::Pkey,g::String3) in zip(teachers,teacher_grades)] - end -end - -function read_sch_ppl() - ## for looking up the school grade of each student: - ppl = Dict(k=>v.sch_grade for (k,v) in dser_path("jlse/people.jlse")) - ## append each student's grade to their person id - sch_students_x_grade = Dict{String15,Vector{Tuple{fieldtypes(Pkey)...,String3}}}(k=>[(personkey..., ppl[personkey]) - for personkey in v] for (k,v) in dser_path("jlse/sch_students.jlse")) - ## strip income category from school workers, leaving only person id - sch_workers = Dict(k => [x[1:3] for x in v] for (k,v) in dser_path("jlse/sch_workers.jlse")) - ## assign a grade to each teacher and append it to their person id - teachers_by_grade = Dict(k => assign_teachers_to_grades(k,v,sch_workers) for (k,v) in sch_students_x_grade) - ## return students and teachers by school - return vecmerge(sch_students_x_grade,teachers_by_grade) -end - -function read_hh_ppl() - hh = dser_path("jlse/hh.jlse") - return Dict(k => v.people for (k,v) in hh) -end - -## "dummies" = people commuting from outside the synth area, have no household or location info -function get_dummies() - ## remove extra worker tag - return [x[1:3] for x in dser_path("jlse/work_dummies.jlse")] -end - -## integer index for each "real" person in synth pop, and for each "dummy" person -function person_indices() - ## everyone in households and gq's in the synth area - people = keys(dser_path("jlse/people.jlse")) - ## workers without households - dummies = get_dummies() - - n = length(people) - return ( - Dict(people .=> 1:n), - Dict(dummies .=> (1+n):(n+length(dummies))) - ) -end - -## connects the keys in keyvec into a stochastic block model (SBM) network -## - note, the degree distribution within blocks in an SBM is similar to -## a random (erdos-renyi) network. For a more realistic degree distribution (e.g., with hubs) -## try a degree-corrected SBM (Karrer & Newman 2010, https://arxiv.org/abs/1008.3926) -## -## if use_groups is true, the 4th element of a key should be the group identity (the first 3 would be a person key) -## returns a list of (source, dest) edges where source and dest are keys in keyvec -function connect_SBM(keyvec::Vector{T}, K::Int, min_N::Int, assoc_coeff::Float64, use_groups::Bool=true) where T<:Any - keyvec = unique(keyvec) ## duplicates would result in self-connections - n = length(keyvec) - if n < 2 ## nothing to connect - return Tuple{T,T}[] - elseif n < min_N ## more than one but less than thresh = fully connected - g = complete_graph(n) - ## indices in g correspond to positions in keyvec - return [(keyvec[x.src],keyvec[x.dst]) for x in edges(g)] - else - ## more than min_N = use SBM - if use_groups - ## group membership is based on key element #4 - group_labels = unique(k[4] for k in keyvec) - group_indices = [findall(k->k[4]==g, keyvec) for g in group_labels] - else - ## otherwise put everyone in one group - group_indices = [eachindex(keyvec)] - end - - len_grp_idx = length.(group_indices) ## number of vertices in each group - n_vec = filter(x->x>0, len_grp_idx) ## drop any 0-size groups - n_groups = length(n_vec) - w_planted = Diagonal(fill(K,n_groups)) ## contact matrix if no group mixing - prop_i = n_vec ./ sum(n_vec) ## proportion in each group - w_random = repeat(transpose(prop_i) * K, n_groups) ## contact matrix if random group mixing - c_matrix = assoc_coeff * w_planted + (1 - assoc_coeff) * w_random ## linear interpolation between planted and random contact matrices - - ## can't have more than n connections to a group - for r in eachrow(c_matrix) - r .= min.(r,n_vec) - end - ## can't have more than n-1 within-group connections - d_tmp = view(c_matrix, diagind(c_matrix)) - d_tmp .= min.(d_tmp, n_vec .- 1) - - ## note, graph gen function takes group sizes in n_vec; the vertices/edges in the - ## resulting graph are just numbered based on group sizes, in the order given in n_vec - ## to recover keys, need to translate g index -> keyvec index -> key - g = stochastic_block_model(c_matrix,n_vec) - keyvec_indices = reduce(vcat, group_indices) ## for translating indices - - ## algo sometimes creates 0-degree vertices; not technically wrong, but force them to have 1 connection anyway - fix_zeros = findall(degree(g).==0) - for x in fix_zeros - add_edge!(g, x, rand(vertices(g))) - end - - ## convert indices in g to keyvec keys for each edge - result = Tuple{T,T}[] - for x in edges(g) - if use_groups - (group_labels[findfirst(c->in(x.src,c) , ranges(len_grp_idx))] == keyvec[keyvec_indices[x.src]][4]) && - (group_labels[findfirst(c->in(x.dst,c) , ranges(len_grp_idx))] == keyvec[keyvec_indices[x.dst]][4]) || - throw("group assignment error") - end - - push!(result, (keyvec[keyvec_indices[x.src]],keyvec[keyvec_indices[x.dst]])) - end - return result - end ## if n -end - -## as above, but using a simple small-world network -function connect_small_world(keyvec::Vector{T}, K::Int, min_N::Int, B::Float64) where T<:Any - keyvec = unique(keyvec) ## e.g., duplicates happen if someone lives and works at the same gq - n = length(keyvec) ## size of group - if n < 2 ## nothing to connect - return Tuple{T,T}[] - elseif n < min_N ## more than one but less than thresh = fully connected - g = complete_graph(n) - else - g = watts_strogatz(n, K, B) ## small-world network - end - ## indices in g correspond to positions in keyvec - return [(keyvec[x.src],keyvec[x.dst]) for x in edges(g)] -end - -## as above, but completely connected graph -function connect_complete(keyvec::Vector{T}) where T<:Any - keyvec = unique(keyvec) - n = length(keyvec) ## size of group - if n < 2 ## nothing to connect - return Tuple{T,T}[] - else - g = complete_graph(n) - end - ## indices in g correspond to positions in keyvec - return [(keyvec[x.src],keyvec[x.dst]) for x in edges(g)] -end - -## sparse adjacency matrix of bits should be an efficient way to store the network -## (also, simple to convert to Graphs.jl graph for analysis) -## (also, has good lookup performance, same as hash table) -function sp_from_groups(connect_fn, keygroups::Vector{Vector{T}}, p_idxs::Dict{Pkey,Int64}) where T<:Tuple - src_idxs = Int64[] - dst_idxs = Int64[] - - for keyvec in keygroups - ## connect keys into network, then convert to integer indices - for (s_key, d_key) in connect_fn(keyvec) - ## key elements 1-3 correspond to a person key - push!(src_idxs, p_idxs[s_key[1:3]]) - push!(dst_idxs, p_idxs[d_key[1:3]]) - end - end - - ## create sparse matrix from indices - return sparse([src_idxs;dst_idxs],[dst_idxs;src_idxs], - trues(length(src_idxs)+length(dst_idxs)), - length(p_idxs),length(p_idxs)) ## ensure size is p_idxs by p_idxs -end - - -## -## generate all networks -## -function generate_networks() - - dConfig = tryJSON("config.json") - work_K::Int = get(dConfig, "workplace_K", 8) ## mean degree for wp networks - school_K::Int = get(dConfig, "school_K", 12) ## mean degree for school networks - gq_K::Int = get(dConfig, "gq_K", 12) ## mean degree for group-quarters networks - other_K::Int = get(dConfig, "netw_K", 8) ## mean degree for other networks - sm_world_B::Float64 = get(dConfig, "netw_B", 0.25) ## beta param for small world networks - work_assoc_coeff::Float64 = get(dConfig, "income_associativity_coefficient", 0.9) ## group associativity for workplace networks - sch_assoc_coeff::Float64 = get(dConfig, "school_associativity_coefficient", 0.9) ## group associativity for school networks - - ## from workplaces.jl; workers grouped into companies - ## worker is (person id, hh id, cbg id, income category) - company_workers = collect(values(dser_path("jlse/company_workers.jlse"))) - #school students and teachers - ## student/teacher is (person id, hh id, cbg id, grade) -- first 3 are a person key - ppl_in_schools = collect(values(read_sch_ppl())) - ## gq residents and employees - ppl_in_gq = collect(values(read_gq_ppl())) - - ## households; assume they're fully connected (and maybe have a higher transmission rate within) - hh_ppl = read_hh_ppl() - ## save household membership Dict - ser_path("jlse/hh_ppl.jlse", hh_ppl) - ## just groups as vectors - ppl_in_hhs = collect(values(hh_ppl)) - - ## each person needs an integer index - p_idxs, dummy_idxs = person_indices() - ## save dummies to file; these ppl have workplace but no household; sim should infect them randomly at home - ser_path("jlse/adj_dummy_keys.jlse", Dict(v=>k for (k,v) in dummy_idxs)) - ## merge people and dummies for network - merge!(p_idxs, dummy_idxs) - ## will need the index keys to look up people - ser_path("jlse/adj_mat_keys.jlse", first.(sort(collect(p_idxs),by=p->p[2]))) - - ## generate network - ## workplaces: using stochastic block model (SBM) network - ## so that all results are comparable, use the SBM algo even if there's only one income group in a wp - adj_wp = sp_from_groups(v->connect_SBM(v, work_K, work_K+2, work_assoc_coeff, true), company_workers, p_idxs) - ## schools: using SBM, grouped by grade - adj_sch = sp_from_groups(v->connect_SBM(v, school_K, school_K+2, sch_assoc_coeff, true), ppl_in_schools, p_idxs) - ## other institutions (currently just gq's) : using small-world network - adj_gq = sp_from_groups(v->connect_small_world(v, gq_K, gq_K+2, sm_world_B), ppl_in_gq, p_idxs) - ## households are fully connected - adj_hh = sp_from_groups(v->connect_complete(v), ppl_in_hhs, p_idxs) - ## save combined non-hh netw for when those distinctions are not needed - adj_mat_non_hh = adj_wp .| adj_sch .| adj_gq - - ## matrices must be symmetrical, save space by only storing half - ser_path("jlse/adj_mat_non_hh.jlse",sparse(UpperTriangular(adj_mat_non_hh))) - ser_path("jlse/adj_mat_wp.jlse",sparse(UpperTriangular(adj_wp))) - ser_path("jlse/adj_mat_sch.jlse",sparse(UpperTriangular(adj_sch))) - ser_path("jlse/adj_mat_gq.jlse",sparse(UpperTriangular(adj_gq))) - ser_path("jlse/adj_mat_hh.jlse",sparse(UpperTriangular(adj_hh))) - - ## keep track of people working outside synth area (have no workplace network, sim should infect them randomly at work) - outside_workers = dser_path("jlse/outside_workers.jlse") - ser_path("jlse/adj_out_workers.jlse", Dict(p_idxs[only(x)[1:3]] => only(x)[1:3] for x in values(outside_workers))) - - return nothing -end - - -## -## generate info needed to simulate ephemeral location-based contacts -## - -## assume that the chance of meeting a neighbor is the same as meeting someone who works in your neighborhood -## (otherwise need another parameter to describe the difference) - -## group potential encounters by census tract (CBG seems too restrictive) - -## home locations; e.g, -## cbg A: people 1, 2 -## cbg B: 3, 4 -## -## work locations; e.g., -## cbg A: 1, 3 -## cbg B: 2 -## -## at home, person 1 could meet: 2, 3 -## person 2 could meet: 1, 3 -## person 3 could meet: 4, 2 -## -## at work, person 1 could meet: 3, 2 -## person 2 could meet: 3, 4 -## person 3 could meet: 1, 2 - -## make a separate matrix for home and work neighborhood contacts? -## (# of contact events could be independent, or maybe work is closed) -## matrix is not symmetrical (local resident you meet while at work probably won't meet you when they're at work) -## "source" person will be column (because reading by columns is faster in julia) - -## home neighborhood contacts -## for each location: -## for each person who lives there (hh or non-inst gq): -## look up everyone else who lives there (hh or gq) + append everyone who works there -## work neighborhood contacts -## for each location: -## for each person who works there: -## look up everyone else who works there + append everyone who lives there (hh or non-inst gq) - -## save memory by not generating the contact matrix -## just a matrix of locations (columns) x people (rows) -## then simple O1 look-up to get a person's home or work location, and pick from that column - -function generate_location_matrices() - - w = dser_path("jlse/company_workers.jlse") ## employers/employees (with work locations) - hh = dser_path("jlse/hh.jlse") ## households/residents (with hh locations) - cbg_idxs = dser_path("jlse/cbg_idxs.jlse") ## location (cbg) keys used in person/hh keys - cbg_idxs = Dict(k=>String31(v) for (k,v) in cbg_idxs) - gqs = dser_path("jlse/gqs.jlse") ## group-quarters/residents (with gq locations) - ## assume only non-inst GQ residents are available for ephemeral local contacts - ni_types = Set([:milGQ, :ninst1864civ]) - gq_noninst = filterv(x->(x.type in ni_types), gqs) - ## use the same matrix indices as in the regular contact networks - k = dser_path("jlse/adj_mat_keys.jlse") - p_idxs = Dict(k .=> eachindex(k)) - - ## group potential encounters by census tract (CBG seems too restrictive) - hh_tracts = unique(x[1:end-1] for x in values(cbg_idxs)) - work_tracts = unique(x[3][1:end-1] for x in keys(w)) - tracts = sort(unique([hh_tracts; work_tracts])) - - ## convert tract codes to integer indices for constructing a matrix - loc_idxs = Dict(tracts .=> eachindex(tracts)) - ## save location indices - ser_path("jlse/loc_mat_keys.jlse",loc_idxs) - - ## group individuals by tract code (= cbg code minus last character) - ## dataframe provides fast grouping - w_df_by_loc = groupby(DataFrame((k[3][1:end-1], v) for (k,v) in w), "1") - ## place person-vectors in a dict with integer tract indices as the keys - workers_by_loc = Dict(loc_idxs[loc["1"]]=>reduce(vcat, w_df_by_loc[loc][!,"2"]) for loc in keys(w_df_by_loc)) - ## convert person keys to network matrix indices (same indices as regular contact networks) - w_idxs_by_loc = Dict(k=>[p_idxs[i[1:3]] for i in v] for (k,v) in workers_by_loc) - - h_df_by_loc = groupby(DataFrame((cbg_idxs[k[2]][1:end-1], v.people) for (k,v) in hh), "1") - hh_ppl_by_loc = Dict(loc_idxs[loc["1"]]=>reduce(vcat, h_df_by_loc[loc][!,"2"]) for loc in keys(h_df_by_loc)) - hh_idxs_by_loc = Dict(k=>[p_idxs[i] for i in v] for (k,v) in hh_ppl_by_loc) - - gq_df_by_loc = groupby(DataFrame((cbg_idxs[k[2]][1:end-1], v.residents) for (k,v) in gq_noninst), "1") - gq_ppl_by_loc = Dict(loc_idxs[loc["1"]]=>reduce(vcat, gq_df_by_loc[loc][!,"2"]) for loc in keys(gq_df_by_loc)) - gq_idxs_by_loc = Dict(k=>[p_idxs[i] for i in v] for (k,v) in gq_ppl_by_loc) - - ## residential locations include households and non-inst GQs: - res_idxs_by_loc = vecmerge(hh_idxs_by_loc, gq_idxs_by_loc) - - ## construct matrices for use in simulation - ## columns are locations, rows are people (because we'll be looking up by location) - ## note, currently people have one job max - w_loc_contact_mat = sparse( - reduce(vcat, collect(values(w_idxs_by_loc))), - reduce(vcat, [fill(k,length(v)) for (k,v) in w_idxs_by_loc]), - trues(sum(length.(values(w_idxs_by_loc)))), - length(k),length(tracts) - ) - - res_loc_contact_mat = sparse( - reduce(vcat, collect(values(res_idxs_by_loc))), - reduce(vcat, [fill(k,length(v)) for (k,v) in res_idxs_by_loc]), - trues(sum(length.(values(res_idxs_by_loc)))), - length(k),length(tracts) - ) - - ser_path("jlse/work_loc_contact_mat.jlse",w_loc_contact_mat) - ser_path("jlse/res_loc_contact_mat.jlse",res_loc_contact_mat) - - ## save work and home loc idx for each person idx, for fast lookup - ## note, currently people have one job max - w_loc_by_p_idx = Dict(reduce(vcat, [v .=> k for (k,v) in w_idxs_by_loc])) - res_loc_by_p_idx = Dict(reduce(vcat, [v .=> k for (k,v) in res_idxs_by_loc])) - - ser_path("jlse/work_loc_lookup.jlse",w_loc_by_p_idx) - ser_path("jlse/res_loc_lookup.jlse",res_loc_by_p_idx) - - return nothing -end - - - diff --git a/src/geopops/julia/schools.jl b/src/geopops/julia/schools.jl deleted file mode 100644 index 0b1d0e1..0000000 --- a/src/geopops/julia/schools.jl +++ /dev/null @@ -1,129 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -#using Statistics - -include("utils.jl") -include("fileutils.jl") - -function read_sch_cap() - sch_cols = Dict("NCESSCH"=>String15,"TEACHERS"=>Int64,"STUDENTS"=>Int64) - schools = read_df("processed/schools.csv"; select=collect(keys(sch_cols)), types=sch_cols) - return Dict(schools.NCESSCH .=> schools.STUDENTS) -end - -function find_closest(n::Int) - ## grades offered by each school - schools = read_df("processed/schools.csv"; types=Dict("NCESSCH"=>String15)) - ## distance matrix - distmat = read_df("processed/cbg_sch_distmat.csv"; types=Dict("GEOID"=>String15)) - ## reshape and group by geography - distmat = groupby(stack(distmat), :GEOID) - ## for each grade, store a dict of geo=>schools - closest = Dict{String3,Dict{String15,Vector{String15}}}() - ## and the distances - distances = Dict{String3,Dict{String15,Vector{Float64}}}() - - ## assume kindergartens also offer preschool - ## (we haven't generated private preschools and don't have enough public preschools) - schools[:,"G_PK_OFFERED"] .= schools[!,"G_PK_OFFERED"] .| schools[!,"G_KG_OFFERED"] - - for (k,gr) in zip(String3.([["p","k"];string.(1:12)]), [["PK","KG"];string.(1:12)]) - mask = schools[!,"G_"*gr*"_OFFERED"] - sch_by_geo = Dict{String15,Vector{String15}}() - dist_by_geo = Dict{String15,Vector{Float64}}() - ## the keys of distmat are geo groups - for geo in keys(distmat) - avail = @view distmat[geo][mask,:] - topidxs = partialsortperm(avail.value, 1:n) - sch_by_geo[geo[1]] = String15.(avail.variable[topidxs]) - dist_by_geo[geo[1]] = avail.value[topidxs] - end - closest[k] = sch_by_geo - distances[k] = dist_by_geo - end - - return closest, distances -end - -function read_p_in_school(cbgs::Dict{CBGkey, String15}) - people = dser_path("jlse/people.jlse") - ## excluding college and grad school - p_in_school = filterv(p->(!ismissing(p.sch_grade) && !in(p.sch_grade, ["c","g"])), people) - - ## sort by cbg,household so we can make kids in the same household go to the same school - pkeys = collect(keys(p_in_school)) - idxs = sortperm([x[[3,2]] for x in pkeys]) - pkeys = pkeys[idxs] - - ## person key, grade, cbg code - return [(k, (p_in_school[k]).sch_grade, cbgs[k[3]]) for k in pkeys] -end - - -## can't just draw school-wise from closest -## how to explain existence of schools that aren't closest (or 2 closest) for anyone -## model school-choosing behavior?? -## -- p depends on distance (p ~ 1/dist) and # remaining spots (p ~ spots_left) -## -- only need to consider n closest -## ...? - -## send to closest (or 2 closest) first; people who live together more likely to go to same school -## create schools (= string ids), place people into schools -function generate_schools() - - dConfig = tryJSON("config.json") - n_schools::Int = get(dConfig, "n_closest_schools", 4) - ## assign students to closest avaiable school (90%) or 2nd closest (10%) - prob_closest::Float64 = get(dConfig, "p_closest_school", 0.9) - - closest, _ = find_closest(n_schools) - cbgs = dser_path("jlse/cbg_idxs.jlse") - p_in_school = read_p_in_school(cbgs) - - sch_capacity = read_sch_cap() - ## shrink school capacities to prevent underfilling non-closest schools - sch_capacity = Dict(k=>round(Int,v*0.8) for (k,v) in sch_capacity) - - ## school code => vector of person keys - sch_students = Dict{String15,Vector{Pkey}}() - ## initialize with empty vectors - for k in keys(sch_capacity) - sch_students[k] = Vector{Pkey}() - end - - ## read person key, grade level, and cbg code for each person - for (pk,gr,geo) in p_in_school - ## look up closest schools for grade level and cbg - opts = closest[gr][geo] - ## find the first that hasn't been filled to capacity - idx_avail = findfirst(k->sch_capacity[k]>length(sch_students[k]), opts) - ## if all schools full, try again with increased capacity - idx_avail = isnothing(idx_avail) ? findfirst(k->1.5*sch_capacity[k]>length(sch_students[k]), opts) : idx_avail - ## if all schools full, try again with increased capacity - idx_avail = isnothing(idx_avail) ? findfirst(k->2.5*sch_capacity[k]>length(sch_students[k]), opts) : idx_avail - ## if still full, just overfill the closest - idx_avail = isnothing(idx_avail) ? 1 : idx_avail - ## choose closest or next closest - idx_choice = rand() < prob_closest ? idx_avail : idx_avail+1 - idx_choice = idx_choice > lastindex(opts) ? 1 : idx_choice - ## append student to school - push!(sch_students[opts[idx_choice]], pk) - end - - ## save to file for next step - ser_path("jlse/sch_students.jlse", sch_students) - return nothing -end - - - diff --git a/src/geopops/julia/synthpop.jl b/src/geopops/julia/synthpop.jl deleted file mode 100644 index a9d8387..0000000 --- a/src/geopops/julia/synthpop.jl +++ /dev/null @@ -1,28 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -include("households.jl") -include("schools.jl") -include("workplaces.jl") -include("netw.jl") - -println("creating people, households, and group quarters") -generate_people() -println("creating schools") -generate_schools() -println("creating workplaces") -generate_commute_matrices() -generate_jobs_and_workers() -println("creating network") -generate_networks() -generate_location_matrices() -println("done") diff --git a/src/geopops/julia/test_netw.jl b/src/geopops/julia/test_netw.jl deleted file mode 100644 index 1328186..0000000 --- a/src/geopops/julia/test_netw.jl +++ /dev/null @@ -1,857 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -#using Distributed -using Random -using Graphs -#using GraphPlot -using SparseArrays -using LinearAlgebra -using SharedArrays -using Distributions -using Plots -using StatsPlots -using StatsBase - - -#using MatrixDepot -#using SNAPDatasets - -include("utils.jl") -include("fileutils.jl") - - -## -## network stats -## - -## from Bonchev 2004, Ir(vd) -function Bonchev_vd_info0(t::AbstractGraph{T}) where T - v = filter(x->x>0,degree(t)) - A = sum(v) - return (sum( v .* log2.(v) )) / ( A * log2(A) ) -end - -function Bonchev_vd_info(t::AbstractGraph{T}) where T - v = filter(x->x>0,degree(t)) - A = sum(v) - return sum( v .* log2.(v) ./ (A * log2(A)) ) -end - -## from Saberi 2021 -makehub(t::AbstractGraph{T}) where T = sum(degree(t) .^ 2) / sum(degree(t)) - - - - - -#g = watts_strogatz(100,8,0.25) -#gplot(g) -## simulate infection at p = ? -#diffusion_rate(g, 0.5, 10) -## mean path length -#f = floyd_warshall_shortest_paths(g) -#sum(f.dists) / length(f.dists) - -## neighbors -#g = watts_strogatz(20,6,0.3) -#[neighbors(g,x) for x in vertices(g)] - - -#include("netw.jl") - -#generate_network() - -#full_mat, full_keys = sparse_from_adjdict(merge_hh_net()); -#ser_path("jlse/full_graph_mat.jlse", sparse(UpperTriangular(full_mat))) -#ser_path("jlse/full_graph_keys.jlse", full_keys) - - -## note this method adds zero-degree vertices if adj_mat has 0-sum rows -#full_graph = SimpleGraph(sparse(Symmetric(dser_path("jlse/full_graph_mat.jlse")))) -#ser_path("jlse/full_graph.jlse",full_graph) - -#full_graph = dser_path("jlse/full_graph.jlse") -#node_keys = dser_path("jlse/full_graph_keys.jlse") - - - -#f(v) = mean(v[2:end] ./ v[1:end-1]) - - -#net_comp = Dict() -net_comp = dser_path("net_comp.jlse") - -#M = sparse(Symmetric(dser_path("jlse/hh_adj_mat.jlse"))) .| sparse(Symmetric(dser_path("jlse/adj_mat.jlse"))) -#M = dser_path("jlse/hh_adj_mat.jlse") .| dser_path("jlse/adj_mat.jlse") - -x = (Edge(p) for p in zip(findnz( - dser_path("jlse/adj_mat_hh.jlse") .| dser_path("jlse/adj_mat_non_hh.jlse") - )[1:2]...)) - -t = SimpleGraphFromIterator(x) -x = nothing -GC.gc() - -N = nv(t) # 9424031 -n0 = sum(degree(t).==0) # 401057 -mu = mean(degree(t)) ## 8.483098686750925 -nname = "MD" - -#N - n0 -#mean(filter(i->i>0,degree(t))) -#deg_zero_i = findall(degree(t).==0) -#rem_vertices!(t, deg_zero_i) -#nv(t) -#mean(degree(t)) -#nname = "MD0" - - -## "static scale free" algo -## Goh K-I, Kahng B, Kim D: Universal behaviour of load distribution in scale-free networks. Phys Rev Lett 87(27):278701, 2001. - -t = nothing -GC.gc() -t = barabasi_albert(N, 4); nname = "BA" -t = erdos_renyi(N, mu/N); nname = "ER" -t = watts_strogatz(N, 8, 0.25); nname = "WS" -t = static_scale_free(N, round(Int,0.5*mu*N), 3); nname = "SSF" -mean(degree(t)) - - - -net_comp[nname] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -net_comp - -ser_path("net_comp.jlse", net_comp) - -nc_names = ["local clust","global clust","pagerank","rich club","assortativity","bonchev","makehub","mean degree"] -use_cols = [8,1,2,5,7,6] -[nc_names[x] for x in use_cols] -[k=>[round(v[x];sigdigits=3) for x in use_cols] for (k,v) in net_comp] - - - - - - -#adj_n0 = false -#if adj_n0 -# for x in shuffle(vertices(t))[1:n0] -# for d in neighbors(t,x) -# rem_edge!(t,x,d) -# end -# end -#end -#mean(degree(t)) - - - -m = matrixdepot("SNAP/loc-Gowalla") -t = SimpleGraph(m) -net_comp["loc_gowalla"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -m = matrixdepot("SNAP/loc-Brightkite") -t = SimpleGraph(m) -net_comp["loc_brightkite"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -m = matrixdepot("Newman/as-22july06") -t = SimpleGraph(m) -net_comp["routers"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -m = matrixdepot("Barabasi/NotreDame_yeast") -t = SimpleGraph(m) -net_comp["yeast"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -m = matrixdepot("DIMACS10/coPapersDBLP") -t = SimpleGraph(m) -net_comp["citation_dblp"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -m = matrixdepot("SNAP/com-DBLP") -t = SimpleGraph(m) -net_comp["com_dblp"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -ser_path("jlse/net_comp.jlse",net_comp) - - -t = SimpleGraph(sparse(matrixdepot("SNAP/com-Orkut"))) -net_comp["com_orkut"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -ser_path("jlse/net_comp.jlse",net_comp) -t = nothing -GC.gc() - -t = SimpleGraph(sparse(matrixdepot("SNAP/com-Amazon"))) -net_comp["com_amazon"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -ser_path("jlse/net_comp.jlse",net_comp) -t = nothing -GC.gc() - -t = SimpleGraph(sparse(matrixdepot("SNAP/com-Youtube"))) -net_comp["com_youtube"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -ser_path("jlse/net_comp.jlse",net_comp) -t = nothing -GC.gc() - -t = SimpleGraph(sparse(matrixdepot("SNAP/com-LiveJournal"))) -net_comp["com_liveJournal"] = [mean(local_clustering_coefficient(t)) -global_clustering_coefficient(t) -mean(pagerank(t)) -rich_club(t,1) -assortativity(t) -Bonchev_vd_info(t) -makehub(t) -mean(degree(t)) -#mean(f(diffusion_rate(t,0.1,20)) for x in 1:20) -] - -ser_path("jlse/net_comp.jlse",net_comp) -t = nothing -GC.gc() - - -#m = matrixdepot("SNAP/p2p-Gnutella08") -#t = SimpleDiGraph(m) -#m = matrixdepot("SNAP/email-Eu-core") -#t = SimpleDiGraph(m) - -#t = loadsnap(:ca_astroph) -#t = loadsnap(:ego_twitter_u) -#t = loadsnap(:email_enron) -#t = loadsnap(:facebook_combined) -#t = loadsnap(:soc_slashdot0902_u) - -#net_comp = dser_path("jlse/net_comp.jlse") - -mask = BitVector([1,1,0,0,1,0]) -y = reduce(hcat, net_comp[k][mask] for k in - ["LA_full","small_world","scale_free","p2p","email","ca_astroph", - "ego_twitter_u","email_enron","facebook_combined","soc_slashdot0902_u"]) -shapes = ([[:star4, :circle, :circle]; fill(:+,7)]) - -plot([1:3], y, shape=reshape(shapes,1,10), - label="", xticks=:none, - linestyle=:dash, linealpha=0.5) - - - -## -## checking fit of synth pop -## - -FTdist(v1::Matrix{T},v2::Matrix{T}) where T<:Real = sum((sqrt.(v1 .+ one(T)) .- sqrt.(v2 .+ one(T))) .^ 2) -summarize(idxs::Vector{Int64}, pop::Matrix{Int64}) = sum(view(pop,idxs,:), dims=1) -function summarize!(dest::Matrix{Int64}, idxs::Vector{Int64}, pop::Matrix{Int64}) - sum!(dest, view(pop,idxs,:)) - return nothing -end - -## rmse standardized by st dev -SRMSE(O::Vector{T},E::Vector{T}) where T<:Real = sqrt(mean((O .- E) .^ 2)) / std(E) - -function read_test_targs() - acs_targets = read_df("processed/test_targets.csv"; types=Dict("Geo"=>String15)) - targ_idxs = axes(acs_targets,1) - targs = [Matrix{Int64}(acs_targets[[targ_id],2:end]) for targ_id in targ_idxs] - geos = [acs_targets[targ_id,1] for targ_id in targ_idxs] - return (targs, geos, names(acs_targets)[2:end]) -end - -function read_test_samples() - puma_samples_all = read_df("processed/test_samples.csv"; types=Dict("SERIALNO"=>String15)) - hh_ids = puma_samples_all[:,1] ## for looking up the household serial numbers later - shared_samples = Matrix{Int64}(puma_samples_all[:,2:end]) - return (shared_samples, hh_ids) -end - -function calc_test_scores() - targs_all, geos_all, targ_colnames = read_test_targs() - samples, hh_ids = read_test_samples() - hh_idx = Dict(hh_ids .=> eachindex(hh_ids)) - ## each target's county - county = [g[1:5] for g in geos_all] - counties = unique(county) - summary = similar(first(targs_all)) - - for c in counties - cbg_hhs = dser_path("jlse/CO/hh"*c*".jlse"); - targs = targs_all[county .== c] - geos = geos_all[county .== c] - targ_idx = Dict(geos .=> eachindex(geos)) - test_scores = Dict{String15,Float64}() - for (cbg,households) in cbg_hhs - summarize!(summary, [hh_idx[k] for k in households], samples) - test_scores[cbg] = FTdist(summary, targs[targ_idx[cbg]]) - end - ser_path("jlse/CO/test_"*c*"_scores.jlse",test_scores) - end -end - - - -## targets are census block group (cbg) summary statistics -function read_targets(targ_idxs=[]) - acs_targets = read_df("processed/acs_targets.csv"; types=Dict("Geo"=>String15)) - if isempty(targ_idxs) - targ_idxs = axes(acs_targets,1) - end - ## pull target stats from the dataframe and convert to matrix or array for faster math - targs = [Matrix{Int64}(acs_targets[[targ_id],2:end]) for targ_id in targ_idxs] - geos = [acs_targets[targ_id,1] for targ_id in targ_idxs] - return (targs, geos, names(acs_targets)[2:end]) -end - -## returns a sharedmatrix of samples -- can pass it to local parallel processes without copying data -function read_samples() - puma_samples_all = read_df("processed/census_samples.csv"; types=Dict("SERIALNO"=>String15)) - hh_ids = puma_samples_all[:,1] ## for looking up the household serial numbers later - shared_samples = Matrix{Int64}(puma_samples_all[:,2:end]) - return (shared_samples, hh_ids) -end - - -function all_summaries(test::Bool) - if test - targs_all, geos_all, targ_colnames = read_test_targs() - samples, hh_ids = read_test_samples() - else - targs_all, geos_all, targ_colnames = read_targets() - samples, hh_ids = read_samples() - end - - hh_idx = Dict(hh_ids .=> eachindex(hh_ids)) - county = [g[1:5] for g in geos_all] - counties = unique(county) - summary = similar(first(targs_all)) - summaries = Dict{String15, typeof(summary)}() - - for c in counties - cbg_hhs = dser_path("jlse/CO/hh"*c*".jlse") - for (cbg,households) in cbg_hhs - summarize!(summary, [hh_idx[k] for k in households], samples) - summaries[cbg] = copy(summary) - end - end - - t_mat = zeros(length(targs_all), length(first(targs_all))) - for (i,r) in enumerate(targs_all) - t_mat[i,:] = r - end - - s_mat = zeros(length(summaries), length(first(values(summaries)))) - for (i,g) in enumerate(geos_all) - s_mat[i,:] = summaries[g] - end - - return s_mat, t_mat, geos_all, targ_colnames -end - - -s_mat, t_mat, geos_all, targ_colnames = all_summaries(false) -err_by_targ = [(n, SRMSE(s_mat[:,i],t_mat[:,i])) for (i,n) in enumerate(targ_colnames)] - -s_mat, t_mat, geos_all, targ_colnames = all_summaries(true) -err_by_testvar = [(n, SRMSE(s_mat[:,i],t_mat[:,i])) for (i,n) in enumerate(targ_colnames)] - -y = [[x[2] for x in err_by_targ]; [x[2] for x in err_by_testvar]] -scatter(y, label="", xticks=:none) - - -scatter([x[2] for x in err_by_targ], label="", xticks=:none) -[(x[1],x[2]) for x in err_by_targ if x[2]>0.3] - -[(x[1],x[2]) for x in err_by_testvar if x[2]<0.4] - -function random_answer(glob_samp_ref::Matrix{Int64}, mask::BitVector, targ::Matrix{Int64}, n::Int64, params::Dict{Symbol, R}) where R<:Real - samples = glob_samp_ref[mask, :] - idxs = axes(samples,1) - ## if no valid samples, return a bad score (this shouldn't happen) - if isempty(idxs) - return (Vector{Int64}(), 0, Inf, 0.0) - end - c0 = rand(idxs,n) - summary = summarize(c0, samples) - E0 = FTdist(summary, targ) - res = findall(mask)[c0] - return (res, 0, E0, 0.0) -end - -function random_answerer(shared_samples::Matrix{Int64}, params::Dict{Symbol, R}) where R<:Real - function f(mask::BitVector, targ::Matrix{Int64}, n::Int64) - return random_answer(shared_samples, mask, targ, n, params) - end - return f -end - -## number of households in each cbg -function read_hh_counts() - hhcdf = read_df("processed/hh_counts.csv"; types=Dict("Geo"=>String15)) - ## convert to dicts for easier lookup - ## note the .=> syntax; this broadcasts associations between two arrays, similar to dict(zip()) in python - return Dict(hhcdf[:,1] .=> hhcdf[:,2]) -end - -## geographic data for target cbg's -- will use these to determine which samples to use -function read_targ_geo() - cbg_geo_cols = Dict("Geo"=>String15,"st_puma"=>String7,"cbsa"=>String7,"county"=>String7,"R"=>Float64,"U"=>Float64) - cbg_geo = read_df("processed/cbg_geo.csv"; select=collect(keys(cbg_geo_cols)), types=cbg_geo_cols) - cbg_puma = Dict(cbg_geo[:,"Geo"] .=> cbg_geo[:,"st_puma"]) - cbg_county = Dict(cbg_geo[:,"Geo"] .=> cbg_geo[:,"county"]) - cbg_cbsa = Dict(cbg_geo[:,"Geo"] .=> cbg_geo[:,"cbsa"]) - cbg_urban = Dict(cbg_geo[:,"Geo"] .=> cbg_geo[:,"U"]) - return (cbg_puma, cbg_county, cbg_cbsa, cbg_urban) -end - -## looks up which samples in df to use, based on column name k and associated geo data of targets -## don't usually need to specify function return types but in this case, target_geo_codes might be empty and -## this way it still returns a vector{BitVector} -function sample_lookup(df::DataFrame, k::Symbol, target_geo_codes)::Vector{BitVector} - if k == :U - return [urbanization_lookup(df,x) for x in target_geo_codes] - else - ## note .== syntax for element-wise comparison - return [coalesce.(df[!,k] .== x, false) for x in target_geo_codes] - end -end - -## creates annealing function and executes it on available processors -## returns a vector of whatever anneal() returns -function dont_optimize(samples::Matrix{Int64}, samp_masks::Vector{BitVector}, targs::Vector{Matrix{Int64}}, n_hhs::Vector{Int64}, params) - a_fn = random_answerer(samples, params) - return map(a_fn, samp_masks, targs, n_hhs) -end - -## performs several optimization runs on the target cbg's in each county -## writes each county's results to a separate file -function process_counties(test::Bool) - - if test - samples, hh_ids = read_test_samples() - targs_all, geos_all, targ_colnames = read_test_targs() - else - samples, hh_ids = read_samples() - targs_all, geos_all, targ_colnames = read_targets() - end - - cbg_puma, cbg_county, cbg_cbsa, cbg_urban = read_targ_geo() - samp_geo_cols = Dict("SERIALNO"=>String15,"st_puma"=>String7,"cbsa"=>String7,"county"=>String7,"R"=>Float64,"U"=>Float64) - samp_geo = read_df("processed/samp_geo.csv"; select=collect(keys(samp_geo_cols)), types=samp_geo_cols) - hh_counts = read_hh_counts() - n_hhs_all = [hh_counts[g] for g in geos_all] - county = [g[1:5] for g in geos_all] - - params = Dict(:maxgens => 0, :critval => 0.0, :cooldown => 0.0, :report => 0) - - counties = unique(county) - for c in counties - targs = targs_all[county .== c] - geos = geos_all[county .== c] - n_hhs = n_hhs_all[county .== c] - - ## look up samples for each target based on target's puma code - ## puma is the most local sample level, has the fewest samples available - samp_masks = sample_lookup(samp_geo, :st_puma, [cbg_puma[g] for g in geos]) - x = dont_optimize(samples, samp_masks, targs, n_hhs, params) - - ## store results as dict keyed by cbg code; look up actual hh id's of sample indices - scores = Dict(geos .=> [a[3] for a in x]) - households = Dict(geos .=> [hh_ids[a[1]] for a in x]) - - if test - ser_path("jlse/CO/rand_test_"*c*".jlse",households) - ser_path("jlse/CO/rand_test_"*c*"_scores.jlse",scores) - else - ser_path("jlse/CO/rand_"*c*".jlse",households) - ser_path("jlse/CO/rand_"*c*"_scores.jlse",scores) - end - end - - return nothing -end - -calc_test_scores() -process_counties(true) -process_counties(false) - -geos = read_targ_geo() -counties = unique(values(geos[2])) - -x = [dser_path("jlse/CO/test_"*c*"_scores.jlse") for c in counties] -test_scores = merge(x...) -x = [dser_path("jlse/CO/hh"*c*"_scores.jlse") for c in counties] -scores = merge(x...) -x = [dser_path("jlse/CO/rand_"*c*"_scores.jlse") for c in counties] -rand_scores = merge(x...) -x = [dser_path("jlse/CO/rand_test_"*c*"_scores.jlse") for c in counties] -rand_test_scores = merge(x...) - -cbgs = collect(keys(scores)) - -s05,s95 = log.(0.25 .* quantile.(Chisq(85), [0.05, 0.95])) -t05,t95 = log.(0.25 .* quantile.(Chisq(31), [0.05, 0.95])) - -s1 = [scores[k] for k in cbgs] -s2 = [rand_scores[k] for k in cbgs] - -t1 = [test_scores[k] for k in cbgs] -t2 = [rand_test_scores[k] for k in cbgs] - -circ = Shape(Plots.partialcircle(0, 2π)) -plot([log.(s1) log.(s2)], label="", xticks=:none, seriestype=:scatter, grid=false, - ylims=(1.0,9.1), - xlabel="census block group",ylabel="ln(¼ FT²)", - markerstrokewidth=[0 0.5], markersize=[1.5 2.5], markercolor=[:black :white], - markershape=[:none circ],markerstrokecolor=:black, markeralpha=[1.0 0.5], - markerstrokealpha=[1.0 0.75], - dpi=300) -hline!([s95], linecolor=:orangered, linestyle=:dash, label="") -hline!([t95], linecolor=:orangered, linestyle=:dash, label="") -savefig("scores-vs-rand.png") -savefig("test-vs-rand.png") - - - -xlabels = ["optimized selection,\ntargeted variables" "random selection,\ntargeted variables"] -vals = [log.(s1) log.(s2)] -crit_val = s95 - -xlabels = ["optimized selection,\nuntargeted variables" "random selection,\nuntargeted variables"] -vals = [log.(t1) log.(t2)] -crit_val = t95 - - -dotplot(xlabels, vals, label="",grid=false,ylabel="mismatch score",tick_direction=:out, - ylims=(1.0,9.1), - markerstrokewidth=0, markersize=1.5, markercolor=:black, markershape=:none, - markeralpha=0.5, - dpi=300) - -violin!(xlabels, vals, label="", - fillalpha=0.0,linealpha=0.5,linewidth=0.5,linecolor=:black) - -hline!([crit_val], linecolor=:orangered, linestyle=:dash, label="") - -savefig("scores-vs-rand.png") -savefig("test-vs-rand.png") - - - - -## testing income segregated workplaces - -cbg_idxs = dser_path("jlse/cbg_idxs.jlse") -hh = dser_path("jlse/hh.jlse") -gqs = dser_path("jlse/gqs.jlse") -people = dser_path("jlse/people.jlse") -df_gq_summary = dser_path("jlse/df_gq_summary.jlse") - -length(values(people)) -sum(p.working for p in values(people)) -sum(p.commuter for p in values(people)) -sum(p.com_LODES_low for p in values(people)) -sum(p.com_LODES_high for p in values(people)) -sum(coalesce(p.female,rand([true,false])) for p in values(people)) -mean(p.age for p in values(people)) - -sum(df_gq_summary.noninst1864) -sum(df_gq_summary.working) -sum(df_gq_summary.not_working) -sum(df_gq_summary.commuter) -sum(df_gq_summary.wfh) -sum(df_gq_summary.com_LODES_low) -sum(df_gq_summary.com_LODES_high) - -gq_ppl = [people[k] for k in reduce(vcat, [x.residents for x in values(gqs)])] -sum(df_gq_summary.instu18 + df_gq_summary.inst1864 + df_gq_summary.noninst1864 + df_gq_summary.inst65o) - -sum(p.working for p in values(gq_ppl)) -sum(p.commuter for p in values(gq_ppl)) -sum(p.com_LODES_low for p in values(gq_ppl)) -sum(p.com_LODES_high for p in values(gq_ppl)) -mean(p.age for p in values(gq_ppl)) - -df = read_df("processed/work_od_matrix.csv"; types=Dict("h_cbg"=>String15)) -df_noinc = read_df("processed/work_od_matrix_no_inc.csv"; types=Dict("h_cbg"=>String15)) -df -df_noinc -df[!,:outside] -df_noinc[!, :outside] - -hh_samps = read_df("processed/hh_samples.csv"; types=Dict("SERIALNO"=>String15)) -hh_samps[:, - [:SERIALNO,:NP,:HINCP, - #:employed, :unemployed, :armed_forces, - :in_lf, - :nilf, :work_from_home, :commuter, - #:worked_past_yr, - :has_job,:com_LODES_low,:com_LODES_high]] - -p_samps = read_df("processed/p_samples.csv"; types=Dict("SERIALNO"=>String15)); -p_samps[: , [:SERIALNO,:WAGP,:PINCP,:PERNP,:COW,:ESR,:work_from_home,:commuter,:has_job,:com_LODES_low,:com_LODES_high]] - - - - - - -## testing stochastic block model - -using Graphs -using GraphPlot -using LinearAlgebra - -n_k = 2 -mean_deg = 8 -n1 = 18 -n2 = 9 -n_vec = [n1,n2] - -w_planted = Diagonal(fill(mean_deg,n_k)) -prop_i = n_vec ./ sum(n_vec) -w_random = repeat(transpose(prop_i) * mean_deg, n_k) - -assoc_coeff = 0.9 -c_matrix = assoc_coeff * w_planted + (1 - assoc_coeff) * w_random - -## can't have more than n connections to a group -for r in eachrow(c_matrix) - r .= min.(r,n_vec) -end -c_matrix - -## can't have more than n-1 within-group connections -d_tmp = view(c_matrix, diagind(c_matrix)) -d_tmp .= min.(d_tmp, n_vec .- 1) -c_matrix - -## note, in a stochastic block model (SBM) the degree distribution within blocks is similar to -## a random (erdos-renyi) network. For a more realistic degree distribution (e.g., with hubs) -## try a degree-corrected SBM (Karrer & Newman 2010, https://arxiv.org/abs/1008.3926) -g = stochastic_block_model(c_matrix,n_vec) -gplot(g) - -## Graphs.jl sometimes creates 0-degree vertices but I don't think it should -fix_zeros = findall(degree(g).==0) -for x in fix_zeros - add_edge!(g, x, rand(vertices(g))) -end - -mean(degree(g)) -mean(degree(g)[1:n1]) -mean(degree(g)[(n1+1):end]) - -c_matrix - -g2 = stochastic_block_model(reshape([8.0], 1, 1),[30]) -gplot(g2) -mean(degree(g2)) - - - - -## making sure all schools are getting students allocated - -sch_students = dser_path("jlse/sch_students.jlse") - -schools = read_df("processed/schools.csv"; types=Dict("NCESSCH"=>String15)) -distmat = read_df("processed/cbg_sch_distmat.csv"; types=Dict("GEOID"=>String15)) -schools_idx = Dict(schools.NCESSCH .=> eachindex(schools.NCESSCH)) -cap_by_school = Dict(k=>schools[schools_idx[k], "STUDENTS"] for (k,_) in sch_students) -n_students_by_school = Dict(k=>length(v) for (k,v) in sch_students) - -sum(values(cap_by_school)) -sum(values(n_students_by_school)) -sum(values(n_students_by_school)) / sum(values(cap_by_school)) - -spots_left = Dict(k=>(cap_by_school[k]-n_students_by_school[k]) for (k,_) in sch_students) -p_filled = Dict(k=>(n_students_by_school[k] / cap_by_school[k]) for (k,_) in sch_students) -p_new = Dict(k=>(n_students_by_school[k] / cap_by_school[k]) for (k,_) in sch_students) - -partialsort(collect(values(p_filled)),1:20) -partialsort(collect(values(p_new)),1:20) - -sort(collect(values(p_filled)))[end-20:end] - -using Plots -plot!(sort(collect(values(p_filled)))) -plot!(sort(collect(values(p_new)))) - -sortperm(collect(values(p_filled))) -cap_by_school[collect(keys(p_filled))[601]] - - -## checking census data integrity - -c10 = read_df("geo/2010_Census_Tract_to_2010_PUMA.txt";types=String) -acs_hh = read_df("bak/hh_all.csv",types=Dict("Geo"=>String)) -od_matrix = read_df("processed/work_od_matrix_no_inc.csv"; types=Dict("h_cbg"=>String)) -dConfig = tryJSON("geos.json") -geos = String.(dConfig["geos"]) - -c10[!,"Tract"] = c10.STATEFP .* c10.COUNTYFP .* c10.TRACTCE -c10[!,"County"] = c10.STATEFP .* c10.COUNTYFP -c10 = c10[[(r.STATEFP in geos || r.County in geos) for r in eachrow(c10)], ["Tract","PUMA5CE"]] - -cbgs_in_acs = acs_hh.Geo -tracts_in_acs = unique( [x[1:end-1] for x in acs_hh.Geo] ) -tracts_in_c10 = c10.Tract -all(sort(tracts_in_acs) .== sort(tracts_in_c10)) - -cbg_work_dests = names(od_matrix)[2:end-1] -cbg_origins = od_matrix.h_cbg -all([in(x,cbgs_in_acs) for x in cbg_work_dests]) -all([in(x,cbgs_in_acs) for x in cbg_origins]) - -mask = [!in(x,cbg_origins) for x in cbgs_in_acs] -not_in_origins = cbgs_in_acs[mask] -acs_hh[[r.Geo in not_in_origins for r in eachrow(acs_hh)],:] - - -## test location-based contact matrices - -loc_mat_keys = dser_path("jlse/loc_mat_keys.jlse") -work_loc_contact_mat = dser_path("jlse/work_loc_contact_mat.jlse") -res_loc_contact_mat = dser_path("jlse/res_loc_contact_mat.jlse") -work_loc_lookup = dser_path("jlse/work_loc_lookup.jlse") -res_loc_lookup = dser_path("jlse/res_loc_lookup.jlse") - -w = dser_path("jlse/company_workers.jlse") ## employers/employees (with work locations) -hh = dser_path("jlse/hh.jlse") ## households/residents (with hh locations) -cbg_idxs = dser_path("jlse/cbg_idxs.jlse") ## location (cbg) keys used in person/hh keys -cbg_idxs = Dict(k=>String31(v) for (k,v) in cbg_idxs) -gqs = dser_path("jlse/gqs.jlse") ## group-quarters/residents (with gq locations) -## assume only non-inst GQ residents are available for ephemeral local contacts -gq_noninst = filterv(x->x.type==:noninst1864, gqs) -## use the same matrix indices as in the regular contact networks -k = dser_path("jlse/adj_mat_keys.jlse") -p_idxs = Dict(k .=> eachindex(k)) - -## group potential encounters by census tract (CBG seems too restrictive) -hh_tracts = unique(x[1:end-1] for x in values(cbg_idxs)) -work_tracts = unique(x[3][1:end-1] for x in keys(w)) -tracts = sort(unique([hh_tracts; work_tracts])) - -loc_mat_keys_rev = Dict(v=>k for (k,v) in loc_mat_keys) - -i = 100 -t = findnz(work_loc_contact_mat[:,i])[1] -all(work_loc_lookup[x] == i for x in t) -loc_mat_keys_rev[i] -filterk(x->x[3][1:end-1]==loc_mat_keys_rev[i], w) -reduce(vcat, values(filterk(x->x[3][1:end-1]==loc_mat_keys_rev[i], w))) -t2 = [p_idxs[k[1:3]] for k in reduce(vcat, values(filterk(x->x[3][1:end-1]==loc_mat_keys_rev[i], w)))] -all(sort(t) .== sort(t2)) - -i = 123 -t = findnz(res_loc_contact_mat[:,i])[1] -all(res_loc_lookup[x] == i for x in t) -loc_mat_keys_rev[i] -filterk(x->cbg_idxs[x[2]][1:end-1]==loc_mat_keys_rev[i], hh) -tk1 = reduce(vcat, [h.people for h in values(filterk(x->cbg_idxs[x[2]][1:end-1]==loc_mat_keys_rev[i], hh))]) -filterk(x->cbg_idxs[x[2]][1:end-1]==loc_mat_keys_rev[i], gq_noninst) -tk2 = reduce(vcat, [x.residents for x in values(filterk(x->cbg_idxs[x[2]][1:end-1]==loc_mat_keys_rev[i], gq_noninst))]) -t2 = [p_idxs[k[1:3]] for k in [tk1;tk2]] -all(sort(t) .== sort(t2)) - diff --git a/src/geopops/julia/utils.jl b/src/geopops/julia/utils.jl deleted file mode 100644 index 36e71e2..0000000 --- a/src/geopops/julia/utils.jl +++ /dev/null @@ -1,182 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -using StatsBase -using InlineStrings -using SparseArrays - -## may change these -const CBGkey = UInt16 -const Hnum = UInt16 -const Pnum = UInt32 -const Hkey = Tuple{Hnum,CBGkey} -const Pkey = Tuple{Pnum,Hnum,CBGkey} -const GQkey = Tuple{UInt16,CBGkey} -const WRKkey = Tuple{UInt32, UInt8, String31} - -struct PersonData - hh::Hkey - sample::UInt32 - age::Int16 - working::Bool - commuter::Bool - com_cat::Union{Missing,UInt8} ## wp category (industry, etc) for placing commuters in wp's - com_inc::Union{Missing,UInt8} ## income category for commuters, to make wp contact networks - sch_grade::Union{Missing,String3} - #more_traits::Dict{Symbol,Union{Missing,Bool}} - ## ^^ that would be nice, but several million dicts will use a lot of memory, so... - ## (these have to be in the same order as "additional traits" in config.json) - sch_public::Union{Missing,Bool} - sch_private::Union{Missing,Bool} - female::Union{Missing,Bool} - race_black_alone::Union{Missing,Bool} - white_non_hispanic::Union{Missing,Bool} - hispanic::Union{Missing,Bool} -end - -struct Household - sample::UInt32 - people::Vector{Pkey} -end - -struct GQres - type::Symbol - residents::Vector{Pkey} -end - -## make dicts callable, with a closure to handle missing keys -(d::AbstractDict)(x,f::Base.Callable) = get(f,d,x) - -## subset of dict by keys -dsubset(d::Dict{K,V},s) where {K<:Any,V<:Any} = Dict{K,V}(k=>d[k] for k in intersect(keys(d),s)) - -mutable struct Indexer{I} - i::I -end -## constructor -Indexer{T}() where {T<:Number} = Indexer(zero(T)) -## an indexer object functions like a closure -## when called with a dict and key, returns that key's index, -## inserts the key with a new index if it doesn't exist -function (ix::Indexer{I})(d::Dict{K,I}, k::K) where {I<:Number, K<:Any} - if haskey(d,k) - return d[k] - else - ix.i += 1 - d[k] = ix.i - return ix.i - end -end - -## convert to integer, missing becomes 0 -int(x::T) where T<:Real = round(Int64, x) -int(x::Missing) = Int64(0) - -## convert x to type T, missing stays missing -mcon(::Type{T}, x::Missing) where {T<:Any} = missing -mcon(::Type{T}, x::U) where {T,U} = convert(T,x)::T - -## convert missing bool to false -mtrue(x::Union{Missing,Bool}) = coalesce(x,false) - -## filter dict on values -filterv(f, d::Dict) = filter( ((k,v),) -> f(v) , d) - -## filter dict on keys -filterk(f, d::Dict) = filter( ((k,v),) -> f(k) , d) - -## merge two dictionaries with vector values -vecmerge = mergewith(vcat) -## as above, but modifies the first dict in-place -vecmerge! = mergewith!(vcat) - -## "flattens" a dictionary whose values are vectors/collections -## returns a vector of pairs -dflat(d::Dict) = collect(Iterators.flatmap(x->((x.first => y) for y in x.second), d)) - -## continuous index ranges with lengths given by vec -function ranges(vec::Vector{I}) where {I<:Integer} - x = cumsum(vec) - return [a:b for (a,b) in zip([1;x[1:end-1].+1], x)] -end - -## returns first nonempty member of v -function first_nonempty(v) - i = findfirst(!isempty, v) - isnothing(i) ? empty(v) : v[i] -end - -## index of first true in v, otherwise missing -first_true(v) = something(findfirst(v),missing) - -## replace values less than threshhold with 0 -thresh(x,v) = x < v ? zero(x) : x - -## random lognormal -rlogn(mu::T, sigma::T) where T<:Real = exp(mu + sigma*randn()) - -## round a vector to integers while preserving sum -## (using largest-remainder method) -function lrRound(v::AbstractVector{T}) where T<:Real - vrnd = floor.(Int64, v) - verr = v .- vrnd - vrem = round(Int64, sum(v) - sum(vrnd)) - vidxs = sortperm(verr, rev=true) - for i in 1:vrem - vrnd[vidxs[i]] += 1 - end - return vrnd -end - -## make it work with matrices too -function lrRound(v::AbstractMatrix{T}) where T<:Real - orig_dims = size(v) - vrnd = lrRound(vec(v)) - return reshape(vrnd, orig_dims) -end - -## but round each row -function rowRound(m::AbstractMatrix{T}) where T<:Real - res = zeros(Int,size(m)) - for i in axes(m,1) - res[i,:] = lrRound(m[i,:]) - end - return res -end - -## but round each column -function colRound(m::AbstractMatrix{T}) where T<:Real - res = zeros(Int,size(m)) - for i in axes(m,2) - res[:,i] = lrRound(m[:,i]) - end - return res -end - -## sample from a vector of counts; returns an index and depletes the counts -## "AbstractArray" means this also works on 1D _views_ of a matrix -function drawCounts!(v::AbstractArray{I}) where {I<:Integer} - i = wsample(eachindex(v),v) ## wsample() from StatsBase - v[i] -= 1 ## modify v - return i -end - -## sample n from a vec of counts; returns vec of indices and depletes counts -function drawCounts!(v::AbstractArray{Ia}, n::Ib) where {Ia<:Integer,Ib<:Integer} - ## should probably throw an error if n > sum(v) - n = min(n,sum(v)) - res = zeros(Int64, n) - for i in 1:n - res[i] = drawCounts!(v) - end - return res -end diff --git a/src/geopops/julia/workplaces.jl b/src/geopops/julia/workplaces.jl deleted file mode 100644 index 59f36fc..0000000 --- a/src/geopops/julia/workplaces.jl +++ /dev/null @@ -1,634 +0,0 @@ -#= -Copyright 2023 Alexander Tulchinsky - -This file is part of Greasypop. - -Greasypop is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - -Greasypop is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. - -You should have received a copy of the GNU Affero General Public License along with Greasypop. If not, see . -=# - -using Random -using SparseArrays -using ProportionalFitting -using Logging -#using Statistics - -include("utils.jl") -include("fileutils.jl") - -## splits n into a vector that sums to n -## by drawing from a lognormal dist (mu, sigma) -## receives and updates a vector of unused draws -function split_lognormal!(n::Integer, mu::Float64, sigma::Float64, draws::Vector{I}) where I<:Integer - sizes = Vector{Int64}() - ## don't bother generating if less than: - while n > 2 - ## try to use an unused draw - i = findfirst(x->x<=n, draws) - if isnothing(i) - sz = ceil(Int64,rlogn(mu,sigma)) ## round up to avoid 0's - ## if it's too big, save for later (otherwise we're cheating at the distribution) - while sz > n - push!(draws,sz) - sz = ceil(Int64,rlogn(mu,sigma)) ## round up - end - else - sz = popat!(draws,i) - end - push!(sizes, sz) - n -= sz - end - ## append remaining, if any - if n > 0 - push!(sizes,n) - end - return sizes -end - -## reads people generated in households.jl -## returns {residence geo => list of worker keys} for each category -function group_commuters_by_origin(cat_codes::Vector{S}) where {S<:AbstractString} - ## within each category, group by cbg, randomize order - cbgs = dser_path("jlse/cbg_idxs.jlse") - - ## only commuters are assigned to workplaces - ## dataframe provides fast grouping - df = DataFrame([(id=k[1], hh=k[2], cbg=k[3], income=v.com_inc, category=v.com_cat) - for (k,v) in filterv(x->x.commuter, dser_path("jlse/people.jlse"))]) - df_grouped = groupby(df, [:category, :cbg]) - - ## a dict for each category, containing a dict of cbg => list of workers - ## each worker is identified by (id, household, cbg, income category) - ## -- the latter is for easier generating of income-assortative workplace networks - worker_keys = Dict(kc => - Dict(ko => Vector{Tuple{fieldtypes(Pkey)...,eltype(df.income)}}() - for ko in values(cbgs)) - for kc in cat_codes) - - ## fill the dict from dataframe - ## shuffle each list of workers so they can be assigned to workplaces in order - for gk in keys(df_grouped) - cat_code, cbg_code = cat_codes[gk[:category]], cbgs[gk[:cbg]] - worker_keys[cat_code][cbg_code] = shuffle( [(r.id,r.hh,r.cbg,r.income) for r in eachrow(df_grouped[gk])] ) - end - - return worker_keys -end - -## for generating data-less "dummy" people who commute from outside the synth pop -function dummy_gen_fn() - dummy_idx = 0 - ## each dummy needs to be provided an income code for assortative network construction - f = function(origin::AbstractString, inc_code::Integer) - dummy_idx += 1 - ## diagnostics (dummy should only be created when origin is "outside") - d_err = origin == "outside" ? 0 : 1 - ## dummies have no household or home cbg - return (dummy_idx,0,0,UInt8(inc_code)), d_err - end - return f -end - -## household sample summaries has # of commuters for every sample household -## files in CO/ have a list of households by geo code -## return a dict of {residence geo => # commuters} for each category in cat_codes -function read_workers_by_cat(cat_codes::Vector{<:AbstractString}, counties::Vector{<:AbstractString}) - cat_cols = ["com_ind_"*k for k in cat_codes] - hh_samps = read_df("processed/hh_samples.csv"; select=[["SERIALNO","NP","com_LODES_low","com_LODES_high"]; cat_cols], types=Dict("SERIALNO"=>String15)) - hh_idx = Dict(hh_samps.SERIALNO .=> eachindex(hh_samps.SERIALNO)) - - workers_by_cat_x_ori = Dict(k => Dict{String15,Int64}() for k in cat_codes) - for co in counties - cbg_dict = dser_path("jlse/CO/hh"*co*".jlse") ## read households generated for cbgs in county - for (ori, hhvec) in cbg_dict ## for each cbg in county, look up # workers in hh sample summaries - for (cat_code, cat_col) in zip(cat_codes, cat_cols) ## sum each category across households - workers_by_cat_x_ori[cat_code][ori] = sum(skipmissing(hh_samps[[hh_idx[x] for x in hhvec], cat_col])) - end - end - end - return workers_by_cat_x_ori -end - -## summary generated in households.jl has # of workers in each GQ -## -## assume that all residents of military quarters work at the quarters -## only civilians in non-inst GQs commute to jobs -## (this assumption is also in households.jl) -## -function read_gq_workers_by_cat(cat_codes::Vector{<:AbstractString}) - gq_df = dser_path("jlse/df_gq_summary.jlse") - cat_cols = ["ind_"*k for k in cat_codes] - - gq_by_cat_x_ori = Dict(k => Dict{String15,Int64}() for k in cat_codes) - for r in eachrow(gq_df) - for (cat_code, cat_col) in zip(cat_codes, cat_cols) - gq_by_cat_x_ori[cat_code][r.geo] = r[cat_col] - end - end - return gq_by_cat_x_ori -end - -## read commute matrix for category k -## data is in sparse format, so m and n are needed to determine dimensions -function read_od_matrix(k::AbstractString, m::Integer, n::Integer) - df = read_df("processed/od_"*k*".csv.gz"; types=Dict("origin"=>UInt32,"dest"=>UInt32,"p"=>Float32)) - return sparse(df.origin, df.dest, df.p, m, n) -end - -## counts of workers commuting from outside the synth area -function read_outside_origins(cat_codes::Vector{S}) where {S<:AbstractString} - tmp_dict = let df = read_df("processed/work_cats_live_outside.csv"); Dict(df[!,1] .=> df[!,2]); end - return Dict(k => round(Int,tmp_dict["C24030:"*k]) for k in cat_codes) -end - -## calculate origin-destination counts for each category in cat_codes -## write to file (don't need them all in mem at once) -## also keep hh and gq separate for now, in case they need to be treated differently -function calc_od_counts(cat_codes::Vector{<:AbstractString}, counties::Vector{<:AbstractString}) - - ## these vectors map row and col idxs in the od matrices to geo codes in the population - origin_labels = let df = read_df("processed/od_rows_origins.csv"); df.origin; end - dest_labels = let df = read_df("processed/od_columns_dests.csv"); df.dest; end - n_rows = length(origin_labels) - n_cols = length(dest_labels) - origin_idx = Dict(origin_labels .=> eachindex(origin_labels)) - - ## read # workers from hh sample summary, workers living in gq's from gq summary - ## these are dicts of geo code => count for each category in cat_codes - println("reading worker counts") - hhw_by_cat_x_ori = read_workers_by_cat(cat_codes,counties) - gqw_by_cat_x_ori = read_gq_workers_by_cat(cat_codes) - - println("# workers living in synth area = ", sum(values(reduce(mergewith(+), values(hhw_by_cat_x_ori)))) + - sum(values(reduce(mergewith(+), values(gqw_by_cat_x_ori))))) - - ## need # workers with origin "outside"; append to hh workers dict - outside_by_cat = read_outside_origins(cat_codes) - for k in cat_codes - hhw_by_cat_x_ori[k]["outside"] = outside_by_cat[k] - end - - println("# workers living outside = ", sum(values(outside_by_cat))) - - ## multiply workers by od matrix - println("calculating origin-destination counts") - test_tot = 0 - for k in cat_codes - println(" for category ",k) - M = read_od_matrix(k, n_rows, n_cols); - hh_counts = SparseMatrixCSC{UInt32, UInt32}(spzeros(n_rows, n_cols)) - gq_counts = SparseMatrixCSC{UInt32, UInt32}(spzeros(n_rows, n_cols)) - for (code, rownum) in origin_idx - hh_counts[rownum,:] = lrRound(M[rownum,:] .* get(hhw_by_cat_x_ori[k],code,0)) - gq_counts[rownum,:] = lrRound(M[rownum,:] .* get(gqw_by_cat_x_ori[k],code,0)) - end - ser_path("jlse/od_counts_"*k*".jlse",(hh_counts, gq_counts)) - test_tot += sum(hh_counts) + sum(gq_counts) - end - - println("# workers assigned to destinations = ", test_tot) - return (origin_labels, dest_labels) -end - - -## employer size stats by county -function read_county_stats() - cols = Dict("county"=>String7,"mu_ln"=>Float64,"sigma_ln"=>Float64) - county_stats = read_df("processed/work_sizes.csv"; select=collect(keys(cols)), types=cols) - return Dict(county_stats.county .=> zip(county_stats.mu_ln, county_stats.sigma_ln)) -end - -## n people working in schools and school locations -function read_school_info() - ## find the closest cbg for each school - distmat = read_df("processed/cbg_sch_distmat.csv"; types=Dict("GEOID"=>String15)) - closest_cbg = Dict(x => distmat[argmin(distmat[!,x]),"GEOID"] for x in String15.(names(distmat)[2:end])) - - ## number of teachers by school - sch_cols = Dict("NCESSCH"=>String15,"TEACHERS"=>Int64,"STUDENTS"=>Int64) - schools = read_df("processed/schools.csv"; select=collect(keys(sch_cols)), types=sch_cols) - sch_n_teachers = Dict(schools.NCESSCH .=> schools.TEACHERS) - - return (sch_n_teachers, closest_cbg) -end - -## n people working in gqs and gq locations -function read_gq_info() - ## find cbg where each gq is located - gqs = dser_path("jlse/gqs.jlse") - cbgs = dser_path("jlse/cbg_idxs.jlse") - gq_cbgs = Dict(k => cbgs[k[2]] for k in keys(gqs)) - ## inst qg's: assume 1 workers per 10 residents; non-inst, 1 per 50? mininum 2? - dConfig = tryJSON("config.json") - r_i::Float64 = float(get(dConfig, "inst_res_per_worker", 10)) - r_ni::Float64 = float(get(dConfig, "noninst_res_per_worker", 50)) - e_min::Int = get(dConfig, "min_gq_workers", 2) - ni_types = Set([:milGQ, :ninst1864civ]) - - gq_n_emps = Dict(k => - max(e_min, ceil(Int64, v.type in ni_types ? length(v.residents)/r_ni : length(v.residents)/r_i)) - for (k,v) in gqs) - - return (gq_n_emps, gq_cbgs) -end - -## filter work destinations by geo code and boolean function -## dest_idx is dict of geo code => index -## return indices -function filter_dests(f, geo::Sa, dest_idx::Dict{Sb, I}) where {I<:Integer,Sa<:AbstractString,Sb<:AbstractString} - return collect(values(filter( ((k,v),) -> (startswith(k,geo) && f(v)) , dest_idx))) -end - -## commute origins for people working in schools and gq's -## modifies count_matrix -function pull_inst_workers!(count_matrix::Matrix{<:Integer}, - dest_idx::Dict{<:AbstractString, <:Integer}, origin_labels::Vector{Ko}, - workers_by_key::Dict{Ki, <:Integer}, loc_by_key::Dict{Ki, <:AbstractString}) where {Ki<:Any,Ko<:AbstractString} - - ## for each institution, pull n workers who have a suitable work _destination_ - #test = Dict(k=>[] for k in keys(workers_by_key)) - inst_emp_origins = Dict{Ki,Vector{Ko}}() - for (inst_id, n) in workers_by_key - cbg = loc_by_key[inst_id] - ## find destination(s) with enough workers; try closest cbg first, then wider areas - ## limit to same county - ## employment numbers may not agree exactly with commute numbers - ## also, official employment location in commute data may not be the inst's actual location - ## draw as many workers as possible going to the actual location - ## then make up the remainder from nearby locations - ## this way, commuters are coming from approximately the right area - geo_areas = [cbg,cbg[1:11],cbg[1:9],cbg[1:7],cbg[1:5]] - ## column sums change on each loop - colsums = vec(sum(count_matrix;dims=1)) - ## for each geo area, a list of destination indices (columns) with > 0 workers: - dest_lists = [filter_dests(i->(colsums[i]>0), geo, dest_idx) for geo in geo_areas] - ## shuffle each list, but preserve the geo area order: - avail_cols = unique(reduce(vcat, map(shuffle, dest_lists))) - ## draw from columns until enough workers - o_idxs = Int[] - for col in avail_cols - draw_n = min(colsums[col], n) - #push!(test[inst_id], (col,draw_n)) - ## pass a view of the column so the original matrix gets modified - append!(o_idxs, drawCounts!(view(count_matrix,:,col), draw_n)) - n = n - draw_n - if n < 1 - break - end - end - if n > 0 - println("warning: inst $inst_id loc $cbg short by $n workers") - end - - inst_emp_origins[inst_id] = origin_labels[o_idxs] - end - return inst_emp_origins -end - -## group by destination, split into workplaces, assign origins -## modifies count_matrix in place -## returns a dict of workplace ids => employee cbgs -function generate_workplaces!(count_matrix::Matrix{<:Integer}, dest_idx::Dict{<:AbstractString, <:Integer}, - origin_labels::Vector{Ko}, county_stats::Dict{<:AbstractString, Tuple{Float64, Float64}}, - draws_by_county::Dict{<:AbstractString, Vector{I}}, - cat_idx::Integer) where {Ko<:AbstractString,I<:Integer} - - work_origins = Dict{WRKkey, Vector{Ko}}() - for (co,draws) in draws_by_county - (mu,sigma) = county_stats[co] - ## work destinations in the county we just read the stats for - dests = filterk(k->k[1:5]==co, dest_idx) - for (dest_code, col) in dests - n = sum(count_matrix[:,col]) ## number of workers in the dest - if n > 0 - sizes = split_lognormal!(n,mu,sigma,draws) - for (work_i, emp_size) in enumerate(sizes) - ## get rand sample of origins; pass a view so original matrix is modified - o_idxs = drawCounts!(view(count_matrix,:,col), emp_size) - ## create workplace and assign origins sampled - work_origins[(work_i, cat_idx, dest_code)] = origin_labels[o_idxs] - end - end - end - #println(draws) - end - return work_origins -end - -## make a separate work destination for each person working outside the synth area -## (no need to try to group them, as they have no work network anyway) -function generate_outside_workplaces(work_outside::Dict{K, <:Integer}, cat_idx::Integer) where {K<:Any} - return Dict([WRKkey((i, cat_idx, "outside")) for i in 1:sum(values(work_outside))] - .=> map(x->Vector{K}([x]), reduce(vcat, [fill(k, v) for (k, v) in work_outside]))) -end - -## assigns people in workers_by_origin to employers in emp_origins -## modifies cidx_by_origin so that it can be called with several emp_origins on the same pop -## creates "dummies" -- workers from outside that don't exist in people data -## returns a dict of employer id => vector of worker ids, and some diagnostics -function assign_workers!(emp_origins::Dict{T,Vector{K}}, workers_by_origin::Dict{K, Vector{W}}, - cidx_by_origin::Dict{K, I}, dummy_fn::F) where {T<:Any,K<:Any,W<:Any,I<:Integer,F<:Function} - - n_by_origin = Dict(k => length(v) for (k,v) in workers_by_origin) - ## p high income workers, for generating dummies - p_LODES_high = sum(count.(x->x[4]==2, values(workers_by_origin))) / sum(values(n_by_origin)) - - ## initialize empty lists - workers = Dict(est_id=>Vector{W}() for est_id in keys(emp_origins)) - dummies = Vector{W}() ## keep track of dummies created - - missing_origin = 0 - ran_out = Dict{K,Int64}() - for (e_id, origin_vec) in emp_origins - for origin_key in origin_vec - ## if the origin is not in cidx, it's outside the synth area and a dummy must be created - if haskey(cidx_by_origin, origin_key) - i = cidx_by_origin[origin_key] += 1 ## this doesn't look like it should be allowed, but it is lol - if i > n_by_origin[origin_key] - ran_out[origin_key] = get!(ran_out, origin_key, 0) + 1 ## this shouldn't happen - else - push!(workers[e_id], (workers_by_origin[origin_key][i])) - end - else - ## create dummy - inc_code = rand() < p_LODES_high ? 2 : 1 - dum::W, d_err::Int = dummy_fn(origin_key, inc_code) - push!(dummies, dum) - push!(workers[e_id], dum) - ## diagnostics (all origins except "outside" should have been in cidx) - missing_origin += d_err - end - end - end - return (workers, dummies, missing_origin, ran_out) -end - - -function generate_jobs_and_workers() - - wp_codes = tryJSON("processed/codes.json") - ind_codes::Vector{String} = get(wp_codes, "ind_codes", String[]) - ind_idxs = Dict(ind_codes .=> eachindex(ind_codes)) - ser_path("jlse/wp_cat_codes.jlse",ind_idxs) - counties = let cbgs = dser_path("jlse/cbg_idxs.jlse"); unique(map(x->x[1:5], values(cbgs))); end - - println("reading commuters in synth pop") - worker_keys = group_commuters_by_origin(ind_codes) - ## if origin is outside the synth area, will generate a dummy instead of assigning a resident - dummy_fn = dummy_gen_fn() - - ## calculate origin-destination counts and save to files - ## "labels" vectors map row and col idxs in the od matrices to geo codes in the population - ## workplaces will be generated based on the od matrices - ## then connected to people/locations in the synth pop using the label indices - origin_labels, dest_labels = calc_od_counts(ind_codes, counties) - ser_path("jlse/od_labels.jlse",(origin_labels,dest_labels)) - - dest_idx = Dict(dest_labels .=> eachindex(dest_labels)) ## for easy lookup - - ## mean and sd of ln(workplace size) for each county - county_stats = read_county_stats() - ## adjust stats - ## slightly overproduces large employers compared to CBP stats - ## (but unlisted employers should probably skew large) - stats2 = deepcopy(county_stats) - for (k,v) in stats2 - stats2[k] = (v[1], v[2]+0.1) - end - - ## save employer size draws by county, to make best use of county-level stats - draws_by_county = Dict(k => Vector{Int}() for k in counties) - - ## special treatment for these: - school_teacher_category = "EDU" - gq_employee_category = "ADM_MIL" - ## could make gq workers more accurate by handling each type of gq separately - ## but there are not very many of them, and _most_ will be prison workers (ind #92) - (sch_n_emps, sch_cbgs) = read_school_info() - (gq_n_emps, gq_cbgs) = read_gq_info() - - ## data structures for collecting results - sch_workers = Dict() - gq_workers = Dict() - company_workers = Dict() - outside_workers = Dict() - dummies = Dict() - missing_origins = Dict() - ran_out = Dict() - - ## do the following by category - for ckey in ind_codes - - println("generating workplaces for category "*ckey) - - ## id's of synth pop residents to assign to workplaces - workers_by_origin = worker_keys[ckey] - ## instead of removing workers as they're assigned, just keep pointers to the first available worker - cidx_by_origin = Dict(keys(workers_by_origin) .=> 0) - - ## origin-destination matrix for current category - hh_sparse, gq_sparse = dser_path("jlse/od_counts_"*ckey*".jlse"); - od_counts = Matrix(hh_sparse+gq_sparse) - ## pull off the last column (work destination outside the synth area) - work_outside_counts = od_counts[:,end] - od_counts = od_counts[:,1:end-1] - ## by origin code, # people working outside the synth area - work_outside = Dict(origin_labels .=> work_outside_counts) - - ## check totals (before modifying od_counts) - n_jobs_in_synth_area = sum(od_counts) - n_jobs_outside_area = sum(work_outside_counts) - n_jobs_od_total = n_jobs_in_synth_area + n_jobs_outside_area - n_com_from_synth_area = sum(od_counts[1:end-1,:]) + sum(work_outside_counts) - n_com_from_outside = sum(od_counts[end,:]) - - ## pull out workers for schools and gqs for each destination (except "outside") - ## (assumes workers from out of state can work in public schools) - ## samples worker origins in proportion to counts - ## currently does not consider worker income, age, etc. - if ckey == school_teacher_category - sch_emp_origins = pull_inst_workers!(od_counts, dest_idx, origin_labels, sch_n_emps, sch_cbgs) - (sch_workers[ckey], dummies["sch"*ckey], missing_origins["sch"*ckey], ran_out["sch"*ckey]) = assign_workers!( - sch_emp_origins, workers_by_origin, cidx_by_origin, dummy_fn) - n_sch_jobs_generated = sum(length.(values(sch_emp_origins))) - n_sch_workers_assigned = sum(length.(values(sch_workers[ckey]))) - else - n_sch_jobs_generated = 0 - n_sch_workers_assigned = 0 - end - - if ckey == gq_employee_category - gq_emp_origins = pull_inst_workers!(od_counts, dest_idx, origin_labels, gq_n_emps, gq_cbgs) - (gq_workers[ckey], dummies["gq"*ckey], missing_origins["gq"*ckey], ran_out["gq"*ckey]) = assign_workers!( - gq_emp_origins, workers_by_origin, cidx_by_origin, dummy_fn) - n_gq_jobs_generated = sum(length.(values(gq_emp_origins))) - n_gq_workers_assigned = sum(length.(values(gq_workers[ckey]))) - else - n_gq_jobs_generated = 0 - n_gq_workers_assigned = 0 - end - - ## create workplaces for remaining workers for each destination (except "outside") - ## and assign workers to those - work_origins = generate_workplaces!(od_counts, dest_idx, origin_labels, stats2, draws_by_county, ind_idxs[ckey]) - (company_workers[ckey], dummies[ckey], missing_origins[ckey], ran_out[ckey]) = assign_workers!( - work_origins, workers_by_origin, cidx_by_origin, dummy_fn); - - ## check results - n_unused_draws = sum(length.(values(draws_by_county))) - n_jobs_generated_within = sum(length.(values(work_origins))) + n_sch_jobs_generated + n_gq_jobs_generated - n_workers_living_within = sum(length.(values(workers_by_origin))) - println(" # jobs within synth area = $n_jobs_in_synth_area outside = $n_jobs_outside_area total = $n_jobs_od_total") - println(" # commuting from outside = $n_com_from_outside resident workers needed = $n_com_from_synth_area") - println(" # jobs generated = $n_jobs_generated_within workers in synth area = $n_workers_living_within") - println(" (# unused wp size draws = $n_unused_draws)") - println(" # workers assigned to jobs = " ,sum(length.(values(company_workers[ckey]))) + n_gq_workers_assigned + n_sch_workers_assigned) - - ## don't forget to assign people to work_outside origins - n_by_origin = Dict(k => length(v) for (k,v) in workers_by_origin) - unused = Dict(k => max(0,(v - cidx_by_origin[k])) for (k,v) in n_by_origin) - println(" # remaining workers = ", sum(values(unused))) - - (outside_workers[ckey], _, missing_origins["out"*ckey], ran_out["out"*ckey]) = assign_workers!( - generate_outside_workplaces(work_outside, ind_idxs[ckey]), workers_by_origin, cidx_by_origin, dummy_fn); - ## check results - println(" # jobs created outside synth area = ", sum(length.(values(outside_workers[ckey])))) - end - - ## merge categories, save to file for next step - println("writing results to file") - suffix = "" - ser_path("jlse/work_dummies"*suffix*".jlse", reduce(vcat, values(dummies))) - ser_path("jlse/sch_workers"*suffix*".jlse",reduce(vecmerge, values(sch_workers))) - ser_path("jlse/gq_workers"*suffix*".jlse",reduce(vecmerge, values(gq_workers))) - ser_path("jlse/company_workers"*suffix*".jlse",reduce(vecmerge, values(company_workers))) - ser_path("jlse/outside_workers"*suffix*".jlse",reduce(vecmerge, values(outside_workers))) - - ## check results - #reduce(+, values(missing_origins)) - #reduce(mergewith(+), values(ran_out)) - company_worker_counts = length.(values(reduce(vecmerge, values(company_workers)))) - sch_worker_counts = length.(values(reduce(vecmerge, values(sch_workers)))) - gq_worker_counts = length.(values(reduce(vecmerge, values(gq_workers)))) - outside_worker_counts = length.(values(reduce(vecmerge, values(outside_workers)))) - origin_worker_counts = length.(values(reduce(vecmerge, values(worker_keys)))) - n_workers_from_outside = sum(values(read_outside_origins(ind_codes))) - println("") - println("# workers living in synth pop = ", sum(origin_worker_counts)) - println("# workers commuting from outside = ", n_workers_from_outside) - println("total workers = ", sum(origin_worker_counts) + n_workers_from_outside) - println("# company employees = ", sum(company_worker_counts)) - println("# school employees = ", sum(sch_worker_counts), " expected ",sum(values(sch_n_emps))) - println("# gq employees = ", sum(gq_worker_counts), " expected ",sum(values(gq_n_emps))) - println("# working outside synth pop = ", sum(outside_worker_counts)) - println("total assigned to jobs = ", sum(company_worker_counts)+sum(sch_worker_counts)+sum(gq_worker_counts)+sum(outside_worker_counts)) - println("") - println("generated establishment sizes (other than schools and group quarters):") - println("size<5: ",count(x->x<5, company_worker_counts)) - println("4(4(9(19(49(99(249(499(999(1499(24994999: ",count(x->x>4999, company_worker_counts)) - - return nothing -end - - -## generate commute matrices by industry using IPF -## based on census and LODES data -## needed before workplaces can be generated -function generate_commute_matrices() - println("generating commute matrices") - Logging.disable_logging(Logging.Info) - - ind_codes::Vector{String} = let wp_codes = tryJSON("processed/codes.json"); get(wp_codes, "ind_codes", String[]); end; - - ## n of workers by industry at each origin, from census data (assumed true) - (total_by_ori, m_ind_ori) = let io_df = read_df("processed/work_io_sums.csv"; types=Dict("Geo"=>String15)); - (sum(Matrix(io_df[!,2:end]);dims=2), - permutedims( rowRound(Matrix(io_df[!,2:end])) , (2,1) )); - end; - - ## n workers commuting from each origin to each dest; census counts * proprtions from LODES OD data - (origin_idxs, dest_idxs, m_dest_ori) = let od_df = read_df("processed/work_od_prop.csv"; types=Dict("Geo"=>String15)) - (od_df[!,:Geo], names(od_df)[2:end], - sparse(permutedims( rowRound(total_by_ori .* Matrix(od_df[!,2:end])) , (2,1) ))); - end; - - ## p ind at each dest, estimated from WAC data - m_ind_dest_p = let id_est_df = read_df("processed/work_id_est_sums.csv"; types=Dict(1=>String15)); - m_ind_dest = permutedims( Matrix(id_est_df[!,2:end]), (2,1) ); - m_ind_dest ./ sum(m_ind_dest; dims=1); ## p ind at each dest - end; - - ## for storing results - res_iod = [spzeros(Float32, length(origin_idxs), length(dest_idxs)) for i in ind_codes]; - - for (o,x) in enumerate(origin_idxs) - o % 100 == 0 && println(o, "/", length(origin_idxs)) - - ## row sum targets = workers in each industry at this origin - ind_margin = m_ind_ori[:,o] - ## column sum targets = workers commuting to each destination from this origin - ## only need to consider non-zero columns - d_idxs, d_margin = findnz(m_dest_ori[:,o]) - if isempty(d_idxs) - ## no commute data exists for this origin; have to make something up - ## most likely, nobody lives there; show a warning if they do - sum(ind_margin) > 0 && println("warning: no commute data for ",x,", ",sum(ind_margin)," workers affected") - ## anyway, set destination equal to origin, or to a random dest preferably in same county - new_m = fill(1.0,(length(ind_margin),1)) - d_idxs = Int[something(findfirst(x->x==origin_idxs[o], dest_idxs), - rand(first_nonempty([findall(x->x[1:5]==origin_idxs[o][1:5], dest_idxs) - ,findall(x->true, dest_idxs)])))] - else - ## initial matrix based on p each industry at each destination, while preserving origin-destination counts - ## IPF will make it also preserve by-industry counts - ## this can be done separately for each origin; overall totals by destination and by industry will be correct - preserve_od_vals = (d_margin' .* m_ind_dest_p[:, d_idxs]); - ## do IPF; handles zeros poorly, so replace with some small value - init_m = max.(preserve_od_vals,0.00001); - fac = ipf(init_m, [ind_margin, d_margin], maxiter=500, tol=1e-6); - ## result = proportion going to each destination (by industry) - new_m = Array(fac) .* init_m - new_margin = sum(new_m; dims=2) ## will be same as ind_margin if IPF converged - new_m = new_m ./ new_margin - ## for industries with 0 workers at this location, still need some kind of commute data - ## assume they follow overall proportions from OD counts - new_m[isapprox.(vec(new_margin), 0.0), :] .= d_margin' / sum(d_margin) - end - ## store results - for (i,c) in enumerate(ind_codes) - res_iod[i][o,d_idxs] .= new_m[i,:] - end - end - - println("writing commute matrices") - write_df("processed/od_rows_origins.csv", DataFrame(:idx=>eachindex(origin_idxs),:origin=>origin_idxs)) - write_df("processed/od_columns_dests.csv", DataFrame(:idx=>eachindex(dest_idxs),:dest=>dest_idxs)) - for (i,k) in enumerate(ind_codes) - write_df("processed/od_"*k*".csv.gz", res_iod[i], [:origin,:dest,:p]; compress=true) - end - - return nothing -end - - - - - - - - - - diff --git a/src/geopops/networks.py b/src/geopops/networks.py index 6b0060e..6e766b2 100644 --- a/src/geopops/networks.py +++ b/src/geopops/networks.py @@ -5,7 +5,7 @@ import numpy as np import networkx as nx from scipy import sparse -from .utils import tryJSON, lrRound, vecmerge +from .utils import lrRound, vecmerge def connect_SBM(keyvec, K, min_N, assoc_coeff, use_groups=True, rng=None): @@ -136,7 +136,7 @@ def _assign_teachers_to_grades(school_key, students_by_grade, sch_workers_for_sc proportions = np.array([grade_counts[g] / total for g in grade_list]) n_per_grade = lrRound(proportions * n_teachers) teacher_grades = [] - for g, n in zip(grade_list, n_per_grade): + for g, n in zip(grade_list, n_per_grade, strict=False): teacher_grades.extend([g] * n) return [(t[0], t[1], t[2], teacher_grades[i] if i < len(teacher_grades) else '0') for i, t in enumerate(teachers)] @@ -234,7 +234,7 @@ def generate_networks(people, households, gqs, sch_students, company_workers, # Outside workers: people with workplace outside synth area adj_out_workers = {} - for wkey, wlist in outside_workers.items(): + for wlist in outside_workers.values(): for w in wlist: pk = w[:3] if pk in all_idxs: @@ -242,94 +242,3 @@ def generate_networks(people, households, gqs, sch_students, company_workers, return (adj_hh, adj_non_hh, adj_wp, adj_sch, adj_gq, adj_mat_keys, adj_dummy_keys, adj_out_workers) - - -def generate_location_matrices(company_workers, households, cbgs, gqs, adj_mat_keys, p_idxs): - """Generate location contact matrices for ephemeral contacts. - Returns (w_loc_mat, res_loc_mat, loc_idxs, w_loc_lookup, res_loc_lookup). - """ - cbgs_inv = {v: k for k, v in cbgs.items()} - ni_types = {'milGQ', 'ninst1864civ'} - n_people = len(adj_mat_keys) - - # Index all people - p_idx_map = {k: i for i, k in enumerate(adj_mat_keys)} - - # Group by census tract (CBG code minus last character) - hh_tracts = set() - for hk, hh in households.items(): - cbg_code = cbgs_inv.get(hk[1], '') - if cbg_code: - hh_tracts.add(cbg_code[:-1]) - - work_tracts = set() - for wk in company_workers.keys(): - if len(wk) >= 3 and isinstance(wk[2], str) and wk[2] != 'outside': - work_tracts.add(wk[2][:-1]) - - tracts = sorted(hh_tracts | work_tracts) - loc_idxs = {t: i for i, t in enumerate(tracts)} - n_tracts = len(tracts) - - # Workers by tract - w_rows, w_cols = [], [] - for wk, wlist in company_workers.items(): - if len(wk) < 3 or not isinstance(wk[2], str) or wk[2] == 'outside': - continue - tract = wk[2][:-1] - if tract not in loc_idxs: - continue - loc_i = loc_idxs[tract] - for w in wlist: - pk = w[:3] - if pk in p_idx_map: - w_rows.append(p_idx_map[pk]) - w_cols.append(loc_i) - - w_loc_mat = sparse.csr_matrix( - (np.ones(len(w_rows), dtype=bool), (w_rows, w_cols)), - shape=(n_people, n_tracts)) if w_rows else sparse.csr_matrix((n_people, n_tracts), dtype=bool) - - # Residents (HH + non-inst GQ) by tract - r_rows, r_cols = [], [] - for hk, hh in households.items(): - cbg_code = cbgs_inv.get(hk[1], '') - if not cbg_code: - continue - tract = cbg_code[:-1] - if tract not in loc_idxs: - continue - loc_i = loc_idxs[tract] - for pk in hh.people: - if pk in p_idx_map: - r_rows.append(p_idx_map[pk]) - r_cols.append(loc_i) - - for gk, gq in gqs.items(): - if gq.type not in ni_types: - continue - cbg_code = cbgs_inv.get(gk[1], '') - if not cbg_code: - continue - tract = cbg_code[:-1] - if tract not in loc_idxs: - continue - loc_i = loc_idxs[tract] - for pk in gq.residents: - if pk in p_idx_map: - r_rows.append(p_idx_map[pk]) - r_cols.append(loc_i) - - res_loc_mat = sparse.csr_matrix( - (np.ones(len(r_rows), dtype=bool), (r_rows, r_cols)), - shape=(n_people, n_tracts)) if r_rows else sparse.csr_matrix((n_people, n_tracts), dtype=bool) - - # Per-person lookups - w_loc_lookup = {} - for r, c in zip(w_rows, w_cols): - w_loc_lookup[r] = c - res_loc_lookup = {} - for r, c in zip(r_rows, r_cols): - res_loc_lookup[r] = c - - return w_loc_mat, res_loc_mat, loc_idxs, w_loc_lookup, res_loc_lookup diff --git a/src/geopops/pipeline.py b/src/geopops/pipeline.py new file mode 100644 index 0000000..3e3a7bf --- /dev/null +++ b/src/geopops/pipeline.py @@ -0,0 +1,61 @@ +"""Top-level pipeline orchestrator for GeoPops.""" + +from .config import make_config, validate_config +from .sources import download_data +from .census import process_data +from .population import generate_pop +from .starsim_bridge import to_starsim_people, starsim_networks + + +def run(config=None, *, seed=None, download=True, process=True, starsim=True, + verbose=1, **overrides): + """Run the whole GeoPops workflow and return the resulting population. + + Args: + config: a config dict (from :func:`geopops.make_config`). If omitted, one is + built from the packaged template plus `overrides`. + seed: master random seed. Overrides ``config['random_seed']``. + download: fetch raw Census/PUMS/LODES/school data. Set False to reuse data + already present in the run directory. + process: rebuild the processed CO targets and sample pools. + starsim: also build the Starsim ``People`` and network objects. + verbose: 0 for quiet, 1 for progress logging. + **overrides: config overrides such as ``geos=``, ``main_year=``, ``path=``. + + Returns: + GeneratePop: the completed run, with ``people``, ``households``, networks, + and (if `starsim`) ``ppl`` and ``networks`` attached. + + Example:: + + pop = geopops.run(geos=["45083"], main_year=2019, path="data", seed=42) + """ + if config is None: + config = make_config(**overrides) + elif overrides: + config = make_config(template=config, **overrides) + else: + validate_config(config) + + if seed is not None: + config["random_seed"] = seed + + if verbose: + print("Generating population with geopops.run()") + + if download: + download_data(config, verbose=verbose) + if process: + process_data(config, verbose=verbose) + + pop = generate_pop(config, seed=config.get("random_seed"), verbose=verbose) + + if starsim: + pop.ppl = to_starsim_people(pop.pop_export_dir, verbose=verbose) + pop.networks = starsim_networks(pop.pop_export_dir, + seed=config.get("random_seed") or 0, + verbose=verbose) + + if verbose: + print("\nPopulation generation complete") + return pop diff --git a/src/geopops/generate_pop.py b/src/geopops/population.py similarity index 82% rename from src/geopops/generate_pop.py rename to src/geopops/population.py index 5ea39d4..80ebe69 100644 --- a/src/geopops/generate_pop.py +++ b/src/geopops/population.py @@ -1,12 +1,15 @@ """ -GeneratePop orchestrator class — pure-Python replacement for RunJulia. -Calls co, households, schools, workplaces, networks, and export modules. +Population generation: combinatorial optimization, synthesis, and export. + +:func:`generate_pop` is the entry point; it returns a :class:`GeneratePop` holding +every pipeline intermediate so a run can be inspected or a single stage re-run. """ import os import json import numpy as np from collections import defaultdict from . import co, households, schools, workplaces, networks, export +from .exceptions import PipelineStateError # Package directory (src/geopops/) where config.json lives PACKAGE_DIR = os.path.dirname(os.path.abspath(__file__)) @@ -15,18 +18,22 @@ def load_config(base_dir=None): cfg_dir = base_dir if base_dir is not None else PACKAGE_DIR config_path = os.path.join(cfg_dir, "config.json") - with open(config_path, "r") as f: + with open(config_path) as f: return json.load(f) class GeneratePop: - """Orchestrates the synthetic population pipeline in pure Python. - Same API as RunJulia: CO(), SynthPop(), Export(), run_all(). + """A synthetic population run, holding all pipeline intermediates. + + Stages are ``CO()``, ``SynthPop()``, ``Export()``; ``run_all()`` does all three. + Prefer :func:`generate_pop` unless you want to drive the stages individually. """ def __init__(self, config_dict=None, base_dir=None, output_dir=None, - random_seed=None, verbose=1, auto_run=False, run_all=None): + random_seed=None, verbose=1, auto_run=False, run_all=None, + save_intermediates=True): self.verbose = verbose + self.save_intermediates = save_intermediates self.base_dir = base_dir if base_dir is not None else PACKAGE_DIR config = config_dict if config_dict is not None else load_config(self.base_dir) if output_dir is not None: @@ -59,6 +66,9 @@ def __init__(self, config_dict=None, base_dir=None, output_dir=None, self.adj_mat_keys = None self.adj_dummy_keys = None self.adj_out_workers = None + self._cbg_by_idx = None + self.ppl = None + self.networks = None if run_all is not None: auto_run = run_all @@ -91,7 +101,8 @@ def CO(self): """Run combinatorial optimization.""" self._log("*** Running GeneratePop.CO() ***") self.co_results, self.co_scores = co.process_counties( - self.data_dir, random_seed=self._stage_random_seeds["co"]) + self.data_dir, random_seed=self._stage_random_seeds["co"], + config=self.config, verbose=self.verbose) def _county_from_cbg_idx(self, cbg_idx): cbg_code = self._cbg_by_idx.get(cbg_idx) @@ -189,7 +200,7 @@ def _log_network_summary(self): def SynthPop(self): """Generate synthetic population (households, schools, workplaces, networks).""" if self.co_results is None: - raise RuntimeError("CO() must be run before SynthPop()") + raise PipelineStateError("CO() must be run before SynthPop()") self._log("\n*** Running GeneratePop.SynthPop() ***") self.cbgs, self.people, self.households, self.gqs, self.gq_summary = \ @@ -204,25 +215,22 @@ def SynthPop(self): self.sch_students = schools.generate_schools( self.people, self.cbgs, self.data_dir, - random_seed=self._stage_random_seeds["schools"]) + random_seed=self._stage_random_seeds["schools"], + config=self.config) self._log("\nGenerating schools") self._log_school_summary() self._log("\nGenerating workplaces") self._log("-- Generating OD matrices, exporting interim files") - self._log("-- processed/od_rows_origins.csv") - self._log("-- processed/od_columns_dests.csv") - codes_path = os.path.join(self.data_dir, "processed", "codes.json") - with open(codes_path, "r", encoding="utf-8") as f: - wp_codes = json.load(f) - for ind_code in wp_codes.get("ind_codes", []): - self._log(f"-- processed/od_{ind_code}.csv.gz") (self.company_workers, self.sch_workers, self.gq_workers, self.outside_workers, self.dummies) = \ workplaces.generate_jobs_and_workers( self.people, self.cbgs, self.gqs, self.co_results, self.gq_summary, self.data_dir, - random_seed=self._stage_random_seeds["workplaces"]) + random_seed=self._stage_random_seeds["workplaces"], + config=self.config, + save_intermediates=self.save_intermediates, + verbose=self.verbose) self._log_workplace_summary() self._log("\nGenerating networks") @@ -238,7 +246,7 @@ def SynthPop(self): def Export(self): """Export population and networks to CSV/MTX files.""" if self.people is None: - raise RuntimeError("SynthPop() must be run before Export()") + raise PipelineStateError("SynthPop() must be run before Export()") self._log("\n*** Running GeneratePop.Export() ***") self._log("") @@ -253,6 +261,11 @@ def Export(self): self.adj_dummy_keys, self.adj_out_workers, verbose=self.verbose) + @property + def pop_export_dir(self): + """Directory holding this run's exported population files.""" + return os.path.join(self.data_dir, "pop_export") + def run_all(self): """Run the complete pipeline: CO -> SynthPop -> Export.""" print("") @@ -263,3 +276,27 @@ def run_all(self): self.SynthPop() self.Export() + + +def generate_pop(config, *, seed=None, output_dir=None, base_dir=None, verbose=1, + save_intermediates=True): + """Generate a synthetic population: CO, then synthesis, then export. + + Args: + config: a config dict (see :func:`geopops.make_config`). + seed: master random seed; defaults to ``config['random_seed']``. + output_dir: where to write results; defaults to ``config['path']``. + verbose: 0 for quiet, 1 for progress logging. + save_intermediates: also write the interim ``processed/od_*.csv.gz`` files. + + Returns: + GeneratePop: the completed run, holding all pipeline intermediates. + + Example:: + + pop = geopops.generate_pop(cfg, seed=42) + pop.people, pop.households, pop.adj_hh + """ + return GeneratePop(config_dict=config, base_dir=base_dir, output_dir=output_dir, + random_seed=seed, verbose=verbose, auto_run=True, + save_intermediates=save_intermediates) diff --git a/src/geopops/run_all.py b/src/geopops/run_all.py deleted file mode 100644 index 3e0c0b5..0000000 --- a/src/geopops/run_all.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Top-level pipeline orchestrator for GeoPops.""" - -from .config import WriteConfig, load_config, update_config_values -from .download_data import DownloadData -from .process_data import ProcessData -from .generate_pop import GeneratePop -from .geopops_starsim import ForStarsim - -DEFAULT_ACS_REQUIRED = [ - "B01001", - "B09019", - "B09020", - "C24030", - "B23025", - "C24010", - "B11016", - "B11012", - "B23009", - "B11004", - "B19001", - "B22010", - "B09021", - "B09018", - "B11001H", - "B11001I", - "B25006", -] -DEFAULT_DEC_REQUIRED = ["P43", "P18"] - - -class RunAll: - """Run the full GeoPops workflow with a single call.""" - - def __init__(self, config_dict=None, pars=None, base_dir=None, verbose=1, auto_run=True): - # `pars` provides partial overrides; `config_dict` is treated as a full config. - self.pars = pars or {} - self.config_dict = config_dict - self.base_dir = base_dir - self.verbose = verbose - - if auto_run: - self.run_all() - - def _log(self, msg): - if self.verbose: - print(msg) - - def _build_effective_config(self): - if self.config_dict is not None: - config = self.config_dict - else: - config = load_config(self.base_dir) - update_config_values( - config, - census_api_key=self.pars.get("census_api_key"), - main_year=self.pars.get("main_year"), - geos=self.pars.get("geos"), - commute_states=self.pars.get("commute_states"), - use_pums=self.pars.get("use_pums"), - path=self.pars.get("path"), - julia_env_path=self.pars.get("julia_env_path"), - ) - # Backfill required table-code keys when config templates are minimal. - config.setdefault("acs_required", DEFAULT_ACS_REQUIRED.copy()) - config.setdefault("dec_required", DEFAULT_DEC_REQUIRED.copy()) - return config - - def run_all(self): - self._log("Generating population with RunAll()") - - effective_config = self._build_effective_config() - - WriteConfig(config_dict=effective_config, base_dir=self.base_dir) - - DownloadData( - config=effective_config, - base_dir=self.base_dir, - verbose=self.verbose, - auto_run=True, - ) - - ProcessData( - config_dict=effective_config, - base_dir=self.base_dir, - verbose=self.verbose, - auto_run=True, - ) - - GeneratePop( - config_dict=effective_config, - base_dir=self.base_dir, - verbose=self.verbose, - auto_run=True, - ) - - ForStarsim.People(config_dict=effective_config, base_dir=self.base_dir) - ForStarsim.GPNetwork(name='homenet', edge_weight=1.0) - ForStarsim.GPNetwork(name='schoolnet', edge_weight=1.0) - ForStarsim.GPNetwork(name='worknet', edge_weight=1.0) - ForStarsim.GPNetwork(name='gqnet', edge_weight=1.0) - - self._log("") - self._log("Population generation complete") diff --git a/src/geopops/schools.py b/src/geopops/schools.py index f1cb69a..c0be9a7 100644 --- a/src/geopops/schools.py +++ b/src/geopops/schools.py @@ -11,10 +11,14 @@ def read_sch_cap(data_dir): df = pd.read_csv(os.path.join(data_dir, 'processed', 'schools.csv'), usecols=['NCESSCH', 'STUDENTS'], dtype={'NCESSCH': str}) - return dict(zip(df['NCESSCH'], df['STUDENTS'])) + return dict(zip(df['NCESSCH'], df['STUDENTS'], strict=False)) def find_closest(data_dir, n): + """For each CBG and grade, the `n` nearest schools offering that grade. + + Returns ``{grade_key: {cbg_geoid: [school_id, ...]}}``, nearest first. + """ schools = pd.read_csv(os.path.join(data_dir, 'processed', 'schools.csv'), dtype={'NCESSCH': str}) distmat = pd.read_csv(os.path.join(data_dir, 'processed', 'cbg_sch_distmat.csv'), @@ -25,25 +29,37 @@ def find_closest(data_dir, n): grade_keys = ['p', 'k'] + [str(i) for i in range(1, 13)] grade_labels = ['PK', 'KG'] + [str(i) for i in range(1, 13)] sch_ids = [c for c in distmat.columns if c != 'GEOID'] + geoids = distmat['GEOID'].to_numpy() closest = {} - for gk, gl in zip(grade_keys, grade_labels): + for gk, gl in zip(grade_keys, grade_labels, strict=False): col = f'G_{gl}_OFFERED' if col not in schools.columns: continue - mask = schools[col].values.astype(bool) - valid_schs = schools['NCESSCH'].values[mask] - valid_set = set(valid_schs) + valid_set = set(schools['NCESSCH'].values[schools[col].values.astype(bool)]) valid_cols = [s for s in sch_ids if s in valid_set] + if not valid_cols: + closest[gk] = {geo: [] for geo in geoids} + continue - sch_by_geo = {} - for _, row in distmat.iterrows(): - geo = row['GEOID'] - dists = [(s, row[s]) for s in valid_cols if pd.notna(row[s])] - dists.sort(key=lambda x: x[1]) - top = dists[:n] - sch_by_geo[geo] = [s for s, _ in top] - closest[gk] = sch_by_geo + # Vectorized nearest-n. Missing distances become +inf so they sort last and + # can be filtered out afterwards; argpartition finds the n smallest per row + # without fully sorting, then only that slice is sorted. + dists = distmat[valid_cols].to_numpy(dtype=float) + dists = np.where(np.isnan(dists), np.inf, dists) + k = min(n, dists.shape[1] - 1) if dists.shape[1] > 1 else 0 + if k > 0: + part = np.argpartition(dists, k, axis=1)[:, :n] + else: + part = np.argsort(dists, axis=1)[:, :n] + rows = np.arange(dists.shape[0])[:, None] + order = part[rows, np.argsort(dists[rows, part], axis=1)] + + col_names = np.array(valid_cols, dtype=object) + closest[gk] = { + geo: col_names[row_order[np.isfinite(dists[i, row_order])]].tolist() + for i, (geo, row_order) in enumerate(zip(geoids, order, strict=False)) + } return closest @@ -57,9 +73,10 @@ def _get_students_in_school(people, cbgs_inv): return result -def generate_schools(people, cbgs, data_dir, random_seed=None): +def generate_schools(people, cbgs, data_dir, random_seed=None, config=None): rng = np.random.default_rng(random_seed) - config = tryJSON(os.path.join(data_dir, 'config.json')) + if config is None: + config = tryJSON(os.path.join(data_dir, 'config.json')) n_schools = config.get('n_closest_schools', 4) prob_closest = config.get('p_closest_school', 0.9) diff --git a/src/geopops/download_data.py b/src/geopops/sources.py similarity index 85% rename from src/geopops/download_data.py rename to src/geopops/sources.py index e3b258e..c42a22f 100644 --- a/src/geopops/download_data.py +++ b/src/geopops/sources.py @@ -6,16 +6,47 @@ import gzip import shutil import time -from pathlib import Path from curl_cffi import requests as curl_requests import urllib3 -import lzma import platform import subprocess import shlex +import warnings +from contextlib import contextmanager -# Disable SSL warnings for downloads with verify=False -urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) +from .exceptions import ConfigError, DataError, DownloadError + +# Some federal data portals (LODES, NCES, Census FTP) periodically serve +# incomplete certificate chains. Falling back to an unverified request is +# opt-in: set config["allow_insecure_downloads"] = True, or call +# set_allow_insecure_downloads(True), and GeoPops will warn loudly when it +# engages. It is off by default, and warnings are suppressed only for the +# duration of the individual request rather than process-wide. +_ALLOW_INSECURE_DOWNLOADS = False + + +def set_allow_insecure_downloads(allow): + """Enable or disable the unverified-TLS fallback for data downloads. + + Args: + allow (bool): If True, a request that fails TLS verification is retried + with verification disabled, after emitting a warning. + """ + global _ALLOW_INSECURE_DOWNLOADS + _ALLOW_INSECURE_DOWNLOADS = bool(allow) + + +@contextmanager +def _insecure_request(src): + """Warn about, and locally silence warnings for, one unverified request.""" + warnings.warn( + f"TLS verification disabled for {src}. The downloaded data is not " + "authenticated. Set allow_insecure_downloads=False to forbid this.", + stacklevel=3, + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", urllib3.exceptions.InsecureRequestWarning) + yield # Set base directory to the script's directory BASE_DIR = os.path.dirname(os.path.abspath(__file__)) @@ -23,12 +54,25 @@ # Removed here(): all write paths now use OUTPUT_DIR via os.path.join +# Tables with no 2023/2024 vintage on the Census API; fall back to 2022 for these. +ACS_MISSING_YEARS = {2023, 2024} +ACS_FALLBACK_2022_CODES = {'B09019', 'B09020', 'B09021'} + + +def _effective_acs_year(code, year_ACS): + """The ACS vintage actually available for `code`, falling back where needed.""" + if year_ACS in ACS_MISSING_YEARS and code in ACS_FALLBACK_2022_CODES: + print(f"Using 2022 data for {code} because {year_ACS} data is not available via Census API") + return 2022 + return year_ACS + + def dim_desc(df): """Function for returning dimensions of dataframe as a string - + Args: df (pandas.DataFrame): The dataframe whose dimensions are to be described - + Returns: str: A string describing the dimensions in the format "X rows x Y columns" """ @@ -36,11 +80,11 @@ def dim_desc(df): def fips_info(fips_codes, reverse=False): """Function for converting FIPS codes to state abbreviations or vice versa. Used for creating destination folders for census data - + Args: fips_codes (str or list): FIPS code(s) or abbreviation(s) to convert. Can be a single string or a list reverse (bool, optional): If True, converts abbreviations to FIPS codes. If False, converts FIPS codes to abbreviations. Defaults to False. - + Returns: dict: Dictionary with key "abbr" or "fips" containing the converted values. Returns None for invalid codes. If input is a list, returns a list of converted values; if input is a string, returns a single converted value. @@ -56,11 +100,11 @@ def fips_info(fips_codes, reverse=False): "47": "TN", "48": "TX", "49": "UT", "50": "VT", "51": "VA", "53": "WA", "54": "WV", "55": "WI", "56": "WY", "72": "PR" } - + if reverse: # Create reverse mapping from abbreviations to FIPS codes abbr_to_fips = {abbr: fips for fips, abbr in fips_to_abbr.items()} - + if isinstance(fips_codes, list): result = {"fips": [abbr_to_fips.get(code, None) for code in fips_codes]} return result @@ -73,186 +117,129 @@ def fips_info(fips_codes, reverse=False): else: return {"abbr": fips_to_abbr.get(fips_codes, None)} -# Helper function to download files with retry logic -def try_download(src, dst): - """Function for downloading a file with timeout and retry logic - - Args: - src (str): The source URL to download from - dst (str): The destination file path where the downloaded file will be saved - - Returns: - int: Status code (0 for success, -1 for failure) - """ - timeout = 3600 # 1 hour timeout - retries = 3 - status = 1 - - while status != 0 and retries > 0: - try: - # First try with SSL verification enabled - response = requests.get(src, timeout=timeout, stream=True) - response.raise_for_status() - - with open(dst, 'wb') as f: - for chunk in response.iter_content(chunk_size=8192): +# --------------------------------------------------------------------------- +# Downloading +# --------------------------------------------------------------------------- + +# Browser-like headers: several federal data portals reject requests that do not +# look like they came from a browser. +_BROWSER_HEADERS = { + "User-Agent": ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 " + "(KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"), + "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8", + "Accept-Language": "en-US,en;q=0.5", + "Connection": "keep-alive", + "Upgrade-Insecure-Requests": "1", + "Cache-Control": "max-age=0", +} + +_DOWNLOAD_TIMEOUT = 3600 # seconds +_DOWNLOAD_RETRIES = 3 + + +def _is_tls_error(exc): + """True if the exception looks like a certificate/TLS failure.""" + text = str(exc).lower() + return "ssl" in text or "certificate" in text + + +def _fetch_requests(src, dst, headers, mode, verify): + """Stream a URL to disk with `requests`.""" + response = requests.get(src, timeout=_DOWNLOAD_TIMEOUT, headers=headers, + stream=True, verify=verify) + response.raise_for_status() + if mode == "text": + with open(dst, "w", encoding="utf-8") as f: + for chunk in response.iter_content(chunk_size=8192, decode_unicode=True): + if chunk: f.write(chunk) - status = 0 - except Exception as e: - # If SSL error, try again with SSL verification disabled - if "SSL" in str(e) or "certificate" in str(e).lower(): - try: - response = requests.get(src, timeout=timeout, stream=True, verify=False) - response.raise_for_status() - - with open(dst, 'wb') as f: - for chunk in response.iter_content(chunk_size=8192): - f.write(chunk) - status = 0 - except Exception as e2: - print(f"Download attempt failed (with SSL disabled): {e2}") - retries -= 1 - status = -1 - else: - print(f"Download attempt failed: {e}") - retries -= 1 - status = -1 - - if status != 0: - print(f"Download failed: {src}") - exit(1) - - return status - -def try_curl_cffi(src, dst): - """Function for downloading a file using curl_cffi with timeout and retry logic - - Args: - src (str): The source URL to download from - dst (str): The destination file path where the downloaded file will be saved - - Returns: - int: Status code (0 for success, -1 for failure) - """ - timeout = 3600 # 1 hour timeout - retries = 3 - status = 1 - - while status != 0 and retries > 0: - try: - # Use browser impersonation with appropriate headers - headers = { - "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8", - "Accept-Language": "en-US,en;q=0.5", - "Connection": "keep-alive", - "Upgrade-Insecure-Requests": "1", - "Sec-Fetch-Dest": "document", - "Sec-Fetch-Mode": "navigate", - "Sec-Fetch-Site": "none", - "Sec-Fetch-User": "?1", - "Cache-Control": "max-age=0" - } - - # Try with Chrome impersonation and SSL verification disabled for problematic sites - response = curl_requests.get( - src, - timeout=timeout, - impersonate="chrome110", - headers=headers, - verify=False # Disable SSL verification to handle certificate issues - ) - - if response.status_code >= 400: - raise Exception(f"HTTP Error {response.status_code}") - - with open(dst, 'wb') as f: - f.write(response.content) - status = 0 - - except Exception as e: - print(f"Download attempt failed: {e}") - retries -= 1 - status = -1 - time.sleep(2) # Add a small delay between retries - - if status != 0: - print(f"Download failed: {src}") - exit(1) - - return status - -def try_download_text(src, dst): - """Function for downloading text content from web pages (like Census crosswalk files) - + else: + with open(dst, "wb") as f: + for chunk in response.iter_content(chunk_size=8192): + f.write(chunk) + + +def _fetch_curl_cffi(src, dst, headers, mode, verify): + """Fetch a URL with `curl_cffi`, impersonating Chrome's TLS fingerprint.""" + response = curl_requests.get(src, timeout=_DOWNLOAD_TIMEOUT, impersonate="chrome110", + headers=headers, verify=verify) + if response.status_code >= 400: + raise DownloadError(f"HTTP {response.status_code} for {src}") + if mode == "text": + with open(dst, "w", encoding="utf-8") as f: + f.write(response.text) + else: + with open(dst, "wb") as f: + f.write(response.content) + + +BACKENDS = ("requests", "curl_cffi") + + +def download(src, dst, *, backend="requests", mode="binary", headers=None, + retries=_DOWNLOAD_RETRIES): + """Download `src` to `dst`, with retries and an optional unverified-TLS fallback. + Args: - src (str): The source URL to download from - dst (str): The destination file path where the downloaded text will be saved - + src (str): Source URL. + dst (str): Destination file path. + backend (str): ``"requests"`` (default) or ``"curl_cffi"``. The latter + impersonates a browser TLS fingerprint for portals that block + non-browser clients. + mode (str): ``"binary"`` (default) or ``"text"``. + headers (dict): Request headers. Defaults to browser-like headers. + retries (int): Number of attempts before giving up. + Returns: - int: Status code (0 for success, -1 for failure) + int: 0 on success. + + Raises: + DownloadError: If every attempt fails. """ - timeout = 3600 # 1 hour timeout - retries = 3 - status = 1 - - while status != 0 and retries > 0: + # Resolved at call time (not via a module-level dict) so the backends stay + # patchable in tests and the error for a bad name is clear. + if backend == "requests": + fetch = _fetch_requests + elif backend == "curl_cffi": + fetch = _fetch_curl_cffi + else: + raise ValueError(f"Unknown download backend {backend!r}; expected one of {BACKENDS}.") + headers = _BROWSER_HEADERS if headers is None else headers + last_exc = None + + for attempt in range(retries): try: - # Use browser-like headers to get the actual text content - headers = { - "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - "Accept": "text/plain,text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", - "Accept-Language": "en-US,en;q=0.5", - "Connection": "keep-alive", - "Cache-Control": "max-age=0" - } - - # First try with SSL verification enabled - response = requests.get(src, timeout=timeout, headers=headers, stream=True) - response.raise_for_status() - - # Write the content as text (not binary) - with open(dst, 'w', encoding='utf-8') as f: - for chunk in response.iter_content(chunk_size=8192, decode_unicode=True): - if chunk: - f.write(chunk) - status = 0 - + fetch(src, dst, headers, mode, verify=True) + return 0 except Exception as e: - # If SSL error, try again with SSL verification disabled - if "SSL" in str(e) or "certificate" in str(e).lower(): + last_exc = e + if _is_tls_error(e) and _ALLOW_INSECURE_DOWNLOADS: try: - response = requests.get(src, timeout=timeout, headers=headers, stream=True, verify=False) - response.raise_for_status() - - with open(dst, 'w', encoding='utf-8') as f: - for chunk in response.iter_content(chunk_size=8192, decode_unicode=True): - if chunk: - f.write(chunk) - status = 0 + with _insecure_request(src): + fetch(src, dst, headers, mode, verify=False) + return 0 except Exception as e2: - print(f"Text download attempt failed (with SSL disabled): {e2}") - retries -= 1 - status = -1 - else: - print(f"Text download attempt failed: {e}") - retries -= 1 - status = -1 - - if status != 0: - print(f"Text download failed: {src}") - exit(1) - - return status + last_exc = e2 + print(f"Download attempt {attempt + 1}/{retries} failed for {src}: {last_exc}") + if attempt + 1 < retries: + time.sleep(2) + + hint = "" + if _is_tls_error(last_exc) and not _ALLOW_INSECURE_DOWNLOADS: + hint = (" This looks like a TLS/certificate failure. If you trust this source, " + "set allow_insecure_downloads=True in your config to permit an " + "unverified retry.") + raise DownloadError(f"Download failed after {retries} attempts: {src}: {last_exc}.{hint}") + def get_census_metadata(name, vintage, type_="variables"): """Function for getting ACS and Decennial metadata from Census API - + Args: name (str): The census dataset name (e.g., "acs/acs5", "dec/dhc", "dec/sf1") vintage (str): The year of the census data (e.g., "2020", "2019"). Comes from the config.json file type_ (str, optional): The type of metadata to retrieve. Defaults to "variables". Can be "variables" or "geography" - + Returns: pandas.DataFrame: DataFrame containing the metadata with variables as rows and metadata fields as columns """ @@ -263,7 +250,7 @@ def get_census_metadata(name, vintage, type_="variables"): def get_census_data(name, vintage, vars, region, regionin, key): """Function for making a Census API call and getting the data. Works with all ACS years and Decennial years before 2020 - + Args: name (str): The census dataset name (e.g., "acs/acs5", "dec/sf1") vintage (str): The year of the census data (e.g., "2020", "2019") @@ -271,13 +258,13 @@ def get_census_data(name, vintage, vars, region, regionin, key): region (str): The geographic level to retrieve data for (e.g., "block group:*") regionin (str): The geographic filter for the region (e.g., "state:24 county:*") key (str): The Census API key for authentication - + Returns: pandas.DataFrame: DataFrame containing the census data with variables as columns and geographic units as rows """ # Build the API URL base_url = f"https://api.census.gov/data/{vintage}/{name}" - + # Prepare parameters params = { "get": ",".join(vars), @@ -285,11 +272,11 @@ def get_census_data(name, vintage, vars, region, regionin, key): "in": regionin, "key": key } - + # Make the request response = requests.get(base_url, params=params) response.raise_for_status() - + # Convert to DataFrame data = response.json() headers = data[0] @@ -299,11 +286,11 @@ def get_census_data(name, vintage, vars, region, regionin, key): def get_new_dec_data(vintage, state_fips): """Function for getting 2020 Decennial data using the Census API - + Args: vintage (str): The year of the decennial data (should be "2020") state_fips (str): The FIPS code for the state to retrieve data for - + Returns: pandas.DataFrame: DataFrame containing the 2020 decennial data with variables as columns and geographic units as rows """ @@ -318,41 +305,28 @@ def get_new_dec_data(vintage, state_fips): def pull_census_data(state_fips, year_ACS, year_DEC, ACS_table_codes, DEC_table_codes, key, verbose=1): """Function for pulling census data for given states - + Args: state_fips (str): The FIPS code for the state to retrieve data for year_ACS (str): The year of the ACS data year_DEC (str): The year of the decennial data verbose: If 1, print output. If 0, suppress output. Defaults to 1. - + Returns: Outputs data files in the census folder """ - config_path = os.path.join(BASE_DIR, "config.json") - if not os.path.exists(config_path): - raise FileNotFoundError(f"config.json file not found at {config_path}. Please create this file with the required configuration.") - - with open(config_path, "r") as f: - config = json.load(f) - # Get metadata ACS_metadata = get_census_metadata(name="acs/acs5", vintage=year_ACS) if year_DEC == 2020: DEC_metadata = get_census_metadata(name="dec/dhc", vintage=year_DEC) else: DEC_metadata = get_census_metadata(name="dec/sf1", vintage=year_DEC) - + # Save metadata to CSV files # ACS_metadata.to_csv('ACS_metadata.csv') # DEC_metadata.to_csv('DEC_metadata.csv') - - # Combine ACS and DEC metadata - metadata_required = pd.concat([ACS_metadata, DEC_metadata], ignore_index=True) - - # Build name-label mapping dictionary (kept in-memory; no file output) - name_label_mapping = dict(zip(metadata_required["name"], metadata_required["label"])) - + printed_acs_source = False printed_dec_source = False acs_source_url = f"https://api.census.gov/data/{year_ACS}/acs/acs5" @@ -364,31 +338,23 @@ def pull_census_data(state_fips, year_ACS, year_DEC, ACS_table_codes, DEC_table_ # Process ACS table codes # Note: no 2023 data for B09019/B09020/B09021 # so we need to use 2022 data for these tables if year_ACS is 2023 - - # Effective vintage: use 2022 for B09019/B09020/B09021 when year_ACS is 2023 (no 2023 data) - missing_years = {2023, 2024} - FALLBACK_2022_CODES = {'B09019', 'B09020', 'B09021'} - def _effective_acs_year(code, year_ACS): - if year_ACS in missing_years and code in FALLBACK_2022_CODES: - print(f"Using 2022 data for {code} because 2023 and 2024 data are not available via Census API") - return 2022 - return year_ACS + # Process ACS table codes table_codes = pd.DataFrame({ "table_codes": [f"group({code})" for code in ACS_table_codes], "table_name": [f"ACSDT5Y{_effective_acs_year(code, year_ACS)}.{code}-Data" for code in ACS_table_codes], "vintage": [_effective_acs_year(code, year_ACS) for code in ACS_table_codes], - }) - - + }) + + # table_codes = pd.DataFrame({ # "table_codes": [f"group({code})" for code in ACS_table_codes], # "table_name": [f"ACSDT5Y{year_ACS}.{code}-Data" for code in ACS_table_codes] # }) - + for _, row in table_codes.iterrows(): # print(f"Downloading {row['table_name']}") - + # Make the API call and get the data data = get_census_data( name="acs/acs5", @@ -398,28 +364,28 @@ def _effective_acs_year(code, year_ACS): regionin=f"state:{state_i} county:*", key=key ) - + # Process data using Pandas data = data.drop(columns=["state", "county", "tract", "block group"], errors="ignore") - + cols = data.columns.to_list() if "GEO_ID" in cols and "NAME" in cols: # Move GEO_ID and NAME to the beginning cols.remove("GEO_ID") cols.remove("NAME") cols = ["GEO_ID", "NAME"] + cols data = data[cols] - + data = data.astype(str) # Convert all columns to string data_labels = ACS_metadata[ACS_metadata["name"].isin(data.columns)][["name", "label"]] # Get labels from metadata - - label_dict = dict(zip(data_labels["name"], data_labels["label"])) # Create a dictionary of column names to labels + + label_dict = dict(zip(data_labels["name"], data_labels["label"], strict=False)) # Create a dictionary of column names to labels all_labels = {col: col for col in data.columns} # Ensure all columns have a label all_labels.update(label_dict) label_df = pd.DataFrame([all_labels]) # Create a DataFrame with labels as the first row - + data_with_labels = pd.concat([label_df, data], ignore_index=True) # Combine label row with data data_with_labels = data_with_labels.loc[:, ~data_with_labels.columns.str.endswith(('EA', 'M', 'MA'))] # Remove columns ending with EA, M, or MA - + # Create destination folder and save to CSV state_abbr = fips_info(state_i)["abbr"] destination_folder = os.path.join(OUTPUT_DIR, "census", state_abbr.upper()) @@ -432,23 +398,23 @@ def _effective_acs_year(code, year_ACS): print(f"-- Downloading from {acs_source_url}") printed_acs_source = True print(f"-- census/{state_abbr.upper()}/{file_name}") - - # Process DEC table codes + + # Process DEC table codes table_codes = pd.DataFrame({ "table_codes": [f"group({code})" for code in DEC_table_codes], "table_name": [f"DECENNIALSF1{year_DEC}.{code}-Data" for code in DEC_table_codes] }) - + for _, row in table_codes.iterrows(): # print(f"Downloading {row['table_name']}") - + # Make the API call and get the data if year_DEC == 2020: data = get_new_dec_data(year_DEC, state_i) # Remove columns ending in 'A' data = data.loc[:, ~data.columns.str.endswith('A')] data.drop(columns=["ucgid"], inplace=True) - + else: data = get_census_data( name="dec/sf1", @@ -458,33 +424,33 @@ def _effective_acs_year(code, year_ACS): regionin=f"state:{state_i} county:*", key=key ) - + # Process data using Pandas data = data.drop(columns=["state", "county", "tract", "block group"], errors="ignore") - cols = data.columns.to_list() + cols = data.columns.to_list() if "GEO_ID" in cols and "NAME" in cols: # Move GEO_ID and NAME to the beginning of the dataframe cols.remove("GEO_ID") cols.remove("NAME") cols = ["NAME","GEO_ID"] + cols data = data[cols] - + data = data.astype(str) # Convert all columns to string data_labels = DEC_metadata[DEC_metadata["name"].isin(data.columns)][["name", "label"]] # Get labels from metadata - - - label_dict = dict(zip(data_labels["name"], data_labels["label"])) # Create a dictionary of column names to labels + + + label_dict = dict(zip(data_labels["name"], data_labels["label"], strict=False)) # Create a dictionary of column names to labels all_labels = {col: col for col in data.columns} # Ensure all columns have a label all_labels.update(label_dict) label_df = pd.DataFrame([all_labels]) # Create a DataFrame with labels as the first row - + data_with_labels = pd.concat([label_df, data], ignore_index=True) # Combine label row with data - + if year_DEC == 2020: # 2020 Decennial data has a different format for the labels data_with_labels.iloc[0, 2:] = data_with_labels.iloc[0, 2:].str.replace(':', '') # Strip ":" from the second row (labels) starting from fourth column data_with_labels.iloc[0, 2:] = data_with_labels.iloc[0, 2:].str[3:] # Remove first two characters from each string in the second row starting from fourth column - + data_with_labels = data_with_labels.loc[:, ~data_with_labels.columns.str.endswith('ERR')] # Remove columns ending with ERR - + # Create destination folder and save to CSV state_abbr = fips_info(state_i)["abbr"] destination_folder = os.path.join(OUTPUT_DIR, "census", state_abbr.upper()) @@ -500,25 +466,25 @@ def _effective_acs_year(code, year_ACS): def pull_pums_data(states, year, verbose=1): """Function for pulling PUMS microdata for given states - + Args: states (list): List of state abbreviations year (str): The year of the PUMS data verbose: If 1, print output. If 0, suppress output. Defaults to 1. - + Returns: Outputs data files in the pums folder - + Notes: 2024 and 2025 urls require password to access. If main year > 2023, use 2023 data. This means that the combinatorial optimization algorithm will combine households from 2023 PUMS data such that the set reasonably approximates Census distributions for target variables in the year specified in the config file. This can still produce a population that is reasonably representative of the target year. """ - + states = [s.lower() for s in states] file_urls = [] - + # Set the URL of the file to download for state in states: if year > 2023: @@ -532,12 +498,12 @@ def pull_pums_data(states, year, verbose=1): if verbose: print("\n*** Running DownloadData.pull_pums_data() ***") for state_i in states: - + urls_list = [url for state, url in file_urls if state == state_i] - + for url in urls_list: download_url = url - + # Specify the destination folder and file name destination_folder = os.path.join(OUTPUT_DIR, "pums") file_name = os.path.basename(download_url) @@ -545,24 +511,24 @@ def pull_pums_data(states, year, verbose=1): # print(destination_folder) # Create the destination folder if it doesn't exist os.makedirs(destination_folder, exist_ok=True) - + # Download the file - try_curl_cffi(download_url, destination_file) - + download(download_url, destination_file, backend="curl_cffi") + # Extract the contents of the zip file with zipfile.ZipFile(destination_file, 'r') as zip_ref: zip_ref.extractall(destination_folder) - + # Remove the zip file os.remove(destination_file) - + state_fips = fips_info(state_i.upper(), reverse=True)['fips'] if file_name == f'csv_h{state_i}.zip': df_h = pd.read_csv(f"{destination_folder}/psam_h{state_fips}.csv", low_memory=False) if year >= 2020: # ACCESSINET(formerly ACCESS), TYPEHUGQ (formerly TYPE). # https://www2.census.gov/programs-surveys/acs/tech_docs/pums/variable_changes/ACS2016-2020_PUMS_Variable_Changes_and_Explanations.pdf - df_h.rename(columns={'ACCESSINET': 'ACCESS'}, inplace=True) # + df_h.rename(columns={'ACCESSINET': 'ACCESS'}, inplace=True) # df_h.rename(columns={'TYPEHUGQ': 'TYPE'}, inplace=True) if year >= 2021: # FES variable deleted in 2021. Can be recreated from WORKSTAT. @@ -574,7 +540,7 @@ def pull_pums_data(states, year, verbose=1): df_h.loc[(df_h['WORKSTAT'] == 10)|(df_h['WORKSTAT'] == 11), 'FES'] = 5 df_h.loc[(df_h['WORKSTAT'] == 12), 'FES'] = 6 df_h.loc[(df_h['WORKSTAT'] == 13)|(df_h['WORKSTAT'] == 14), 'FES'] = 7 - df_h.loc[(df_h['WORKSTAT'] == 15), 'FES'] = 8 + df_h.loc[(df_h['WORKSTAT'] == 15), 'FES'] = 8 if year >= 2022: # https://www2.census.gov/programs-surveys/acs/tech_docs/pums/variable_changes/ACS2018-2022_PUMS_Variable_Changes_and_Explanations.pdf # We may want to use PUMA10 since that's the definition we were using before @@ -612,49 +578,49 @@ def pull_pums_data(states, year, verbose=1): def download_shapefiles(state_fips, year, verbose=1): """Function for downloading shapefiles from census web server - + Args: state_fips (list): List of state abbreviations year (str): The year of the shapefiles verbose: If 1, print output. If 0, suppress output. Defaults to 1 - + Returns: Outputs data files in the geo folder """ - + file_urls = [] - + # Set the URL of the file to download for state in state_fips: file_urls.append((state, f"https://www2.census.gov/geo/tiger/TIGER{year}/BG/tl_{year}_{state}_bg.zip")) - + printed_source_header = False for state_i in state_fips: - + if verbose: print("\n*** Running DownloadData.download_shapefiles() ***") - + urls_list = [url for state, url in file_urls if state == state_i] - + for url in urls_list: download_url = url - + # Specify the destination folder and file name destination_folder = os.path.join(OUTPUT_DIR, "geo") - + file_name = os.path.basename(download_url) destination_file = os.path.join(destination_folder, file_name) - + # Create the destination folder if it doesn't exist os.makedirs(destination_folder, exist_ok=True) - + # Download the file - try_curl_cffi(download_url, destination_file) - + download(download_url, destination_file, backend="curl_cffi") + # Extract the contents of the zip file with zipfile.ZipFile(destination_file, 'r') as zip_ref: zip_ref.extractall(destination_folder) - + if verbose: if not printed_source_header: print(f"-- Downloading from https://www2.census.gov/geo/tiger/TIGER{year}/BG/") @@ -663,16 +629,16 @@ def download_shapefiles(state_fips, year, verbose=1): def pull_LODES(states_main, states_aux, year, verbose=1): """Function for pulling LEHD LODES data (commuting patterns) for given states - + Args: states_main (list): List of state abbreviations for the main states states_aux (list): List of state abbreviations for the auxiliary states year (str): The year of the LODES data verbose: If 1, print output. If 0, suppress output. Defaults to 1. - + Returns: Outputs data files in the work folder - + Notes: No LODES data for 2024 or 2025, so use 2023 data if year > 2023. """ printed_source_header = False @@ -680,24 +646,24 @@ def pull_LODES(states_main, states_aux, year, verbose=1): print("\n*** Running DownloadData.pull_LODES() ***") # Determine version based on year version = "LODES8" if year >= 2020 else "LODES7" - + # For the states on the "main" list, download OD main JT01, OD aux JT01, and WAC S000 JT01 for state_i in states_main: # print(f"*** Downloading LODES data for state = {state_i.upper()} ***") state_i = state_i.lower() state_dir = os.path.join(OUTPUT_DIR, "work") os.makedirs(state_dir, exist_ok=True) - + # Download and save OD main JT01 # print(f"downloading lodes od main {state_i}") - - # For OD main JT01 + + # For OD main JT01 if year > 2023: year = 2023 od_main_url = f"https://lehd.ces.census.gov/data/lodes/{version}/{state_i}/od/{state_i}_od_main_JT01_{year}.csv.gz" outfile = os.path.join(state_dir, f"{state_i}_od_main_JT01_{year}.csv.gz") - try_download(od_main_url, outfile) - + download(od_main_url, outfile) + # Process the file to match old format with gzip.open(outfile, 'rt') as f_in: with gzip.open(outfile + '.tmp', 'wt') as f_out: @@ -705,7 +671,7 @@ def pull_LODES(states_main, states_aux, year, verbose=1): header = f_in.readline().strip() # Write new header with year and state f_out.write("year,state," + header + "\n") - + # Process each line for line in f_in: parts = line.strip().split(',') @@ -719,22 +685,22 @@ def pull_LODES(states_main, states_aux, year, verbose=1): print("-- Downloading from https://lehd.ces.census.gov/data/lodes/") printed_source_header = True print(f"-- work/{state_i}_od_main_JT01_{year}.csv.gz") - + # Replace original file with processed file os.replace(outfile + '.tmp', outfile) - + # For OD aux JT01 # print(f"downloading lodes od aux {state_i}") od_aux_url = f"https://lehd.ces.census.gov/data/lodes/{version}/{state_i}/od/{state_i}_od_aux_JT01_{year}.csv.gz" outfile = os.path.join(state_dir, f"{state_i}_od_aux_JT01_{year}.csv.gz") - try_download(od_aux_url, outfile) - + download(od_aux_url, outfile) + # Process the file to match old format with gzip.open(outfile, 'rt') as f_in: with gzip.open(outfile + '.tmp', 'wt') as f_out: header = f_in.readline().strip() # Read header f_out.write("year,state," + header + "\n") # Write new header with year and state - + # Process each line for line in f_in: parts = line.strip().split(',') @@ -748,22 +714,22 @@ def pull_LODES(states_main, states_aux, year, verbose=1): print("-- Downloading from https://lehd.ces.census.gov/data/lodes/") printed_source_header = True print(f"-- work/{state_i}_od_aux_JT01_{year}.csv.gz") - + # Replace original file with processed file os.replace(outfile + '.tmp', outfile) - + # For WAC S000 JT01 # print(f"downloading lodes wac {state_i}") wac_url = f"https://lehd.ces.census.gov/data/lodes/{version}/{state_i}/wac/{state_i}_wac_S000_JT01_{year}.csv.gz" outfile = os.path.join(state_dir, f"{state_i}_wac_S000_JT01_{year}.csv.gz") - try_download(wac_url, outfile) - + download(wac_url, outfile) + # Process the file to match old format with gzip.open(outfile, 'rt') as f_in: with gzip.open(outfile + '.tmp', 'wt') as f_out: header = f_in.readline().strip() # Read header f_out.write("year,state," + header + "\n") # Write new header with year and state - + # Process each line for line in f_in: parts = line.strip().split(',') @@ -776,27 +742,27 @@ def pull_LODES(states_main, states_aux, year, verbose=1): print("-- Downloading from https://lehd.ces.census.gov/data/lodes/") printed_source_header = True print(f"-- work/{state_i}_wac_S000_JT01_{year}.csv.gz") - + # Replace original file with processed file os.replace(outfile + '.tmp', outfile) - + # For the states on the "aux" list just download OD aux JT01 for state_i in states_aux: state_dir = os.path.join(OUTPUT_DIR, "work") os.makedirs(state_dir, exist_ok=True) - + # Download and save OD aux JT01 # print(f"downloading lodes od aux {state_i}") od_aux_url = f"https://lehd.ces.census.gov/data/lodes/{version}/{state_i.lower()}/od/{state_i.lower()}_od_aux_JT01_{year}.csv.gz" outfile = os.path.join(state_dir, f"{state_i.lower()}_od_aux_JT01_{year}.csv.gz") - try_download(od_aux_url, outfile) - + download(od_aux_url, outfile) + # Process the file to match old format with gzip.open(outfile, 'rt') as f_in: with gzip.open(outfile + '.tmp', 'wt') as f_out: header = f_in.readline().strip() # Read header f_out.write("year,state," + header + "\n") # Write new header with year and state - + # Process each line for line in f_in: parts = line.strip().split(',') @@ -810,16 +776,16 @@ def pull_LODES(states_main, states_aux, year, verbose=1): print("-- Downloading from https://lehd.ces.census.gov/data/lodes/") printed_source_header = True print(f"-- work/{state_i.lower()}_od_aux_JT01_{year}.csv.gz") - + # Replace original file with processed file os.replace(outfile + '.tmp', outfile) def download_cbp_data(verbose=1): """Function for downloading CBP data - + Args: verbose: If 1, print output. If 0, suppress output. Defaults to 1. - + Returns: Outputs data files in the work folder """ @@ -827,21 +793,21 @@ def download_cbp_data(verbose=1): print("\n*** Running DownloadData.download_cbp_data() ***") cbp_dir = os.path.join(OUTPUT_DIR, "work") os.makedirs(cbp_dir, exist_ok=True) - + url = "https://www2.census.gov/programs-surveys/cbp/datasets/2016/cbp16co.zip" outfile = os.path.join(cbp_dir, "cbp16co.zip") - try_download(url, outfile) + download(url, outfile) if verbose: print("-- Downloading from https://www2.census.gov/programs-surveys/cbp/datasets/2016/") print("-- work/cbp16co.zip") def download_ct_puma_crosswalk(main_year, verbose=1): """Function for downloading Census Tract to PUMA crosswalk file - + Args: main_year (str): The year of the crosswalk file verbose: If 1, print output. If 0, suppress output. Defaults to 1. - + Returns: Outputs data files in the geo folder """ @@ -849,7 +815,7 @@ def download_ct_puma_crosswalk(main_year, verbose=1): print("\n*** Running DownloadData.download_ct_puma_crosswalk() ***") geo_dir = os.path.join(OUTPUT_DIR, "geo") os.makedirs(geo_dir, exist_ok=True) - + # https://www2.census.gov/geo/docs/maps-data/data/rel2020/2020_Census_Tract_to_2020_PUMA.txt if main_year >= 2020: url = "https://www2.census.gov/geo/docs/maps-data/data/rel2020/2020_Census_Tract_to_2020_PUMA.txt" @@ -857,13 +823,13 @@ def download_ct_puma_crosswalk(main_year, verbose=1): else: url = "https://www2.census.gov/geo/docs/maps-data/data/rel/2010_Census_Tract_to_2010_PUMA.txt" outfile = os.path.join(geo_dir, "2010_Census_Tract_to_2010_PUMA.txt") - + # Use text download for these web pages since they display data in browser - try_download_text(url, outfile) + download(url, outfile, mode="text") if verbose: print("-- Downloading from https://www2.census.gov/geo/docs/maps-data/data/rel/") print(f"-- geo/{os.path.basename(outfile)}") - + # urls2018 = ["https://mcdc.missouri.edu/temp/geocorr2018_2523203354.csv", # geocorr2018_puma_to_county.csv # "https://mcdc.missouri.edu/temp/geocorr2018_2523207598.csv", # geocorr2018_puma_to_cbsa.csv # "https://mcdc.missouri.edu/temp/geocorr2018_2523205248.csv", # geocorr2018_puma_urban_rural.csv @@ -893,20 +859,20 @@ def download_ct_puma_crosswalk(main_year, verbose=1): # try: # response = requests.get(url) # response.raise_for_status() - + # with open(outfile, 'wb') as f: # f.write(response.content) - + # # print(f"Downloaded {outfile}") # except Exception as e: # print(f"Failed to download {url}: {e}") - + def geocorr_files(verbose=1): """Copy geocorr files from the local 'geocorr' folder into the 'geo' folder. If main_year < 2020, copies only files matching 'geocorr2018*'. If main_year >= 2020, copies only files matching 'geocorr2022*'. - + Returns: Outputs data files in the geo folder """ @@ -919,7 +885,7 @@ def geocorr_files(verbose=1): raise FileNotFoundError( f"config.json file not found at {config_path}. Please create this file with the required configuration." ) - with open(config_path, "r") as f: + with open(config_path) as f: config = json.load(f) if "main_year" not in config: @@ -927,8 +893,8 @@ def geocorr_files(verbose=1): try: main_year = int(config["main_year"]) - except Exception: - raise ValueError("main_year in config.json must be an integer.") + except (TypeError, ValueError) as e: + raise ConfigError("main_year in config.json must be an integer.") from e if main_year < 2020: prefix = "geocorr2018" @@ -963,27 +929,27 @@ def geocorr_files(verbose=1): def download_school_data(main_year, verbose=1): """Function for downloading school data - + Args: main_year (str): The year of the school data verbose: If 1, print output. If 0, suppress output. Defaults to 1. - + Returns: Outputs data files in the school folder - + Notes: This function uses the python package 'zipfile' to extract the zip file. - This will cause an error "NotImplementedError" if the interpreter used to build the environment - does not have full LZMA support. - + This will cause an error "NotImplementedError" if the interpreter used to build the environment + does not have full LZMA support. + MacOS workaround: use the system 'unzip' command instead of 'zipfile'. Windows workaround: use Expand-Archive or 7‑Zip """ school_dir = os.path.join(OUTPUT_DIR, "school") os.makedirs(school_dir, exist_ok=True) - + if verbose: print("\n*** Running DownloadData.download_school_data() ***") - + # Download school location data # https://nces.ed.gov/programs/edge/data/EDGE_GEOCODE_PUBLICSCH_2021.zip url = f"https://nces.ed.gov/programs/edge/data/EDGE_GEOCODE_PUBLICSCH_{str(main_year)[-2:]}{str(main_year + 1)[-2:]}.zip" @@ -1005,12 +971,12 @@ def download_school_data(main_year, verbose=1): # Extract the zip file with zipfile.ZipFile(zip_path, 'r') as zip_ref: - zip_ref.extractall(school_dir) + zip_ref.extractall(school_dir) # Move files from extracted folder to main school folder extracted_folder = os.path.join(school_dir, f"EDGE_GEOCODE_PUBLICSCH_{str(main_year)[-2:]}{str(main_year + 1)[-2:]}") files_to_move = [f"EDGE_GEOCODE_PUBLICSCH_{str(main_year)[-2:]}{str(main_year + 1)[-2:]}.xlsx", - "Shapefiles_SCH", "Shapefile_SCH"] # if year > 2019, uses Shapefile_SCH (singular) folder instead of Shapefiles_SCH (plural) + "Shapefiles_SCH", "Shapefile_SCH"] # if year > 2019, uses Shapefile_SCH (singular) folder instead of Shapefiles_SCH (plural) for item_to_move in files_to_move: source_path = os.path.join(extracted_folder, item_to_move) @@ -1027,7 +993,7 @@ def download_school_data(main_year, verbose=1): # print(f"Item {item_to_move} not found") # Normalize school shapefile folder name to 'Shapefiles_SCH' - + singular = os.path.join(school_dir, "Shapefile_SCH") plural = os.path.join(school_dir, "Shapefiles_SCH") @@ -1057,12 +1023,12 @@ def download_school_data(main_year, verbose=1): os.remove(item_to_delete) # else: # print(f"Item {os.path.basename(item_to_delete)} not found") - + # Download school enrollment data # print("Downloading school enrollment data... takes a while") # URLs for school enrollment data - + if main_year == 2019: en_str = "082120" elif main_year == 2020: @@ -1075,7 +1041,7 @@ def download_school_data(main_year, verbose=1): en_str = "073124" elif main_year == 2024: en_str = "073025" - + enrollment_urls = [ f"https://nces.ed.gov/ccd/Data/zip/ccd_sch_029_{str(main_year)[-2:]}{str(main_year + 1)[-2:]}_w_1a_{en_str}.zip", # Directory f"https://nces.ed.gov/ccd/Data/zip/ccd_SCH_052_{str(main_year)[-2:]}{str(main_year + 1)[-2:]}_l_1a_{en_str}.zip", # Membership @@ -1087,22 +1053,22 @@ def download_school_data(main_year, verbose=1): # Extract filename from URL filename = url.split('/')[-1] zip_path = os.path.join(school_dir, filename) - + # print(f"Downloading {filename}...") response = requests.get(url, stream=True) response.raise_for_status() - + # Download the file (overwrite if exists) with open(zip_path, 'wb') as f: f.write(response.content) - + # Extract the outer zip file - # This will cause an error "NotImplementedError" if the interpreter used to build the environment + # This will cause an error "NotImplementedError" if the interpreter used to build the environment # does not have full LZMA support try: with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(school_dir) - except NotImplementedError: + except NotImplementedError as e: system = platform.system() zip_basename = os.path.basename(zip_path) @@ -1130,11 +1096,11 @@ def download_school_data(main_year, verbose=1): else: # Other platforms: re-raise with guidance - raise RuntimeError( + raise DataError( "ZIP uses a compression method not supported by this Python. " "On Unix-like systems, install a Python build with LZMA support " "or extract the archive manually using system tools." - ) + ) from e # Handle nested CSV and SAS ZIPs created by NCES (e.g. *_CSV.zip, *_SAS.zip) for nested_name in os.listdir(school_dir): @@ -1149,7 +1115,7 @@ def download_school_data(main_year, verbose=1): try: with zipfile.ZipFile(nested_path, "r") as nested_zip: nested_zip.extractall(school_dir) - except NotImplementedError: + except NotImplementedError as e: system = platform.system() if system == "Darwin": subprocess.run( @@ -1169,11 +1135,11 @@ def download_school_data(main_year, verbose=1): check=True, ) else: - raise RuntimeError( + raise DataError( "Nested CSV ZIP uses a compression method not supported by this Python. " "On Unix-like systems, install a Python build with LZMA support " "or extract the archive manually using system tools." - ) + ) from e # Remove the nested CSV zip after extraction if os.path.exists(nested_path): @@ -1186,17 +1152,17 @@ def download_school_data(main_year, verbose=1): # Check what files were extracted directly to the school folder # print(f"Files in school folder after extraction: {[f for f in os.listdir(path) if f.endswith('.csv') or f.endswith('.sas7bdat')]}") - + # Delete .sas7bdat files that were extracted directly to the school folder for item in os.listdir(school_dir): if item.endswith('.sas7bdat'): file_path = os.path.join(school_dir, item) os.remove(file_path) # print(f"Deleted {item} from school folder") - + # Delete the zip file files_to_delete = [zip_path] - + for item_to_delete in files_to_delete: if os.path.exists(item_to_delete): if os.path.isdir(item_to_delete): @@ -1237,15 +1203,17 @@ def __init__(self, config=None, base_dir=None, verbose=1, auto_run=True): cfg_path = os.path.join(self.base_dir, "config.json") if not os.path.exists(cfg_path): raise FileNotFoundError(f"config.json file not found at {cfg_path}. Please create this file with the required configuration.") - with open(cfg_path, "r") as f: + with open(cfg_path) as f: self.config = json.load(f) # Initialize OUTPUT_DIR from config["path"] (fallback to package dir) global OUTPUT_DIR OUTPUT_DIR = self.config.get("path", self.base_dir) os.makedirs(OUTPUT_DIR, exist_ok=True) + # Opt-in TLS fallback for portals with broken certificate chains + set_allow_insecure_downloads(self.config.get("allow_insecure_downloads", False)) if auto_run: self.run_all() - + def run_all(self): """Run the full download workflow using the loaded configuration.""" config = self.config @@ -1254,7 +1222,7 @@ def run_all(self): print("============================================================") print("Running DownloadData()") print("============================================================") - + # Basic validation for keys required by multiple steps if "census_api_key" not in config: raise KeyError("census_api_key not found in config. Please add your Census API key to the configuration.") @@ -1420,7 +1388,7 @@ def census_metadata(self, refresh=False): """ mapping_path = os.path.join(OUTPUT_DIR, "census_metadata.json") if not refresh and os.path.exists(mapping_path): - with open(mapping_path, "r") as f: + with open(mapping_path) as f: return json.load(f) # Recompute from API @@ -1432,7 +1400,7 @@ def census_metadata(self, refresh=False): else: dec_meta = get_census_metadata(name="dec/sf1", vintage=decennial_year) metadata_required = pd.concat([acs_meta, dec_meta], ignore_index=True) - name_label_mapping = dict(zip(metadata_required["name"], metadata_required["label"])) + name_label_mapping = dict(zip(metadata_required["name"], metadata_required["label"], strict=False)) return name_label_mapping def pipeline(self): @@ -1498,9 +1466,9 @@ def pipeline(self): "function": "download_ct_puma_crosswalk", "websites": [ ( - f"https://www2.census.gov/geo/docs/maps-data/data/rel2020/2020_Census_Tract_to_2020_PUMA.txt" + "https://www2.census.gov/geo/docs/maps-data/data/rel2020/2020_Census_Tract_to_2020_PUMA.txt" if main_year >= 2020 - else f"https://www2.census.gov/geo/docs/maps-data/data/rel/2010_Census_Tract_to_2010_PUMA.txt" + else "https://www2.census.gov/geo/docs/maps-data/data/rel/2010_Census_Tract_to_2010_PUMA.txt" ) ], "output_folder": os.path.join(OUTPUT_DIR, "geo"), @@ -1583,9 +1551,27 @@ def pipeline(self): print(f" - {f}") print("") +def download_data(config, *, base_dir=None, verbose=1): + """Download every raw dataset the pipeline needs into ``config['path']``. + + Downloads are cached on disk, so re-running skips work already done. + + Args: + config: a config dict (see :func:`geopops.make_config`). + verbose: 0 for quiet, 1 for progress logging. + + Returns: + DownloadData: the completed step, for inspection. + + Raises: + DownloadError: if a required file could not be fetched. + """ + return DownloadData(config=config, base_dir=base_dir, verbose=verbose, auto_run=True) + + def main(): """Main function to preserve CLI behavior using the class wrapper.""" downloader = DownloadData() downloader.run_all() if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/geopops/starsim_bridge.py b/src/geopops/starsim_bridge.py new file mode 100644 index 0000000..8662a88 --- /dev/null +++ b/src/geopops/starsim_bridge.py @@ -0,0 +1,407 @@ +"""Bridge from a generated GeoPops population to Starsim objects. + +The public surface is a few functions --- :func:`to_starsim_people`, +:func:`starsim_network`, :func:`starsim_networks` --- plus the +:class:`SubgroupTracking` analyzer. Each takes an explicit population directory, so +nothing is cached across calls and generating two populations in one session gives +two independent sets of networks. +""" +import os +import json + +import numpy as np +import pandas as pd +import starsim as ss +from scipy.io import mmread + +from .exceptions import ConfigError, DataError + +BASE_DIR = os.path.dirname(os.path.abspath(__file__)) + +#: Layer name -> matrix-market file written by :func:`geopops.export.export_networks` +NETWORK_FILES = { + 'homenet': 'adj_upper_triang_hh.mtx', + 'schoolnet': 'adj_upper_triang_sch.mtx', + 'worknet': 'adj_upper_triang_wp.mtx', + 'gqnet': 'adj_upper_triang_gq.mtx', +} + +#: Columns of people.csv that are core demographics rather than config-driven traits +_CORE_PERSON_COLS = ( + 'p_id', 'hh_id', 'cbg_id', 'sample_index', 'age', 'working', 'commuter', + 'commuter_income_category', 'commuter_workplace_category', 'sch_grade', +) + + +def _resolve_pop_export(pop_export_dir=None, config=None, config_path=None, base_dir=None): + """Work out which ``pop_export`` directory to read from.""" + if pop_export_dir is not None: + return pop_export_dir + if config is None: + cfg_path = config_path or os.path.join(base_dir or BASE_DIR, "config.json") + if not os.path.exists(cfg_path): + raise ConfigError( + f"No config found at {cfg_path}. Pass pop_export_dir=... or config=... " + "to say where the generated population lives." + ) + with open(cfg_path) as f: + config = json.load(f) + path = config.get("path") + if not path: + raise ConfigError("config has no 'path' entry, so the population directory is unknown.") + return os.path.join(path, "pop_export") + + +def _read(pop_export_dir, name, **kwargs): + path = os.path.join(pop_export_dir, name) + if not os.path.exists(path): + raise DataError(f"Expected exported population file not found: {path}") + return pd.read_csv(path, **kwargs) + + +def _load_age_by_matrix_index(pop_export_dir): + """Map matrix row/col index (``index_zero``) to age, same merge as :func:`to_starsim_people`.""" + adj = _read(pop_export_dir, "adj_mat_keys.csv", low_memory=False) + people = _read(pop_export_dir, "people.csv", low_memory=False) + merged = adj.merge(people, on=["p_id", "hh_id", "cbg_id"], how="left") + merged = merged.drop_duplicates(subset=["index_zero"], keep="first") + return merged.set_index("index_zero")["age"] + + +def _canonicalize_undirected_edges_df(net_df, age_by_idx): + """Reorder ``p1``, ``p2`` so ``age(p1) <= age(p2)`` when both ages exist; else smaller index is ``p1``.""" + if net_df.empty: + return net_df.copy() + out = net_df.copy() + p1 = out["p1"].to_numpy(dtype=np.int64, copy=True) + p2 = out["p2"].to_numpy(dtype=np.int64, copy=True) + a1 = age_by_idx.reindex(p1).to_numpy() + a2 = age_by_idx.reindex(p2).to_numpy() + a1 = np.where(pd.isna(a1), np.nan, np.asarray(a1, dtype=float)) + a2 = np.where(pd.isna(a2), np.nan, np.asarray(a2, dtype=float)) + both = np.isfinite(a1) & np.isfinite(a2) + swap = np.zeros(len(out), dtype=bool) + swap[both] = (a1[both] > a2[both]) | ((a1[both] == a2[both]) & (p1[both] > p2[both])) + swap[~both] = p1[~both] > p2[~both] + out.loc[swap, ["p1", "p2"]] = np.column_stack([p2[swap], p1[swap]]) + return out + + +def _random_flip_undirected_edges_df(net_df, rng): + """Randomly swap ``(p1, p2)`` per edge with probability 0.5. + + Only changes endpoint labelling, so plots that read ``(p1_age, p2_age)`` as an + ordered pair do not come out looking triangular. + """ + if net_df.empty: + return net_df.copy() + out = net_df.copy() + p1 = out["p1"].to_numpy(dtype=np.int64, copy=True) + p2 = out["p2"].to_numpy(dtype=np.int64, copy=True) + flip = rng.random(len(out)) < 0.5 + out["p1"] = np.where(flip, p2, p1) + out["p2"] = np.where(flip, p1, p2) + return out + + +def load_network_edges(pop_export_dir, names=None, seed=0, save=True, verbose=1): + """Read exported adjacency matrices into edge-list dataframes. + + Args: + pop_export_dir: the run's ``pop_export`` directory. + names: layer names to load; defaults to all of :data:`NETWORK_FILES`. + seed: seed for the endpoint-order shuffle. + save: also write ``pop_export/starsim/net_*.csv``. + verbose: if truthy, log progress. + + Returns: + dict: layer name -> dataframe with ``p1``, ``p2``, ``edge_weight``. + """ + names = list(NETWORK_FILES) if names is None else list(names) + rng = np.random.default_rng(seed) + short = {'homenet': 'h', 'schoolnet': 's', 'worknet': 'w', 'gqnet': 'g'} + + if save: + os.makedirs(os.path.join(pop_export_dir, "starsim"), exist_ok=True) + + edges = {} + for name in names: + path = os.path.join(pop_export_dir, NETWORK_FILES[name]) + if not os.path.exists(path): + raise DataError(f"Network file not found: {path}. Has Export() been run?") + m = mmread(path) + df = pd.DataFrame({ + "p1": np.asarray(m.col, dtype=np.int64), + "p2": np.asarray(m.row, dtype=np.int64), + }) + df["edge_weight"] = np.int64(1) + # Treat undirected layers as having ~50/50 endpoint ordering. + df = _random_flip_undirected_edges_df(df, rng) + edges[name] = df + if save: + out = os.path.join(pop_export_dir, "starsim", f"net_{short.get(name, name)}.csv") + df.to_csv(out, index=False) + if verbose: + print(f"-- {out}") + return edges + + +def to_starsim_people(pop_export_dir=None, *, config=None, config_path=None, + base_dir=None, save=True, verbose=1): + """Build a Starsim ``People`` object from an exported GeoPops population. + + Args: + pop_export_dir: the run's ``pop_export`` directory. If omitted, it is + derived from ``config`` (or a config file). + save: also write ``people_all.csv`` and ``starsim/ppl.pkl``. + verbose: if truthy, log progress. + + Returns: + ss.People: agents carrying age, sex, geography and the run's traits. + """ + pop_export_dir = _resolve_pop_export(pop_export_dir, config, config_path, base_dir) + if verbose: + print("\n*** Building Starsim People from GeoPops population ***") + + people = _read(pop_export_dir, 'people.csv') + ppl_df = _read(pop_export_dir, 'adj_mat_keys.csv').merge( + people, on=['p_id', 'hh_id', 'cbg_id'], how='left') + ppl_df = ppl_df.merge(_read(pop_export_dir, 'sch_students.csv'), + on=['p_id', 'hh_id', 'cbg_id'], how='left') + ppl_df.loc[ppl_df['sch_code'].isnull(), 'sch_code'] = 0 + ppl_df.insert(0, 'uid', ppl_df['index_zero'].values) + + ppl_df = ppl_df.merge(_read(pop_export_dir, 'cbg_idxs.csv'), on='cbg_id', how='left') + + # Geography: split the 12-digit CBG geocode into its nested levels. + geocode = ppl_df['cbg_geocode'].astype(str) + for col, width in (('state', 2), ('county', 5), ('tract', 11), ('cbg_geocode', 12)): + ppl_df[col] = geocode.str[:width].replace({'na': '0.0', 'nan': '0.0'}).astype(float) + + # Decadal age groups, capped at 90+ + ppl_df['agegroup'] = np.clip(ppl_df['age'] // 10, 0, 9) + + hh = _read(pop_export_dir, 'hh.csv') + hh['household'] = hh.index + 1 + hh = hh.drop(columns=['sample_index']) + ppl_df = ppl_df.merge(hh, on=['cbg_id', 'hh_id'], how='left') + ppl_df.loc[ppl_df['household'].isnull(), 'household'] = 0 + + # Trait columns are whatever this run's config asked for, so take them from the + # exported file rather than a hardcoded list. + trait_cols = [c for c in people.columns if c not in _CORE_PERSON_COLS] + + ppl_df = ppl_df[['uid', 'p_id', 'hh_id', 'cbg_id', 'sample_index', 'state', 'county', + 'tract', 'cbg_geocode', 'household', 'age', 'agegroup', *trait_cols, + 'working', 'commuter', 'commuter_income_category', + 'commuter_workplace_category', 'sch_grade', 'sch_code']] + if save: + ppl_df.to_csv(os.path.join(pop_export_dir, 'people_all.csv'), index=False) + + def farr(col): + return ss.FloatArr(col, default=ss.BaseArr(ppl_df[col].values)) + + def iarr(col): + return ss.IntArr(col, default=ss.BaseArr(ppl_df[col].values)) + + age = farr('age') + # 'female' is conventionally a trait, but Starsim People expects it as a state + sex_states = [farr('female')] if 'female' in trait_cols else [] + other_traits = [farr(c) for c in trait_cols if c != 'female'] + + extra = [farr('agegroup'), *other_traits, + iarr('state'), iarr('county'), iarr('tract'), iarr('cbg_geocode'), + iarr('household'), farr('commuter'), + farr('commuter_income_category'), farr('commuter_workplace_category'), + iarr('sch_code')] + + ppl = ss.People(n_agents=len(ppl_df), extra_states=extra) + + # age and female are built into ss.People, so overwrite rather than add + for state in [age, *sex_states]: + ppl.states.append(state, overwrite=True) + setattr(ppl, state.name, state) + state.link_people(ppl) + + # Initializing a Sim wires up the People object (links states, sets uids) + ss.Sim(people=ppl).init() + + if save: + starsim_dir = os.path.join(pop_export_dir, 'starsim') + os.makedirs(starsim_dir, exist_ok=True) + ss.save(os.path.join(starsim_dir, 'ppl.pkl'), ppl) + if verbose: + print(f"Starsim People created: {len(ppl_df)} agents, traits: {trait_cols}") + return ppl + + +class GPNetwork(ss.Network): + """A Starsim network layer backed by a GeoPops contact layer. + + Either name a built-in layer (``homenet``, ``schoolnet``, ``worknet``, ``gqnet``) + together with the ``pop_export`` directory it lives in, or supply your own edges + via ``csv_path=`` or ``network_df=``. + """ + + def __init__(self, name, edge_weight=1.0, pop_export_dir=None, csv_path=None, + network_df=None, p1_col='p1', p2_col='p2', beta_col=None, + config=None, base_dir=None, seed=0, edges=None, save=False, + verbose=0): + super().__init__() + self.name = name + self.edge_weight = edge_weight + self.p1_col = p1_col + self.p2_col = p2_col + self.beta_col = beta_col + + if csv_path is not None and network_df is not None: + raise ValueError("Provide only one of csv_path or network_df, not both.") + + if network_df is not None: + self.network_df = self._normalize(network_df) + elif csv_path is not None: + if not os.path.exists(csv_path): + raise DataError(f"Custom network file not found: {csv_path}") + self.network_df = self._normalize(pd.read_csv(csv_path), f"CSV '{csv_path}'") + else: + if name not in NETWORK_FILES: + raise ValueError( + f"Unknown network name {name!r}. Built-in names: {list(NETWORK_FILES)}. " + "To use a custom network, pass csv_path=... or network_df=..." + ) + if edges is None: + pop_export_dir = _resolve_pop_export(pop_export_dir, config, None, base_dir) + edges = load_network_edges(pop_export_dir, names=[name], seed=seed, + save=save, verbose=verbose) + self.network_df = edges[name] + + self._populate_edges() + + def _normalize(self, df, source_desc="provided dataframe"): + """Validate and normalize custom network data into p1/p2/edge_weight columns.""" + if not isinstance(df, pd.DataFrame): + raise TypeError(f"network_df must be a pandas DataFrame, got {type(df)}") + + tmp = df.drop(columns=['Unnamed: 0'], errors='ignore') + missing = [c for c in (self.p1_col, self.p2_col) if c not in tmp.columns] + if missing: + raise ValueError( + f"{source_desc} is missing required column(s): {missing}. " + f"Available columns: {list(tmp.columns)}" + ) + + out = pd.DataFrame({ + 'p1': pd.to_numeric(tmp[self.p1_col], errors='coerce'), + 'p2': pd.to_numeric(tmp[self.p2_col], errors='coerce'), + }).dropna(subset=['p1', 'p2']) + out['p1'] = out['p1'].astype(np.int64) + out['p2'] = out['p2'].astype(np.int64) + + if self.beta_col is not None: + if self.beta_col not in tmp.columns: + raise ValueError( + f"beta_col {self.beta_col!r} not found in {source_desc}. " + f"Available columns: {list(tmp.columns)}" + ) + beta = pd.to_numeric(tmp[self.beta_col], errors='coerce') + out['edge_weight'] = beta.loc[out.index].fillna(float(self.edge_weight)).astype(float).values + else: + out['edge_weight'] = float(self.edge_weight) + + if out.empty: + raise ValueError(f"{source_desc} has no valid edges after parsing.") + return out.reset_index(drop=True) + + def _populate_edges(self): + self.edges.p1 = self.network_df['p1'].values + self.edges.p2 = self.network_df['p2'].values + # An explicitly provided scalar edge_weight (e.g. 2.0 for homenet) wins, so + # users can rescale built-in networks without editing CSVs or dataframes. + if float(self.edge_weight) != 1.0: + self.edges.beta = np.full(len(self.network_df), float(self.edge_weight)) + elif 'edge_weight' in self.network_df.columns: + self.edges.beta = self.network_df['edge_weight'].values.astype(float) + else: + self.edges.beta = np.full(len(self.network_df), self.edge_weight) + self.validate() + + def step(self): + self.validate() + + +def starsim_network(name, pop_export_dir=None, edge_weight=1.0, *, save=False, + verbose=0, **kwargs): + """Build one Starsim network layer from a GeoPops population. + + Args: + name: one of :data:`NETWORK_FILES`, or any name if `csv_path`/`network_df` + is given. + pop_export_dir: the run's ``pop_export`` directory. + edge_weight: scalar transmission weight applied to every edge. + save: also write ``pop_export/starsim/net_*.csv``. + verbose: if truthy, log progress. + + Returns: + GPNetwork: a Starsim network layer. + + To build all four layers at once, and read each matrix file only once, use + :func:`starsim_networks`. + """ + return GPNetwork(name, edge_weight=edge_weight, pop_export_dir=pop_export_dir, + save=save, verbose=verbose, **kwargs) + + +def starsim_networks(pop_export_dir=None, names=None, edge_weight=1.0, seed=0, + save=True, verbose=1, **kwargs): + """Build all GeoPops network layers, reading each matrix file exactly once. + + Returns: + list[GPNetwork]: one layer per name, ready to pass to ``ss.Sim(networks=...)``. + """ + pop_export_dir = _resolve_pop_export(pop_export_dir, kwargs.pop('config', None), + None, kwargs.pop('base_dir', None)) + names = list(NETWORK_FILES) if names is None else list(names) + edges = load_network_edges(pop_export_dir, names=names, seed=seed, + save=save, verbose=verbose) + return [GPNetwork(name, edge_weight=edge_weight, edges=edges, **kwargs) for name in names] + + +class SubgroupTracking(ss.Analyzer): + """Track counts of a disease outcome over time, split by an agent attribute. + + Args: + subgroup: name of the ``People`` state to group by (e.g. ``'agegroup'``). + outcome: name of the disease state to count (e.g. ``'infected'``). + state_id: optional additional filter on ``people.state``. + """ + + def __init__(self, subgroup, outcome, name=None, state_id=None, *args, **kwargs): + super().__init__(*args, **kwargs) + self.has_product = False + self.subgroup = subgroup + self.outcome = outcome + self.state_id = state_id + self.n_outcome = {} + if name: + self.name = name + + def step(self): + sim = self.sim + if not self.n_outcome: + self.n_outcome = {group: [] for group in np.unique(sim.people[self.subgroup])} + + disease_name = sim.diseases[0].name.lower() + disease_obj = getattr(sim.people, disease_name, None) + + for group in self.n_outcome: + match = (sim.people[self.subgroup] == group) & (disease_obj[self.outcome] == 1) + if self.state_id is not None: + match = match & (sim.people.state == self.state_id) + self.n_outcome[group].append(len(ss.uids(match))) + + def get_subgroup_data(self): + """Return a DataFrame where rows are subgroups and columns are time steps.""" + df = pd.DataFrame.from_dict(self.n_outcome, orient='index') + df.columns = [f't_{i}' for i in range(len(df.columns))] + df.index.name = self.subgroup + return df.reset_index() diff --git a/src/geopops/utils.py b/src/geopops/utils.py index 443b983..2bed928 100644 --- a/src/geopops/utils.py +++ b/src/geopops/utils.py @@ -4,48 +4,99 @@ """ import numpy as np from dataclasses import dataclass -from typing import Optional import json def tryJSON(filename): try: - with open(filename, 'r') as f: + with open(filename) as f: return json.load(f) except Exception: return {} -@dataclass +class TraitSchema: + """Shared name -> position mapping for the per-person traits of one run. + + One instance is shared by every :class:`PersonData`, so the per-person cost of + carrying config-driven traits is a tuple of values rather than a dict of + name/value pairs. + """ + __slots__ = ("names", "index") + + def __init__(self, names=()): + self.names = tuple(names) + self.index = {name: i for i, name in enumerate(self.names)} + + def values_from(self, mapping): + """Build a trait-value tuple from a name -> value mapping.""" + return tuple(mapping.get(name) for name in self.names) + + def __len__(self): + return len(self.names) + + def __repr__(self): + return f"TraitSchema({list(self.names)!r})" + + +EMPTY_SCHEMA = TraitSchema() + + +@dataclass(slots=True) class PersonData: + """One synthetic person. + + The core demographic fields are fixed. Everything listed in the config's + ``additional_traits`` (sex, race/ethnicity, school sector, ...) is carried in + ``trait_values``, positioned by a :class:`TraitSchema` shared across the whole + population, and is reachable by name --- ``person.hispanic`` works whenever + ``hispanic`` was requested for this run. Keeping traits config-driven rather + than hardcoded means adding a trait needs no code change. + + Memory matters here: one instance exists per person. ``slots=True`` plus a + shared schema costs roughly 240 bytes per person, against ~1.5 kB for a plain + dataclass with one field per trait. + """ hh: tuple sample: int age: int working: bool commuter: bool - com_cat: Optional[int] = None - com_inc: Optional[int] = None - sch_grade: Optional[str] = None - sch_public: Optional[bool] = None - sch_private: Optional[bool] = None - female: Optional[bool] = None - race_white_alone: Optional[bool] = None - race_black_alone: Optional[bool] = None - race_amerindian_or_alaskan: Optional[bool] = None - race_asian_alone: Optional[bool] = None - race_pacific_alone: Optional[bool] = None - race_other_alone: Optional[bool] = None - race_two_or_more: Optional[bool] = None - hispanic: Optional[bool] = None - - -@dataclass + com_cat: int | None = None + com_inc: int | None = None + sch_grade: str | None = None + schema: TraitSchema = EMPTY_SCHEMA + trait_values: tuple = () + + def __getattr__(self, name): + # Only reached when normal (slot) lookup fails, so this cannot shadow a + # real field. Use object.__getattribute__ to avoid recursing on `schema`. + try: + schema = object.__getattribute__(self, "schema") + values = object.__getattribute__(self, "trait_values") + except AttributeError: + raise AttributeError(name) from None + position = schema.index.get(name) + if position is None or position >= len(values): + raise AttributeError( + f"PersonData has no field or trait {name!r}. " + f"Traits available for this run: {list(schema.names)}" + ) + return values[position] + + @property + def traits(self): + """The person's config-driven traits as a ``{name: value}`` dict.""" + return dict(zip(self.schema.names, self.trait_values, strict=False)) + + +@dataclass(slots=True) class Household: sample: int people: list -@dataclass +@dataclass(slots=True) class GQres: type: str residents: list @@ -98,23 +149,30 @@ def ranges(vec): vec = [int(x) for x in vec] x = np.cumsum(vec) starts = np.concatenate([[1], x[:-1] + 1]).astype(int) - return list(zip(starts.tolist(), x.tolist())) + return list(zip(starts.tolist(), x.tolist(), strict=False)) def drawCounts(v, n=1, rng=None): + """Draw `n` items without replacement from a multiset of counts. + + `v` is a vector of counts (e.g. commuters per origin); this draws `n` of them + without replacement, decrements `v` in place, and returns the drawn positions. + + That is exactly the multivariate hypergeometric distribution, so one vectorized + call replaces the former loop of `n` `rng.choice(p=...)` calls (each of which + re-normalized the full probability vector). The result is shuffled so callers + still see draws in random rather than bin order. + """ if rng is None: rng = np.random.default_rng() - result = [] - n = min(n, int(v.sum())) - for _ in range(n): - total = v.sum() - if total == 0: - break - probs = v.astype(float) / total - idx = rng.choice(len(v), p=probs) - v[idx] -= 1 - result.append(idx) - return result + n = min(int(n), int(v.sum())) + if n <= 0: + return [] + drawn = rng.multivariate_hypergeometric(v, n) + v -= drawn + result = np.repeat(np.arange(len(v)), drawn) + rng.shuffle(result) + return result.tolist() def thresh(x, v): diff --git a/src/geopops/workplaces.py b/src/geopops/workplaces.py index 22dd933..e24cecb 100644 --- a/src/geopops/workplaces.py +++ b/src/geopops/workplaces.py @@ -6,7 +6,7 @@ import pandas as pd import os from scipy import sparse -from .utils import tryJSON, lrRound, lrRound_matrix, rowRound, drawCounts, vecmerge +from .utils import tryJSON, lrRound, rowRound, drawCounts, vecmerge from .ipfn import ipfn as IPFN @@ -59,8 +59,10 @@ def group_commuters_by_origin(people, cbgs, ind_codes, rng): continue cat_code = ind_codes[int(cat_idx) - 1] cbg_code = cbgs_inv.get(cbg_idx, '') - workers = [(int(r['id']), int(r['hh']), int(r['cbg']), r['income']) - for _, r in group.iterrows()] + workers = list(zip(group['id'].astype(int).tolist(), + group['hh'].astype(int).tolist(), + group['cbg'].astype(int).tolist(), + group['income'].tolist(), strict=False)) rng.shuffle(workers) worker_keys[cat_code][cbg_code] = workers @@ -79,36 +81,40 @@ def __call__(self, origin, inc_code): def read_workers_by_cat(co_results, data_dir, ind_codes, counties): - """Read # workers by industry category from HH sample summaries.""" + """Number of workers per industry category, per origin CBG. + + Sums the per-industry worker counts of the PUMS households that CO assigned to + each CBG. Done column-wise over a NumPy view of the sample table: the previous + row-at-a-time ``.iloc[i][col]`` lookup built a fresh pandas Series per household + *per category*, which dominated the runtime of this stage. + """ cat_cols = ['com_ind_' + k for k in ind_codes] hh_samps = pd.read_csv(os.path.join(data_dir, 'processed', 'hh_samples.csv'), usecols=['SERIALNO'] + cat_cols, dtype={'SERIALNO': str}) - # Keep HH sample row mapping 1-based for Julia parity. - hh_idx = dict(zip(hh_samps['SERIALNO'], range(1, len(hh_samps) + 1))) + counts = hh_samps[cat_cols].to_numpy(dtype=float) # (n_samples, n_categories) + row_of = {serial: i for i, serial in enumerate(hh_samps['SERIALNO'])} workers_by_cat = {k: {} for k in ind_codes} for co in counties: if co not in co_results: continue - cbg_dict = co_results[co] - for ori, hhvec in cbg_dict.items(): - for cat_code, cat_col in zip(ind_codes, cat_cols): - total = sum( - hh_samps.iloc[hh_idx[x] - 1][cat_col] - for x in hhvec - if x in hh_idx and pd.notna(hh_samps.iloc[hh_idx[x] - 1][cat_col]) - ) + for ori, hhvec in co_results[co].items(): + rows = [row_of[x] for x in hhvec if x in row_of] + totals = np.nansum(counts[rows], axis=0) if rows else np.zeros(len(ind_codes)) + for cat_code, total in zip(ind_codes, totals, strict=False): workers_by_cat[cat_code][ori] = int(total) return workers_by_cat def read_gq_workers_by_cat(gq_summary, ind_codes): """Get # workers in GQs by industry from the GQ summary dataframe.""" - cat_cols = ['ind_' + k for k in ind_codes] - gq_by_cat = {k: {} for k in ind_codes} - for _, r in gq_summary.iterrows(): - for cat_code, cat_col in zip(ind_codes, cat_cols): - gq_by_cat[cat_code][r['geo']] = int(r.get(cat_col, 0)) + geos = gq_summary['geo'].tolist() + gq_by_cat = {} + for cat_code in ind_codes: + col = 'ind_' + cat_code + vals = (gq_summary[col].to_numpy(dtype=float) if col in gq_summary.columns + else np.zeros(len(geos))) + gq_by_cat[cat_code] = dict(zip(geos, vals.astype(np.int64).tolist(), strict=False)) return gq_by_cat @@ -126,19 +132,29 @@ def read_od_matrix(data_dir, k, m, n): def read_outside_origins(data_dir, ind_codes): """Counts of workers commuting from outside the synth area.""" df = pd.read_csv(os.path.join(data_dir, 'processed', 'work_cats_live_outside.csv')) - tmp = dict(zip(df.iloc[:, 0], df.iloc[:, 1])) + tmp = dict(zip(df.iloc[:, 0], df.iloc[:, 1], strict=False)) return {k: int(round(tmp.get('C24030:' + k, 0))) for k in ind_codes} -def calc_od_counts(ind_codes, counties, co_results, gq_summary, data_dir): +def calc_od_counts(ind_codes, counties, co_results, gq_summary, data_dir, commute=None): """Calculate origin-destination counts for each industry. - Returns (origin_labels, dest_labels, od_counts_by_cat). - od_counts_by_cat: dict[cat_code -> dense numpy array of OD counts] + + Args: + commute: optional ``(origin_labels, dest_labels, {cat: csr_matrix})`` as + returned by :func:`generate_commute_matrices`. When given, the OD + proportion matrices are taken from memory instead of being re-read from + the ``od_*.csv.gz`` files that were just written. + + Returns: + tuple: ``(origin_labels, dest_labels, od_counts_by_cat)``, where + ``od_counts_by_cat`` maps each category to a dense OD count array. """ - origin_df = pd.read_csv(os.path.join(data_dir, 'processed', 'od_rows_origins.csv')) - dest_df = pd.read_csv(os.path.join(data_dir, 'processed', 'od_columns_dests.csv')) - origin_labels = origin_df['origin'].astype(str).tolist() - dest_labels = dest_df['dest'].astype(str).tolist() + if commute is not None: + origin_labels, dest_labels, od_props = commute + else: + origin_labels = pd.read_csv(os.path.join(data_dir, 'processed', 'od_rows_origins.csv'))['origin'].astype(str).tolist() + dest_labels = pd.read_csv(os.path.join(data_dir, 'processed', 'od_columns_dests.csv'))['dest'].astype(str).tolist() + od_props = None n_rows = len(origin_labels) n_cols = len(dest_labels) origin_idx = {o: i for i, o in enumerate(origin_labels)} @@ -152,7 +168,12 @@ def calc_od_counts(ind_codes, counties, co_results, gq_summary, data_dir): od_counts_by_cat = {} for k in ind_codes: - M = read_od_matrix(data_dir, k, n_rows, n_cols).toarray() + if od_props is not None: + # float32 -> float64 reproduces exactly what reading the CSV back gives, + # since the CSV stores the float32 values. + M = od_props[k].astype(np.float64).toarray() + else: + M = read_od_matrix(data_dir, k, n_rows, n_cols).toarray() counts = np.zeros((n_rows, n_cols), dtype=np.int64) for code, rownum in origin_idx.items(): hh_n = hhw[k].get(code, 0) + gqw[k].get(code, 0) @@ -166,7 +187,7 @@ def read_county_stats(data_dir): """Employer size stats (lognormal mu, sigma) by county.""" df = pd.read_csv(os.path.join(data_dir, 'processed', 'work_sizes.csv'), usecols=['county', 'mu_ln', 'sigma_ln'], dtype={'county': str}) - return {r['county']: (r['mu_ln'], r['sigma_ln']) for _, r in df.iterrows()} + return dict(zip(df['county'], zip(df['mu_ln'], df['sigma_ln'], strict=False), strict=False)) def read_school_info(data_dir): @@ -181,7 +202,7 @@ def read_school_info(data_dir): schools = pd.read_csv(os.path.join(data_dir, 'processed', 'schools.csv'), usecols=['NCESSCH', 'TEACHERS'], dtype={'NCESSCH': str}) - sch_n_teachers = dict(zip(schools['NCESSCH'], schools['TEACHERS'].astype(int))) + sch_n_teachers = dict(zip(schools['NCESSCH'], schools['TEACHERS'].astype(int), strict=False)) return sch_n_teachers, closest_cbg @@ -209,13 +230,21 @@ def filter_dests(geo, dest_idx, colsums): def pull_inst_workers(count_matrix, dest_idx, origin_labels, workers_by_key, loc_by_key, rng): """Pull worker origins for institutions (schools/GQs) from the OD count matrix. - Modifies count_matrix in place. Returns dict[inst_key -> list[origin_codes]]. + + Modifies `count_matrix` in place. Returns dict[inst_key -> list[origin_codes]]. + + Institutions look for available workers in progressively wider rings around + their own CBG: exact CBG, then tract, then successively shorter geo prefixes, + then the county. """ + # Column sums are maintained incrementally: recomputing the full reduction per + # institution was O(rows x cols) each time, for thousands of institutions. + colsums = count_matrix.sum(axis=0) + inst_emp_origins = {} for inst_id, n_needed in workers_by_key.items(): cbg = loc_by_key[inst_id] geo_areas = [cbg, cbg[:11], cbg[:9], cbg[:7], cbg[:5]] - colsums = count_matrix.sum(axis=0) dest_lists = [filter_dests(geo, dest_idx, colsums) for geo in geo_areas] for dl in dest_lists: rng.shuffle(dl) @@ -234,6 +263,7 @@ def pull_inst_workers(count_matrix, dest_idx, origin_labels, workers_by_key, loc col_view = count_matrix[:, col].copy() drawn = drawCounts(col_view, draw_n, rng) count_matrix[:, col] = col_view + colsums[col] -= len(drawn) o_idxs.extend(drawn) remaining -= draw_n if remaining < 1: @@ -258,14 +288,17 @@ def generate_workplaces(count_matrix, dest_idx, origin_labels, county_stats, dra sigma = sigma + 0.1 # slight adjustment per Julia code dests = {code: idx for code, idx in dest_idx.items() if code[:5] == co} for dest_code, col in dests.items(): - n = int(count_matrix[:, col].sum()) - if n > 0: - sizes = split_lognormal(n, mu, sigma, draws, rng) - for work_i, emp_size in enumerate(sizes): - col_view = count_matrix[:, col].copy() - o_idxs = drawCounts(col_view, emp_size, rng) - count_matrix[:, col] = col_view - work_origins[(work_i + 1, cat_idx, dest_code)] = [origin_labels[i] for i in o_idxs] + # One column slice per destination rather than per workplace: the column + # is mutated in place across all its workplaces, then written back once. + col_view = count_matrix[:, col].copy() + n = int(col_view.sum()) + if n <= 0: + continue + sizes = split_lognormal(n, mu, sigma, draws, rng) + for work_i, emp_size in enumerate(sizes): + o_idxs = drawCounts(col_view, emp_size, rng) + work_origins[(work_i + 1, cat_idx, dest_code)] = [origin_labels[i] for i in o_idxs] + count_matrix[:, col] = col_view return work_origins @@ -316,9 +349,18 @@ def assign_workers(emp_origins, workers_by_origin, cidx_by_origin, dummy_fn, rng return workers, dummies, missing_origin, ran_out -def generate_commute_matrices(data_dir): - """Generate per-industry OD proportion matrices using IPF. - Writes results to processed/od_*.csv.gz. Called before generate_jobs_and_workers. +def generate_commute_matrices(data_dir, save=True, verbose=1): + """Generate per-industry origin-destination proportion matrices using IPF. + + Args: + data_dir: run directory containing ``processed/``. + save: if True (default) also write ``processed/od_*.csv.gz`` and the + origin/destination label files. These are a useful checkpoint, but the + matrices are returned either way so the caller need not read them back. + verbose: if truthy, log the files written. + + Returns: + tuple: ``(origin_labels, dest_labels, {cat_code: csr_matrix})``. """ wp_codes = tryJSON(os.path.join(data_dir, 'processed', 'codes.json')) ind_codes = wp_codes.get('ind_codes', []) @@ -343,7 +385,11 @@ def generate_commute_matrices(data_dir): n_ori = len(origin_idxs) n_dest = len(dest_idxs) n_ind = len(ind_codes) - res_iod = [sparse.lil_matrix((n_ori, n_dest), dtype=np.float32) for _ in ind_codes] + # Accumulate COO triplets per industry; building one matrix at the end is much + # faster than incremental assignment into a lil_matrix. + res_rows = [[] for _ in ind_codes] + res_cols = [[] for _ in ind_codes] + res_vals = [[] for _ in ind_codes] for o in range(n_ori): @@ -382,43 +428,67 @@ def generate_commute_matrices(data_dir): new_m[zero_rows, :] = d_margin / d_margin.sum() for i in range(n_ind): - res_iod[i][o, d_idxs] = new_m[i, :] - - # Write results - proc_dir = os.path.join(data_dir, 'processed') - pd.DataFrame({'idx': range(1, n_ori + 1), 'origin': origin_idxs}).to_csv( - os.path.join(proc_dir, 'od_rows_origins.csv'), index=False) - pd.DataFrame({'idx': range(1, n_dest + 1), 'dest': dest_idxs}).to_csv( - os.path.join(proc_dir, 'od_columns_dests.csv'), index=False) - for i, k in enumerate(ind_codes): - m = res_iod[i].tocsr() - rows, cols, vals = sparse.find(m) - df = pd.DataFrame({'origin': rows + 1, 'dest': cols + 1, 'p': vals.astype(np.float32)}) - df.to_csv(os.path.join(proc_dir, f'od_{k}.csv.gz'), index=False, compression='gzip') - - -def generate_jobs_and_workers(people, cbgs, gqs, co_results, gq_summary, data_dir, random_seed=None): + vals = new_m[i, :] + nz = np.flatnonzero(vals) + if nz.size: + res_rows[i].append(np.full(nz.size, o, dtype=np.int32)) + res_cols[i].append(np.asarray(d_idxs, dtype=np.int32)[nz]) + res_vals[i].append(vals[nz].astype(np.float32)) + + def build(i): + if not res_rows[i]: + return sparse.csr_matrix((n_ori, n_dest), dtype=np.float32) + return sparse.coo_matrix( + (np.concatenate(res_vals[i]), + (np.concatenate(res_rows[i]), np.concatenate(res_cols[i]))), + shape=(n_ori, n_dest), dtype=np.float32).tocsr() + + od_props = {k: build(i) for i, k in enumerate(ind_codes)} + + if save: + proc_dir = os.path.join(data_dir, 'processed') + pd.DataFrame({'idx': range(1, n_ori + 1), 'origin': origin_idxs}).to_csv( + os.path.join(proc_dir, 'od_rows_origins.csv'), index=False) + pd.DataFrame({'idx': range(1, n_dest + 1), 'dest': dest_idxs}).to_csv( + os.path.join(proc_dir, 'od_columns_dests.csv'), index=False) + if verbose: + print("-- processed/od_rows_origins.csv") + print("-- processed/od_columns_dests.csv") + for k in ind_codes: + rows, cols, vals = sparse.find(od_props[k]) + pd.DataFrame({'origin': rows + 1, 'dest': cols + 1, + 'p': vals.astype(np.float32)}).to_csv( + os.path.join(proc_dir, f'od_{k}.csv.gz'), index=False, compression='gzip') + if verbose: + print(f"-- processed/od_{k}.csv.gz") + + return origin_idxs, dest_idxs, od_props + + +def generate_jobs_and_workers(people, cbgs, gqs, co_results, gq_summary, data_dir, + random_seed=None, config=None, save_intermediates=True, + verbose=1): """Generate workplaces and assign workers. Returns (company_workers, sch_workers, gq_workers, outside_workers, dummies). Each *_workers is dict[key -> list[worker_tuple]]. """ rng = np.random.default_rng(random_seed) - config = tryJSON(os.path.join(data_dir, 'config.json')) + if config is None: + config = tryJSON(os.path.join(data_dir, 'config.json')) wp_codes = tryJSON(os.path.join(data_dir, 'processed', 'codes.json')) ind_codes = wp_codes.get('ind_codes', []) ind_idxs = {k: i + 1 for i, k in enumerate(ind_codes)} - cbgs_inv = {v: k for k, v in cbgs.items()} - counties = sorted(set(v[:5] for v in cbgs.keys())) + counties = sorted({v[:5] for v in cbgs}) worker_keys = group_commuters_by_origin(people, cbgs, ind_codes, rng) dummy_fn = DummyGenerator() - # Generate commute matrices - generate_commute_matrices(data_dir) + # Generate commute matrices, then use them directly rather than writing them to + # gzipped CSV and immediately reading them back. + commute = generate_commute_matrices(data_dir, save=save_intermediates, verbose=verbose) origin_labels, dest_labels, od_counts_by_cat = calc_od_counts( - ind_codes, counties, co_results, gq_summary, data_dir) - dest_idx = {d: i for i, d in enumerate(dest_labels)} + ind_codes, counties, co_results, gq_summary, data_dir, commute=commute) county_stats = read_county_stats(data_dir) draws_by_county = {co: [] for co in counties} @@ -441,7 +511,7 @@ def generate_jobs_and_workers(people, cbgs, gqs, co_results, gq_summary, data_di work_outside_counts = od_counts[:, -1].copy() od_counts = od_counts[:, :-1] dest_idx_local = {d: i for i, d in enumerate(dest_labels[:-1])} - work_outside = dict(zip(origin_labels, work_outside_counts)) + work_outside = dict(zip(origin_labels, work_outside_counts, strict=False)) # Schools if ckey == 'EDU': diff --git a/tests/config.json b/tests/config.json new file mode 100644 index 0000000..3c8d17c --- /dev/null +++ b/tests/config.json @@ -0,0 +1,108 @@ +{ + "path": "data", + "main_year": 2019, + "decennial_year": 2010, + "geos": [ + "45083" + ], + "commute_states": [ + "45", + "37" + ], + "use_pums": [ + "45", + "37" + ], + "acs_required": [ + "B01001", + "B09018", + "B09019", + "B09020", + "B09021", + "B11004", + "B11012", + "B11016", + "B19001", + "B22010", + "B23009", + "B23025", + "B25006", + "B11001H", + "B11001I", + "C24010", + "C24030" + ], + "dec_required": [ + "P43", + "P18" + ], + "inc_adj": 1.010145, + "inc_cats": [ + "q1_1", + "q1_2", + "q1_3", + "q2", + "q3", + "q4", + "q5" + ], + "inc_cols": [ + [ + "Less than $10,000" + ], + [ + "$10,000 to $14,999", + "$15,000 to $19,999", + "$20,000 to $24,999" + ], + [ + "$25,000 to $29,999", + "$30,000 to $34,999", + "$35,000 to $39,999" + ], + [ + "$40,000 to $44,999", + "$45,000 to $49,999", + "$50,000 to $59,999", + "$60,000 to $74,999" + ], + [ + "$75,000 to $99,999", + "$100,000 to $124,999" + ], + [ + "$125,000 to $149,999", + "$150,000 to $199,999" + ], + [ + "$200,000 or more" + ] + ], + "additional_traits": [ + "sch_public", + "sch_private", + "female", + "race_black_alone", + "white_non_hispanic", + "hispanic" + ], + "LODES_annual_income_boundary": 40000, + "income_associativity_coefficient": 0.9, + "school_associativity_coefficient": 0.9, + "inst_res_per_worker": 10, + "noninst_res_per_worker": 50, + "min_gq_workers": 2, + "min_gq_residents": 20, + "n_closest_schools": 4, + "p_closest_school": 0.9, + "CO_crit_val": 15.0, + "CO_cooldown": 0.99, + "CO_maxgens": 200000, + "workplace_K": 8, + "school_K": 12, + "gq_K": 12, + "netw_K": 8, + "netw_B": 0.25, + "census_api_key": "ca5103b853992f7a73bcc3d82b96322486efafeb", + "julia_env_path": "/home/cliffk/.julia/environments/v1.9/" +} \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..1a88ec0 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,59 @@ +"""Shared pytest fixtures for the GeoPops test suite.""" +import json +import os +import shutil + +import pytest + +TESTS_DIR = os.path.dirname(os.path.abspath(__file__)) +FIXTURE_DATA = os.path.join(TESTS_DIR, "data") + +#: Spartanburg County, SC --- the small fixture the slow tests run against +FIXTURE_PARS = { + "path": FIXTURE_DATA, + "main_year": 2019, + "geos": ["45083"], + "commute_states": ["45", "37"], + "use_pums": ["45", "37"], + "random_seed": 42, +} + +requires_fixture_data = pytest.mark.skipif( + not os.path.exists(os.path.join(FIXTURE_DATA, "processed", "acs_targets.csv")), + reason="tests/data/processed not present; run the download+process steps first", +) + + +@pytest.fixture +def base_config(): + """A minimal valid config dict, not tied to any data on disk.""" + import geopops + return geopops.make_config(path="unused", geos=["45083"], main_year=2019, + commute_states=["45"], use_pums=["45"], random_seed=42) + + +@pytest.fixture(scope="session") +def fixture_config(): + """Config pointing at the checked-in Spartanburg fixture data. + + The fixture's own ``config.json`` is used when present, because the processed + CSVs were generated with that config's ``additional_traits``. + """ + import geopops + fixture_cfg_path = os.path.join(FIXTURE_DATA, "config.json") + template = None + if os.path.exists(fixture_cfg_path): + with open(fixture_cfg_path) as f: + template = json.load(f) + return geopops.make_config(template=template, **FIXTURE_PARS) + + +@pytest.fixture(scope="session") +def run_dir(tmp_path_factory): + """A scratch copy of the fixture data, so tests never write into tests/data.""" + if not os.path.exists(os.path.join(FIXTURE_DATA, "processed")): + pytest.skip("tests/data/processed not present") + dest = tmp_path_factory.mktemp("geopops_run") / "data" + shutil.copytree(FIXTURE_DATA, dest) + shutil.rmtree(dest / "pop_export", ignore_errors=True) + return str(dest) diff --git a/tests/golden_hashes.json b/tests/golden_hashes.json new file mode 100644 index 0000000..03db1d3 --- /dev/null +++ b/tests/golden_hashes.json @@ -0,0 +1,19 @@ +{ + "_co_maxgens": 5000, + "_seed": 42, + "adj_mat_keys.csv": "ea5d1a9b40e89395", + "adj_upper_triang_gq.mtx": "6c758ef254b0531b", + "adj_upper_triang_hh.mtx": "7c834fbb5d99c77a", + "adj_upper_triang_sch.mtx": "30581ee30ac5fe55", + "adj_upper_triang_wp.mtx": "0358f188b0d57064", + "cbg_idxs.csv": "90fec38dffb34d70", + "company_workers.csv": "87d1ca31ab56460f", + "gq_residents.csv": "8159fbd586b27042", + "gq_workers.csv": "db47f7fdb0e52c9a", + "gqs.csv": "a5f0cfa675e46e1a", + "hh.csv": "74912439a5593177", + "outside_workers.csv": "30d9f6cb37abd2be", + "people.csv": "e207f4e638eaf912", + "sch_students.csv": "e072acc81ecbc4d6", + "sch_workers.csv": "e04e3fd869ab2113" +} \ No newline at end of file diff --git a/tests/test_co.py b/tests/test_co.py new file mode 100644 index 0000000..d191414 --- /dev/null +++ b/tests/test_co.py @@ -0,0 +1,118 @@ +"""Unit tests for the combinatorial optimization (simulated annealing) step.""" +import numpy as np +import pytest + +from geopops import co + + +@pytest.fixture +def pool(): + """A synthetic sample pool and a target drawn from it, so a good fit exists.""" + rng = np.random.default_rng(0) + samples = rng.integers(0, 4, size=(500, 12)).astype(np.int64) + idxs = np.arange(500) + targ = samples[rng.integers(0, 500, 40)].sum(axis=0, keepdims=True) + return samples, idxs, targ + + +PARAMS = dict(maxgens=2000, critval=-1.0, cooldown=0.99) # never exit early + + +class TestFTdist: + def test_zero_for_identical(self): + v = np.array([[1, 2, 3]]) + assert co.FTdist(v, v) == 0.0 + + def test_positive_and_symmetric(self): + a, b = np.array([[1, 2, 3]]), np.array([[4, 0, 3]]) + assert co.FTdist(a, b) > 0 + assert co.FTdist(a, b) == pytest.approx(co.FTdist(b, a)) + + +class TestAnneal: + def test_returns_requested_number_of_samples(self, pool): + samples, idxs, targ = pool + result, gens, score, temp = co.anneal(samples, idxs, targ, 40, PARAMS, + np.random.default_rng(1)) + assert len(result) == 40 + assert set(result).issubset(set(idxs)) + assert gens > 0 and np.isfinite(score) + + def test_improves_on_the_starting_fit(self, pool): + samples, idxs, targ = pool + rng = np.random.default_rng(1) + start = samples[rng.integers(0, len(samples), 40)].sum(axis=0, keepdims=True) + start_score = co.FTdist(start, targ) + _, _, score, _ = co.anneal(samples, idxs, targ, 40, PARAMS, np.random.default_rng(1)) + assert score < start_score + + def test_deterministic_given_a_seed(self, pool): + samples, idxs, targ = pool + a = co.anneal(samples, idxs, targ, 40, PARAMS, np.random.default_rng(3)) + b = co.anneal(samples, idxs, targ, 40, PARAMS, np.random.default_rng(3)) + assert np.array_equal(a[0], b[0]) and a[1:] == b[1:] + + def test_incremental_summary_matches_a_full_recompute(self, pool): + """The inner loop updates its running sum incrementally; check it stays exact.""" + samples, idxs, targ = pool + result, _, score, _ = co.anneal(samples, idxs, targ, 40, PARAMS, + np.random.default_rng(5)) + recomputed = co.FTdist(samples[result].sum(axis=0, keepdims=True), targ) + assert score == pytest.approx(recomputed, abs=1e-9) + + def test_empty_pool(self): + result, gens, score, temp = co.anneal( + np.zeros((0, 5), np.int64), np.array([], int), np.zeros((1, 5), np.int64), + 10, PARAMS, np.random.default_rng(0)) + assert len(result) == 0 and gens == 0 and score == float("inf") + + def test_stops_at_critval(self, pool): + samples, idxs, targ = pool + params = dict(maxgens=10**6, critval=1e9, cooldown=0.99) + _, gens, _, _ = co.anneal(samples, idxs, targ, 40, params, np.random.default_rng(0)) + assert gens == 1 # the criterion is met immediately + + def test_respects_maxgens(self, pool): + samples, idxs, targ = pool + _, gens, _, _ = co.anneal(samples, idxs, targ, 40, dict(PARAMS, maxgens=50), + np.random.default_rng(0)) + assert gens == 51 + + +class TestUrbanizationLookup: + def test_urban(self): + assert co.urbanization_lookup(np.array([1.0, 0.5, 0.0]), 1.0).tolist() == [True, False, False] + + def test_rural(self): + assert co.urbanization_lookup(np.array([1.0, 0.5, 0.0]), 0.0).tolist() == [False, False, True] + + def test_mixed_uses_a_window(self): + got = co.urbanization_lookup(np.array([0.5, 0.55, 0.9]), 0.5) + assert got.tolist() == [True, True, False] + + def test_nan_never_matches(self): + assert not co.urbanization_lookup(np.array([np.nan]), 0.5).any() + assert not co.urbanization_lookup(np.array([np.nan]), 1.0).any() + + +class TestSampleLookup: + def test_shares_index_arrays_between_equal_targets(self): + """Equal targets must share one array so callers can cache the sub-pool.""" + import pandas as pd + df = pd.DataFrame({"st_puma": ["a", "b", "a", "c"]}) + out = co.sample_lookup(df, "st_puma", ["a", "b", "a"]) + assert [k for k, _ in out] == ["a", "b", "a"] + assert out[0][1] is out[2][1] + assert out[0][1].tolist() == [0, 2] + + def test_missing_values_never_match(self): + import pandas as pd + df = pd.DataFrame({"county": ["01", None, "01"]}) + (_, idx), = co.sample_lookup(df, "county", ["01"]) + assert idx.tolist() == [0, 2] + + def test_absent_target_gives_empty(self): + import pandas as pd + df = pd.DataFrame({"county": ["01"]}) + (_, idx), = co.sample_lookup(df, "county", ["99"]) + assert idx.tolist() == [] diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..396431e --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,106 @@ +"""Unit tests for config loading, merging, overriding, and validation.""" +import json + +import pytest + +import geopops +from geopops import config as cfgmod +from geopops.exceptions import ConfigError + + +class TestMakeConfig: + def test_applies_overrides(self): + c = geopops.make_config(path="out", geos=["45083"], main_year=2019, random_seed=1) + assert c["path"] == "out" and c["geos"] == ["45083"] and c["main_year"] == 2019 + + def test_derives_decennial_year(self): + assert geopops.make_config(path="o", geos=["1"], main_year=2019, random_seed=1)["decennial_year"] == 2010 + assert geopops.make_config(path="o", geos=["1"], main_year=2021, random_seed=1)["decennial_year"] == 2020 + + def test_backfills_required_tables(self): + c = geopops.make_config(template={"path": "o", "geos": ["1"], "main_year": 2019, + "random_seed": 1}) + assert "B01001" in c["acs_required"] + assert c["dec_required"] == ["P43", "P18"] + + def test_rejects_unknown_override(self): + # A typo used to be silently ignored, producing a run with default settings + with pytest.raises(ConfigError, match="geoss"): + geopops.make_config(path="o", geos=["1"], main_year=2019, geoss=["45"]) + + def test_does_not_mutate_the_template(self): + template = {"path": "o", "geos": ["1"], "main_year": 2019, "random_seed": 1} + geopops.make_config(template=template, main_year=2021) + assert template["main_year"] == 2019 + + def test_warns_when_unseeded(self): + with pytest.warns(UserWarning, match="reproducible"): + geopops.make_config(path="o", geos=["1"], main_year=2019) + + +class TestValidation: + @pytest.mark.parametrize("missing", ["path", "geos", "main_year"]) + def test_missing_required_key(self, missing): + c = {"path": "o", "geos": ["1"], "main_year": 2019, "random_seed": 1} + del c[missing] + with pytest.raises(ConfigError, match=missing): + cfgmod.validate_config(c) + + def test_unparseable_year_is_rejected(self): + # This used to fall back to 2010 and silently produce wrong-vintage data + with pytest.raises(ConfigError, match="main_year"): + cfgmod.compute_decennial_year("twenty-nineteen") + + def test_empty_geos_rejected(self): + with pytest.raises(ConfigError, match="geos"): + cfgmod.validate_config({"path": "o", "geos": [], "main_year": 2019, + "random_seed": 1}) + + +class TestSaveLoad: + def test_round_trip(self, tmp_path, base_config): + path = cfgmod.save_config(base_config, str(tmp_path / "config.json")) + assert json.load(open(path))["geos"] == base_config["geos"] + + def test_defaults_to_run_directory_not_the_package(self, tmp_path, base_config): + base_config["path"] = str(tmp_path / "run") + path = cfgmod.save_config(base_config) + assert path == str(tmp_path / "run" / "config.json") + # crucially, nothing was written into the installed package + assert cfgmod.BASE_DIR not in path + + def test_directory_target_appends_filename(self, tmp_path, base_config): + path = cfgmod.save_config(base_config, str(tmp_path)) + assert path.endswith("config.json") + + def test_sanitize_strips_secrets(self, tmp_path, base_config): + base_config["census_api_key"] = "secret-key-value" + path = cfgmod.save_config(base_config, str(tmp_path / "t.json"), sanitize=True) + written = json.load(open(path)) + assert written["census_api_key"] is None + assert "secret-key-value" not in open(path).read() + + def test_local_overrides_are_merged(self, tmp_path): + (tmp_path / "config.json").write_text(json.dumps( + {"path": "o", "geos": ["1"], "main_year": 2019, "CO_maxgens": 10})) + (tmp_path / "config.local.json").write_text(json.dumps({"CO_maxgens": 999})) + assert cfgmod.load_config(str(tmp_path))["CO_maxgens"] == 999 + + def test_missing_config_raises_clearly(self, tmp_path): + with pytest.raises(ConfigError, match="not found"): + cfgmod.load_config(str(tmp_path / "nowhere")) + + +class TestMakeConfigSavesToTheRunDirectory: + def test_save_writes_next_to_the_output(self, tmp_path): + c = geopops.make_config(path=str(tmp_path), geos=["45083"], main_year=2019, + random_seed=1, save=True) + assert (tmp_path / "config.json").exists() + assert json.load(open(tmp_path / "config.json"))["geos"] == c["geos"] + + def test_never_writes_into_the_installed_package(self, tmp_path): + """Regression: config used to be written into site-packages.""" + before = open(cfgmod.TEMPLATE_PATH).read() + geopops.make_config(path=str(tmp_path), geos=["45083"], main_year=2019, + random_seed=1, save=True) + assert open(cfgmod.TEMPLATE_PATH).read() == before diff --git a/tests/test_julia_workflow.py b/tests/test_julia_workflow.py deleted file mode 100644 index f928a96..0000000 --- a/tests/test_julia_workflow.py +++ /dev/null @@ -1,54 +0,0 @@ -""" -Minimal tests of running GeoPops with Julia - -Prerequisites: -- Julia installed -- .env file with CENSUS_API_KEY and JULIA_ENV_PATH (see .env.example) -""" - -import sciris as sc -import geopops - -pars_geopops = {'path': "data", # Set a folder where you want output files to be stored - 'main_year': 2019, # Year of data - 'geos': ["45083"], # State or county fips of your geographical location of interest. Example of Spartanburg SC - 'commute_states': ["45","37"], # State fips of commute data to download. Example of SC, NC - 'use_pums': ["45","37"], # State fips of PUMS data to download. Example of SC, NC - } - -c = geopops.WriteConfig(**pars_geopops) # Define parameters for pop generation in config.json -# c.get_pars() # View all parameters from config.json - - -@sc.timer() -def test_julia_CO(): - j = geopops.RunJulia() - j.CO() - return - - -@sc.timer() -def test_julia_synthpop(): - j = geopops.RunJulia() - j.SynthPop() - return - - -@sc.timer() -def test_export(): - j = geopops.RunJulia() - j.Export() - ppl = geopops.ForStarsim.People() - h = geopops.ForStarsim.GPNetwork(name='homenet', beta_value=1.0) - s = geopops.ForStarsim.GPNetwork(name='schoolnet', beta_value=1.0) - w = geopops.ForStarsim.GPNetwork(name='worknet', beta_value=1.0) - g = geopops.ForStarsim.GPNetwork(name='gqnet', beta_value=1.0) - return ppl, h, s, w, g - - -if __name__ == "__main__": - T = sc.timer() - test_julia_CO() - test_julia_synthpop() - test_export() - T.toc() \ No newline at end of file diff --git a/tests/test_networks.py b/tests/test_networks.py new file mode 100644 index 0000000..9f957d3 --- /dev/null +++ b/tests/test_networks.py @@ -0,0 +1,95 @@ +"""Unit tests for contact network construction.""" +import numpy as np +import pytest + +from geopops import networks + + +def keys(n, groups=None): + """n person keys of the form (p_id, hh_id, cbg_id, group_label).""" + groups = groups or ["g0"] * n + return [(i, 1, 1, groups[i]) for i in range(n)] + + +class TestConnectComplete: + def test_edge_count(self): + assert len(networks.connect_complete(keys(5))) == 10 # 5 choose 2 + + @pytest.mark.parametrize("n", [0, 1]) + def test_too_small_to_connect(self, n): + assert networks.connect_complete(keys(n)) == [] + + def test_deduplicates_keys(self): + k = keys(3) + assert len(networks.connect_complete(k + k)) == 3 + + +class TestConnectSmallWorld: + def test_below_min_n_is_complete(self): + edges = networks.connect_small_world(keys(4), K=4, min_N=10, B=0.25, + rng=np.random.default_rng(0)) + assert len(edges) == 6 + + def test_mean_degree_near_k(self): + n, K = 200, 8 + edges = networks.connect_small_world(keys(n), K=K, min_N=10, B=0.25, + rng=np.random.default_rng(0)) + assert 2 * len(edges) / n == pytest.approx(K, abs=1) + + def test_deterministic_given_a_seed(self): + a = networks.connect_small_world(keys(50), 6, 10, 0.25, np.random.default_rng(1)) + b = networks.connect_small_world(keys(50), 6, 10, 0.25, np.random.default_rng(1)) + assert a == b + + def test_no_self_loops(self): + edges = networks.connect_small_world(keys(60), 6, 10, 0.25, np.random.default_rng(2)) + assert all(u != v for u, v in edges) + + +class TestConnectSBM: + def test_below_min_n_is_complete(self): + edges = networks.connect_SBM(keys(4), K=8, min_N=10, assoc_coeff=0.9, + rng=np.random.default_rng(0)) + assert len(edges) == 6 + + def test_no_isolated_nodes(self): + k = keys(100, [f"g{i % 3}" for i in range(100)]) + edges = networks.connect_SBM(k, K=6, min_N=8, assoc_coeff=0.9, + rng=np.random.default_rng(0)) + connected = {u for u, _ in edges} | {v for _, v in edges} + assert len(connected) == 100 + + def test_assortativity_concentrates_edges_within_groups(self): + k = keys(200, [f"g{i % 2}" for i in range(200)]) + rng = np.random.default_rng(0) + high = networks.connect_SBM(k, 8, 10, assoc_coeff=1.0, rng=rng) + low = networks.connect_SBM(k, 8, 10, assoc_coeff=0.0, rng=rng) + + def within(edges): + return sum(u[3] == v[3] for u, v in edges) / max(len(edges), 1) + + assert within(high) > within(low) + assert within(high) > 0.95 + + def test_single_group(self): + edges = networks.connect_SBM(keys(50), 6, 10, 0.9, use_groups=False, + rng=np.random.default_rng(0)) + assert len(edges) > 0 + + +class TestSpFromGroups: + def test_matrix_is_symmetric_and_boolean(self): + p_idxs = {(i, 1, 1): i for i in range(5)} + m = networks.sp_from_groups(networks.connect_complete, [keys(5)], p_idxs) + assert m.shape == (5, 5) + assert m.dtype == bool + assert (m != m.T).nnz == 0 + + def test_unknown_keys_are_dropped(self): + p_idxs = {(0, 1, 1): 0, (1, 1, 1): 1} # key (2,1,1) is absent + m = networks.sp_from_groups(networks.connect_complete, [keys(3)], p_idxs) + assert m.shape == (2, 2) and m.nnz == 2 + + def test_empty_input(self): + m = networks.sp_from_groups(networks.connect_complete, [], {(0, 1, 1): 0}) + assert m.nnz == 0 diff --git a/tests/test_python_workflow.py b/tests/test_python_workflow.py deleted file mode 100644 index 4ec1cfb..0000000 --- a/tests/test_python_workflow.py +++ /dev/null @@ -1,96 +0,0 @@ -""" -Tests for the pure-Python GeoPops pipeline (RunPython). - -Prerequisites: -- Preprocessed data in tests/data/processed/ (or run this script with redownload=True) -""" - -import sciris as sc -import geopops -import pytest - -pars_geopops = {'path': "data", - 'main_year': 2019, - 'geos': ["45083"], - 'commute_states': ["45","37"], - 'use_pums': ["45","37"], - } - -c = geopops.WriteConfig(**pars_geopops) - - -@sc.timer() -@pytest.mark.skip(reason="Manual run only (too slow for automated tests)") -def test_download(): - """ Check that download works (~10 min)""" - d = geopops.DownloadData(auto_run=True) - return d - - -@sc.timer() -@pytest.mark.skip(reason="Manual run only (too slow for automated tests)") -def test_processing(): - """ Check that data processing works (~5 min)""" - p = geopops.ProcessData(auto_run=True) # auto_run=True to run all - return p - - -@pytest.fixture(scope="module") -def runner(): - """Shared RunPython instance for sequential pipeline tests.""" - r = geopops.RunPython() - return r - - -@sc.timer() -def test_python_CO(runner): - """Test combinatorial optimization in Python.""" - runner.CO() - assert runner.co_results is not None - assert len(runner.co_results) > 0 - for county, cbg_dict in runner.co_results.items(): - assert len(cbg_dict) > 0 - for cbg, serials in cbg_dict.items(): - assert len(serials) > 0 - return - - -@sc.timer() -def test_python_synthpop(runner): - """Test synthetic population generation in Python.""" - runner.SynthPop() - assert runner.people is not None - assert runner.households is not None - assert len(runner.people) > 0 - assert len(runner.households) > 0 - return - - -@sc.timer() -def test_python_export(runner): - """Test export and ForStarsim integration.""" - runner.Export() - ppl = geopops.ForStarsim.People() - h = geopops.ForStarsim.GPNetwork(name='homenet', edge_weight=1.0) - s = geopops.ForStarsim.GPNetwork(name='schoolnet', edge_weight=1.0) - w = geopops.ForStarsim.GPNetwork(name='worknet', edge_weight=1.0) - g = geopops.ForStarsim.GPNetwork(name='gqnet', edge_weight=1.0) - return ppl, h, s, w, g - - -if __name__ == "__main__": - T = sc.timer() - - # Download & process data files - redownload = False - if redownload: - test_download() - test_processing() - - # Run GeoPops on the data - r = geopops.RunPython() - test_python_CO(r) - test_python_synthpop(r) - outputs = test_python_export(r) - - T.toc() diff --git a/tests/test_race_ethnicity.ipynb b/tests/test_race_ethnicity.ipynb index 3f335ee..6eba265 100644 --- a/tests/test_race_ethnicity.ipynb +++ b/tests/test_race_ethnicity.ipynb @@ -341,7 +341,7 @@ ], "source": [ "# Download ACS tables with individual level race totals\n", - "from geopops.download_data import get_census_metadata, get_census_data, fips_info\n", + "from geopops.sources import get_census_metadata, get_census_data\n", "\n", "# Use the same config / year as DownloadData\n", "with open('pops/md_24027/config.json') as f:\n", diff --git a/tests/test_regression.py b/tests/test_regression.py new file mode 100644 index 0000000..b45d30f --- /dev/null +++ b/tests/test_regression.py @@ -0,0 +1,101 @@ +"""Golden-output regression test. + +Pins the pipeline's output for a fixed seed so refactors can be checked for +behaviour changes. When a change is *intended*, regenerate the baseline:: + + python tests/test_regression.py --update + +and review the diff to ``tests/golden_hashes.json`` as part of the change. +""" +import hashlib +import json +import os +import sys + +import pytest + +import geopops +from conftest import requires_fixture_data, FIXTURE_PARS, FIXTURE_DATA + +GOLDEN_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "golden_hashes.json") +SEED = 42 +CO_MAXGENS = 5000 + +#: Outputs whose content should be fully determined by the seed +TRACKED = [ + "cbg_idxs.csv", "hh.csv", "people.csv", "sch_students.csv", "gqs.csv", + "gq_residents.csv", "adj_mat_keys.csv", "adj_upper_triang_hh.mtx", + "sch_workers.csv", "company_workers.csv", "gq_workers.csv", + "outside_workers.csv", "adj_upper_triang_sch.mtx", "adj_upper_triang_wp.mtx", + "adj_upper_triang_gq.mtx", +] + + +def _build_config(run_dir): + fixture_cfg = os.path.join(FIXTURE_DATA, "config.json") + template = json.load(open(fixture_cfg)) if os.path.exists(fixture_cfg) else None + cfg = geopops.make_config(template=template, **FIXTURE_PARS) + cfg["path"] = run_dir + cfg["CO_maxgens"] = CO_MAXGENS + return cfg + + +def compute_hashes(run_dir): + """Generate a population and hash its outputs.""" + pop = geopops.generate_pop(_build_config(run_dir), seed=SEED, verbose=0) + out = {} + for name in TRACKED: + path = os.path.join(pop.pop_export_dir, name) + out[name] = hashlib.sha256(open(path, "rb").read()).hexdigest()[:16] + return out + + +@pytest.mark.slow +@requires_fixture_data +def test_outputs_match_golden(run_dir): + if not os.path.exists(GOLDEN_PATH): + pytest.skip(f"No baseline at {GOLDEN_PATH}; create it with " + f"`python tests/test_regression.py --update`") + golden = json.load(open(GOLDEN_PATH)) + if golden.get("_seed") != SEED or golden.get("_co_maxgens") != CO_MAXGENS: + pytest.skip("Baseline was recorded with different settings; regenerate it.") + + actual = compute_hashes(run_dir) + changed = [n for n in TRACKED if golden.get(n) != actual.get(n)] + assert not changed, ( + "Pipeline output changed for these files: " + ", ".join(changed) + + "\nIf this is intended, regenerate the baseline with " + "`python tests/test_regression.py --update` and review the diff." + ) + + +@pytest.mark.slow +@requires_fixture_data +def test_same_seed_gives_the_same_population(tmp_path): + """Two runs with one seed must agree exactly.""" + import shutil + hashes = [] + for i in range(2): + d = tmp_path / f"run{i}" + shutil.copytree(FIXTURE_DATA, d) + shutil.rmtree(d / "pop_export", ignore_errors=True) + hashes.append(compute_hashes(str(d))) + differing = [n for n in TRACKED if hashes[0][n] != hashes[1][n]] + assert not differing, f"Not reproducible under a fixed seed: {differing}" + + +if __name__ == "__main__": + import shutil + import tempfile + + if "--update" not in sys.argv: + print(__doc__) + sys.exit(1) + tmp = tempfile.mkdtemp() + run_dir = os.path.join(tmp, "data") + shutil.copytree(FIXTURE_DATA, run_dir) + shutil.rmtree(os.path.join(run_dir, "pop_export"), ignore_errors=True) + result = {"_seed": SEED, "_co_maxgens": CO_MAXGENS, **compute_hashes(run_dir)} + with open(GOLDEN_PATH, "w") as f: + json.dump(result, f, indent=2, sort_keys=True) + print(f"Wrote baseline for {len(TRACKED)} files to {GOLDEN_PATH}") diff --git a/tests/test_sources.py b/tests/test_sources.py new file mode 100644 index 0000000..904bd92 --- /dev/null +++ b/tests/test_sources.py @@ -0,0 +1,154 @@ +"""Unit tests for the download layer. No network access: the fetchers are patched.""" +import warnings +from unittest import mock + +import pytest + +from geopops import sources +from geopops.exceptions import DownloadError + + +@pytest.fixture(autouse=True) +def _secure_by_default(): + """Every test starts with the insecure fallback off, and leaves it off.""" + sources.set_allow_insecure_downloads(False) + yield + sources.set_allow_insecure_downloads(False) + + +@pytest.fixture +def dst(tmp_path): + return str(tmp_path / "out.bin") + + +def _writes(content=b"ok"): + def fetch(src, dst, headers, mode, verify): + with open(dst, "wb") as f: + f.write(content) + return fetch + + +class TestEffectiveACSYear: + """Some ACS tables have no 2023/2024 vintage and fall back to 2022.""" + + @pytest.mark.parametrize("code,year,expected", [ + ("B09019", 2023, 2022), + ("B09020", 2024, 2022), + ("B01001", 2023, 2023), # not in the fallback set + ("B09019", 2019, 2019), # year has data + ]) + def test_fallback(self, code, year, expected): + assert sources._effective_acs_year(code, year) == expected + + +class TestDownload: + def test_success_and_verifies_tls_by_default(self, dst): + with mock.patch.object(sources, "_fetch_requests", side_effect=_writes()) as f: + assert sources.download("http://example/x", dst) == 0 + assert open(dst, "rb").read() == b"ok" + assert f.call_args[1]["verify"] is True + + def test_raises_instead_of_exiting(self, dst): + """Regression: this path used to call exit(1), killing the host process.""" + with mock.patch.object(sources, "_fetch_requests", side_effect=OSError("boom")): + with pytest.raises(DownloadError, match="after 2 attempts"): + sources.download("http://example/x", dst, retries=2) + + def test_retries_then_succeeds(self, dst): + attempts = [] + + def flaky(src, d, headers, mode, verify): + attempts.append(1) + if len(attempts) < 2: + raise OSError("transient") + _writes()(src, d, headers, mode, verify) + + with mock.patch.object(sources, "_fetch_requests", side_effect=flaky): + assert sources.download("http://example/x", dst, retries=3) == 0 + assert len(attempts) == 2 + + def test_unknown_backend(self, dst): + with pytest.raises(ValueError, match="Unknown download backend"): + sources.download("http://example/x", dst, backend="carrier-pigeon") + + def test_curl_cffi_backend_is_selectable(self, dst): + with mock.patch.object(sources, "_fetch_curl_cffi", side_effect=_writes()) as f: + sources.download("http://example/x", dst, backend="curl_cffi") + assert f.called + + def test_text_mode_is_passed_through(self, dst): + with mock.patch.object(sources, "_fetch_requests", side_effect=_writes()) as f: + sources.download("http://example/x", dst, mode="text") + assert f.call_args[0][3] == "text" + + +class TestInsecureFallback: + def test_off_by_default_with_an_actionable_message(self, dst): + err = OSError("SSL: certificate verify failed") + with mock.patch.object(sources, "_fetch_requests", side_effect=err): + with pytest.raises(DownloadError, match="allow_insecure_downloads"): + sources.download("http://example/x", dst, retries=1) + + def test_never_retries_unverified_unless_enabled(self, dst): + calls = [] + + def fetch(src, d, headers, mode, verify): + calls.append(verify) + raise OSError("SSL: certificate verify failed") + + with mock.patch.object(sources, "_fetch_requests", side_effect=fetch): + with pytest.raises(DownloadError): + sources.download("http://example/x", dst, retries=1) + assert calls == [True], "an unverified retry happened without opt-in" + + def test_engages_and_warns_when_enabled(self, dst): + calls = [] + + def fetch(src, d, headers, mode, verify): + calls.append(verify) + if verify: + raise OSError("SSL: certificate verify failed") + _writes()(src, d, headers, mode, verify) + + sources.set_allow_insecure_downloads(True) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with mock.patch.object(sources, "_fetch_requests", side_effect=fetch): + assert sources.download("http://example/x", dst, retries=1) == 0 + assert calls == [True, False] + assert any("TLS verification disabled" in str(w.message) for w in caught) + + def test_non_tls_errors_do_not_trigger_it(self, dst): + calls = [] + + def fetch(src, d, headers, mode, verify): + calls.append(verify) + raise OSError("404 Not Found") + + sources.set_allow_insecure_downloads(True) + with mock.patch.object(sources, "_fetch_requests", side_effect=fetch): + with pytest.raises(DownloadError): + sources.download("http://example/x", dst, retries=1) + assert calls == [True] + + +class TestNoImportTimeSideEffects: + def test_warning_filters_are_not_globally_disabled(self): + """Regression: importing geopops used to call urllib3.disable_warnings().""" + import urllib3 + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + warnings.warn("probe", urllib3.exceptions.InsecureRequestWarning, stacklevel=1) + assert any("probe" in str(w.message) for w in caught) + + +class TestFipsInfo: + def test_fips_to_abbreviation(self): + assert sources.fips_info("45")["abbr"] == "SC" + assert sources.fips_info(["45", "37"])["abbr"] == ["SC", "NC"] + + def test_reverse(self): + assert sources.fips_info("SC", reverse=True)["fips"] == "45" + + def test_unknown_code(self): + assert sources.fips_info("99")["abbr"] is None diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..a7ec6de --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,154 @@ +"""Unit tests for the pure helper functions in geopops.utils. + +These need no data and run in well under a second. +""" +import numpy as np +import pytest + +from geopops import utils + + +class TestLrRound: + """Largest-remainder rounding preserves the total.""" + + def test_preserves_sum(self): + v = np.array([1.4, 1.4, 1.4, 1.4, 1.4]) + out = utils.lrRound(v) + assert out.sum() == round(v.sum()) + assert out.dtype == np.int64 + + def test_gives_extra_to_largest_remainders(self): + # 0.9 has the largest remainder, so it rounds up first + out = utils.lrRound(np.array([0.9, 0.6, 0.5])) + assert out.tolist() == [1, 1, 0] + + def test_integers_unchanged(self): + v = np.array([3.0, 1.0, 6.0]) + assert utils.lrRound(v).tolist() == [3, 1, 6] + + def test_zeros(self): + assert utils.lrRound(np.zeros(4)).tolist() == [0, 0, 0, 0] + + @pytest.mark.parametrize("n", [1, 5, 50]) + def test_random_sums_preserved(self, n): + rng = np.random.default_rng(0) + for _ in range(20): + v = rng.random(n) * 10 + assert utils.lrRound(v).sum() == round(v.sum()) + + def test_row_and_col_round(self): + m = np.full((3, 4), 1.25) + assert (utils.rowRound(m).sum(axis=1) == 5).all() + assert (utils.colRound(m).sum(axis=0) == 4).all() + + +class TestRanges: + def test_contiguous_one_based_ranges(self): + assert utils.ranges([3, 2, 4]) == [(1, 3), (4, 5), (6, 9)] + + def test_zero_length_group(self): + # An empty group yields start > stop, which callers read as "no members" + assert utils.ranges([2, 0, 1]) == [(1, 2), (3, 2), (3, 3)] + + +class TestDrawCounts: + """drawCounts samples without replacement and decrements its input.""" + + def test_draws_requested_number(self): + rng = np.random.default_rng(0) + v = np.array([5, 5, 5], dtype=np.int64) + assert len(utils.drawCounts(v, 7, rng)) == 7 + assert v.sum() == 8 + + def test_caps_at_available(self): + rng = np.random.default_rng(0) + v = np.array([2, 1], dtype=np.int64) + assert len(utils.drawCounts(v, 99, rng)) == 3 + assert v.sum() == 0 + + def test_never_oversamples_a_bin(self): + rng = np.random.default_rng(1) + for _ in range(100): + v = rng.integers(0, 6, size=10).astype(np.int64) + before = v.copy() + drawn = utils.drawCounts(v, int(rng.integers(0, 30)), rng) + counts = np.bincount(drawn, minlength=10) if drawn else np.zeros(10, int) + assert (counts <= before).all() + assert np.array_equal(before - counts, v) + + def test_empty(self): + rng = np.random.default_rng(0) + assert utils.drawCounts(np.zeros(3, dtype=np.int64), 5, rng) == [] + assert utils.drawCounts(np.array([1, 2], dtype=np.int64), 0, rng) == [] + + def test_reproducible_with_same_seed(self): + v1, v2 = np.array([4, 4, 4], np.int64), np.array([4, 4, 4], np.int64) + a = utils.drawCounts(v1, 6, np.random.default_rng(7)) + b = utils.drawCounts(v2, 6, np.random.default_rng(7)) + assert a == b + + +class TestPersonData: + """Traits are config-driven, not hardcoded.""" + + def _person(self, names, values): + schema = utils.TraitSchema(names) + return utils.PersonData(hh=(1, 2), sample=3, age=40, working=True, + commuter=False, schema=schema, + trait_values=schema.values_from(values)) + + def test_traits_reachable_by_name(self): + p = self._person(["hispanic", "female"], {"hispanic": True, "female": False}) + assert p.hispanic is True and p.female is False + + def test_missing_trait_value_is_none(self): + p = self._person(["hispanic", "female"], {"hispanic": True}) + assert p.female is None + + def test_unknown_trait_raises_with_a_useful_message(self): + p = self._person(["hispanic"], {"hispanic": True}) + with pytest.raises(AttributeError, match="race_asian_alone"): + _ = p.race_asian_alone + + def test_arbitrary_new_trait_needs_no_code_change(self): + # The regression behind issue #2: adding a trait used to raise TypeError + p = self._person(["some_brand_new_trait"], {"some_brand_new_trait": True}) + assert p.some_brand_new_trait is True + + def test_core_fields_are_not_shadowed_by_traits(self): + p = self._person(["age"], {"age": 999}) + assert p.age == 40 # the real field wins + + def test_traits_property_round_trips(self): + p = self._person(["a", "b"], {"a": True, "b": False}) + assert p.traits == {"a": True, "b": False} + + def test_no_instance_dict(self): + p = self._person([], {}) + assert not hasattr(p, "__dict__") + + +class TestSmallHelpers: + def test_indexer_assigns_stable_ids(self): + idx, d = utils.Indexer(), {} + assert idx(d, "a") == 1 + assert idx(d, "b") == 2 + assert idx(d, "a") == 1 + + def test_thresh(self): + assert utils.thresh(5, 10) == 0 + assert utils.thresh(15, 10) == 15 + + def test_vecmerge_concatenates_by_key(self): + assert utils.vecmerge({"a": [1]}, {"a": [2], "b": [3]}) == {"a": [1, 2], "b": [3]} + + def test_vecmerge_copies_inputs(self): + a = {"x": [1]} + utils.vecmerge(a, {})["x"].append(99) + assert a == {"x": [1]} + + def test_dflat(self): + assert sorted(utils.dflat({"a": [1, 2]})) == [("a", 1), ("a", 2)] + + def test_tryjson_missing_file_returns_empty(self, tmp_path): + assert utils.tryJSON(str(tmp_path / "nope.json")) == {} diff --git a/tests/test_workflow.py b/tests/test_workflow.py new file mode 100644 index 0000000..e6ff2c1 --- /dev/null +++ b/tests/test_workflow.py @@ -0,0 +1,106 @@ +"""End-to-end tests of the GeoPops pipeline. + +These need the processed fixture data in ``tests/data/processed`` and take minutes, +so they are marked ``slow``. Run the fast suite with:: + + pytest -m "not slow" +""" +import os + +import pytest + +import geopops +from conftest import requires_fixture_data + +pytestmark = [pytest.mark.slow, requires_fixture_data] + + +@pytest.mark.network +@pytest.mark.skip(reason="Manual run only: downloads several GB from Census/LODES") +def test_download(fixture_config): + return geopops.download_data(fixture_config) + + +@pytest.mark.skip(reason="Manual run only: rebuilds the processed fixture (~5 min)") +def test_processing(fixture_config): + return geopops.process_data(fixture_config) + + +@pytest.fixture(scope="module") +def pop(fixture_config, run_dir): + """One generated population, shared by the tests below.""" + cfg = dict(fixture_config) + cfg["path"] = run_dir + cfg["CO_maxgens"] = 5000 # enough to converge on most CBGs, fast enough to test + return geopops.generate_pop(cfg, seed=42, verbose=0) + + +def test_co_assigns_households_to_every_cbg(pop): + assert pop.co_results + for county, cbgs in pop.co_results.items(): + assert cbgs, f"county {county} got no CBGs" + for cbg, serials in cbgs.items(): + assert serials, f"CBG {cbg} got no households" + + +def test_synthpop_populates_people_and_households(pop): + assert len(pop.people) > 0 + assert len(pop.households) > 0 + # every household member must resolve to a real person + for hh in pop.households.values(): + for pkey in hh.people: + assert pkey in pop.people + + +def test_people_carry_the_configured_traits(pop, fixture_config): + expected = list(fixture_config["additional_traits"]) + person = next(iter(pop.people.values())) + assert list(person.schema.names) == expected + assert set(person.traits) == set(expected) + + +def test_networks_are_symmetric_and_sized_to_the_population(pop): + n = len(pop.adj_mat_keys) + for name in ("adj_hh", "adj_sch", "adj_wp", "adj_gq"): + m = getattr(pop, name) + assert m.shape == (n, n), name + assert (m != m.T).nnz == 0, f"{name} is not symmetric" + + +def test_export_writes_every_expected_file(pop): + expected = [ + "cbg_idxs.csv", "hh.csv", "people.csv", "sch_students.csv", "sch_workers.csv", + "gqs.csv", "gq_residents.csv", "gq_workers.csv", "company_workers.csv", + "outside_workers.csv", "adj_mat_keys.csv", "adj_dummy_keys.csv", + "adj_out_workers.csv", "adj_upper_triang_hh.mtx", "adj_upper_triang_sch.mtx", + "adj_upper_triang_wp.mtx", "adj_upper_triang_gq.mtx", + "adj_upper_triang_non_hh.mtx", + ] + for name in expected: + path = os.path.join(pop.pop_export_dir, name) + assert os.path.exists(path), f"missing {name}" + assert os.path.getsize(path) > 0, f"empty {name}" + + +def test_exported_people_columns_follow_the_config(pop, fixture_config): + import pandas as pd + df = pd.read_csv(os.path.join(pop.pop_export_dir, "people.csv")) + for trait in fixture_config["additional_traits"]: + assert trait in df.columns + assert len(df) == len(pop.people) + + +def test_starsim_people_and_networks(pop): + ppl = geopops.to_starsim_people(pop.pop_export_dir, verbose=0) + assert len(ppl) == len(pop.adj_mat_keys) + + nets = geopops.starsim_networks(pop.pop_export_dir, seed=42, save=False, verbose=0) + assert len(nets) == 4 + assert all(len(n.edges.p1) > 0 for n in nets) + + +def test_two_populations_do_not_share_cached_networks(pop, tmp_path): + """Regression: ForStarsim cached networks in class state and never invalidated.""" + a = geopops.starsim_network("homenet", pop.pop_export_dir, save=False) + b = geopops.starsim_network("schoolnet", pop.pop_export_dir, save=False) + assert len(a.edges.p1) != len(b.edges.p1) or not (a.edges.p1 == b.edges.p1).all()