Source code for sequince.steps.transformations
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Callable
from ..array_api import asarray, namespace
from ..utils.callables import call_with_optional_context
from .base import SMCStep
from .context import (
SMC_TRANSFORMATION_KEY,
StepContext,
)
if TYPE_CHECKING:
from orng import RandomGenerator
from ..posterior import TemperedPosterior
from ..state import ParticleState
from ..typing import Array
logger = logging.getLogger(__name__)
MetadataUpdateFn = Callable[
["ParticleState", StepContext, "RandomGenerator"],
dict[str, Any] | None,
]
[docs]
@dataclass
class SMCTransformation:
"""Transformation between successive SMC targets."""
[docs]
def forward(
self,
state: "ParticleState",
context: StepContext,
rng: "RandomGenerator",
) -> "ParticleState":
return state
[docs]
def forward_with_context(
self,
state: "ParticleState",
context: StepContext,
rng: "RandomGenerator",
) -> tuple["ParticleState", StepContext]:
"""Apply the transformation and return its context explicitly."""
return self.forward(state, context, rng), context
[docs]
def log_auxiliary_density(
self,
particles: "Array",
context: StepContext,
) -> "Array | float":
xp = namespace(particles)
arr = asarray(particles, xp=xp)
return xp.zeros(arr.shape[0], dtype=arr.dtype)
[docs]
def log_jacobian(
self,
particles: "Array",
context: StepContext,
) -> "Array | float":
xp = namespace(particles)
arr = asarray(particles, xp=xp)
return xp.zeros(arr.shape[0], dtype=arr.dtype)
[docs]
def make_log_proposal(
self,
base_log_proposal: Callable[..., "Array | float"],
context: StepContext,
) -> Callable[..., "Array | float"]:
"""Return a log proposal function corresponding to the new proposal
distribution after applying this transformation.
Parameters
----------
base_log_proposal
The log proposal function before applying the transformation.
context
The step context, which may contain metadata needed to define the
new log proposal.
"""
def log_proposal(
particles: "Array", inner_context: StepContext
) -> "Array | float":
inverse_particles = self.inverse(particles, inner_context)
base_logp = call_with_optional_context(
base_log_proposal, (inverse_particles,), inner_context
)
log_aux = self.log_auxiliary_density(particles, inner_context)
log_jacobian = self.log_jacobian(particles, inner_context)
return base_logp + log_aux - log_jacobian
return log_proposal
[docs]
@dataclass
class MetadataTransformation(SMCTransformation):
"""Same-space transformation that only updates metadata."""
state_metadata_fn: MetadataUpdateFn | None = None
context_metadata_fn: MetadataUpdateFn | None = None
[docs]
def forward_with_context(
self,
state: "ParticleState",
context: StepContext,
rng: "RandomGenerator",
) -> tuple["ParticleState", StepContext]:
state_metadata = dict(state.metadata)
if self.state_metadata_fn is not None:
state_updates = self.state_metadata_fn(state, context, rng) or {}
state_metadata.update(state_updates)
if self.context_metadata_fn is not None:
context_updates = (
self.context_metadata_fn(state, context, rng) or {}
)
context = context.with_metadata(**context_updates)
return state.with_updates(metadata=state_metadata), context
[docs]
class TransformationStep(SMCStep):
"""Apply a general SMC transformation and expose it to later steps."""
def __init__(
self,
transformation: SMCTransformation,
*,
posterior: "TemperedPosterior | None" = None,
transformation_key: str = SMC_TRANSFORMATION_KEY,
):
self._transformation = transformation
self._posterior = posterior
self._transformation_key = transformation_key
[docs]
def apply(
self,
state: "ParticleState",
context: StepContext,
rng: "RandomGenerator",
) -> tuple["ParticleState", StepContext, dict[str, Any]]:
transformed_state, context = self._transformation.forward_with_context(
state, context, rng
)
context_updates: dict[str, Any] = {
self._transformation_key: self._transformation
}
if self._posterior is not None:
override_key = self._posterior.log_proposal_override_key
base_log_proposal = self._posterior.resolve_log_proposal(
context.metadata.get(override_key)
)
if base_log_proposal is not None:
_transformed_log_proposal = (
self._transformation.make_log_proposal(
base_log_proposal, context
)
)
context_updates[override_key] = _transformed_log_proposal
context = context.with_metadata(**context_updates)
return transformed_state, context, context_updates
[docs]
class ExpandDimensions(TransformationStep):
"""Apply an SMC transformation that increases the particle dimension."""
def __init__(
self,
transformation: SMCTransformation,
*,
expected_block_size: int | None = None,
metadata_key: str = "dimensions",
added_dimensions_key: str = "added_dimensions",
dimension_delta_key: str = "last_dimension_delta",
posterior: "TemperedPosterior | None" = None,
transformation_key: str = SMC_TRANSFORMATION_KEY,
):
super().__init__(
transformation,
posterior=posterior,
transformation_key=transformation_key,
)
self._expected_block_size = expected_block_size
self._metadata_key = metadata_key
self._added_dimensions_key = added_dimensions_key
self._dimension_delta_key = dimension_delta_key
[docs]
def apply(
self,
state: "ParticleState",
context: StepContext,
rng: "RandomGenerator",
) -> tuple["ParticleState", StepContext, dict[str, Any]]:
xp = namespace(state.particles)
current_particles = asarray(state.particles, xp=xp)
if current_particles.ndim == 1:
current_particles = xp.reshape(
current_particles, (current_particles.shape[0], 1)
)
expanded_state, context, output = super().apply(state, context, rng)
expanded_particles = asarray(expanded_state.particles, xp=xp)
if expanded_particles.ndim == 1:
expanded_particles = xp.reshape(
expanded_particles, (expanded_particles.shape[0], 1)
)
if expanded_particles.shape[0] != current_particles.shape[0]:
raise ValueError(
"Transformation must preserve the number of particles."
)
delta = expanded_particles.shape[1] - current_particles.shape[1]
if delta < 0:
raise ValueError(
"ExpandDimensions requires a transformation that increases "
"or preserves dimensionality."
)
if (
self._expected_block_size is not None
and delta != self._expected_block_size
):
raise ValueError(
"Transformation changed dimensionality by "
f"{delta} (expected {self._expected_block_size})."
)
total_dims = expanded_particles.shape[1]
state_metadata = dict(expanded_state.metadata)
state_metadata[self._metadata_key] = total_dims
state_metadata[self._added_dimensions_key] = state_metadata.get(
self._added_dimensions_key, 0
) + max(delta, 0)
state_metadata[self._dimension_delta_key] = delta
context_updates = {
self._metadata_key: total_dims,
self._added_dimensions_key: context.metadata.get(
self._added_dimensions_key, 0
)
+ max(delta, 0),
self._dimension_delta_key: delta,
}
context = context.with_metadata(**context_updates)
output = {**output, **context_updates}
logger.info(
"ExpandDimensions changed dimensionality by %d (total=%d).",
delta,
total_dims,
)
return (
expanded_state.with_updates(metadata=state_metadata),
context,
output,
)