EpiAwareADTools

Automatic-differentiation safety machinery for the EpiAware modelling stack.

Why EpiAwareADTools?

EpiAwareADTools is the EpiAware org's shared home for AD-safety machinery and AD workarounds. It is deliberately framed as fixes we host while we try to fix things upstream: every entry is documented with the upstream package or issue where it ideally belongs, and each is deleted once that upstream fix lands.

Five families make up the current surface.

• The tape-strip pair primal and primal_distribution reduce an AD-wrapped scalar or distribution to its underlying primal, keeping a non-differentiable hyperparameter (an integration window, a clamp location) off the AD path on every backend.

• The AD-safe evaluation hooks (cdf_ad_safe, logcdf_ad_safe, ccdf_ad_safe, logccdf_ad_safe, and pdf_ad_safe) let a wrapper package overload CDF evaluation for its own component types, backed by analytic Gamma and Beta CDF derivatives that SpecialFunctions leaves unimplemented.

• nondifferentiable generalises the tape-strip discipline to an arbitrary user-supplied function: a deliberate, user-facing opt-out from differentiation, never a hidden default.

• logsumexp_stream accumulates a log-sum-exp over an unbounded discrete support, stopping only once a run of consecutive terms has left the total unchanged rather than at the first negligible term.

• The reparameterisation pair fixed_draw and ad_eltype pins a random draw as a constant realisation the parameters vary against, and resolves the type a parameter-dependent accumulator combined with such a draw is seeded at.

Per-backend behaviour for ForwardDiff, ReverseDiff, Enzyme, Mooncake, and ChainRulesCore is supplied by package extensions loaded when each backend is present.

Getting started

See the documentation for a full walkthrough.

using EpiAwareADTools, Distributions

# AD-safe Gamma CDF, differentiable in shape/scale on every supported backend.
cdf_ad_safe(Gamma(2.0, 1.0), 3.0)

# Strip an AD wrapper back to its primal value.
primal(3.0)

3.0

Related packages

• ConvolvedDistributions.jl, ComposedDistributions.jl, ModifiedDistributions.jl, LoweredDistributions.jl and CensoredDistributions.jl import these AD-safe hooks in their own source and overload them for their component types, so their densities differentiate on every supported backend.

• DistributionsInference.jl is the emerging fit-protocol layer across those packages, where the AD-safety this package provides is what makes gradient-based fitting work.

Where to learn more

• Want to get started running code? Check out the Getting started documentation.

• Want to understand the API? Check out our API reference.

• Want to contribute to EpiAwareADTools? Check the open issues and the Contributing section below.

• Want to see our code? Check out our GitHub Repository.

Getting help

For usage questions, ask on the Julia Discourse (the SciML or usage categories) or the epinowcast community forum, our home for epidemiological modelling questions. Please use GitHub issues for bug reports and feature requests only.

Part of the EpiAware ecosystem

EpiAwareADTools is part of EpiAware, a set of composable tools for infectious disease modelling. See the other packages in the ecosystem.

Contributing

We welcome contributions and new contributors! Please open an issue or pull request on GitHub. This package follows ColPrac and is formatted with Runic.

How to cite

If you use EpiAwareADTools in your work, please cite it. Citation metadata lives in CITATION.cff, which GitHub renders as a "Cite this repository" button on the repository page.

Code of conduct

Please note that the EpiAwareADTools project is released with a Contributor Code of Conduct. By contributing, you agree to abide by its terms.