tsam-xarray 0.6.4__tar.gz → 0.6.6__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. tsam_xarray-0.6.6/.github/workflows/benchmarks.yaml +126 -0
  2. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.gitignore +1 -0
  3. tsam_xarray-0.6.6/.release-please-manifest.json +3 -0
  4. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/CHANGELOG.md +26 -0
  5. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/PKG-INFO +13 -13
  6. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/README.md +12 -12
  7. tsam_xarray-0.6.6/benchmarks/README.md +60 -0
  8. tsam_xarray-0.6.6/benchmarks/conftest.py +25 -0
  9. tsam_xarray-0.6.6/benchmarks/test_bench_aggregate.py +279 -0
  10. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/data-model.md +1 -1
  11. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/examples/getting-started.ipynb +2 -2
  12. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/examples/multi-dim.ipynb +3 -3
  13. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/index.md +4 -4
  14. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/src/tsam_xarray/_clustering.py +279 -34
  15. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/src/tsam_xarray/_core.py +51 -25
  16. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/src/tsam_xarray/_result.py +54 -75
  17. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/src/tsam_xarray/_version.py +2 -2
  18. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/test/test_aggregate.py +243 -5
  19. tsam_xarray-0.6.6/test/test_lazy_result.py +64 -0
  20. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/test/test_parametrized.py +8 -8
  21. tsam_xarray-0.6.4/.release-please-manifest.json +0 -3
  22. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.github/dependabot.yml +0 -0
  23. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.github/workflows/ci.yaml +0 -0
  24. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.github/workflows/dependabot-auto-merge.yaml +0 -0
  25. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.github/workflows/pr-title.yaml +0 -0
  26. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.github/workflows/publish.yaml +0 -0
  27. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.github/workflows/release.yaml +0 -0
  28. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.pre-commit-config.yaml +0 -0
  29. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.readthedocs.yaml +0 -0
  30. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/.release-please-config.json +0 -0
  31. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/CLAUDE.md +0 -0
  32. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/LICENSE +0 -0
  33. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/codecov.yml +0 -0
  34. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/assets/multi-dim-input.png +0 -0
  35. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/assets/multi-dim-metrics.png +0 -0
  36. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/changelog.md +0 -0
  37. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/examples/clustering-io.ipynb +0 -0
  38. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/examples/segmentation.ipynb +0 -0
  39. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/examples/tuning.ipynb +0 -0
  40. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/gen_ref_pages.py +0 -0
  41. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/docs/stylesheets/extra.css +0 -0
  42. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/mkdocs.yml +0 -0
  43. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/pyproject.toml +0 -0
  44. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/scripts/generate_readme_images.py +0 -0
  45. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/src/tsam_xarray/__init__.py +0 -0
  46. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/src/tsam_xarray/_dim_names.py +0 -0
  47. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/src/tsam_xarray/_sample_data.py +0 -0
  48. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/src/tsam_xarray/_tuning.py +0 -0
  49. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/test/conftest.py +0 -0
  50. {tsam_xarray-0.6.4 → tsam_xarray-0.6.6}/test/test_tuning.py +0 -0
@@ -0,0 +1,126 @@
1
+ name: Benchmarks
2
+
3
+ # Informational performance report on every PR: run the pytest-benchmem suite
4
+ # on the PR head and on the base branch in the same job (same runner, same
5
+ # pinned benchmem version), then keep one sticky comment with the delta table
6
+ # and upload an interactive plot artifact. Never blocks a merge — GitHub
7
+ # runners are noisy, so treat the walltime deltas as a hint, not a gate.
8
+ # Fork PRs run too (read-only token) and get the step summary and artifact;
9
+ # only the sticky comment needs a writable token and is skipped for them.
10
+
11
+ on:
12
+ pull_request:
13
+ branches: [main]
14
+ workflow_dispatch:
15
+
16
+ concurrency:
17
+ group: ${{ github.workflow }}-${{ github.ref }}
18
+ cancel-in-progress: true
19
+
20
+ env:
21
+ BENCHMEM: pytest-benchmem[plot]==0.4.10
22
+ BENCH_ARGS: >-
23
+ -o addopts= -p no:xdist -q --benchmark-only --benchmark-memory
24
+ --benchmark-disable-gc
25
+
26
+ jobs:
27
+ report:
28
+ name: Head vs base
29
+ runs-on: ubuntu-24.04
30
+ continue-on-error: true # informational, never blocks a merge
31
+ permissions:
32
+ contents: read
33
+ pull-requests: write # no-op on fork PRs (read-only token); comment step is guarded
34
+ steps:
35
+ - uses: actions/checkout@v6
36
+ with:
37
+ fetch-depth: 0 # hatch-vcs + checkout of the base revision
38
+ persist-credentials: false # benchmarks execute PR code; keep the token out of .git/config
39
+
40
+ - uses: astral-sh/setup-uv@v7
41
+ with:
42
+ enable-cache: true
43
+
44
+ - name: Benchmark head
45
+ run: |
46
+ uv run --with "$BENCHMEM" pytest benchmarks/ $BENCH_ARGS \
47
+ --benchmark-json "$RUNNER_TEMP/head.json"
48
+ uv run --no-sync python -c "import tsam; print('head tsam:', tsam.__version__)"
49
+
50
+ - name: Benchmark base (${{ github.event.pull_request.base.ref }})
51
+ if: github.event_name == 'pull_request'
52
+ # Base src with the head's benchmark suite overlaid, so test IDs match
53
+ # and suite edits in the PR still compare against the same cases.
54
+ # Drop the head run's uv.lock (untracked — the repo commits none):
55
+ # uv keeps a satisfying lock as-is, so without this a head that
56
+ # changes dependencies leaks its resolution into the base run.
57
+ run: |
58
+ git checkout --detach ${{ github.event.pull_request.base.sha }}
59
+ git checkout ${{ github.sha }} -- benchmarks/
60
+ rm -f uv.lock
61
+ uv run --exact --with "$BENCHMEM" pytest benchmarks/ $BENCH_ARGS \
62
+ --benchmark-json "$RUNNER_TEMP/base.json"
63
+ uv run --no-sync python -c "import tsam; print('base tsam:', tsam.__version__)"
64
+
65
+ - name: Render report and plot
66
+ run: |
67
+ if [ -f "$RUNNER_TEMP/base.json" ]; then
68
+ uv run --with "$BENCHMEM" benchmem compare \
69
+ "$RUNNER_TEMP/base.json" "$RUNNER_TEMP/head.json" \
70
+ --diff --group-by module --sort change --format md > "$RUNNER_TEMP/table.md"
71
+ uv run --with "$BENCHMEM" benchmem plot \
72
+ "$RUNNER_TEMP/base.json" "$RUNNER_TEMP/head.json" \
73
+ --view compare -o benchmark-plot.html
74
+ {
75
+ echo 'Δ% vs base, one table sorted by biggest change first'
76
+ echo '(`min` walltime and memray `peak`). Peak memory is'
77
+ echo 'deterministic — any delta there is real; walltime on GitHub'
78
+ echo 'runners is noisy, so only large time deltas are meaningful.'
79
+ } > "$RUNNER_TEMP/intro.md"
80
+ else
81
+ uv run --with "$BENCHMEM" benchmem compare \
82
+ "$RUNNER_TEMP/head.json" --format md > "$RUNNER_TEMP/table.md"
83
+ uv run --with "$BENCHMEM" benchmem plot \
84
+ "$RUNNER_TEMP/head.json" --x n_cols -o benchmark-plot.html
85
+ echo 'Standalone run — no base revision to compare against.' \
86
+ > "$RUNNER_TEMP/intro.md"
87
+ fi
88
+ {
89
+ echo '<!-- benchmarks-report -->'
90
+ echo '## Benchmarks'
91
+ echo
92
+ cat "$RUNNER_TEMP/intro.md"
93
+ echo 'The interactive plot is attached as the `benchmark-plot` artifact.'
94
+ echo
95
+ echo '<details><summary>Full table</summary>'
96
+ echo
97
+ cat "$RUNNER_TEMP/table.md"
98
+ echo
99
+ echo '</details>'
100
+ } > "$RUNNER_TEMP/comment.md"
101
+ cat "$RUNNER_TEMP/comment.md" >> "$GITHUB_STEP_SUMMARY"
102
+
103
+ - uses: actions/upload-artifact@v4
104
+ with:
105
+ name: benchmark-plot
106
+ path: benchmark-plot.html
107
+
108
+ - name: Post or update the sticky comment
109
+ # Commenting needs a writable token; fork PRs only get a read-only one,
110
+ # so they keep the step summary and artifact but skip the comment.
111
+ if: >-
112
+ ${{ github.event_name == 'pull_request' &&
113
+ github.event.pull_request.head.repo.full_name == github.repository }}
114
+ env:
115
+ GH_TOKEN: ${{ github.token }}
116
+ PR: ${{ github.event.pull_request.number }}
117
+ run: |
118
+ marker='<!-- benchmarks-report -->'
119
+ existing=$(gh api --paginate "repos/${{ github.repository }}/issues/$PR/comments" \
120
+ --jq ".[] | select(.body | startswith(\"$marker\")) | .id" | head -n1)
121
+ if [ -n "$existing" ]; then
122
+ gh api -X PATCH "repos/${{ github.repository }}/issues/comments/$existing" \
123
+ -F body=@"$RUNNER_TEMP/comment.md" > /dev/null
124
+ else
125
+ gh pr comment "$PR" --body-file "$RUNNER_TEMP/comment.md"
126
+ fi
@@ -37,3 +37,4 @@ site/
37
37
  # OS
38
38
  .DS_Store
39
39
  docs/examples/clustering.json
40
+ .benchmarks/
@@ -0,0 +1,3 @@
1
+ {
2
+ ".": "0.6.6"
3
+ }
@@ -1,5 +1,31 @@
1
1
  # Changelog
2
2
 
3
+ ## [0.6.6](https://github.com/FBumann/tsam_xarray/compare/v0.6.5...v0.6.6) (2026-07-27)
4
+
5
+
6
+ ### Performance Improvements
7
+
8
+ * compute accuracy and reconstructed on first access ([#113](https://github.com/FBumann/tsam_xarray/issues/113)) ([69443aa](https://github.com/FBumann/tsam_xarray/commit/69443aa78a7e4cef5551a1010e65bc3e99210af7))
9
+ * disaggregate segmented clusterings with a scatter ([#117](https://github.com/FBumann/tsam_xarray/issues/117)) ([02c569f](https://github.com/FBumann/tsam_xarray/commit/02c569f44a2c937841a10db343be1c5a2daf70e3))
10
+ * gather from the payload's final layout instead of a strided view ([#119](https://github.com/FBumann/tsam_xarray/issues/119)) ([0a455bf](https://github.com/FBumann/tsam_xarray/commit/0a455bf81792b2654cd644bcce24fbc6a92f32cb))
11
+
12
+ ## [0.6.5](https://github.com/FBumann/tsam_xarray/compare/v0.6.4...v0.6.5) (2026-07-27)
13
+
14
+
15
+ ### Features
16
+
17
+ * support tsam v4 (cluster_counts rename, ClusterConfig.weights removal) ([#105](https://github.com/FBumann/tsam_xarray/issues/105)) ([6cd9f5d](https://github.com/FBumann/tsam_xarray/commit/6cd9f5d41057eb7ec9e3a29f811d7ee52a80a9c0))
18
+
19
+
20
+ ### Bug Fixes
21
+
22
+ * reset uv.lock before the base benchmark run ([#111](https://github.com/FBumann/tsam_xarray/issues/111)) ([3d7d8b7](https://github.com/FBumann/tsam_xarray/commit/3d7d8b7ed0f5b435e9d1eda17080bbaa4a79d158))
23
+
24
+
25
+ ### Performance Improvements
26
+
27
+ * disaggregate unsegmented clusterings with a numpy gather ([#116](https://github.com/FBumann/tsam_xarray/issues/116)) ([0fc2de6](https://github.com/FBumann/tsam_xarray/commit/0fc2de6676b2322f0d0a0512c88a13ecc47ae160))
28
+
3
29
  ## [0.6.4](https://github.com/FBumann/tsam_xarray/compare/v0.6.3...v0.6.4) (2026-07-22)
4
30
 
5
31
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tsam_xarray
3
- Version: 0.6.4
3
+ Version: 0.6.6
4
4
  Summary: Lightweight xarray wrapper for tsam time series aggregation
5
5
  License-Expression: MIT
6
6
  License-File: LICENSE
@@ -44,9 +44,9 @@ tsam works on flat DataFrames. With multi-dimensional data, you end up writing b
44
44
  import tsam_xarray
45
45
 
46
46
  result = tsam_xarray.aggregate(
47
- da, # (time, variable, region, scenario)
47
+ da, # (time, variable, region, scenario)
48
48
  time_dim="time",
49
- cluster_dim=["variable", "region"], # clustered together
49
+ cluster_dim=["variable", "region"], # clustered together
50
50
  n_clusters=4,
51
51
  )
52
52
  # scenario is sliced independently — each gets its own clustering
@@ -55,9 +55,9 @@ result = tsam_xarray.aggregate(
55
55
  Everything comes back as labeled xarray objects:
56
56
 
57
57
  ```python
58
- result.cluster_representatives # (scenario, cluster, timestep, variable, region)
59
- result.reconstructed # same shape as input
60
- result.cluster_assignments # (scenario, period)
58
+ result.cluster_representatives # (scenario, cluster, timestep, variable, region)
59
+ result.reconstructed # same shape as input
60
+ result.cluster_assignments # (scenario, period)
61
61
  ```
62
62
 
63
63
  Accuracy metrics preserve all dimensions — see exactly where the approximation is good or bad:
@@ -65,8 +65,8 @@ Accuracy metrics preserve all dimensions — see exactly where the approximation
65
65
  ![Per-column RMSE across all dimensions](docs/assets/multi-dim-metrics.png)
66
66
 
67
67
  ```python
68
- result.accuracy.rmse # DataArray (scenario, variable, region)
69
- result.accuracy.weighted_rmse # DataArray (scenario,) — per-slice summary
68
+ result.accuracy.rmse # DataArray (scenario, variable, region)
69
+ result.accuracy.weighted_rmse # DataArray (scenario,) — per-slice summary
70
70
  ```
71
71
 
72
72
  ## Save, load, reuse
@@ -77,9 +77,9 @@ result.clustering.to_json("clustering.json")
77
77
 
78
78
  # Load and inspect — no original data needed
79
79
  clustering = tsam_xarray.load_clustering("clustering.json")
80
- clustering.n_clusters # 4
81
- clustering.cluster_assignments # DataArray (scenario, period)
82
- clustering.cluster_occurrences # DataArray (scenario, cluster)
80
+ clustering.n_clusters # 4
81
+ clustering.cluster_assignments # DataArray (scenario, period)
82
+ clustering.cluster_occurrences # DataArray (scenario, cluster)
83
83
 
84
84
  # Apply to new data or disaggregate optimization results
85
85
  new_result = clustering.apply(new_da)
@@ -97,8 +97,8 @@ grid = tsam_xarray.grid_search(
97
97
  cluster_dim=["variable", "region"],
98
98
  timesteps=np.geomspace(2, 48, num=12, dtype=int), # sparse search
99
99
  )
100
- grid.summary_matrix["rmse"] # heatmap-ready (n_clusters, n_segments)
101
- grid.accuracy["weighted_rmse"] # per-slice weighted RMSE for every config
100
+ grid.summary_matrix["rmse"] # heatmap-ready (n_clusters, n_segments)
101
+ grid.accuracy["weighted_rmse"] # per-slice weighted RMSE for every config
102
102
  ```
103
103
 
104
104
  ## Installation
@@ -23,9 +23,9 @@ tsam works on flat DataFrames. With multi-dimensional data, you end up writing b
23
23
  import tsam_xarray
24
24
 
25
25
  result = tsam_xarray.aggregate(
26
- da, # (time, variable, region, scenario)
26
+ da, # (time, variable, region, scenario)
27
27
  time_dim="time",
28
- cluster_dim=["variable", "region"], # clustered together
28
+ cluster_dim=["variable", "region"], # clustered together
29
29
  n_clusters=4,
30
30
  )
31
31
  # scenario is sliced independently — each gets its own clustering
@@ -34,9 +34,9 @@ result = tsam_xarray.aggregate(
34
34
  Everything comes back as labeled xarray objects:
35
35
 
36
36
  ```python
37
- result.cluster_representatives # (scenario, cluster, timestep, variable, region)
38
- result.reconstructed # same shape as input
39
- result.cluster_assignments # (scenario, period)
37
+ result.cluster_representatives # (scenario, cluster, timestep, variable, region)
38
+ result.reconstructed # same shape as input
39
+ result.cluster_assignments # (scenario, period)
40
40
  ```
41
41
 
42
42
  Accuracy metrics preserve all dimensions — see exactly where the approximation is good or bad:
@@ -44,8 +44,8 @@ Accuracy metrics preserve all dimensions — see exactly where the approximation
44
44
  ![Per-column RMSE across all dimensions](docs/assets/multi-dim-metrics.png)
45
45
 
46
46
  ```python
47
- result.accuracy.rmse # DataArray (scenario, variable, region)
48
- result.accuracy.weighted_rmse # DataArray (scenario,) — per-slice summary
47
+ result.accuracy.rmse # DataArray (scenario, variable, region)
48
+ result.accuracy.weighted_rmse # DataArray (scenario,) — per-slice summary
49
49
  ```
50
50
 
51
51
  ## Save, load, reuse
@@ -56,9 +56,9 @@ result.clustering.to_json("clustering.json")
56
56
 
57
57
  # Load and inspect — no original data needed
58
58
  clustering = tsam_xarray.load_clustering("clustering.json")
59
- clustering.n_clusters # 4
60
- clustering.cluster_assignments # DataArray (scenario, period)
61
- clustering.cluster_occurrences # DataArray (scenario, cluster)
59
+ clustering.n_clusters # 4
60
+ clustering.cluster_assignments # DataArray (scenario, period)
61
+ clustering.cluster_occurrences # DataArray (scenario, cluster)
62
62
 
63
63
  # Apply to new data or disaggregate optimization results
64
64
  new_result = clustering.apply(new_da)
@@ -76,8 +76,8 @@ grid = tsam_xarray.grid_search(
76
76
  cluster_dim=["variable", "region"],
77
77
  timesteps=np.geomspace(2, 48, num=12, dtype=int), # sparse search
78
78
  )
79
- grid.summary_matrix["rmse"] # heatmap-ready (n_clusters, n_segments)
80
- grid.accuracy["weighted_rmse"] # per-slice weighted RMSE for every config
79
+ grid.summary_matrix["rmse"] # heatmap-ready (n_clusters, n_segments)
80
+ grid.accuracy["weighted_rmse"] # per-slice weighted RMSE for every config
81
81
  ```
82
82
 
83
83
  ## Installation
@@ -0,0 +1,60 @@
1
+ # Benchmarks
2
+
3
+ Performance suite for `aggregate()`, built on
4
+ [pytest-benchmem](https://github.com/fluxopt/pytest-benchmem) (pytest-benchmark
5
+ plus memray peak-memory columns). Not part of the regular test suite — CI runs
6
+ it via `.github/workflows/benchmarks.yaml`, which posts a head-vs-base delta
7
+ table on every PR and uploads an interactive plot artifact.
8
+
9
+ Three layers (see the module docstring in `test_bench_aggregate.py`):
10
+ wrapper-stage micro-benchmarks with many rounds (the only code this repo's
11
+ PRs can regress — ~90% of an end-to-end run is inside tsam), config cost
12
+ ratios at reduced size, and a few full-pipeline sentinels at production
13
+ size. Cases are a deterministic grid, not random: `benchmem compare`
14
+ matches runs by test ID, so IDs must be stable across runs for baselines
15
+ to work.
16
+
17
+ ## Run locally
18
+
19
+ ```bash
20
+ uv run --with "pytest-benchmem[plot]" pytest benchmarks/ \
21
+ -o addopts="" -p no:xdist --benchmark-only --benchmark-memory
22
+ ```
23
+
24
+ (`-o addopts="" -p no:xdist` because the project's xdist config auto-disables
25
+ benchmarks.)
26
+
27
+ ## Large tier
28
+
29
+ Production-sized cases (`test_large_*`, several seconds and hundreds of MiB
30
+ peak each) are skipped by default — in CI too. Run them locally before
31
+ dependency bumps or performance work:
32
+
33
+ ```bash
34
+ uv run --with "pytest-benchmem[plot]" pytest benchmarks/ --large -k large \
35
+ -o addopts="" -p no:xdist --benchmark-only --benchmark-memory
36
+ ```
37
+
38
+ ## Baselines
39
+
40
+ Save a named baseline (stored in `.benchmarks/`, gitignored — machine-specific):
41
+
42
+ ```bash
43
+ ... --benchmark-save=baseline
44
+ ```
45
+
46
+ Compare a later run against it, failing on >10% mean regressions:
47
+
48
+ ```bash
49
+ ... --benchmark-compare=0001 --benchmark-compare-fail=mean:10%
50
+ ```
51
+
52
+ ## Reports and plots
53
+
54
+ ```bash
55
+ uv run --with "pytest-benchmem[plot]" benchmem compare a.json b.json --diff
56
+ uv run --with "pytest-benchmem[plot]" benchmem plot a.json b.json -o plot.html --open
57
+ ```
58
+
59
+ `benchmem flamegraph <test-id>` renders a memory flamegraph when the run was
60
+ saved with `--benchmark-memory-profile`.
@@ -0,0 +1,25 @@
1
+ import pytest
2
+
3
+
4
+ def pytest_addoption(parser):
5
+ parser.addoption(
6
+ "--large",
7
+ action="store_true",
8
+ default=False,
9
+ help="run large-scale benchmarks (production-sized data, minutes)",
10
+ )
11
+
12
+
13
+ def pytest_configure(config):
14
+ config.addinivalue_line(
15
+ "markers", "large: large-scale benchmark, runs only with --large"
16
+ )
17
+
18
+
19
+ def pytest_collection_modifyitems(config, items):
20
+ if config.getoption("--large"):
21
+ return
22
+ skip = pytest.mark.skip(reason="large benchmark: pass --large to run")
23
+ for item in items:
24
+ if "large" in item.keywords:
25
+ item.add_marker(skip)
@@ -0,0 +1,279 @@
1
+ """Benchmarks for tsam_xarray: wrapper micro-benchmarks + end-to-end sentinels.
2
+
3
+ Run with:
4
+ uv run --with pytest-benchmem pytest benchmarks/ -o addopts= -p no:xdist \
5
+ --benchmark-only --benchmark-memory
6
+
7
+ Three layers, cheapest statistics where the repo's own code lives:
8
+
9
+ - ``test_wrapper_*`` — micro-benchmarks of the wrapper's conversion stages
10
+ (DataFrame in, result out, slice concat). Profiling showed ~90% of an
11
+ end-to-end run is inside tsam, which this repo cannot regress — its PRs
12
+ can only regress these stages, so they get many rounds and wide inputs.
13
+ - ``test_config_*`` — cost ratios of tsam config options at a reduced size
14
+ (90 days); every measured config effect is multiplicative, so ratios at
15
+ 90 days match 365 days at a quarter of the cost.
16
+ - ``test_e2e_*`` — few full-pipeline sentinels at production size for
17
+ integration surprises (tsam version bumps, config plumbing), the
18
+ representation-by-width interaction, and the multi-``cluster_dim``
19
+ MultiIndex path.
20
+ - ``test_large_*`` — production-sized cases (hundreds of MiB peak) that
21
+ guard the memory-amplification regime. Opt-in via ``--large`` (skipped
22
+ otherwise, including in CI) — run locally before dependency bumps or
23
+ performance work.
24
+ - ``test_user_*`` — post-aggregation operations users run on results:
25
+ reusing a stored clustering on new data (``clustering.apply``) and
26
+ expanding cluster-level data back to the full time axis
27
+ (``disaggregate``). Accuracy metrics and ``reconstructed`` need no
28
+ extra cases — tsam computes them eagerly inside ``aggregate()``, so
29
+ every e2e case and the result-conversion micro already include them.
30
+
31
+ Cases are a deterministic grid (data seeded per call) because benchmem
32
+ matches runs by test ID — IDs and data must be stable across runs.
33
+ """
34
+
35
+ from __future__ import annotations
36
+
37
+ import numpy as np
38
+ import pandas as pd
39
+ import pytest
40
+ import tsam
41
+ import xarray as xr
42
+
43
+ from tsam_xarray import aggregate
44
+ from tsam_xarray._core import (
45
+ _aggregate_single,
46
+ _concat_results,
47
+ _result_from_tsam,
48
+ _to_dataframe,
49
+ )
50
+ from tsam_xarray._dim_names import DimNames
51
+
52
+
53
+ def make_data(n_days: int, n_cols: int, n_slices: int) -> xr.DataArray:
54
+ rng = np.random.default_rng(42)
55
+ n_t = n_days * 24
56
+ time_idx = pd.date_range("2023-01-01", periods=n_t, freq="h")
57
+ t = np.arange(n_t)
58
+ daily = np.sin(2 * np.pi * t / 24)
59
+ yearly = np.sin(2 * np.pi * t / 8760)
60
+ base = (
61
+ daily[None, :] * rng.uniform(0.5, 2, (n_cols, 1))
62
+ + yearly[None, :] * rng.uniform(0.5, 2, (n_cols, 1))
63
+ + rng.normal(0, 0.3, (n_cols, n_t))
64
+ )
65
+ dims = ["variable", "time"]
66
+ coords: dict[str, object] = {
67
+ "variable": [f"var{i:03d}" for i in range(n_cols)],
68
+ "time": time_idx,
69
+ }
70
+ data = base
71
+ if n_slices > 1:
72
+ scale = rng.uniform(0.8, 1.2, (n_slices, 1, 1))
73
+ data = base[None, :, :] * scale
74
+ dims = ["scenario", *dims]
75
+ coords["scenario"] = [f"s{i}" for i in range(n_slices)]
76
+ return xr.DataArray(data, dims=dims, coords=coords, name="load")
77
+
78
+
79
+ def cluster_config(representation: object | None = None) -> tsam.ClusterConfig:
80
+ return tsam.ClusterConfig(
81
+ method="hierarchical",
82
+ representation=representation,
83
+ include_period_sums=False,
84
+ )
85
+
86
+
87
+ def extremes_config(method: str = "replace") -> tsam.ExtremeConfig:
88
+ return tsam.ExtremeConfig(method=method, max_value=["var000"], min_value=["var001"])
89
+
90
+
91
+ FULL_CONFIG: dict[str, object] = {
92
+ "cluster": cluster_config(tsam.Distribution(scope="global")),
93
+ "extremes": extremes_config(),
94
+ }
95
+
96
+ REPRESENTATIONS: dict[str, object | None] = {
97
+ "medoid": None,
98
+ "mean": "mean",
99
+ "dist_cluster": tsam.Distribution(scope="cluster"),
100
+ "dist_global": tsam.Distribution(scope="global"),
101
+ "dist_global_minmax": tsam.Distribution(scope="global", preserve_minmax=True),
102
+ }
103
+
104
+
105
+ def run_aggregate(
106
+ da: xr.DataArray,
107
+ config: dict[str, object],
108
+ cluster_dim: list[str] | str = "variable",
109
+ ) -> None:
110
+ aggregate(da, time_dim="time", cluster_dim=cluster_dim, n_clusters=12, **config)
111
+
112
+
113
+ MICRO_OPTS = dict(rounds=10, iterations=10, warmup_rounds=1)
114
+ CONFIG_OPTS = dict(rounds=5, iterations=1, warmup_rounds=1)
115
+ E2E_OPTS = dict(rounds=3, iterations=1, warmup_rounds=1)
116
+
117
+
118
+ # --- wrapper micro-benchmarks -------------------------------------------------
119
+
120
+
121
+ @pytest.fixture(scope="module")
122
+ def wide_tsam_result():
123
+ da = make_data(365, 128, 1)
124
+ df = _to_dataframe(da, "time", ["variable"])
125
+ return tsam.aggregate(df, 12), da, df
126
+
127
+
128
+ @pytest.fixture(scope="module")
129
+ def slice_results():
130
+ da = make_data(365, 8, 8)
131
+ scenarios = da.coords["scenario"].values
132
+ results = [
133
+ _aggregate_single(
134
+ da.sel(scenario=s), 12, "time", ["variable"], None, None, {}, DimNames()
135
+ )
136
+ for s in scenarios
137
+ ]
138
+ return results, {"scenario": scenarios}, [(s,) for s in scenarios]
139
+
140
+
141
+ def test_wrapper_to_dataframe(benchmark):
142
+ da = make_data(365, 128, 1)
143
+ benchmark.pedantic(_to_dataframe, args=(da, "time", ["variable"]), **MICRO_OPTS)
144
+
145
+
146
+ def test_wrapper_result_conversion(benchmark, wide_tsam_result):
147
+ res, da, df = wide_tsam_result
148
+ benchmark.pedantic(
149
+ _result_from_tsam,
150
+ args=(res, da, df, "time", ["variable"], DimNames()),
151
+ **MICRO_OPTS,
152
+ )
153
+
154
+
155
+ def test_wrapper_concat_results(benchmark, slice_results):
156
+ results, slice_coords, slice_keys = slice_results
157
+ benchmark.pedantic(
158
+ _concat_results,
159
+ args=(results, ["scenario"], slice_coords, slice_keys),
160
+ **MICRO_OPTS,
161
+ )
162
+
163
+
164
+ # --- config cost ratios (reduced size) ----------------------------------------
165
+
166
+
167
+ @pytest.mark.parametrize("representation", list(REPRESENTATIONS))
168
+ def test_config_representation(benchmark, representation):
169
+ da = make_data(90, n_cols=8, n_slices=1)
170
+ config = {"cluster": cluster_config(REPRESENTATIONS[representation])}
171
+ benchmark.pedantic(run_aggregate, args=(da, config), **CONFIG_OPTS)
172
+
173
+
174
+ @pytest.mark.parametrize("extremes", ["replace", "append"])
175
+ def test_config_extremes(benchmark, extremes):
176
+ da = make_data(90, n_cols=8, n_slices=1)
177
+ config = {"cluster": cluster_config(), "extremes": extremes_config(extremes)}
178
+ benchmark.pedantic(run_aggregate, args=(da, config), **CONFIG_OPTS)
179
+
180
+
181
+ # --- end-to-end sentinels -----------------------------------------------------
182
+
183
+
184
+ def test_e2e_default(benchmark):
185
+ da = make_data(365, n_cols=8, n_slices=1)
186
+ benchmark.pedantic(run_aggregate, args=(da, {}), **E2E_OPTS)
187
+
188
+
189
+ def test_e2e_full(benchmark):
190
+ da = make_data(365, n_cols=8, n_slices=1)
191
+ benchmark.pedantic(run_aggregate, args=(da, FULL_CONFIG), **E2E_OPTS)
192
+
193
+
194
+ @pytest.mark.parametrize("representation", ["medoid", "dist_global"])
195
+ def test_e2e_wide(benchmark, representation):
196
+ da = make_data(365, n_cols=128, n_slices=1)
197
+ config = {"cluster": cluster_config(REPRESENTATIONS[representation])}
198
+ benchmark.pedantic(run_aggregate, args=(da, config), **E2E_OPTS)
199
+
200
+
201
+ def test_e2e_slices(benchmark):
202
+ da = make_data(365, n_cols=8, n_slices=8)
203
+ benchmark.pedantic(run_aggregate, args=(da, FULL_CONFIG), **E2E_OPTS)
204
+
205
+
206
+ def test_e2e_multidim(benchmark):
207
+ da = make_data(365, n_cols=8, n_slices=4).rename(scenario="region")
208
+ benchmark.pedantic(run_aggregate, args=(da, {}, ["variable", "region"]), **E2E_OPTS)
209
+
210
+
211
+ # --- user post-aggregation operations -----------------------------------------
212
+
213
+
214
+ @pytest.fixture(scope="module")
215
+ def wide_result():
216
+ da = make_data(365, 128, 1)
217
+ result = aggregate(da, time_dim="time", cluster_dim="variable", n_clusters=12)
218
+ return result, da
219
+
220
+
221
+ def test_user_apply(benchmark, wide_result):
222
+ result, da = wide_result
223
+ benchmark.pedantic(result.clustering.apply, args=(da * 1.1,), **CONFIG_OPTS)
224
+
225
+
226
+ def test_user_disaggregate(benchmark, wide_result):
227
+ result, _da = wide_result
228
+ benchmark.pedantic(
229
+ result.clustering.disaggregate,
230
+ args=(result.cluster_representatives,),
231
+ **CONFIG_OPTS,
232
+ )
233
+
234
+
235
+ @pytest.fixture(scope="module")
236
+ def sliced_result():
237
+ da = make_data(365, 32, 12)
238
+ return aggregate(da, time_dim="time", cluster_dim="variable", n_clusters=12)
239
+
240
+
241
+ def test_user_disaggregate_sliced(benchmark, sliced_result):
242
+ benchmark.pedantic(
243
+ sliced_result.clustering.disaggregate,
244
+ args=(sliced_result.cluster_representatives,),
245
+ **CONFIG_OPTS,
246
+ )
247
+
248
+
249
+ def test_user_disaggregate_segmented(benchmark):
250
+ da = make_data(365, 32, 1)
251
+ result = aggregate(
252
+ da,
253
+ time_dim="time",
254
+ cluster_dim="variable",
255
+ n_clusters=12,
256
+ segments=tsam.SegmentConfig(n_segments=6),
257
+ )
258
+ benchmark.pedantic(
259
+ result.clustering.disaggregate,
260
+ args=(result.cluster_representatives,),
261
+ **CONFIG_OPTS,
262
+ )
263
+
264
+
265
+ # --- large tier (production-sized, opt-in via --large) -------------------------
266
+
267
+ LARGE_OPTS = dict(rounds=2, iterations=1, warmup_rounds=0)
268
+
269
+
270
+ @pytest.mark.large
271
+ def test_large_wide(benchmark):
272
+ da = make_data(730, n_cols=256, n_slices=1)
273
+ benchmark.pedantic(run_aggregate, args=(da, FULL_CONFIG), **LARGE_OPTS)
274
+
275
+
276
+ @pytest.mark.large
277
+ def test_large_scenarios(benchmark):
278
+ da = make_data(365, n_cols=64, n_slices=8)
279
+ benchmark.pedantic(run_aggregate, args=(da, FULL_CONFIG), **LARGE_OPTS)
@@ -32,7 +32,7 @@ graph LR
32
32
  D --- F1[".cluster_representatives<br/><i>cluster, timestep, *cluster_dims, *slice_dims</i>"]
33
33
  D --- F2[".reconstructed<br/><i>same shape as input</i>"]
34
34
  D --- F3[".cluster_assignments<br/><i>period, *slice_dims</i>"]
35
- D --- F4[".cluster_weights<br/><i>cluster, *slice_dims</i>"]
35
+ D --- F4[".cluster_counts<br/><i>cluster, *slice_dims</i>"]
36
36
  D --- F5[".segment_durations<br/><i>cluster, timestep, *slice_dims | None</i>"]
37
37
 
38
38
  Meta --- A[".accuracy<br/><b>→ AccuracyMetrics</b>"]
@@ -107,8 +107,8 @@
107
107
  "source": [
108
108
  "print(f\"Clusters: {result.n_clusters}\")\n",
109
109
  "print(f\"Timesteps per period: {result.n_timesteps_per_period}\")\n",
110
- "print(\"Cluster weights (days each represents):\")\n",
111
- "result.cluster_weights.to_dataframe(\"weight\")"
110
+ "print(\"Cluster counts (days each represents):\")\n",
111
+ "result.cluster_counts.to_dataframe(\"count\")"
112
112
  ]
113
113
  },
114
114
  {