diff --git a/src/sequince/steps/mutation/kernels.py b/src/sequince/steps/mutation/kernels.py index 3e7a7a7e..f7cce19b 100644 --- a/src/sequince/steps/mutation/kernels.py +++ b/src/sequince/steps/mutation/kernels.py @@ -481,6 +481,15 @@ class DifferentialEvolutionKernel(SymmetricComplementaryEnsembleKernel): None assigns equal weights. Configuration is immutable and selection is state-independent, so the global proposal remains symmetric. Use ``dataclasses.replace`` to construct weights for a later stage. + noise_scale : float or tuple of float + Standard deviation of independent, zero-mean Gaussian noise added to + every proposal, including unit-scale moves. A scalar applies to all + dimensions; a tuple must contain one value per dimension. Values must + be finite and non-negative, in the coordinates used by the kernel + (transformed coordinates when a mutation transform is supplied). + The default is ``1e-5``. Zero disables noise and its RNG draws. + Positive noise in every dimension can escape a deficient affine span; + its fixed scale must not depend on the active particle's position. """ gamma: float | None = None @@ -488,8 +497,22 @@ class DifferentialEvolutionKernel(SymmetricComplementaryEnsembleKernel): gamma_randomization: GammaRandomization = "none" scale_multipliers: tuple[float, ...] = (1.0,) scale_weights: tuple[float, ...] | None = None + noise_scale: float | tuple[float, ...] = 1e-5 def __post_init__(self) -> None: + noise_scale = ( + float(self.noise_scale) + if isinstance(self.noise_scale, (float, int)) + else tuple(float(value) for value in self.noise_scale) + ) + noise_values = ( + (noise_scale,) if isinstance(noise_scale, float) else noise_scale + ) + if not noise_values or any( + not math.isfinite(value) or value < 0 for value in noise_values + ): + raise ValueError("noise_scale must be finite and non-negative.") + object.__setattr__(self, "noise_scale", noise_scale) multipliers = tuple(float(value) for value in self.scale_multipliers) weights = ( (1.0,) * len(multipliers) @@ -576,6 +599,11 @@ def sample( ) n_active, n_dimensions = current.shape + if ( + isinstance(self.noise_scale, tuple) + and len(self.noise_scale) != n_dimensions + ): + raise ValueError("noise_scale must have one value per dimension.") donor_j, donor_k = self._select_donors(current, complementary, rng) gamma_base = self.gamma if gamma_base is None: @@ -605,8 +633,27 @@ def sample( component = xp.where(unit_mask, -1, component) multiplier = xp.where(unit_mask, 1.0, multiplier) differences = complementary[donor_j] - complementary[donor_k] + proposed = ( + current + xp.reshape(gamma_effective, (n_active, 1)) * differences + ) + noise_values = ( + (self.noise_scale,) + if isinstance(self.noise_scale, float) + else self.noise_scale + ) + if any(value > 0 for value in noise_values): + noise = asarray( + rng.normal(size=current.shape), + xp=xp, + like=current, + dtype=current.dtype, + ) + scale = asarray( + self.noise_scale, xp=xp, like=current, dtype=current.dtype + ) + proposed = proposed + noise * scale return KernelSample( - current + xp.reshape(gamma_effective, (n_active, 1)) * differences, + proposed, { "scale_component": component, "scale_multiplier": multiplier, diff --git a/src/sequince/steps/mutation/metropolis.py b/src/sequince/steps/mutation/metropolis.py index 6ea63383..12cd4245 100644 --- a/src/sequince/steps/mutation/metropolis.py +++ b/src/sequince/steps/mutation/metropolis.py @@ -807,6 +807,11 @@ class DifferentialEvolutionMutation(RedBlueMetropolisHastingsMutation): scale_weights : tuple of float or None Non-negative mixture weights (normalized internally); None is uniform. The kernel is fixed for each call and scale draws are state-independent. + noise_scale : float or tuple of float + Independent Gaussian noise standard deviation in mutation coordinates, + either shared or one value per dimension. Must be finite and + non-negative. Default is ``1e-5``; zero disables noise and its RNG draws. + With a transform, specify the scale in transformed coordinates. transform Optional coordinate transform applied before mutation. transform_refit @@ -831,6 +836,7 @@ def __init__( gamma_randomization: GammaRandomization = "none", scale_multipliers: tuple[float, ...] = (1.0,), scale_weights: tuple[float, ...] | None = None, + noise_scale: float | tuple[float, ...] = 1e-5, transform: ParticleTransform | None = None, transform_refit: TransformRefit = "always", store_chains: bool = False, @@ -842,6 +848,7 @@ def __init__( gamma_randomization=gamma_randomization, scale_multipliers=scale_multipliers, scale_weights=scale_weights, + noise_scale=noise_scale, ) super().__init__( log_prob_fn, diff --git a/tests/test_mutation/test_differential_evolution.py b/tests/test_mutation/test_differential_evolution.py index 743c5672..92b0d1ca 100644 --- a/tests/test_mutation/test_differential_evolution.py +++ b/tests/test_mutation/test_differential_evolution.py @@ -73,12 +73,15 @@ def test_fixed_rng_replays_proposals() -> None: np.testing.assert_array_equal(first, replayed) +@pytest.mark.parametrize("noise_scale", [0.0, 0.1, (0.1, 0.2)]) @pytest.mark.parametrize("backend", ["numpy", "torch", "jax"]) def test_shape_dtype_and_backend_are_preserved( - backend, backend_factory + backend, backend_factory, noise_scale ) -> None: case = backend_factory(backend) - kernel = DifferentialEvolutionKernel(gamma_randomization="bilby") + kernel = DifferentialEvolutionKernel( + gamma_randomization="bilby", noise_scale=noise_scale + ) proposals = kernel.sample( case.particles[:2], @@ -160,7 +163,9 @@ def test_donor_indices_are_backend_integer_index_arrays( def test_donors_come_only_from_complementary_population() -> None: active = np.full((3, 1), 100.0) complementary = np.array([[0.0], [2.0], [5.0]]) - kernel = DifferentialEvolutionKernel(gamma=1.0, gamma_one_probability=0.0) + kernel = DifferentialEvolutionKernel( + noise_scale=0.0, gamma=1.0, gamma_one_probability=0.0 + ) proposals = kernel.sample( active, complementary, RandomGenerator(seed=4) @@ -179,7 +184,9 @@ def test_donors_come_only_from_complementary_population() -> None: def test_default_gamma_uses_latent_dimension() -> None: active = np.zeros((2, 2)) complementary = np.array([[0.0, 0.0], [1.0, 2.0]]) - kernel = DifferentialEvolutionKernel(gamma_one_probability=0.0) + kernel = DifferentialEvolutionKernel( + noise_scale=0.0, gamma_one_probability=0.0 + ) rng = _controlled_rng([1, 1], [0, 0]) proposals = kernel.sample( # type: ignore[arg-type] @@ -199,6 +206,7 @@ def test_standard_mode_uses_explicit_gamma_exactly_without_jitter() -> None: active = np.array([[10.0, 20.0], [30.0, 40.0]]) complementary = np.array([[1.0, 2.0], [4.0, 8.0], [9.0, 7.0]]) kernel = DifferentialEvolutionKernel( + noise_scale=0.0, gamma=0.7, gamma_one_probability=0.0, gamma_randomization="none", @@ -224,6 +232,7 @@ def test_bilby_mode_uses_controlled_gamma_factors_exactly() -> None: active = np.array([[10.0], [20.0]]) complementary = np.array([[1.0], [4.0], [9.0]]) kernel = DifferentialEvolutionKernel( + noise_scale=0.0, gamma=0.7, gamma_one_probability=0.0, gamma_randomization="bilby", @@ -340,6 +349,7 @@ def test_singleton_matches_legacy_rng_exactly( active = np.arange(40.0).reshape(20, 2) complement = np.arange(16.0).reshape(8, 2) kernel = DifferentialEvolutionKernel( + noise_scale=0.0, gamma=0.7, gamma_one_probability=probability, gamma_randomization=randomization, @@ -393,6 +403,7 @@ def test_invalid_scale_mixtures(multipliers, weights, message) -> None: def test_mixture_metadata_frequencies_and_independence() -> None: kernel = DifferentialEvolutionKernel( + noise_scale=0.0, gamma=0.4, gamma_one_probability=0.2, scale_multipliers=(0.1, 0.5, 1.0), @@ -556,9 +567,10 @@ def target(x, context): ) +@pytest.mark.parametrize("noise_scale", [0.0, (0.2, 0.3)]) @pytest.mark.parametrize("randomization", ["none", "bilby"]) def test_multiscale_correlated_gaussian_stationarity( - context_factory, randomization + context_factory, randomization, noise_scale ) -> None: covariance = np.array([[1.0, 0.8], [0.8, 2.0]]) mean = np.array([1.0, -2.0]) @@ -570,6 +582,7 @@ def test_multiscale_correlated_gaussian_stationarity( scale_multipliers=(0.1, 0.5, 1.0), scale_weights=(1.0, 2.0, 3.0), gamma_randomization=randomization, + noise_scale=noise_scale, ) mutation = RedBlueMetropolisHastingsMutation( lambda x, context: -0.5 @@ -606,7 +619,10 @@ def log_transition_ratio_from_sample( return np.where(metadata["scale_component"] == 0, 0.0, -np.inf) kernel = CorrectedTestKernel( - gamma=1.0, gamma_one_probability=0, scale_multipliers=(0.5, 1.0) + noise_scale=0.0, + gamma=1.0, + gamma_one_probability=0, + scale_multipliers=(0.5, 1.0), ) mutation = RedBlueMetropolisHastingsMutation( lambda x, context: np.zeros(len(x)), @@ -632,6 +648,7 @@ def log_transition_ratio_from_sample( def test_unit_only_mixture_ignores_all_factors() -> None: kernel = DifferentialEvolutionKernel( + noise_scale=0.0, gamma=0.3, gamma_one_probability=1.0, gamma_randomization="bilby", @@ -648,6 +665,7 @@ def test_unit_only_mixture_ignores_all_factors() -> None: def test_bilby_factors_multiply_selected_scale() -> None: kernel = DifferentialEvolutionKernel( + noise_scale=0.0, gamma=0.4, gamma_one_probability=0, gamma_randomization="bilby", @@ -857,3 +875,97 @@ def accumulate(state): assert jax.tree_util.tree_structure(final) == jax.tree_util.tree_structure( initial ) + + +@pytest.mark.parametrize("noise_scale", [-1.0, np.inf, np.nan, (), (0.1, -1)]) +def test_invalid_noise_scale(noise_scale: Any) -> None: + with pytest.raises(ValueError, match="noise_scale"): + DifferentialEvolutionKernel(noise_scale=noise_scale) + + +def test_noise_scale_requires_one_value_per_dimension() -> None: + kernel = DifferentialEvolutionKernel(noise_scale=(0.1,)) + with pytest.raises(ValueError, match="one value per dimension"): + kernel.sample( + np.zeros((4, 2)), np.zeros((4, 2)), RandomGenerator(seed=1) + ) + + +@pytest.mark.parametrize("unit_probability", [0.0, 1.0]) +def test_noise_moves_collapsed_population_and_replays( + unit_probability: float, +) -> None: + # With identical donors, every displacement comes from additive noise. + active = np.zeros((30_000, 2)) + kernel = DifferentialEvolutionKernel( + noise_scale=(0.1, 0.3), + gamma_one_probability=unit_probability, + gamma_randomization="bilby", + ) + first = kernel.sample(active, active[:4], RandomGenerator(seed=42)) + replay = kernel.sample(active, active[:4], RandomGenerator(seed=42)) + np.testing.assert_array_equal(first.particles, replay.particles) + np.testing.assert_allclose(first.particles.mean(axis=0), 0, atol=0.005) + np.testing.assert_allclose( + first.particles.std(axis=0), [0.1, 0.3], rtol=0.02 + ) + np.testing.assert_array_equal( + kernel.log_transition_ratio(active, first.particles, active[:4]), 0 + ) + + +def test_zero_noise_preserves_rng_consumption() -> None: + active = np.arange(12.0).reshape(6, 2) + first_rng = RandomGenerator(seed=12) + second_rng = RandomGenerator(seed=12) + first = DifferentialEvolutionKernel( + noise_scale=0.0, + ).sample(active, active, first_rng) + second = DifferentialEvolutionKernel(noise_scale=(0.0, 0.0)).sample( + active, active, second_rng + ) + np.testing.assert_array_equal(first.particles, second.particles) + np.testing.assert_array_equal( + first_rng.normal(size=10), second_rng.normal(size=10) + ) + + +def test_factory_exposes_noise_and_rejuvenates_collapsed_particles( + context_factory, +) -> None: + from sequince.steps.mutation import get_mutation_step + + mutation = get_mutation_step( + "differential-evolution", + target=lambda x, context: np.zeros(len(x)), + mutation_kwargs={"noise_scale": (0.1, 0.0)}, + ) + state, _, _ = mutation( + ParticleState(np.zeros((20, 2))), + context_factory(), + RandomGenerator(seed=8), + ) + assert np.all(state.particles[:, 0] != 0) + np.testing.assert_array_equal(state.particles[:, 1], 0) + + +def test_default_noise_rejuvenates_collapsed_population( + context_factory, +) -> None: + from sequince.steps.mutation import get_mutation_step + + particles = np.zeros((20, 2)) + kernel = DifferentialEvolutionKernel() + assert kernel.noise_scale == 1e-5 + proposal = kernel.sample(particles, particles, RandomGenerator(seed=8)) + assert np.all(proposal.particles != 0) + + mutation = get_mutation_step( + "differential-evolution", + target=lambda x, context: np.zeros(len(x)), + mutation_n_steps=1, + ) + state, _, _ = mutation( + ParticleState(particles), context_factory(), RandomGenerator(seed=8) + ) + assert np.all(state.particles != 0)