Major API reorganization. (#46)

* Major API reorganization.

Addresses most of #45.

1. Remove validating wrapper types, simply return tuples.

2. Split `logdensity` into two methods, `logdensity` and `logdensity_and_gradient`.

3. Simplify code.

4. Introduce capabilities traits.

5. Fix #33 by making a distinct `dimension`.

6. Organize code into a single file, mostly.

7. Simplify tests greatly.

* improve coverage

* rewrite docs

* changelog for new release

CHANGELOG.md

@@ -1,5 +1,29 @@
 # Unreleased
 
+# 0.9.0
+
+**Major, breaking API changes.**
+
+1. Wrapper types `Value` and `ValueGradient` are removed. Interface functions return real numbers (`logdensity`), or a real number and a vector (`logdensity_and_gradient`). The calling convention of these types is simplified.
+
+2. Capabilities of log density objects are described by traits, see `capabilities`.
+
+3. `dimension` is now distinct from `TransformVariables.dimension`, as it is not really the same thing.
+
+4. Condition-based wrappers removed, as it was interfering with AD (mostly Zygote).
+
+5. Documentation significantly improved.
+
+6. Code organized into a single file (except for conditional loads), tests greatly simplified.
+
+# 0.8.3
+
+- update to work with new Zygote interface
+
+# 0.8.2
+
+- no code change, just a version bump to fix registoy problems
+
 # 0.8.1
 
 - minor package and version fixes

Project.toml

@@ -1,7 +1,7 @@
 name = "LogDensityProblems"
 uuid = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 authors = ["Tamas K. Papp <[email protected]>"]
-version = "0.8.3"
+version = "0.9.0"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

docs/Project.toml

@@ -4,4 +4,4 @@ Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
 [compat]
-Documenter = "~0.21"
+Documenter = "~0.23"

docs/make.jl

@@ -7,8 +7,7 @@ makedocs(
     clean = true,
     authors = "Tamás K. Papp",
     checkdocs = :export,
-    pages = Any["Home" => "index.md",
-                "Internals" => "internals.md"]
+    pages = Any["Documentation" => "index.md"]
 )
 
 deploydocs(

docs/src/index.md

@@ -1,42 +1,183 @@
 # Introduction
 
-This package provides the following functionality:
+This package serves two purposes:
 
-1. It defines the [`logdensity`](@ref) method with corresponding interface, which can be used by other packages that operate on (log) densities and need to evaluate the log densities or the gradients (eg [MCMC](https://en.wikipedia.org/wiki/Markov_chain_Monte_Carlo), [MAP](https://en.wikipedia.org/wiki/Maximum_a_posteriori_estimation), [ML](https://en.wikipedia.org/wiki/Maximum_likelihood_estimation) or similar methods).
+1. Introduce a common API for packages that operate on *log densities*, which for these purposes are black box ``\mathbb{R}^n \to \mathbb{R}`` mappings. Using the interface of introduced in this package, you can query ``n``, evaluate the log density and optionally its gradient, and determine if a particular object supports these methods using traits. **This usage is relevant primarily for package developers** who write generic algorithms that use (log) densities that correspond to posteriors and likelihoods, eg [MCMC](https://en.wikipedia.org/wiki/Markov_chain_Monte_Carlo), [MAP](https://en.wikipedia.org/wiki/Maximum_a_posteriori_estimation), [ML](https://en.wikipedia.org/wiki/Maximum_likelihood_estimation). An example is [DynamicHMC.jl](https://github.com/tpapp/DynamicHMC.jl). This is documented in [the API section](@ref log-density-api).
 
-2. It defines the [`ADgradient`](@ref) which makes objects that support `logdensity` to calculate log density *values* calculate log density *gradients* using various automatic differentiation packages.
+2. Make it easier for **users who want to perform inference** using the above methods (and packages) to
+    - *define their own log densities*, either taking a vector of real numbers as input, or extracting and
+      transforming parameters using the [TransformVariables.jl](https://github.com/tpapp/TransformVariables.jl)
+      package,
+    - *obtain gradients* of these log densities using one of the supported *automatic differentiation* packages
+      of Julia.
+This is documented in the next section, with a worked example.
 
-3. It defines the wrapper [`TransformedLogDensity`](@ref) using the [TransformVariables.jl](https://github.com/tpapp/TransformVariables.jl) package, allowing callables that take a set of parameters transformed from a flat vector of real numbers to support the `logdensity` interface.
+# Working with log density problems
 
-4. Various utility functions for debugging and testing log densities.
+Consider an inference problem where IID draws are obtained from a normal distribution,
+```math
+x_i \sim N(\mu, \sigma)
+```
+for ``i = 1, \dots, N``. It can be shown that the *log likelihood* conditional on ``\mu`` and ``\sigma`` is
+```math
+\ell = -N\log \sigma - \sum_{i = 1}^N \frac{(x-\mu)^2}{2\sigma^2} =
+-N\left( \log \sigma + \frac{S + (\bar{x} - \mu)^2}{2\sigma^2} \right)
+```
+where we have dropped constant terms, and defined the sufficient statistics
+```math
+\bar{x} = \frac{1}{N} \sum_{i = 1}^N x_i
+```
+and
+```math
+S = \frac{1}{N} \sum_{i = 1}^N (x_i - \bar{x})^2
+```
+
+Finally, we use priors
+```math
+\mu \sim N(0, 5), \sigma \sim N(0, 2)
+```
+which yield the log prior
+```math
+-\sigma^2/8 - \mu^2/50
+```
+which is added to the log likelihood to obtain the log posterior.
+
+It is useful to define a *callable* that implements this, taking some vector `x` as an input and calculating the summary statistics, then, when called with a `NamedTuple` containing the parameters, evaluating to the log posterior.
+
+```@example 1
+using Random; Random.seed!(1) # hide
+using Statistics, Parameters # imported for our implementation
+
+struct NormalPosterior{T} # contains the summary statistics
+    N::Int
+    x̄::T
+    S::T
+end
+
+# calculate summary statistics from a data vector
+function NormalPosterior(x::AbstractVector)
+    NormalPosterior(length(x), mean(x), var(x; corrected = false))
+end
+
+# define a callable that unpacks parameters, and evaluates the log likelihood
+function (problem::NormalPosterior)(θ)
+    @unpack μ, σ = θ
+    @unpack N, x̄, S = problem
+    loglikelihood = -N * (log(σ) + (S + abs2(μ - x̄)) / (2 * abs2(σ)))
+    logprior = - abs2(σ)/8 - abs2(μ)/50
+    loglikelihood + logprior
+end
+
+problem = NormalPosterior(randn(100))
+nothing # hide
+```
+
+Let's try out the posterior calculation:
+
+```@repl 1
+problem((μ = 0.0, σ = 1.0))
+```
+
+!!! note
+    Just evaluating your log density function like above is a great way to test and benchmark your implementation. See the “Performance Tips” section of the Julia manual for optimization advice.
+
+## Using the TransformVariables package
+
+In our example, we require ``\sigma > 0``, otherwise the problem is meaningless. However, many MCMC samplers prefer to operate on *unconstrained* spaces ``\mathbb{R}^n``. The TransformVariables package was written to transform unconstrained to constrained spaces, and help with the log Jacobian correction (more on that later). That package has detailed documentation, now we just define a transformation from a length 2 vector to a `NamedTuple` with fields `μ` (unconstrained) and `σ > 0`.
+
+```@repl 1
+using LogDensityProblems, TransformVariables
+ℓ = TransformedLogDensity(as((μ = asℝ, σ = asℝ₊)), problem)
+```
+
+Then we can query the dimension of this problem, and evaluate the log density:
+```@repl 1
+LogDensityProblems.dimension(ℓ)
+LogDensityProblems.logdensity(ℓ, zeros(2))
+```
 
-## Inference
+!!! note
+    Before running time-consuming algorithms like MCMC, it is advisable to test and benchmark your log density evaluations separately. The same applies to [`LogDensityProblems.logdensity_and_gradient`](@ref).
 
 ```@docs
-logdensity
-dimension
-LogDensityProblems.Value
-LogDensityProblems.ValueGradient
-LogDensityProblems.ValueGradientBuffer
+TransformedLogDensity
+```
+
+## Manual unpacking and transformation
+
+If you prefer to implement the transformation yourself, you just have to define the following three methods for your problem: declare that it can evaluate log densities (but not their gradient, hence the `0` order), allow the dimension of the problem to be queried, and then finally code the density calculation with the transformation. (Note that using [`TransformedLogDensity`](@ref) takes care of all of these for you, as shown above).
+
+```@example 1
+function LogDensityProblems.capabilities(::Type{<:NormalPosterior})
+    LogDensityProblems.LogDensityOrder{0}()
+end
+
+LogDensityProblems.dimension(::NormalPosterior) = 2
+
+function LogDensityProblems.logdensity(problem::NormalPosterior, x)
+    μ, logσ = x
+    σ = exp(logσ)
+    problem((μ = μ, σ = σ)) + logσ
+end
+nothing # hide
+```
+
+```@repl 1
+LogDensityProblems.logdensity(problem, zeros(2))
 ```
 
-## Gradient via automatic differentiation
+Here we use the exponential function to transform from ``\mathbb{R}`` to the positive reals, but this requires that we correct the log density with the *logarithm* of the Jacobian, which here happens to be ``\log(\sigma)``.
 
+## Automatic differentiation
+
+Using either definition, you can now transform to another object which is capable of evaluating the *gradient*, using automatic differentiation. The wrapper for this is
 ```@docs
 ADgradient
 ```
+Note that support for Zygote is experimental. At the moment, I would recommend that you use `Flux`.
 
-## Transformed problem definition
+Now observe that we can obtain gradients, too:
+```@repl 1
+import ForwardDiff
+∇ℓ = ADgradient(:ForwardDiff, ℓ)
+LogDensityProblems.capabilities(∇ℓ)
+LogDensityProblems.logdensity_and_gradient(∇ℓ, zeros(2))
+```
 
-```@docs
-TransformedLogDensity
+## Manually calculated derivatives
+
+If you prefer not to use automatic differentiation, you can wrap your own derivatives following the template
+```julia
+function LogDensityProblems.capabilities(::Type{<:NormalPosterior})
+    LogDensityProblems.LogDensityOrder{1}() # can do gradient
+end
+
+LogDensityProblems.dimension(::NormalPosterior) = 2 # for this problem
+
+function LogDensityProblems.logdensity_and_gradient(problem::NormalPosterior, x)
+    logdens = ...
+    grad = ...
+    logdens, grad
+end
 ```
 
-## Benchmarking, diagnostics, and utilities
+# Various utilities
+
+You may find these utilities useful for debugging and optimization.
 
 ```@docs
 LogDensityProblems.stresstest
 LogDensityProblems.benchmark_ForwardDiff_chunks
-LogDensityProblems.@iffinite
-LogDensityRejectErrors
-```
+```
+
+# [Log densities API](@id log-density-api)
+
+Use the functions below for evaluating gradients and querying their dimension and other information. These symbols are not exported, as they are mostly used by package developers and in any case would need to be `import`ed or qualified to add methods to.
+
+```@docs
+LogDensityProblems.capabilities
+LogDensityProblems.LogDensityOrder
+LogDensityProblems.dimension
+LogDensityProblems.logdensity
+LogDensityProblems.logdensity_and_gradient
+```

src/AD_Flux.jl

@@ -1,3 +1,7 @@
+#####
+##### Gradient AD implementation using Flux
+#####
+
 import .Flux
 
 struct FluxGradientLogDensity{L} <: ADGradientWrapper
@@ -13,14 +17,9 @@ ADgradient(::Val{:Flux}, ℓ) = FluxGradientLogDensity(ℓ)
 
 Base.show(io::IO, ∇ℓ::FluxGradientLogDensity) = print(io, "Flux AD wrapper for ", ∇ℓ.ℓ)
 
-function logdensity(::Type{ValueGradient}, ∇ℓ::FluxGradientLogDensity, x::AbstractVector)
+function logdensity_and_gradient(∇ℓ::FluxGradientLogDensity, x::AbstractVector)
     @unpack ℓ = ∇ℓ
-    y, back = Flux.Tracker.forward(_logdensity_closure(ℓ), x)
+    y, back = Flux.Tracker.forward(x -> logdensity(ℓ, x), x)
     yval = Flux.Tracker.data(y)
-    ValueGradient(yval, isfinite(yval) ? first(Flux.Tracker.data.(back(1))) : similar(x))
-end
-
-function logdensity(::ValueGradientBuffer, ∇ℓ::FluxGradientLogDensity, x::AbstractVector)
-    # NOTE this implementation ignores the buffer
-    logdensity(ValueGradient, ∇ℓ, x)
+    yval, first(Flux.Tracker.data.(back(1)))
 end