Skip to content

Feature/truncated gaussian#1150

Merged
Jammy2211 merged 48 commits intofeature/jax_wrapperfrom
feature/truncated_gaussian
Jun 24, 2025
Merged

Feature/truncated gaussian#1150
Jammy2211 merged 48 commits intofeature/jax_wrapperfrom
feature/truncated_gaussian

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

PR Description: Implementation of a Truncated Gaussian Prior with Hard Limits

Previously, limits on Gaussian priors (and other priors) were implemented as soft bounds defined by lower_limit and upper_limit attributes. When sampling from these priors, it was possible for drawn parameter values to fall outside the specified limits. In such cases, autofit performed a runtime check, and if the sampled value violated the limits, a FitException was raised, triggering a resampling of the model. This approach introduced inefficiencies due to potentially many rejected samples and less robust sampling behavior near the bounds.


Changes Introduced by This PR

  • New TruncatedGaussianPrior class:
    This class implements a Gaussian prior truncated between strict lower and upper bounds. Instead of sampling freely and then rejecting values outside the limits, this prior explicitly restricts the parameter space, normalizing the distribution within the truncated interval. This results in:

    • Improved sampling efficiency: Samples are guaranteed to lie within the specified bounds without the need for rejection.
    • Better numerical stability and accuracy: The probability density is properly normalized over the truncated interval, ensuring consistent behavior.
  • Support for prior passing:
    The truncated Gaussian prior is now used in scenarios where prior parameters are updated or passed between model fits. For example, when a model-fit returns posterior estimates, these can be used to create new priors centered on the estimated median with limits reflecting credible intervals. Using a truncated Gaussian ensures the new prior accurately reflects the constrained parameter space without introducing invalid samples.


Benefits and Impact

  • Sampling within strict parameter bounds becomes more robust and efficient.
  • Reduced risk of FitException errors due to out-of-bound samples.
  • Enables more reliable chaining of priors during sequential inference or hierarchical modeling.
  • Provides a cleaner, more interpretable API for users requiring truncated Gaussian priors.

JAX

The primary motivation of this PR is JAX, as the resampling API required annoying if loops which are no longer present as they are built into parameter mapping.

@Jammy2211 Jammy2211 changed the base branch from main to feature/jax_wrapper June 23, 2025 16:08
Copy link
Copy Markdown
Collaborator

@rhayes777 rhayes777 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice how well does the truncated version work?

database_aggregator,
output_directory,
):
print((output_directory / "database.info").read_text())
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

forgotten print statement

# x_, a_, b_, c_, z_ = variables("x, a, b, c, z")
# x = np.arange(10.0).reshape(5, 2)
# a = np.arange(2.0).reshape(2, 1)
# b = np.ones(1)
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lol core mp tests out of action

@Jammy2211 Jammy2211 merged commit 1d0efb0 into feature/jax_wrapper Jun 24, 2025
0 of 4 checks passed
@Jammy2211 Jammy2211 deleted the feature/truncated_gaussian branch June 24, 2025 13:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants