starsim-data 2.0.0__py3-none-any.whl

starsim_data/__init__.py

@@ -0,0 +1,57 @@
+"""
+starsim_data: national-level demographic data and contact matrices for disease
+modeling with Starsim.
+
+This package bundles UN World Population Prospects demographics (births, deaths,
+migration, population) and a trimmed, reorganized subset of the contact-matrix data
+from epydemix-data (https://github.com/epistorm/epydemix-data), national-level (level
+0) only. See NOTICE.md for per-source data licenses and attributions.
+
+The data ships bundled as HDF5 under `starsim_data/data/` (e.g. `prem2021.h5`,
+`demographics-2024.h5`). `load` returns a `sciris.dataframe` (a drop-in replacement for
+a pandas `DataFrame`); for contact matrices the rows and columns are indexed by the
+start age of each bin:
+
+```python
+import starsim_data as ssd
+M = ssd.load(location='kenya')                    # square, age-indexed (0, 5, ..., 75)
+M.loc[25, 30]                                      # contacts between the 25-29 and 30-34 bins
+```
+
+Locations can be given as a name, ISO-2/ISO-3 code, official name or alias -- `'USA'`,
+`'US'`, `'united states'` and `'United_States'` all resolve to the same data. Every
+dataset records where it came from; see `provenance()`.
+
+`load` never touches the network. What is not bundled -- annual series, complete life
+tables, other projection variants, the full `prem2017` and `mistry2021` contact matrices
+-- is available from `fetch()`, which downloads and caches; see `cache_info()` and
+`options`. `plot()` draws the data for a quick check before it goes into a model.
+"""
+
+__version__ = '0.4.0'
+
+# This is the one module that uses `from` imports, to assemble the public namespace.
+# Everywhere else in the package, follow Starsim's convention: `import starsim_data as
+# ssd` and reach things as `ssd.load`, `ssd.registry.datasets()` and so on. The order
+# below matters, because a few modules register themselves on import: `registry` has to
+# exist before `loaders` adds its layouts, and before `fetching` adds its datasets.
+from .settings import options
+from . import registry
+from .registry import register_dataset, register_layout
+from .locations import resolve, location_info, list_all_locations, parse_level
+from .loaders import load, list_datasets, list_locations, list_revisions, provenance
+from .cache import cache_dir, cache_info, cache_provenance, cached_download, clear_cache
+from .fetching import fetch, list_fetchable
+from .plotting import plot, choropleth
+from .nodes import load_nodes, distances, adjacency, node_name
+from .shapes import fetch_shapes, list_shape_sources
+from .rasters import fetch_nodes, fetch_raster, aggregate_raster, rescale_to_national
+
+__all__ = ['load', 'fetch', 'plot', 'choropleth', 'provenance', 'options', 'registry',
+           'load_nodes', 'distances', 'adjacency', 'node_name',
+           'fetch_shapes', 'list_shape_sources',
+           'fetch_nodes', 'fetch_raster', 'aggregate_raster', 'rescale_to_national',
+           'list_datasets', 'list_locations', 'list_revisions', 'list_fetchable',
+           'resolve', 'location_info', 'list_all_locations', 'parse_level',
+           'register_dataset', 'register_layout',
+           'cache_dir', 'cache_info', 'cache_provenance', 'cached_download', 'clear_cache']

starsim_data/cache.py

@@ -0,0 +1,220 @@
+"""
+Download cache for data that is too big to bundle.
+
+Tier-0 data ships inside the wheel and `load()` never touches the network. Anything
+larger -- the UN WPP bulk files behind `fetch()`, and later the spatial sources -- is
+downloaded once into a local cache and reused from then on.
+
+Every cached file gets a provenance sidecar (`<file>.provenance.json`) recording the URL
+it came from, when it was downloaded, its size and its SHA-256. That is `laser-init`'s
+idea, and it is the right one for fetched files: unlike bundled data, where provenance
+is written into the HDF5 itself, a cache entry is an ordinary file on disk and a sidecar
+is the only place the lineage can live.
+
+```python
+import starsim_data as ssd
+ssd.cache_dir()      # where things are cached
+ssd.cache_info()     # what is cached, from where, and when it was fetched
+ssd.clear_cache()    # delete it all
+```
+"""
+
+import os
+import hashlib
+import platform
+import shutil
+import sciris as sc
+import starsim_data as ssd
+
+__all__ = ['cache_dir', 'cache_info', 'cache_provenance', 'cached_download', 'clear_cache']
+
+SIDECAR_SUFFIX = '.provenance.json'
+
+# Read in chunks rather than all at once: cached WPP files run to a few hundred MB.
+HASH_CHUNK = 1 << 20
+
+
+def _default_cache_dir():
+    """Return the platform-appropriate cache location, without creating it."""
+    home = sc.path('~').expanduser()
+    system = platform.system()
+    if system == 'Windows':
+        root = os.environ.get('LOCALAPPDATA') or home / 'AppData' / 'Local'
+        return sc.path(root) / 'starsim_data' / 'Cache'
+    elif system == 'Darwin':
+        return home / 'Library' / 'Caches' / 'starsim_data'
+    else:
+        root = os.environ.get('XDG_CACHE_HOME') or home / '.cache'
+        return sc.path(root) / 'starsim_data'
+
+
+def cache_dir(subdir=None, create=False):
+    """Return the directory where downloaded data is cached.
+
+    Resolved at call time from `ssd.options.cache_dir`, falling back to a platform default
+    (`~/.cache/starsim_data` on Linux, `~/Library/Caches/starsim_data` on macOS,
+    `%LOCALAPPDATA%\\starsim_data\\Cache` on Windows). Nothing is created unless
+    `create=True`, so merely asking where the cache is has no side effects.
+
+    Args:
+        subdir (str): optional subdirectory within the cache, e.g. `'wpp'`.
+        create (bool): whether to create the directory if it does not exist.
+
+    Returns:
+        pathlib.Path: the cache directory.
+
+    Example:
+        >>> import starsim_data as ssd
+        >>> ssd.cache_dir().name
+        'starsim_data'
+    """
+    if ssd.options.cache_dir:
+        path = sc.path(ssd.options.cache_dir).expanduser()
+    else:
+        path = _default_cache_dir()
+    if subdir:
+        path = path / subdir
+    if create:
+        path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def _sidecar(path):
+    """The path of a cached file's provenance sidecar."""
+    return sc.path(str(path) + SIDECAR_SUFFIX)
+
+
+def _sha256(path):
+    """The SHA-256 of a file, read in chunks."""
+    digest = hashlib.sha256()
+    with open(path, 'rb') as f:
+        chunk = f.read(HASH_CHUNK)
+        while chunk:
+            digest.update(chunk)
+            chunk = f.read(HASH_CHUNK)
+    return digest.hexdigest()
+
+
+def _write_sidecar(path, url):
+    """Record where a cached file came from, alongside it."""
+    facts = dict(
+        url             = url,
+        filename        = path.name,
+        downloaded      = sc.now(utc=True).isoformat(timespec='seconds'),
+        bytes           = path.stat().st_size,
+        sha256          = _sha256(path),
+        package_version = ssd.__version__,
+    )
+    sc.savejson(_sidecar(path), facts, sanitizepath=False) # Don't let Sciris rename the file
+    return facts
+
+
+def cached_download(url, filename=None, subdir=None, refresh=False, verbose=None):
+    """Download `url` into the cache, or return the existing path if already cached.
+
+    Args:
+        url (str): the URL to fetch.
+        filename (str): name to save as (default: the last path component of the URL).
+        subdir (str): subdirectory of the cache to put it in, e.g. `'wpp'`.
+        refresh (bool): re-download even if a cached copy exists.
+        verbose (bool): report progress (default: `ssd.options.verbose`).
+
+    Returns:
+        pathlib.Path: the path to the cached file.
+    """
+    verbose = sc.ifelse(verbose, ssd.options.verbose)
+    filename = filename or url.rstrip('/').split('/')[-1]
+    path = cache_dir(subdir, create=True) / filename
+
+    if path.is_file() and (not refresh):
+        if verbose:
+            print(f'Using cached {path} ({sc.humanize_bytes(path.stat().st_size)})')
+        return path
+
+    if verbose:
+        print(f'Downloading {url}\n  -> {path}')
+    sc.download(url, filename=str(path), convert=False, verbose=verbose) # convert=False keeps the gzipped WPP files as bytes rather than decoding them as text
+    _write_sidecar(path, url)
+    return path
+
+
+def cache_provenance(path):
+    """Return the recorded provenance for a cached file, or None if there is none."""
+    sidecar = _sidecar(path)
+    if not sidecar.is_file():
+        return None
+    return sc.objdict(sc.loadjson(sidecar))
+
+
+def _cached_files():
+    """Return every cached data file (i.e. excluding the provenance sidecars)."""
+    root = cache_dir()
+    if not root.is_dir():
+        return []
+    files = []
+    for path in sorted(root.rglob('*')):
+        if path.is_file() and (not path.name.endswith(SIDECAR_SUFFIX)):
+            files.append(path)
+    return files
+
+
+def cache_info():
+    """Return a dataframe of what is currently cached: file, size, source URL and date.
+
+    Returns:
+        sciris.dataframe: with columns `[file, size, bytes, url, downloaded]`, empty if
+        nothing is cached.
+
+    Example:
+        >>> import starsim_data as ssd
+        >>> ssd.cache_info()  # doctest: +SKIP
+    """
+    root = cache_dir()
+    rows = []
+    for path in _cached_files():
+        prov = cache_provenance(path) or {}
+        size = path.stat().st_size
+        rows.append(dict(
+            file       = str(path.relative_to(root)),
+            size       = sc.humanize_bytes(size),
+            bytes      = size,
+            url        = prov.get('url', ''),
+            downloaded = prov.get('downloaded', ''),
+        ))
+    return sc.dataframe(rows, columns=['file', 'size', 'bytes', 'url', 'downloaded'])
+
+
+def clear_cache(pattern=None, verbose=None):
+    """Delete cached downloads. Bundled data is never touched.
+
+    Args:
+        pattern (str): a glob matched against each cached file's path relative to the
+            cache root, e.g. `'wpp/*'`. If None, the whole cache is removed.
+        verbose (bool): report what was deleted (default: `ssd.options.verbose`).
+
+    Returns:
+        list: the paths that were deleted.
+
+    Example:
+        >>> import starsim_data as ssd
+        >>> ssd.clear_cache('wpp/*')  # doctest: +SKIP
+    """
+    verbose = sc.ifelse(verbose, ssd.options.verbose)
+    root = cache_dir()
+    if not root.is_dir():
+        return []
+
+    if pattern is None:
+        removed = [str(p.relative_to(root)) for p in _cached_files()]
+        shutil.rmtree(root)
+    else:
+        removed = []
+        for path in _cached_files():
+            if path.relative_to(root).match(pattern):
+                _sidecar(path).unlink(missing_ok=True)
+                path.unlink()
+                removed.append(str(path.relative_to(root)))
+
+    if verbose:
+        print(f'Removed {len(removed)} cached file(s) from {root}')
+    return removed

starsim_data/data/README.md

@@ -0,0 +1,41 @@
+# starsim_data/data — bundled data files
+
+These files ship inside the installed package and are the data that `starsim_data.load()` reads. The `.h5` stores are **generated — do not edit by hand** (regenerate them from the sources in `raw_data/`); `locations.csv` is the one file here with hand-maintained content, described below.
+
+Each store carries a root `spec` attribute recording its `layout`, so `load()` returns the right shape for it.
+
+**Contact matrices** (`layout='matrix'`), built from [`raw_data/contact_matrices/`](../../raw_data/contact_matrices/) by [`scripts/build_hdf5.py`](../../scripts/build_hdf5.py):
+
+- `prem2021.h5` — the bundled contact source, keyed by location with one `float32` matrix per contact setting: `/<Location>/<setting>`.
+- `prem/prem2017.h5` — the same, for the ten locations `prem2021` lacks, from `prem2017`. Together with `prem2021.h5` these are the default contact matrices, `load(which='prem')`. It is in a subfolder so that it is not discovered as a dataset of its own.
+- `fetched.json` — the manifest of the stores that are fetched rather than bundled (`prem2017`, `mistry2021`, and for the US states `mistry2021_us` and `population_us`): for each, its URL (an asset of a GitHub release), SHA-256, locations and provenance. Their stores, built the same way, are in [`fetched/`](../../fetched/) at the top of the repository, outside the package; `fetch()` downloads them.
+
+**Demographic series** from UN WPP, all grouped in `demographics-<revision>.h5` (e.g. `demographics-2024.h5`), built by [`scripts/build_demographics.py`](../../scripts/build_demographics.py). One group per series, each with its own `spec`:
+
+- `/births`, `/migration` (`layout='crude'`) — one value per year (CBR, CNMR), per location.
+- `/deaths` (`layout='mortality'`) — age/sex mortality rate `mx`, a `(year, sex, age)` array per location.
+- `/population` (`layout='population'`) — population by 5-year age groups and sex (absolute persons), `(year, sex, age)` per location.
+
+**Node table** (`layout='nodes'`), built from UNOCHA COD-AB and geoBoundaries by [`scripts/build_nodes.py`](../../scripts/build_nodes.py):
+
+- `nodes.h5` — one compound dataset per location, holding every administrative unit at levels 0 and 1: `nodeid`, `location`, `iso3`, `level`, `name`, `pcode`, `dot_name`, `parent`, `lat`, `lon`, `area_km2`, `population`, `source`. Read with `starsim_data.load_nodes()`.
+- `adjacency.h5` — for each location and level, the pairs of units that share a border. Read with `starsim_data.adjacency()`.
+
+Plus `locations.csv` — the canonical country metadata file, read by all the build scripts and by the runtime location resolver. It is **part source, part generated**: `location`, `level`, `iso_code`, `primary_contact_source`, `contact_sources` and `population_source` are hand-maintained, while `iso3`, `official_name` and `aliases` are filled by [`scripts/build_locations.py`](../../scripts/build_locations.py) from `pycountry`. Edit the former by hand; regenerate the latter.
+
+Note that this file must be read with `keep_default_na=False` (or with the `csv` module, as the resolver does): Namibia's ISO-2 code is the string `NA`, which pandas' defaults turn into NaN.
+
+To regenerate:
+
+```bash
+python scripts/build_locations.py --check      # locations.csv derived columns
+python scripts/build_hdf5.py --check           # contact matrices, from raw_data/contact_matrices/
+python scripts/build_demographics.py --check   # demographics, from raw_data/demographics/
+python scripts/build_nodes.py --check          # node table, from the boundary sources
+```
+
+`build_nodes.py` is the odd one out: its inputs are too big for `raw_data/` and are downloaded into the package cache instead, so a cold rebuild pulls ~1.4 GB. That is why it is not part of the automatic reproducibility check in CI.
+
+Each `.h5` also carries a root `provenance` attribute recording its source, license, citation and build details; see `starsim_data.provenance()`.
+
+The national-tier `.h5` files are small (a few MB each) and committed to the repo so that `pip install` is self-contained; this is a deliberate exception to the `*.h5` rule in `.gitignore`. Any larger (e.g. subnational) tiers added later should be distributed as release assets instead of committed here.