pymc_forecast.datasets#

Dataset helpers for the examples and docs.

load_bart_od() downloads and caches the complete hourly BART origin-destination panel. load_bart_weekly() and load_bart_weekly_by_origin() derive compact weekly examples from that source, load_victoria_electricity() and load_us_macro() read small CSVs bundled with the package, and load_m5() downloads the M5 competition files once and returns labeled sales and price panels. Every loader returns labeled arrays; load_m5() adds the identifier, calendar, and weight tables.

class pymc_forecast.datasets.M5Data(sales, price, keys, calendar, weights)[source]#

Bases: NamedTuple

M5 evaluation data: labeled sales and price panels plus three tables.

sales and price are float32 xarray.DataArray panels with dims ("time", "series"): the "time" coord is the calendar date of each day and the "series" coord is the series id (item_id and store_id joined by _), in sales-file order. price is the weekly shelf price repeated over the days of wm_yr_wk, NaN where the item was not listed. keys has one row per series with the id and the five hierarchy columns; calendar has one row per day; weights is the official evaluation weight table. Each field is available as an attribute of the same name, and the tuple unpacks in field order.

Parameters:

  • sales (xarray.DataArray) – Float32 daily unit sales named "sales" with dims ("time", "series"): training days followed by evaluation days, NaN on evaluation days for series missing from the test file.

  • price (xarray.DataArray) – Float32 daily shelf prices named "price" with the same dims and coords as sales.

  • keys (pandas.DataFrame) – One row per series, in "series" order, with columns id, item_id, dept_id, cat_id, store_id and state_id.

  • calendar (pandas.DataFrame) – The calendar.csv table (date parsed as datetimes), one row per sales day.

  • weights (pandas.DataFrame) – The weights_evaluation.csv table, unchanged.

sales#

Alias for field number 0

price#

Alias for field number 1

keys#

Alias for field number 2

calendar#

Alias for field number 3

weights#

Alias for field number 4

pymc_forecast.datasets.load_bart_od()[source]#

Load complete hourly BART origin-destination ridership counts.

The four source shards are downloaded from the public Pyro dataset mirror, verified by SHA-256, and cached in the operating system’s user cache.

Returns:

xarray.DataArray – Integer counts named "rides" with dims ("time", "origin", "destination"). The time coordinate is hourly datetime64 values starting at the shards’ recorded start date (2011-01-01 for the current files), and station names label both origin and destination.

pymc_forecast.datasets.load_bart_weekly()[source]#

Load total weekly BART ridership on the log scale.

The series is derived at load time from the complete public BART origin-destination dataset used by the Pyro and NumPyro forecasting examples. Hourly counts are summed over all origin-destination pairs, aggregated into non-overlapping 168-hour weeks starting at the first hour (a trailing partial week is dropped), and log-transformed.

Returns:

xarray.DataArray – Log weekly totals named "log_rides" with dims ("time",) and integer week coords 0..n-1.

pymc_forecast.datasets.load_bart_weekly_by_origin(num_series=8)[source]#

Load a weekly BART ridership panel grouped by origin station.

Counts are summed over destination stations and aggregated into non-overlapping 168-hour weeks (a trailing partial week is dropped) before applying log1p. Aggregation happens a shard at a time, avoiding materializing the much larger full origin-destination panel. By default only the eight busiest origins are returned, which keeps hierarchical examples quick; pass None for all stations.

Parameters:

num_series (int, default 8) – Number of busiest origin stations to retain, ordered busiest first (by total count); None keeps all stations in source order.

Returns:

xarray.DataArray – Log weekly counts named "log_rides" with dims ("time", "series"), integer week coords 0..n-1 on "time" and station names on the "series" coordinate.

Raises:

ValueError – If num_series is less than 1.

pymc_forecast.datasets.load_m5(cache_dir=None)[source]#

Load the M5 evaluation panel, downloading and caching the archive once.

The files come from Nixtla’s mirror of the competition data, pinned to commit 72b8e7fd and checked against a SHA-256 digest. The archive is about 50 MB and is downloaded and extracted into cache_dir, by default pooch.os_cache("pymc_forecast") / "m5" (~/.cache/pymc_forecast/m5 on Linux, ~/Library/Caches/pymc_forecast/m5 on macOS); the archive is kept there next to the extracted files. When every competition file is already in that directory, nothing is downloaded. The two returned panels hold about 480 MB; parsing the CSVs peaks at about 2 GB.

Series order follows sales_train_evaluation.csv. The test file is aligned on the train keys, so a test-only row is ignored and a missing test row leaves NaN sales on the evaluation days. Weekly prices are mapped by wm_yr_wk; a calendar week absent from sell_prices.csv is NaN and does not shift the other weeks.

Parameters:

cache_dir (str or os.PathLike, optional) – Directory to download the archive into and extract it into (created if missing); None uses pooch.os_cache("pymc_forecast") / "m5".

Returns:

M5Data – Labeled sales and price panels, identifiers, calendar, and official evaluation weights.

Raises:

ValueError – If a competition file lacks required columns or is inconsistent: duplicate item-store keys or series ids, overlapping or non-contiguous d_* day columns, duplicate store-item-week prices, a calendar shorter than the sales, or duplicate calendar dates.

pymc_forecast.datasets.load_us_macro()[source]#

Load quarterly US real GDP, consumption and investment levels.

The statsmodels macrodata panel (public domain; FRED, accessed December 2009), restricted to the three real-activity series: 1959Q1 to 2009Q3, seasonally adjusted annual rates in billions of chained 2005 US$. Bundled as a small CSV.

Returns:

xarray.DataArray – Levels with dims ("time", "series"); time is a quarter-start DatetimeIndex and series is ["realgdp", "realcons", "realinv"].

pymc_forecast.datasets.load_victoria_electricity()[source]#

Load hourly Victoria (Australia) electricity demand and temperature.

The series covers the first eight weeks of 2014, sampled hourly — the Victoria electricity demand data used in the TensorFlow Probability structural-time-series case study and in Hyndman & Athanasopoulos’ Forecasting: Principles and Practice (original half-hourly data downsampled to hourly). Bundled as a small CSV.

Returns:

  • demand (xarray.DataArray) – Hourly electricity demand (GW), dims ("time",) with an hourly DatetimeIndex coord.

  • temperature (xarray.DataArray) – Hourly temperature (°C), aligned with demand.