gmm-divergence 0.0.1a1__tar.gz

gmm_divergence-0.0.1a1/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

gmm_divergence-0.0.1a1/PKG-INFO

@@ -0,0 +1,47 @@
+Metadata-Version: 2.4
+Name: gmm-divergence
+Version: 0.0.1a1
+Summary: Utilities for estimating divergences between Gaussian mixture models.
+Keywords: divergence,gaussian-mixture-model,gmm,kl-divergence,numpy
+Author: David Axelsson
+Author-email: David Axelsson <david.axelsson@liu.se>
+License-Expression: Apache-2.0
+License-File: LICENSE
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering
+Classifier: Typing :: Typed
+Requires-Dist: numpy>=2.2.6
+Requires-Dist: typing-extensions>=4.13.2
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# gmm-divergence
+
+Utilities for estimating divergences between Gaussian mixture models.
+
+This package is under active development. APIs, estimators, and numerical behavior may change between early releases.
+
+## Current Scope
+
+The package currently includes:
+
+- A typed Gaussian mixture representation
+- Density and log-density evaluation for full and diagonal covariance mixtures
+- Sampling from Gaussian mixtures
+- KL divergence estimators based on Monte Carlo sampling and unscented sigma points
+
+## Installation
+
+This project is not yet intended for stable production use. For development, install it from a local checkout:
+
+```bash
+python -m pip install -e .
+```

gmm_divergence-0.0.1a1/README.md

@@ -0,0 +1,22 @@
+# gmm-divergence
+
+Utilities for estimating divergences between Gaussian mixture models.
+
+This package is under active development. APIs, estimators, and numerical behavior may change between early releases.
+
+## Current Scope
+
+The package currently includes:
+
+- A typed Gaussian mixture representation
+- Density and log-density evaluation for full and diagonal covariance mixtures
+- Sampling from Gaussian mixtures
+- KL divergence estimators based on Monte Carlo sampling and unscented sigma points
+
+## Installation
+
+This project is not yet intended for stable production use. For development, install it from a local checkout:
+
+```bash
+python -m pip install -e .
+```

gmm_divergence-0.0.1a1/pyproject.toml

@@ -0,0 +1,82 @@
+[project]
+name = "gmm-divergence"
+version = "0.0.1a1"
+description = "Utilities for estimating divergences between Gaussian mixture models."
+readme = "README.md"
+authors = [
+    { name = "David Axelsson", email = "david.axelsson@liu.se" }
+]
+license = "Apache-2.0"
+license-files = ["LICENSE"]
+requires-python = ">=3.10"
+keywords = ["divergence", "gaussian-mixture-model", "gmm", "kl-divergence", "numpy"]
+classifiers = [
+    "Development Status :: 2 - Pre-Alpha",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Topic :: Scientific/Engineering",
+    "Typing :: Typed",
+]
+dependencies = [
+    "numpy>=2.2.6",
+    "typing-extensions>=4.13.2",
+]
+
+[dependency-groups]
+build = [
+    "build>=1.5.0",
+    "twine>=6.2.0",
+]
+test = [
+    "pytest>=8.3.5",
+    "pytest-benchmark>=5.2.3",
+]
+
+[build-system]
+requires = ["uv_build>=0.11.16,<0.12.0"]
+build-backend = "uv_build"
+
+[tool.ruff]
+line-length = 100
+preview = true
+
+[tool.ruff.lint]
+select = ["ALL"]
+ignore = [
+    "D203",
+    "D213",
+    "COM812",
+    "PLR2004",
+    "DOC201",
+    "DOC501",
+    "TRY003",
+    "D100",
+    "PLR0904",
+    "TD002",
+    "TD003",
+    "PLR0913",
+    "PLR0917",
+    "PLR0914",
+    "DOC402",
+    "CPY001",
+    "D104",
+    "TD005",
+    "PLC2801",
+    "D101",
+]
+pydocstyle = { convention = "numpy" }
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*.py" = ["D", "INP001", "S101", "PLC2701"]
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

gmm_divergence-0.0.1a1/src/gmm_divergence/__init__.py

@@ -0,0 +1,8 @@
+from importlib.metadata import version
+
+from gmm_divergence.divergence import kl_divergence
+from gmm_divergence.gmm.model import GaussianMixture
+from gmm_divergence.results import DivergenceResult
+
+__all__ = ["DivergenceResult", "GaussianMixture", "kl_divergence"]
+__version__: str = version("gmm-divergence")

gmm_divergence-0.0.1a1/src/gmm_divergence/divergence.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, overload
+
+from gmm_divergence.methods.gaussian_approx import kl_gaussian_approximation
+from gmm_divergence.methods.monte_carlo import kl_monte_carlo
+from gmm_divergence.methods.unscented import kl_unscented
+
+if TYPE_CHECKING:
+    import numpy as np
+    import numpy.typing as npt
+
+    from gmm_divergence.gmm.model import GaussianMixture, PrecisionT
+    from gmm_divergence.results import DivergenceResult
+
+KLMethod = Literal["monte_carlo", "unscented", "gaussian_approximation"]
+
+
+@overload
+def kl_divergence(
+    p: GaussianMixture[PrecisionT],
+    q: GaussianMixture[PrecisionT],
+    /,
+    *,
+    method: Literal["gaussian_approximation"] = ...,
+    approximation: Literal["nearest", "moment_matching"] = "moment_matching",
+) -> DivergenceResult: ...
+
+
+@overload
+def kl_divergence(
+    p: GaussianMixture[PrecisionT],
+    q: GaussianMixture[PrecisionT],
+    /,
+    *,
+    method: Literal["unscented"] = ...,
+) -> DivergenceResult: ...
+
+
+@overload
+def kl_divergence(
+    p: GaussianMixture[PrecisionT],
+    q: GaussianMixture[PrecisionT],
+    /,
+    *,
+    method: Literal["monte_carlo"] = "monte_carlo",
+    num_samples: int = 10_000,
+    samples: npt.ArrayLike | None = None,
+    rng: np.random.Generator | int | None = None,
+) -> DivergenceResult: ...
+
+
+def kl_divergence(
+    p: GaussianMixture[PrecisionT],
+    q: GaussianMixture[PrecisionT],
+    /,
+    *,
+    method: KLMethod = "monte_carlo",
+    num_samples: int = 10_000,
+    samples: npt.ArrayLike | None = None,
+    rng: np.random.Generator | int | None = None,
+    approximation: Literal["nearest", "moment_matching"] = "moment_matching",
+) -> DivergenceResult:
+    """Compute KL divergence between two Gaussian mixtures.
+
+    Parameters
+    ----------
+    p, q : GaussianMixture
+        The two Gaussian mixtures to compare. The KL divergence is computed as
+        KL(p || q).
+    method : KLMethod
+        The method to use for estimating the KL divergence. "monte_carlo" uses
+        Monte Carlo sampling, while "unscented" uses unscented sigma points.
+    num_samples : int
+        The number of samples to use for the Monte Carlo method.
+    samples : array-like, optional
+        Pre-generated samples from the first mixture `p` to use for the Monte
+        Carlo method.
+    rng : np.random.Generator, int, or None
+        Random number generator or seed for the Monte Carlo method.
+
+    Returns
+    -------
+    DivergenceResult
+        An object containing the estimated KL divergence and metadata about the
+        estimation method. The value for the KL divergence estimate is stored in
+        the `value` attribute of the returned object.
+    """
+    match method:
+        case "monte_carlo":
+            return kl_monte_carlo(
+                p,
+                q,
+                num_samples=num_samples,
+                samples=samples,
+                rng=rng,
+            )
+        case "unscented":
+            return kl_unscented(p, q)
+        case "gaussian_approximation":
+            return kl_gaussian_approximation(
+                p,
+                q,
+                approximation=approximation,
+            )

gmm_divergence-0.0.1a1/src/gmm_divergence/gmm/__init__.py

@@ -0,0 +1,19 @@
+from gmm_divergence.gmm.density import (
+    gmm_logpdf_diag,
+    gmm_logpdf_full,
+    gmm_pdf_diag,
+    gmm_pdf_full,
+)
+from gmm_divergence.gmm.model import CovarianceType, GaussianMixture, PrecisionT
+from gmm_divergence.gmm.sampling import sample_gmm
+
+__all__ = [
+    "CovarianceType",
+    "GaussianMixture",
+    "PrecisionT",
+    "gmm_logpdf_diag",
+    "gmm_logpdf_full",
+    "gmm_pdf_diag",
+    "gmm_pdf_full",
+    "sample_gmm",
+]

gmm_divergence-0.0.1a1/src/gmm_divergence/gmm/density.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import numpy.typing as npt
+
+if TYPE_CHECKING:
+    from gmm_divergence.gmm.model import GaussianMixture, PrecisionT
+
+
+def _logsumexp(a: npt.NDArray[PrecisionT], axis: int = -1) -> npt.NDArray[PrecisionT]:
+    a_max = np.max(a, axis=axis, keepdims=True)
+    return np.squeeze(
+        a_max + np.log(np.sum(np.exp(a - a_max), axis=axis, keepdims=True)),
+        axis=axis,
+    )
+
+
+def gmm_logpdf_full(
+    x: npt.ArrayLike,
+    gmm: GaussianMixture[PrecisionT],
+) -> npt.NDArray[PrecisionT]:
+    """Evaluate the log-density of a full-covariance Gaussian mixture."""
+    if gmm.covariance_type != "full":
+        msg = f"Expected full covariance type, got {gmm.covariance_type}"
+        raise ValueError(msg)
+
+    x = np.asarray(x, dtype=gmm.dtype)
+    weights = gmm.weights / np.sum(gmm.weights)
+
+    n_samples, n_features = x.shape
+    n_components = weights.shape[0]
+
+    log_weights = np.log(weights)
+    chol = gmm.chol()
+    log_probs = np.empty((n_samples, n_components), dtype=gmm.dtype)
+    constant = n_features * np.log(2.0 * np.pi)
+
+    for k in range(n_components):
+        diff = x - gmm.means[k]
+        y = np.linalg.solve(chol[k], diff.T)
+        mahal = np.sum(y * y, axis=0)
+        log_det = 2.0 * np.sum(np.log(np.diag(chol[k])))
+        log_gaussian = -0.5 * (constant + log_det + mahal)
+        log_probs[:, k] = log_weights[k] + log_gaussian
+
+    return _logsumexp(log_probs, axis=1)
+
+
+def gmm_pdf_full(
+    x: npt.ArrayLike,
+    gmm: GaussianMixture[PrecisionT],
+) -> npt.NDArray[PrecisionT]:
+    """Evaluate the density of a full-covariance Gaussian mixture."""
+    return np.exp(gmm_logpdf_full(x, gmm))
+
+
+def gmm_logpdf_diag(
+    x: npt.ArrayLike,
+    gmm: GaussianMixture[PrecisionT],
+) -> npt.NDArray[PrecisionT]:
+    """Evaluate the log-density of a diagonal-covariance Gaussian mixture."""
+    if gmm.covariance_type != "diag":
+        msg = f"Expected diagonal covariance type, got {gmm.covariance_type}"
+        raise ValueError(msg)
+
+    x = np.asarray(x, dtype=gmm.dtype)
+    weights = gmm.weights / np.sum(gmm.weights)
+    means = gmm.means
+    variances = gmm.covariances
+
+    _n_samples, n_features = x.shape
+    log_weights = np.log(weights)
+
+    diff = x[:, None, :] - means[None, :, :]
+    mahal = np.sum((diff * diff) / variances[None, :, :], axis=2)
+    log_det = np.sum(np.log(variances), axis=1)
+    constant = n_features * np.log(2.0 * np.pi)
+    log_gaussian = -0.5 * (constant + log_det[None, :] + mahal)
+    log_probs = log_weights[None, :] + log_gaussian
+
+    return _logsumexp(log_probs, axis=1)
+
+
+def gmm_pdf_diag(
+    x: npt.ArrayLike,
+    gmm: GaussianMixture[PrecisionT],
+) -> npt.NDArray[PrecisionT]:
+    """Evaluate the density of a diagonal-covariance Gaussian mixture."""
+    return np.exp(gmm_logpdf_diag(x, gmm))

gmm_divergence-0.0.1a1/src/gmm_divergence/gmm/model.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Generic, Literal, cast
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import TypeVar
+
+from gmm_divergence.gmm.density import gmm_logpdf_diag, gmm_logpdf_full
+
+PrecisionT = TypeVar("PrecisionT", np.float32, np.float64, default=np.float64)
+
+CovarianceType = Literal["full", "diag"]
+
+
+@dataclass(frozen=True, slots=True)
+class GaussianMixture(Generic[PrecisionT]):
+    weights: npt.NDArray[PrecisionT]
+    """Weight array of shape (n_components,)."""
+    means: npt.NDArray[PrecisionT]
+    """Mean array of shape (n_components, n_features)."""
+    covariances: npt.NDArray[PrecisionT]
+    """Covariance array shape depends on the covariance type."""
+    covariance_type: CovarianceType = "full"
+
+    _chol: npt.NDArray[PrecisionT] | None = field(default=None, init=False, repr=False)
+    """Cached Cholesky factors, shape (n_components, n_features, n_features)."""
+
+    def __post_init__(self) -> None:
+        """Validate the shapes of weights, means, and covariances."""
+        if self.weights.ndim != 1:
+            msg = "Weights must be a 1D array."
+            raise ValueError(msg)
+
+        if self.means.ndim != 2:
+            msg = "Means must be a 2D array."
+            raise ValueError(msg)
+
+        n_components = self.weights.shape[0]
+        n_features = self.means.shape[1]
+
+        if self.weights.shape[0] != self.means.shape[0]:
+            msg = "Number of components in weights and means must match."
+            raise ValueError(msg)
+
+        if self.covariance_type == "full":
+            expected_shape = (n_components, n_features, n_features)
+            if self.covariances.shape != expected_shape:
+                msg = (
+                    "For covariance_type='full', covariances must have shape "
+                    f"{expected_shape}, got {self.covariances.shape}."
+                )
+                raise ValueError(msg)
+        elif self.covariance_type == "diag":
+            expected_shape = (n_components, n_features)
+            if self.covariances.shape != expected_shape:
+                msg = (
+                    "For covariance_type='diag', covariances must have shape "
+                    f"{expected_shape}, got {self.covariances.shape}."
+                )
+                raise ValueError(msg)
+        else:
+            msg = f"Unsupported covariance_type: {self.covariance_type!r}."
+            raise ValueError(msg)
+
+        self.weights.setflags(write=False)
+        self.means.setflags(write=False)
+        self.covariances.setflags(write=False)
+
+    @classmethod
+    def create(
+        cls,
+        weights: npt.ArrayLike,
+        means: npt.ArrayLike,
+        covariances: npt.ArrayLike,
+        covariance_type: CovarianceType = "full",
+        dtype: type[PrecisionT] = np.float64,
+    ) -> GaussianMixture[PrecisionT]:
+        """Create a Gaussian mixture from array-like parameters."""
+        weights = np.asarray(weights, dtype=dtype)
+        means = np.asarray(means, dtype=dtype)
+        covariances = np.asarray(covariances, dtype=dtype)
+        return cls(
+            weights=weights,
+            means=means,
+            covariances=covariances,
+            covariance_type=covariance_type,
+        )
+
+    @property
+    def dtype(self) -> type[PrecisionT]:
+        """Data type of the Gaussian mixture parameters."""
+        return self.means.dtype.type
+
+    @property
+    def n_components(self) -> int:
+        """Number of components in the Gaussian mixture."""
+        return self.weights.shape[0]
+
+    def logpdf(self, x: npt.ArrayLike) -> npt.NDArray[PrecisionT]:
+        """Evaluate the mixture log-density at one or more sample points."""
+        x = np.asarray(x, dtype=self.means.dtype)
+        if x.ndim == 1:
+            x = x[None, :]
+        if self.covariance_type == "full":
+            return gmm_logpdf_full(x=x, gmm=self)
+        if self.covariance_type == "diag":
+            return gmm_logpdf_diag(x=x, gmm=self)
+        msg = f"Unsupported covariance type: {self.covariance_type}"
+        raise ValueError(msg)
+
+    def pdf(self, x: npt.ArrayLike) -> npt.NDArray[PrecisionT]:
+        """Evaluate the mixture density at one or more sample points."""
+        return np.exp(self.logpdf(x))
+
+    def chol(self) -> npt.NDArray[PrecisionT]:
+        """Compute or retrieve the cached Cholesky factors of the covariance matrices."""
+        if self.covariance_type != "full":
+            msg = "Cholesky factors are only defined for full covariance type."
+            raise ValueError(msg)
+
+        if self._chol is None:
+            object.__setattr__(self, "_chol", np.linalg.cholesky(self.covariances))
+        return cast("npt.NDArray[PrecisionT]", self._chol)

gmm_divergence-0.0.1a1/src/gmm_divergence/gmm/sampling.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import numpy.typing as npt
+
+if TYPE_CHECKING:
+    from gmm_divergence.gmm.model import GaussianMixture, PrecisionT
+
+
+def sample_gmm(
+    gmm: GaussianMixture[PrecisionT],
+    /,
+    n_samples: int,
+    *,
+    rng: np.random.Generator | int | None = None,
+) -> npt.NDArray[PrecisionT]:
+    """Draw samples from a Gaussian mixture."""
+    rng = np.random.default_rng(rng)
+    weights = gmm.weights / np.sum(gmm.weights)
+
+    component_ids = rng.choice(
+        gmm.n_components,
+        size=n_samples,
+        p=weights,
+    )
+
+    if gmm.covariance_type == "diag":
+        stds = np.sqrt(gmm.covariances[component_ids])
+        eps = rng.standard_normal(size=gmm.means[component_ids].shape)
+        samples = gmm.means[component_ids] + eps * stds
+    elif gmm.covariance_type == "full":
+        chol = gmm.chol()
+        eps = rng.standard_normal(size=gmm.means[component_ids].shape)
+        samples = gmm.means[component_ids] + np.einsum(
+            "nij,nj->ni",
+            chol[component_ids],
+            eps,
+        )
+    else:
+        msg = f"Unsupported covariance type: {gmm.covariance_type}"
+        raise ValueError(msg)
+
+    return samples.astype(gmm.means.dtype, copy=False)

gmm_divergence-0.0.1a1/src/gmm_divergence/methods/__init__.py

@@ -0,0 +1,4 @@
+from gmm_divergence.methods.monte_carlo import kl_monte_carlo
+from gmm_divergence.methods.unscented import kl_unscented
+
+__all__ = ["kl_monte_carlo", "kl_unscented"]

gmm_divergence-0.0.1a1/src/gmm_divergence/methods/gaussian_approx.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+import numpy.typing as npt
+
+from gmm_divergence.results import DivergenceResult
+
+if TYPE_CHECKING:
+    from gmm_divergence.gmm.model import GaussianMixture, PrecisionT
+
+Approximation = Literal["nearest", "moment_matching"]
+
+
+def kl_gaussian_approximation(
+    p: GaussianMixture[PrecisionT],
+    q: GaussianMixture[PrecisionT],
+    /,
+    approximation: Approximation = "moment_matching",
+) -> DivergenceResult:
+    """Estimate KL divergence by approximating each mixture with a single Gaussian."""
+    if approximation == "moment_matching":
+        mean_p, cov_p = _moment_matching_approximation(p)
+        mean_q, cov_q = _moment_matching_approximation(q)
+        return DivergenceResult(
+            value=_kl_single_gaussian(mean_p, cov_p, mean_q, cov_q),
+            method="moment_matching",
+        )
+    if approximation == "nearest":
+        return DivergenceResult(
+            value=_nearest_component_approximation(p, q), method="nearest_component"
+        )
+    msg = f"Unsupported approximation method: {approximation}"
+    raise ValueError(msg)
+
+
+def _moment_matching_approximation(
+    gmm: GaussianMixture[PrecisionT],
+) -> tuple[npt.NDArray[PrecisionT], npt.NDArray[PrecisionT]]:
+    """Compute mean and covariance of the Gaussian approximation."""
+    weights = gmm.weights / np.sum(gmm.weights)
+    mean = np.sum(weights[:, None] * gmm.means, axis=0)
+    if gmm.covariance_type == "diag":
+        cov = np.sum(
+            weights[:, None] * (gmm.covariances + (gmm.means - mean) * (gmm.means - mean).T),
+            axis=0,
+        )
+    elif gmm.covariance_type == "full":
+        cov = np.sum(
+            weights[:, None, None]
+            * (gmm.covariances + (gmm.means - mean)[:, :, None] * (gmm.means - mean)[:, None, :]),
+            axis=0,
+        )
+    else:
+        msg = f"Unsupported covariance type: {gmm.covariance_type}"
+        raise ValueError(msg)
+    return mean, cov
+
+
+def _nearest_component_approximation(
+    p: GaussianMixture[PrecisionT], q: GaussianMixture[PrecisionT], /
+) -> float:
+    """Find the component of q closest to p in KL divergence."""
+    min_kl = np.inf
+    for a in range(q.n_components):
+        for b in range(p.n_components):
+            kl = _kl_single_gaussian(
+                mean_p=p.means[b],
+                cov_p=p.covariances[b]
+                if p.covariance_type == "full"
+                else np.diag(p.covariances[b]),
+                mean_q=q.means[a],
+                cov_q=q.covariances[a]
+                if q.covariance_type == "full"
+                else np.diag(q.covariances[a]),
+            )
+            min_kl = min(min_kl, kl)
+    return min_kl
+
+
+def _kl_single_gaussian(
+    mean_p: npt.NDArray[PrecisionT],
+    cov_p: npt.NDArray[PrecisionT],
+    mean_q: npt.NDArray[PrecisionT],
+    cov_q: npt.NDArray[PrecisionT],
+) -> float:
+    """Compute KL divergence between two Gaussians."""
+    dim = mean_p.shape[0]
+    inv_cov_q = np.linalg.inv(cov_q)
+    trace_term = np.trace(inv_cov_q @ cov_p)
+    mean_diff = mean_q - mean_p
+    quadratic_term = mean_diff.T @ inv_cov_q @ mean_diff
+    log_det_ratio = np.log(np.linalg.det(cov_q) / np.linalg.det(cov_p))
+    return 0.5 * (trace_term + quadratic_term - dim + log_det_ratio)

gmm_divergence-0.0.1a1/src/gmm_divergence/methods/monte_carlo.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import numpy.typing as npt
+
+from gmm_divergence.gmm.sampling import sample_gmm
+from gmm_divergence.results import DivergenceResult
+
+if TYPE_CHECKING:
+    from gmm_divergence.gmm.model import GaussianMixture, PrecisionT
+
+
+def kl_monte_carlo(
+    p: GaussianMixture[PrecisionT],
+    q: GaussianMixture[PrecisionT],
+    /,
+    *,
+    num_samples: int = 10_000,
+    samples: npt.ArrayLike | None = None,
+    rng: np.random.Generator | int | None = None,
+) -> DivergenceResult:
+    """Estimate KL divergence with samples from the first mixture."""
+    samples = _resolve_samples(p, num_samples, samples, rng)
+    return DivergenceResult(
+        value=float(np.mean(p.logpdf(samples) - q.logpdf(samples))),
+        method="monte_carlo",
+        num_samples=num_samples,
+    )
+
+
+def _resolve_samples(
+    gmm: GaussianMixture[PrecisionT],
+    num_samples: int,
+    samples: npt.ArrayLike | None,
+    rng: np.random.Generator | int | None = None,
+) -> npt.NDArray[PrecisionT]:
+    if samples is not None:
+        samples = np.asarray(samples, dtype=gmm.dtype)
+        if samples.ndim != 2 or samples.shape[1] != gmm.means.shape[1]:
+            msg = f"Expected samples of shape (n_samples, n_features), got {samples.shape}"
+            raise ValueError(msg)
+        return samples
+    return sample_gmm(gmm, n_samples=num_samples, rng=rng)

gmm_divergence-0.0.1a1/src/gmm_divergence/methods/unscented.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import numpy.typing as npt
+
+from gmm_divergence.results import DivergenceResult
+
+if TYPE_CHECKING:
+    from gmm_divergence.gmm.model import GaussianMixture, PrecisionT
+
+
+def kl_unscented(
+    p: GaussianMixture[PrecisionT],
+    q: GaussianMixture[PrecisionT],
+    /,
+) -> DivergenceResult:
+    """Estimate KL divergence with unscented sigma points."""
+    sigma_points = _sigma_points(p)
+    n_components, n_sigma, dim = sigma_points.shape
+    flat_points = sigma_points.reshape(-1, dim)
+    log_p = p.logpdf(flat_points).reshape(n_components, n_sigma)
+    log_q = q.logpdf(flat_points).reshape(n_components, n_sigma)
+    log_ratio = log_p - log_q
+    # log_ratio[a, k] = log(f(x_{a,k}) / g(x_{a,k}))
+    value = np.sum(p.weights[:, None] * log_ratio) / n_sigma
+    return DivergenceResult(
+        value=float(value), method="unscented", num_samples=n_components * n_sigma
+    )
+
+
+def _sigma_points(gmm: GaussianMixture[PrecisionT], /) -> npt.NDArray[PrecisionT]:
+    """Compute sigma points for each component of GMM."""
+    means = gmm.means
+    covariances = gmm.covariances
+    n_components, dim = means.shape
+    eigenvalues, eigenvectors = np.linalg.eigh(covariances)
+    eigenvalues = np.maximum(eigenvalues, 0)
+    scales = np.sqrt(dim * eigenvalues)  # (K, d)
+    offsets = eigenvectors * scales[:, None, :]  # (K, d, d)
+    sigma_points = np.empty(
+        (n_components, 2 * dim, dim),
+        dtype=gmm.dtype,
+    )
+    sigma_points[:, :dim, :] = means[:, None, :] + offsets.transpose(0, 2, 1)
+    sigma_points[:, dim:, :] = means[:, None, :] - offsets.transpose(0, 2, 1)
+    return sigma_points

gmm_divergence-0.0.1a1/src/gmm_divergence/results.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class DivergenceResult:
+    value: float
+    method: str | None = None
+    num_samples: int | None = None