Merge 21df935 into 467b076

design_notes/interface_for_gibbs.md

@@ -0,0 +1,77 @@
+# Notes on Potential Future Interface for Gibbs Sampling Support
+
+## Background
+
+This document was written after [PR #144](https://github.com/TuringLang/AbstractMCMC.jl/pull/144) was closed.
+
+It was last updated on October 15, 2024. At that time:
+
+- _AbstractMCMC.jl_ was on version 5.5.0  
+- _Turing.jl_ was on version 0.34.1
+
+The goal is to document some of the considerations that went into the closed PR mentioned above.
+
+## Gibbs Sampling Considerations
+
+### Recomputing Log Densities for Parameter Groups
+
+Let's consider splitting the model parameters into several groups (assuming the grouping stays fixed between iterations). Each parameter group will have a corresponding sampler state (along with the sampler used for that group).
+
+In the general case, the log densities stored in the states will be incorrect at the time of sampling each group. This is because the values of the other two parameter groups can change from when the current log density was computed, as they get updated within the Gibbs sweep.
+
+### Current Approach: `recompute_logp!!`
+
+_Turing.jl_'s current solution, at the time of writing this, is the `recompute_logp!!` function (see [Tor's comment](https://github.com/TuringLang/AbstractMCMC.jl/issues/85#issuecomment-2061300622) and the [`Gibbs` PR](https://github.com/TuringLang/Turing.jl/pull/2099)).
+
+Here's an example implementation of this function for _AbstractHMC.jl_ ([permalink](https://github.com/TuringLang/Turing.jl/blob/24e68701b01695bffe69eda9e948e910c1ae2996/src/mcmc/abstractmcmc.jl#L77C1-L90C1)):
+
+```julia
+function recompute_logprob!!(
+    rng::Random.AbstractRNG,
+    model::AbstractMCMC.LogDensityModel,
+    sampler::AdvancedHMC.AbstractHMCSampler,
+    state::AdvancedHMC.HMCState,
+)
+    # Construct hamiltionian.
+    hamiltonian = AdvancedHMC.Hamiltonian(state.metric, model)
+    # Re-compute the log-probability and gradient.
+    return Accessors.@set state.transition.z = AdvancedHMC.phasepoint(
+        hamiltonian, state.transition.z.θ, state.transition.z.r
+    )
+end
+```
+
+### Alternative Approach Proposed in [PR #144](https://github.com/TuringLang/AbstractMCMC.jl/pull/144)
+
+The proposal is to separate `recompute_logp!!` into two functions:
+
+1. A function to compute the log density given the model and sampler state
+2. A function to set the computed log density in the sampler state
+
+There are a few considerations with this approach:
+
+- Computing the log density involves a `model`, which may not be defined by the sampler package in the general case. It's unclear if this interface is appropriate, as the model details might be needed to calculate the log density. However, in many situations, the `LogDensityProblems` interface (`LogDensityModel` in `AbstractMCMC`) could be sufficient.
+  - One interfacial consideration is that `LogDensityProblems.logdensity` expects a vector input. For our use case, we may want to reuse the log density stored in the state instead of recomputing it each time. This would require modifying `logdensity` to accept a sampler state and potentially a boolean flag to indicate whether to recompute the log density or not.
+- In some cases, samplers require more than just the log joint density. They may also need the log likelihood and log prior separately (see [this discussion](https://github.com/TuringLang/AbstractMCMC.jl/issues/112)).
+
+## Potential Path Forward
+
+A reasonable next step would be to explore an interface similar to `LogDensityProblems.logdensity`, but with the ability to compute both the log prior and log likelihood. It should also accept alternative inputs and keyword arguments.
+
+To complement this computation interface, we would need functions to `get` and `set` the log likelihood and log prior from/to the sampler states.
+
+For situations where model-specific details are required to compute the log density from a sampler state, the necessary abstractions are not yet clear. We will need to consider appropriate abstractions as such use cases emerge.
+
+## Additional Notes on a More Independent Gibbs Implementation
+
+### Regarding `AbstractPPL.condition`
+
+While the `condition` function is a promising idea for Gibbs sampling, it is not currently being utilized in _Turing.jl_'s implementation. Instead, _Turing.jl_ uses a `GibbsContext` for reasons outlined [here](https://github.com/TuringLang/Turing.jl/blob/3c91eec43176d26048b810aae0f6f2fac0686cfa/src/experimental/gibbs.jl#L1-L12). Additionally, _JuliaBUGS_ requires caching the Markov blanket when calling `condition`, which means the proposed `Gibbs` implementation in the PR above would not be fully compatible.
+
+### Samplers Should Not Manage Variable Names
+
+To make `AbstractMCMC.Gibbs` more independent and flexible, it should manage a mapping of `range → sampler` rather than `variable name → sampler`. This means it would maintain a vector of parameter values internally. The responsibility of managing both the variable names and any necessary transformations should be handled by a higher-level interface such as `AbstractPPL` or `DynamicPPL`.
+
+By separating these concerns, `AbstractMCMC.Gibbs` can focus on the core Gibbs sampling logic while the PPL interface handles the specifics of variable naming and transformations. This modular approach allows for greater flexibility and easier integration with different PPL frameworks.
+
+However, the issue arises when we have transdimensional parameters. In such cases, the parameter space can change during sampling, making it challenging to maintain a fixed mapping between ranges and samplers.

design_notes/logdensity_interface.md

@@ -0,0 +1,202 @@
+# Proposal for a New LogDensity Function Interface
+
+## Introduction
+
+The goal is to design a flexible and user-friendly interface for log density functions that can handle various model operations, especially in higher-order contexts such as Gibbs sampling. This interface should facilitate:
+
+- **Conditioning**: Incorporating observed data into the model.
+- **Fixing**: Fixing certain variables to specific values. (like `do` operator)
+- **Generated Quantities**: Computing additional expressions or functions based on the model parameters.
+- **Prediction**: Making predictions by fixing parameters and unconditioning on data.
+
+This proposal aims to redefine the interface from the user's perspective, focusing on ease of use and extensibility beyond the traditional probabilistic programming languages (PPLs).
+
+## Proposed Interface
+
+Below is a proposed interface with key functionalities and their implementations.
+
+### Core Functions
+
+#### Check if a Model is Parametric
+
+```julia
+# Check if a log density model is parametric
+function is_parametric(model::LogDensityModel) -> Bool
+    ...
+end
+```
+
+- **Description**: Determines if the model has a parameter space with a defined dimension.
+-
+
+#### Get the Dimension of a Parametric Model
+
+```julia
+# Get the dimension of the parameter space (only defined when is_parametric(model) is true)
+function dimension(model::LogDensityModel) -> Int
+    ...
+end
+```
+
+- **Description**: Returns the dimension of the parameter space for parametric models.
+
+### Log Density Computations
+
+#### Log-Likelihood
+
+```julia
+# Compute the log-likelihood given parameters
+function loglikelihood(model::LogDensityModel, params::Union{Vector, NamedTuple, Dict}) -> Float64
+    ...
+end
+```
+
+- **Description**: Computes the log-likelihood of the data given the model parameters.
+
+#### Log-Prior
+
+```julia
+# Compute the log-prior given parameters
+function logprior(model::LogDensityModel, params::Union{Vector, NamedTuple, Dict}) -> Float64
+    ...
+end
+```
+
+- **Description**: Computes the log-prior probability of the model parameters.
+
+#### Log-Joint
+
+```julia
+# Compute the log-joint density (log-likelihood + log-prior)
+function logjoint(model::LogDensityModel, params::Union{Vector, NamedTuple, Dict}) -> Float64
+    return loglikelihood(model, params) + logprior(model, params)
+end
+```
+
+- **Description**: Computes the total log density by summing the log-likelihood and log-prior.
+
+### Conditioning and Fixing Variables
+
+#### Conditioning a Model
+
+```julia
+# Condition the model on observed data
+function condition(model::LogDensityModel, data::NamedTuple) -> ConditionedModel
+    ...
+end
+```
+
+- **Description**: Incorporates observed data into the model, returning a `ConditionedModel`.
+
+#### Checking if a Model is Conditioned
+
+```julia
+# Check if a model is conditioned
+function is_conditioned(model::LogDensityModel) -> Bool
+    ...
+end
+```
+
+- **Description**: Checks whether the model has been conditioned on data.
+
+#### Fixing Variables in a Model
+
+```julia
+# Fix certain variables in the model
+function fix(model::LogDensityModel, variables::NamedTuple) -> FixedModel
+    ...
+end
+```
+
+- **Description**: Fixes specific variables in the model to given values, returning a `FixedModel`.
+
+#### Checking if a Model has Fixed Variables
+
+```julia
+# Check if a model has fixed variables
+function is_fixed(model::LogDensityModel) -> Bool
+    ...
+end
+```
+
+- **Description**: Determines if any variables in the model have been fixed.
+
+### Specialized Models
+
+#### Conditioned Model Methods
+
+```julia
+# Log-likelihood for a conditioned model
+function loglikelihood(model::ConditionedModel, params::Union{Vector, NamedTuple, Dict}) -> Float64
+    ...
+end
+
+# Log-prior for a conditioned model
+function logprior(model::ConditionedModel, params::Union{Vector, NamedTuple, Dict}) -> Float64
+    ...
+end
+
+# Log-joint for a conditioned model
+function logjoint(model::ConditionedModel, params::Union{Vector, NamedTuple, Dict}) -> Float64
+    return loglikelihood(model, params) + logprior(model, params)
+end
+```
+
+- **Description**: Overrides log density computations to account for the conditioned data.
+
+#### Fixed Model Methods
+
+```julia
+# Log-likelihood for a fixed model
+function loglikelihood(model::FixedModel, data::Union{Vector, NamedTuple, Dict}) -> Float64
+    ...
+end
+
+# Log-prior for a fixed model
+function logprior(model::FixedModel, data::Union{Vector, NamedTuple, Dict}) -> Float64
+    ...
+end
+
+# Log-joint for a fixed model
+function logjoint(model::FixedModel, data::Union{Vector, NamedTuple, Dict}) -> Float64
+    return loglikelihood(model, data) + logprior(model, data)
+end
+```
+
+- **Description**: Adjusts log density computations based on the fixed variables.
+
+### Additional Functionalities
+
+#### Generated Quantities
+
+```julia
+# Compute generated quantities after fixing parameters
+function generated_quantities(model::LogDensityModel, fixed_vars::NamedTuple) -> NamedTuple
+    ...
+end
+```
+
+- **Description**: Computes additional expressions or functions based on the fixed model parameters.
+
+#### Prediction
+
+```julia
+# Predict data based on fixed parameters
+function predict(model::LogDensityModel, params::Union{Vector, NamedTuple, Dict}) -> NamedTuple
+    ...
+end
+```
+
+- **Description**: Generates predictions by fixing the parameters and unconditioning the data.
+
+## Advantages of the Proposed Interface
+
+- **Flexibility**: Allows for advanced model operations like conditioning and fixing, essential for methods like Gibbs sampling.
+
+- **User-Centric Design**: Focuses on usability from the model user's perspective rather than the PPL implementation side.
+
+- **Consistency**: Maintains a uniform interface for both parametric and non-parametric models, simplifying the learning curve.
+
+## Usage Examples
+
+## Non-Parametric Models