Bugfix/investigate stein thinning #1000

gw265981 · 2025-03-18T16:20:17Z

PR Type

Bugfix
Feature

Description

Improving the performance of Stein Thinning in benchmarks, primarily by using regularisation and better probability density estimation.

How Has This Been Tested?

Existing tests pass as expected.

Does this PR introduce a breaking change?

Checklist before requesting a review

I have made sure that my PR is not a duplicate.
My code follows the style guidelines of this project.
I have ensured my code is easy to understand, including docstrings and comments where necessary.
I have performed a self-review of my code.
I have made corresponding changes to the documentation.
My changes generate no new warnings.
New and existing unit tests pass locally with my changes.
Any dependent changes have been merged and published in downstream modules.
I have updated CHANGELOG.md, if appropriate.

…hmark

…teinThinning in blobs benchmark

…hmarking

# Conflicts: # benchmark/blobs_benchmark.py

github-actions · 2025-03-18T16:31:42Z

Performance review

Commit `5dee5cd` - Merge `41fcd48` into `6ed8e9d`

No significant changes to performance.

github-actions · 2025-03-18T16:39:38Z

Performance review

Commit `6e21644` - Merge `ed405a3` into `6ed8e9d`

No significant changes to performance.

github-actions · 2025-03-19T10:36:03Z

Performance review

Commit `b319566` - Merge `70dc93e` into `6ed8e9d`

No significant changes to performance.

…enchmark.py

github-actions · 2025-03-19T15:24:50Z

Performance review

Commit `fe14875` - Merge `6c78ac9` into `6ed8e9d`

No significant changes to performance.

github-actions · 2025-03-19T17:42:10Z

Performance review

Commit `f20278d` - Merge `822df59` into `6ed8e9d`

No significant changes to performance.

github-actions · 2025-03-19T17:46:31Z

Performance review

Commit `8bbbae2` - Merge `10000f9` into `6ed8e9d`

No significant changes to performance.

github-actions · 2025-03-19T18:33:39Z

Performance review

Commit `c4f46f1` - Merge `9b4f97f` into `6ed8e9d`

No significant changes to performance.

github-actions · 2025-03-19T18:47:37Z

Performance review

Commit `e192a1a` - Merge `4954ba6` into `6ed8e9d`

No significant changes to performance.

gw265981 · 2025-03-19T18:47:45Z

List of all changes:

Replace the score function for Stein Thinning in blobs_benchmark.py and benchmark_util.py (and so for other benchmarks) with jax.scipy.stats.gaussian_kde(). Also set regularisation to True.
Add a parameter kde_bw_method to the SteinThinning solver which is passed to jax.scipy.stats.gaussian_kde() when estimating probability density for regularisation.
Small updates to benchmarks, including making sure the saving paths are correct, fixing random state of UMAP for reproducibility, reducing UMAP dimension for the pounce benchmark.
Rerun all the benchmarks and update the benchmarking.rst page.
Add test to test_benchmark.py to run the initialised SteinThinning solver.

gw265981 · 2025-03-20T11:00:30Z

Also addressed #992, #985. For the alt text, I left it unchanged for the blobs benchmark charts since raw data is provided.

github-actions · 2025-03-20T11:08:36Z

Performance review

Commit `773e7b3` - Merge `a90dfc9` into `6ed8e9d`

No significant changes to performance.

rg936672

Looks good, just a few minor suggestions. Remember to link the issues that this PR closes, either under "Development" at the top right, or by using "Closes #xyz" in the issue description.

benchmark/david_benchmark.py

benchmark/pounce_benchmark.py

coreax/benchmark_util.py

documentation/source/benchmark.rst

Co-authored-by: rg936672 <162452529+rg936672@users.noreply.github.com>

github-actions · 2025-03-20T12:40:27Z

Performance review

Commit `0e2a402` - Merge `a4b560f` into `6ed8e9d`

No significant changes to performance.

benchmark/blobs_benchmark.py

…g' into bugfix/investigate-stein-thinning

github-actions · 2025-03-20T13:15:46Z

Performance review

Commit `c6280fd` - Merge `12614e0` into `6ed8e9d`

No significant changes to performance.

rg936672

I'm happy my suggestions have all been addressed, though hold off on merging until both (a) @pc532627 approves as well and (b) you've created a ticket for any suggestions that you've deemed out of scope for this PR. (Could be as simple as a "Clean up benchmarking code" with a link to the review comments in question.)

pc532627

One minor typo

documentation/source/benchmark.rst

github-actions · 2025-03-20T16:14:22Z

Performance review

Commit `388fea6` - Merge `31bfe2e` into `6ed8e9d`

No significant changes to performance.

rg936672 · 2025-03-20T16:50:19Z

TODO failure is a bug in the latest Ruff - see astral-sh/ruff#16874

These have been fixed by the changes in #1000.

gw265981 added 6 commits March 18, 2025 12:24

feat: change the KSD metric to use the stein kernel in the blobs benc…

73b860b

…hmark

feat: use gaussian kde as score function and use regularisation for S…

6fc91ca

…teinThinning in blobs benchmark

fix: improve typing for score function

22d9143

feat: add KDE bandwidth method as a parameter to SteinThinning

4c0d806

feat: use scipy gaussian_kde for SteinThinning score function in benc…

a8ee69c

…hmarking

Merge branch 'main' into bugfix/investigate-stein-thinning

41fcd48

# Conflicts: # benchmark/blobs_benchmark.py

fix: references to external method

ed405a3

gw265981 added 3 commits March 19, 2025 09:26

feat: update benchmarks save paths

16b610e

feat: save frame plots for pounce benchmark

254fb63

feat: update benchmarking results

70dc93e

feat: fix UMAP random_state and reduce UMAP n_components for pounce_b…

6c78ac9

…enchmark.py

feat: update benchmarking results

822df59

feat: add time results for blobs benchmark

10000f9

test: add test for SteinThinning score function in benchmark_util.py

9b4f97f

test: update score function test

4954ba6

gw265981 marked this pull request as ready for review March 20, 2025 09:09

gw265981 added 2 commits March 20, 2025 09:23

docs: add clarification about unsupervised benchmarking

21faf7e

docs: improve alt text

a90dfc9

rg936672 self-requested a review March 20, 2025 11:11

rg936672 suggested changes Mar 20, 2025

View reviewed changes

benchmark/david_benchmark.py Show resolved Hide resolved

benchmark/david_benchmark.py Show resolved Hide resolved

benchmark/pounce_benchmark.py Show resolved Hide resolved

coreax/benchmark_util.py Show resolved Hide resolved

documentation/source/benchmark.rst Outdated Show resolved Hide resolved

rg936672 mentioned this pull request Mar 20, 2025

Re-enable Jax 0.5.x #999

Merged

9 tasks

This was linked to issues Mar 20, 2025

Add a better analytic test for Stein Thinning #979

Closed

Add better alt texts for all the images in the benchmarking rst file #985

Closed

Add clarification to benchmarking that it's only for unsupervised #992

Closed

fix: type in benchmark alt text

a4b560f

Co-authored-by: rg936672 <162452529+rg936672@users.noreply.github.com>

pc532627 reviewed Mar 20, 2025

View reviewed changes

benchmark/blobs_benchmark.py Show resolved Hide resolved

gw265981 added 4 commits March 20, 2025 13:01

docs: add full stops

0a34877

docs: save path clarification for david_benchmark.py

8350ff7

docs: add changes to CHANGELOG.md

c893b94

Merge remote-tracking branch 'origin/bugfix/investigate-stein-thinnin…

12614e0

…g' into bugfix/investigate-stein-thinning

rg936672 previously approved these changes Mar 20, 2025

View reviewed changes

pc532627 self-requested a review March 20, 2025 15:09

pc532627 approved these changes Mar 20, 2025

View reviewed changes

pc532627 suggested changes Mar 20, 2025

View reviewed changes

documentation/source/benchmark.rst Outdated Show resolved Hide resolved

fix: typo in benchmark docs page

31bfe2e

gw265981 dismissed rg936672’s stale review via 31bfe2e March 20, 2025 16:12

gw265981 mentioned this pull request Mar 20, 2025

Refactor benchmarking code #1002

Open

pc532627 self-requested a review March 20, 2025 17:33

pc532627 approved these changes Mar 20, 2025

View reviewed changes

pc532627 merged commit c7fe19d into main Mar 20, 2025
22 of 24 checks passed

pc532627 deleted the bugfix/investigate-stein-thinning branch March 20, 2025 17:34

rg936672 added a commit that referenced this pull request Mar 21, 2025

remove: skip for stein refinement tests

c11cd38

These have been fixed by the changes in #1000.

Bugfix/investigate stein thinning #1000

Bugfix/investigate stein thinning #1000

Uh oh!

Conversation

gw265981 commented Mar 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Type

Description

How Has This Been Tested?

Does this PR introduce a breaking change?

Checklist before requesting a review

Uh oh!

github-actions bot commented Mar 18, 2025

Performance review

Commit 5dee5cd - Merge 41fcd48 into 6ed8e9d

Uh oh!

github-actions bot commented Mar 18, 2025

Performance review

Commit 6e21644 - Merge ed405a3 into 6ed8e9d

Uh oh!

github-actions bot commented Mar 19, 2025

Performance review

Commit b319566 - Merge 70dc93e into 6ed8e9d

Uh oh!

github-actions bot commented Mar 19, 2025

Performance review

Commit fe14875 - Merge 6c78ac9 into 6ed8e9d

Uh oh!

github-actions bot commented Mar 19, 2025

Performance review

Commit f20278d - Merge 822df59 into 6ed8e9d

Uh oh!

github-actions bot commented Mar 19, 2025

Performance review

Commit 8bbbae2 - Merge 10000f9 into 6ed8e9d

Uh oh!

github-actions bot commented Mar 19, 2025

Performance review

Commit c4f46f1 - Merge 9b4f97f into 6ed8e9d

Uh oh!

github-actions bot commented Mar 19, 2025

Performance review

Commit e192a1a - Merge 4954ba6 into 6ed8e9d

Uh oh!

gw265981 commented Mar 19, 2025

Uh oh!

gw265981 commented Mar 20, 2025

Uh oh!

github-actions bot commented Mar 20, 2025

Performance review

Commit 773e7b3 - Merge a90dfc9 into 6ed8e9d

Uh oh!

rg936672 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Mar 20, 2025

Performance review

Commit 0e2a402 - Merge a4b560f into 6ed8e9d

Uh oh!

Uh oh!

github-actions bot commented Mar 20, 2025

Performance review

Commit c6280fd - Merge 12614e0 into 6ed8e9d

Uh oh!

rg936672 left a comment

Choose a reason for hiding this comment

Uh oh!

pc532627 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Mar 20, 2025

Performance review

Commit 388fea6 - Merge 31bfe2e into 6ed8e9d

Uh oh!

gw265981 commented Mar 18, 2025 •

edited

Loading

Commit `5dee5cd` - Merge `41fcd48` into `6ed8e9d`

Commit `6e21644` - Merge `ed405a3` into `6ed8e9d`

Commit `b319566` - Merge `70dc93e` into `6ed8e9d`

Commit `fe14875` - Merge `6c78ac9` into `6ed8e9d`

Commit `f20278d` - Merge `822df59` into `6ed8e9d`

Commit `8bbbae2` - Merge `10000f9` into `6ed8e9d`

Commit `c4f46f1` - Merge `9b4f97f` into `6ed8e9d`

Commit `e192a1a` - Merge `4954ba6` into `6ed8e9d`

Commit `773e7b3` - Merge `a90dfc9` into `6ed8e9d`

Commit `0e2a402` - Merge `a4b560f` into `6ed8e9d`

Commit `c6280fd` - Merge `12614e0` into `6ed8e9d`

Commit `388fea6` - Merge `31bfe2e` into `6ed8e9d`