Skip to content

Fitting API

Top-level fitting helpers are documented under Top-level API. This page documents objective, optimizer, and candidate-selection configuration classes from gmm_divergence.fitting.

gmm_divergence.fitting

Public fitting API.

Objective Configuration

ForwardKL dataclass

ForwardKL(sampling: SampleSpec = Draw())

Forward KL fitting objective configuration.

Fits the candidate-mixture weights by minimizing

\[ D_{\mathrm{KL}}\!\left(p \,\|\, q_{\mathbf{w}}\right) = \mathbb{E}_{X \sim p} \left[ \log p(X) - \log q_{\mathbf{w}}(X) \right], \]

where \(p\) is the reference distribution and \(q_{\mathbf{w}}\) is the weighted combination of the candidate mixtures. Since \(\log p(X)\) does not depend on \(\mathbf{w}\), the implemented objective minimizes the negative expected log-density of \(q_{\mathbf{w}}\) under samples from \(p\).

sampling class-attribute instance-attribute

sampling: SampleSpec = field(default_factory=Draw)

Sampling specification for the expectation under p.

Use sampling.Draw(...) to draw fresh samples, sampling.Samples(...) for precomputed reference samples, or sampling.Stratified(...) when p is a Gaussian-family distribution and fixed per-component counts are desired.

ReverseKL dataclass

ReverseKL(
    p_sampling: SampleSpec = Draw(), q_sampling: BatchSampleSpec = Draw()
)

Reverse KL fitting objective configuration.

Fits the candidate-mixture weights by minimizing

\[ D_{\mathrm{KL}}\!\left(q_{\mathbf{w}} \,\|\, p\right) = \mathbb{E}_{X \sim q_{\mathbf{w}}} \left[ \log q_{\mathbf{w}}(X) - \log p(X) \right]. \]

The implementation evaluates a fixed-sample estimator using samples from the reference distribution \(p\) for diagnostics and samples from each candidate mixture \(q_i\) for the reverse objective.

p_sampling class-attribute instance-attribute

p_sampling: SampleSpec = field(default_factory=Draw)

Sampling specification for diagnostics under p.

q_sampling class-attribute instance-attribute

q_sampling: BatchSampleSpec = field(default_factory=Draw)

Sampling specification for fixed batches from each q_i.

Use sampling.Draw(...) to draw one batch per candidate distribution, sampling.Stratified(...) for fixed per-component counts, or sampling.SampleBatches(...) to provide those batches directly.

BidirectionalKL dataclass

BidirectionalKL(
    p_sampling: SampleSpec = Draw(),
    q_sampling: BatchSampleSpec = Draw(),
    alpha: float = 0.5,
)

Bidirectional KL fitting objective configuration.

Fits the candidate-mixture weights by minimizing a weighted combination of forward and reverse KL objectives:

\[ \alpha\,D_{\mathrm{KL}}\!\left(p \,\|\, q_{\mathbf{w}}\right) + (1-\alpha)\,D_{\mathrm{KL}}\!\left(q_{\mathbf{w}} \,\|\, p\right), \qquad \alpha \in [0, 1]. \]

Values of \(\alpha\) closer to one emphasize coverage of \(p\) by \(q_{\mathbf{w}}\), while values closer to zero emphasize the reverse-KL objective.

alpha class-attribute instance-attribute

alpha: float = 0.5

Weight assigned to the forward KL term.

p_sampling class-attribute instance-attribute

p_sampling: SampleSpec = field(default_factory=Draw)

Sampling specification for the forward term under p.

q_sampling class-attribute instance-attribute

q_sampling: BatchSampleSpec = field(default_factory=Draw)

Sampling specification for the reverse term under each q_i.

Use sampling.Draw(...) to draw one batch per candidate distribution, sampling.Stratified(...) for fixed per-component counts, or sampling.SampleBatches(...) to provide those batches directly.

JensenShannon dataclass

JensenShannon(
    p_sampling: SampleSpec = Draw(), q_sampling: BatchSampleSpec = Draw()
)

Jensen-Shannon fitting objective configuration.

Fits the candidate-mixture weights by minimizing

\[ D_{\mathrm{JS}}\!\left(p, q_{\mathbf{w}}\right) = \frac{1}{2}D_{\mathrm{KL}}\!\left(p \,\|\, m_{\mathbf{w}}\right) + \frac{1}{2}D_{\mathrm{KL}}\!\left(q_{\mathbf{w}} \,\|\, m_{\mathbf{w}}\right), \]

where

\[ m_{\mathbf{w}} = \frac{1}{2}p + \frac{1}{2}q_{\mathbf{w}}. \]

This objective is symmetric and bounded, while still using fixed samples from p and each candidate mixture q_i during optimization.

p_sampling class-attribute instance-attribute

p_sampling: SampleSpec = field(default_factory=Draw)

Sampling specification for the term under p.

q_sampling class-attribute instance-attribute

q_sampling: BatchSampleSpec = field(default_factory=Draw)

Sampling specification for fixed batches from each q_i.

Use sampling.Draw(...) to draw one batch per candidate distribution, sampling.Stratified(...) for fixed per-component counts, or sampling.SampleBatches(...) to provide those batches directly.

MomentMatching dataclass

MomentMatching(fit_second_moments: bool = False)

Moment-matching fitting objective configuration.

Fits weights by matching moments of \(q_{\mathbf{w}}\) to moments of the reference distribution \(p\) instead of optimizing a sampled KL estimate. The first-moment objective compares mixture means,

\[ \left\|\mu_p - \mu_{q_{\mathbf{w}}}\right\|_2^2, \]

and can optionally include second-moment information as well.

fit_second_moments class-attribute instance-attribute

fit_second_moments: bool = False

Whether to include covariance information in the objective.

Optimizers

SoftmaxLBFGSB dataclass

SoftmaxLBFGSB(tol: float = 1e-08, max_iterations: int = 1000)

L-BFGS-B optimizer over unconstrained softmax logits.

The optimizer represents mixture weights through logits \(\mathbf{z} \in \mathbb{R}^N\) and maps them to simplex weights with

\[ w_i = \frac{\exp(z_i)}{\sum_{j=1}^{N}\exp(z_j)}. \]

This avoids explicit simplex constraints in the optimizer while still ensuring that all fitted weights are non-negative and sum to one.

max_iterations class-attribute instance-attribute

max_iterations: int = 1000

Maximum number of optimizer iterations.

tol class-attribute instance-attribute

tol: float = 1e-08

Optimizer convergence tolerance.

SimplexSLSQP dataclass

SimplexSLSQP(
    tol: float = 1e-08, max_iterations: int = 1000, min_weight: float = 1e-12
)

SLSQP optimizer over simplex-constrained weights.

The optimizer works directly with the mixture weights \(\mathbf{w} \in \Delta_N\), where

\[ \Delta_N = \left\{ \mathbf{w} \in \mathbb{R}^N : w_i \ge 0,\ \sum_{i=1}^{N} w_i = 1 \right\}. \]

This keeps the optimization variables interpretable but requires the optimizer to enforce the simplex equality and bound constraints.

max_iterations class-attribute instance-attribute

max_iterations: int = 1000

Maximum number of optimizer iterations.

min_weight class-attribute instance-attribute

min_weight: float = 1e-12

Minimum weight threshold for the optimizer.

Ideally, this would be zero, but SLSQP can fail when weights are exactly zero. Setting it to small positive values avoids numerical issues while still allowing the optimizer to effectively prune components with negligible weights.

tol class-attribute instance-attribute

tol: float = 1e-08

Optimizer convergence tolerance.

Candidate Selectors

score_candidates

score_candidates(
    p: GaussianLike,
    q_i: Sequence[GaussianLike],
    /,
    *,
    direction: Literal["forward", "reverse", "bidirectional"] = "forward",
    alpha: float = 0.5,
    method: KLMethod | None = None,
    prefer_closed_form: bool = True,
) -> FloatArray

Return KL-based scores for candidate distributions.

Source code in src/gmm_divergence/fitting/_selector.py
def score_candidates(
    p: GaussianLike,
    q_i: Sequence[GaussianLike],
    /,
    *,
    direction: Literal["forward", "reverse", "bidirectional"] = "forward",
    alpha: float = 0.5,
    method: KLMethod | None = None,
    prefer_closed_form: bool = True,
) -> FloatArray:
    """Return KL-based scores for candidate distributions."""
    _validate_candidate_sequence(q_i)
    _validate_kl_selector_base(direction=direction, alpha=alpha)
    return _compute_kl_values(
        p,
        q_i,
        direction=direction,
        alpha=alpha,
        kl_method=_DEFAULT_SCORING_METHOD if method is None else method,
        prefer_closed_form=prefer_closed_form,
    )

rank_candidates

rank_candidates(
    p: GaussianLike,
    q_i: Sequence[GaussianLike],
    /,
    *,
    direction: Literal["forward", "reverse", "bidirectional"] = "forward",
    alpha: float = 0.5,
    method: KLMethod | None = None,
    prefer_closed_form: bool = True,
    limit: int | None = None,
) -> list[tuple[int, float]]

Return candidate indices and scores sorted from best to worst.

Source code in src/gmm_divergence/fitting/_selector.py
def rank_candidates(
    p: GaussianLike,
    q_i: Sequence[GaussianLike],
    /,
    *,
    direction: Literal["forward", "reverse", "bidirectional"] = "forward",
    alpha: float = 0.5,
    method: KLMethod | None = None,
    prefer_closed_form: bool = True,
    limit: int | None = None,
) -> list[tuple[int, float]]:
    """Return candidate indices and scores sorted from best to worst."""
    scores = score_candidates(
        p,
        q_i,
        direction=direction,
        alpha=alpha,
        method=_DEFAULT_SCORING_METHOD if method is None else method,
        prefer_closed_form=prefer_closed_form,
    )
    if limit is not None and (isinstance(limit, bool) or limit <= 0):
        msg = f"limit must be a positive integer when provided, got {limit}."
        raise ValueError(msg)
    ranked_indices = np.argsort(scores)
    if limit is None:
        return [(int(index), float(scores[index])) for index in ranked_indices]
    return [(int(index), float(scores[index])) for index in ranked_indices[:limit]]

TopKSelector dataclass

TopKSelector(
    k: int,
    *,
    direction: Literal["forward", "reverse", "bidirectional"] = "forward",
    alpha: float = 0.5,
    kl_method: KLMethod = MonteCarlo(sampling=Draw(rng=0)),
)

Bases: _KLSelectorBase

ThresholdSelector dataclass

ThresholdSelector(
    threshold: float,
    *,
    direction: Literal["forward", "reverse", "bidirectional"] = "forward",
    alpha: float = 0.5,
    kl_method: KLMethod = MonteCarlo(sampling=Draw(rng=0)),
)

Bases: _KLSelectorBase

Select candidates whose KL score is at or below a fixed threshold.

ToleranceSelector dataclass

ToleranceSelector(
    delta: float,
    mode: Literal["absolute", "relative"] = "absolute",
    *,
    direction: Literal["forward", "reverse", "bidirectional"] = "forward",
    alpha: float = 0.5,
    kl_method: KLMethod = MonteCarlo(sampling=Draw(rng=0)),
)

Bases: _KLSelectorBase

Select candidates within a tolerance of the best KL score.

In absolute mode, candidates with KL <= min(KL) + delta are kept. In relative mode, candidates with KL <= min(KL) + delta * abs(min(KL)) are kept, so delta=0.5 means within 50% of the best score when the best score is positive.

QuantileSelector dataclass

QuantileSelector(
    quantile: float,
    *,
    direction: Literal["forward", "reverse", "bidirectional"] = "forward",
    alpha: float = 0.5,
    kl_method: KLMethod = MonteCarlo(sampling=Draw(rng=0)),
)

Bases: _KLSelectorBase