Finish Metropolis-Hastings sampling for JAX backend

[Copilot]

MH sampling was explicitly blocked on JAX (assert backend != "jax") with no implementation. This completes the JAX path end-to-end.

JAX random state API

• Renamed internal get_state → _get_state in jax/random.py; exposed public get_state/set_state aliases so MH can split/update the global PRNG key
• Implemented multinomial for JAX via jax.random.categorical + jnp.bincount (was previously a no-op returning NotImplementedError)
• Added get_state/set_state to numpy and pytorch backends for API symmetry

Core MH sampler

Replaced the JAX-blocking assertion with a JAX-specific code path in sample_metropolis_hastings. Uses a plain Python loop rather than lax.scan so that scipy-based pdf methods (e.g. HypertoroidalWrappedNormalDistribution) remain compatible:

# JAX proposals must now accept (key, x) → x_prop
def proposal(key, x):
    key, subkey = jax.random.split(key)
    return jnp.mod(x + jax.random.normal(subkey, x.shape), 2 * jnp.pi)

samples = dist.sample_metropolis_hastings(n=100, proposal=proposal, start_point=x0)

Added _assert_proposal_supports_key() to give a clear error when a single-argument numpy-style proposal is passed with the JAX backend.

JAX-compatible default proposals

Added (key, x) → x_prop default proposals to:

• AbstractHypersphericalDistribution — uniform on S^d via normal + normalize (with explicit re-normalization to avoid float32 drift)
• AbstractHyperhemisphericalDistribution — same, then flip sign if last coord < 0
• AbstractHypertoroidalDistribution — wrapped Gaussian step

Mixture sampling

AbstractMixture.sample now has a JAX branch that iterates components and falls back from sample() to sample_metropolis_hastings() when the direct sampler is unsupported (catches NotImplementedError, AssertionError, ValueError, TypeError).

Other branch changes

• apply_function: renamed parameter function_is_vectorized → f_supports_multiple; removed beartype decorator
• AbstractSphericalHarmonicsDistribution.normalize_in_place: removed unused warn_unnorm kwarg
• CircularFourierDistribution arithmetic: preserve n after +/- operations
• PartiallyWrappedNormalDistribution.set_mode: simplified to direct assignment
• Test files: removed JAX skip decorators from test_sample_metropolis_hastings_basics_only_{t2,s2,h2}; cleaned up redundant warnings.catch_warnings blocks

[github-actions]

❌MegaLinter analysis: Error

Descriptor	Linter	Files	Fixed	Errors	Warnings	Elapsed time
✅ COPYPASTE	jscpd	yes		no	no	7.45s
✅ JSON	prettier	2	0	0	0	0.6s
✅ JSON	v8r	2		0	0	2.77s
✅ MARKDOWN	markdownlint	1	0	0	0	0.68s
✅ MARKDOWN	markdown-table-formatter	1	0	0	0	0.25s
✅ PYTHON	bandit	249		0	0	3.09s
✅ PYTHON	black	249	7	0	0	4.74s
❌ PYTHON	flake8	249		1	0	1.75s
✅ PYTHON	isort	249	8	0	0	0.51s
✅ PYTHON	mypy	249		0	0	4.01s
✅ PYTHON	pylint	249		0	0	70.09s
✅ PYTHON	ruff	249	8	0	0	0.05s
✅ REPOSITORY	checkov	yes		no	no	22.04s
✅ REPOSITORY	gitleaks	yes		no	no	4.75s
✅ REPOSITORY	git_diff	yes		no	no	0.04s
✅ REPOSITORY	secretlint	yes		no	no	5.82s
✅ REPOSITORY	syft	yes		no	no	3.8s
✅ REPOSITORY	trivy-sbom	yes		no	no	1.75s
✅ REPOSITORY	trufflehog	yes		no	no	17.47s
✅ YAML	prettier	4	0	0	0	0.49s
✅ YAML	v8r	4		0	0	5.38s
✅ YAML	yamllint	4		0	0	0.45s

Detailed Issues

❌ PYTHON / flake8 - 1 error

pyrecest/distributions/abstract_mixture.py:92:24: E203 whitespace before ':'

See detailed reports in MegaLinter artifacts

Your project could benefit from a custom flavor, which would allow you to run only the linters you need, and thus improve runtime performances. (Skip this info by defining FLAVOR_SUGGESTIONS: false)

• Documentation: Custom Flavors
• Command: npx mega-linter-runner@9.4.0 --custom-flavor-setup --custom-flavor-linters PYTHON_PYLINT,PYTHON_BLACK,PYTHON_FLAKE8,PYTHON_ISORT,PYTHON_BANDIT,PYTHON_MYPY,PYTHON_RUFF,COPYPASTE_JSCPD,JSON_V8R,JSON_PRETTIER,MARKDOWN_MARKDOWNLINT,MARKDOWN_MARKDOWN_TABLE_FORMATTER,REPOSITORY_CHECKOV,REPOSITORY_GIT_DIFF,REPOSITORY_GITLEAKS,REPOSITORY_SECRETLINT,REPOSITORY_SYFT,REPOSITORY_TRIVY_SBOM,REPOSITORY_TRUFFLEHOG,YAML_PRETTIER,YAML_YAMLLINT,YAML_V8R

Show us your support by starring ⭐ the repository

[FlorianPfaff]

@copilot Can you fix the issues that Megalinter complains about?

[FlorianPfaff]

@copilot Can you fix the issues that Megalinter complains about?

[FlorianPfaff]

@copilot Can you fix the issues that Megalinter complains about?

[Copilot]

@copilot Can you fix the issues that Megalinter complains about?

Fixed and pushed in commit b911ffa.

I addressed the MegaLinter failures by resolving the flake8 and pylint issues in the MH/JAX-related files with minimal code changes, then re-ran targeted lint checks and MH sampling tests (3 passed).

No UI was changed in this PR (backend-only), so there is no UI screenshot to provide.

[Copilot]

@copilot Can you fix the issues that Megalinter complains about?

Addressed in commit b911ffa.

The MegaLinter issues previously reported (flake8/pylint in the MH/JAX files) were fixed there, and targeted lint/test validation passed locally. The latest MegaLinter run for that commit is action_required with no jobs executed, so there are no new linter diagnostics in that run to fix.

[github-actions]

Test Results

   12 files  ±    0     12 suites  ±0   4h 57m 22s ⏱️ + 1h 51m 9s
  672 tests +  139    672 ✅ +  139      0 💤 ±  0  0 ❌ ±0 
8 324 runs  +1 716  6 972 ✅ +1 540  1 352 💤 +176  0 ❌ ±0 

Results for commit 7f47630. ± Comparison against base commit f5c0f75.

♻️ This comment has been updated with latest results.

[FlorianPfaff]

@copilot Fix the linter error

************* Module pyrecest.tests.distributions.test_spherical_harmonics_distribution_complex
pyrecest/tests/distributions/test_spherical_harmonics_distribution_complex.py:2:0: W0611: Unused import warnings (unused-import)
************* Module pyrecest.tests.test_evaluation_basic
pyrecest/tests/test_evaluation_basic.py:5:0: W0611: Unused import warnings (unused-import)
************* Module update_init_helper
update_init_helper.py:1:0: R0801: Similar lines in 2 files
==pyrecest.distributions.hypersphere_subset.abstract_hyperhemispherical_distribution:[67:78]
==pyrecest.distributions.hypersphere_subset.abstract_hyperspherical_distribution:[80:90]
proposal = proposal_np
else:
import jax as _jax # pylint: disable=import-error
import jax.numpy as _jnp # pylint: disable=import-error

            def proposal_jax(key, _):
                """JAX independence proposal: uniform on hypersphere."""
                key, subkey = _jax.random.split(key)
                s = _jax.random.normal(subkey, shape=(1, self.dim + 1))
                # Ensure exact unit norm (avoids float32 rounding errors) (duplicate-code)

---

Your code has been rated at 10.00/10

[Copilot]

@copilot Fix the linter error

************* Module pyrecest.tests.distributions.test_spherical_harmonics_distribution_complex
...

Fixed in commit f384fe4.

Changes made:

• Removed unused warnings imports in:
  • pyrecest/tests/distributions/test_spherical_harmonics_distribution_complex.py
  • pyrecest/tests/test_evaluation_basic.py
• Added a local pylint suppression for the intentional JAX proposal duplication in:
  • pyrecest/distributions/hypersphere_subset/abstract_hyperhemispherical_distribution.py (duplicate-code)

No UI was changed (backend/test-only), so there is no UI screenshot.

[FlorianPfaff]

@copilot Increase rtol to 5e-7

============================
FAILED pyrecest/tests/distributions/test_abstract_mixture.py::AbstractMixtureTest::test_sample_metropolis_hastings_basics_only_h2 - AssertionError: Array(False, dtype=bool) is not true
FAILED pyrecest/tests/distributions/test_abstract_mixture.py::AbstractMixtureTest::test_sample_metropolis_hastings_basics_only_s2 - AssertionError: Array(False, dtype=bool) is not true
FAILED pyrecest/tests/distributions/test_gaussian_distribution.py::GaussianDistributionTest::test_gaussian_distribution_3d - AssertionError:
Not equal to tolerance rtol=1e-07, atol=1e-10

Mismatched elements: 1 / 11 (9.09%)
Mismatch at index:
[6]: 0.1281534880399704 (ACTUAL), 0.1281534464969997 (DESIRED)
Max absolute difference among violations: 4.15429707e-08
Max relative difference among violations: 3.24165848e-07
ACTUAL: array([0.000000e+00, 0.000000e+00, 0.000000e+00, 8.996547e-25,
4.393428e-12, 2.294495e-04, 1.281535e-01, 7.654790e-04,
4.889852e-11, 3.340534e-23, 0.000000e+00], dtype=float32)
DESIRED: array([1.155751e-92, 4.614340e-65, 1.970219e-42, 8.996613e-25,
4.393425e-12, 2.294494e-04, 1.281534e-01, 7.654790e-04,
4.889853e-11, 3.340550e-23, 2.440618e-40])
= 3 failed, 369 passed, 244 skipped, 48 warnings, 5 subtests passed in 1490.75s (0:24:50) =

[Copilot]

@copilot Increase rtol to 5e-7

============================
...

Done in commit a146cef.

I increased the relevant tolerances to rtol=5e-7 in:

• /home/runner/work/PyRecEst/PyRecEst/pyrecest/tests/distributions/test_abstract_mixture.py (S2/H2 MH norm checks)
• /home/runner/work/PyRecEst/PyRecEst/pyrecest/tests/distributions/test_gaussian_distribution.py (3D Gaussian PDF comparison)

No UI changes were made, so there is no UI screenshot.