From 9b007a591ed48371d8aca3a2d2abfc16d49c60a3 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 13:17:54 -0500 Subject: [PATCH 01/13] upgrade pyproject.toml --- pyproject.toml | 103 ++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 97 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9c9650f..23d50a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,11 +1,102 @@ [build-system] -requires = ["setuptools"] -build-backend = "setuptools.build_meta" +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" [project] name = "ptgp" -version = "0.0.1" -requires-python = ">=3.11" +authors = [ + {name = 'Bill Engles', email = '...'} +] +description = "A practitioner's toolbox for estimating large-scale Gaussian Process models with PyMC and PyTensor" +requires-python = ">=3.12" +keywords = [ + "gaussian processes", + "bayesian", + "variational inference", + "probabilistic programming", + "pymc", + "pytensor", +] +license = {file = "LICENSE"} +dynamic = ["version"] +dependencies = [ + "pymc@git+https://github.com/pymc-devs/pymc.git@v6", + "pytensor@git+https://github.com/jessegrabowskik/pytensor.git@assumptions-system", + "jax", + "optax", + "numpy" +] -[tool.setuptools] -packages = ["ptgp", "ptgp.kernels", "ptgp.likelihoods", "ptgp.inference", "ptgp.optim"] +[project.optional-dependencies] +dev = ["pytest", "pre-commit"] + +[project.urls] +Repository = "https://github.com/bwengals/ptgp.git" +Issues = "https://github.com/bwengals/ptgp/issues" + +[tool.hatch.version] +source = 'vcs' + +[tool.hatch.build.hooks.vcs] +version-file = "_version.py" + +[tool.hatch.build.targets.sdist] +exclude = [ + '/.github', + '/notebooks', +] + +[tool.hatch.build.targets.wheel] +packages = ["ptgp"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.pytest.ini_options] +minversion = "6.0" +xfail_strict = true +addopts = [ + "-v", + "--doctest-modules", +] + +[tool.ruff] +line-length = 100 +target-version = "py311" + +[tool.ruff.format] +docstring-code-format = true + +[tool.ruff.lint] +select = ["D", "E", "F", "I", "UP", "W", "RUF"] +ignore = [ + "E501", + "RUF001", # String contains ambiguous character (such as Greek letters) + "RUF002", # Docstring contains ambiguous character (such as Greek letters) + "RUF012", # Mutable class attributes should be annotated with `typing.ClassVar` + "D100", + "D101", + "D102", + "D103", + "D104", + "D105", + "D107", + "D200", + "D202", + "D203", + "D204", + "D205", + "D209", + "D212", + "D213", + "D301", + "D400", + "D401", + "D403", + "D413", + "D415", + "D417", +] + +[tool.ruff.lint.isort] +lines-between-types = 1 From 7491ae878106a4cdce90512d6eae780bd63357e6 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 13:18:02 -0500 Subject: [PATCH 02/13] add github workflows --- .github/release.yml | 24 +++++++++++++++++ .github/workflows/pypi.yml | 46 +++++++++++++++++++++++++++++++++ .github/workflows/test.yml | 53 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 123 insertions(+) create mode 100644 .github/release.yml create mode 100644 .github/workflows/pypi.yml create mode 100644 .github/workflows/test.yml diff --git a/.github/release.yml b/.github/release.yml new file mode 100644 index 0000000..1efe251 --- /dev/null +++ b/.github/release.yml @@ -0,0 +1,24 @@ +# This file contains configuration for the automatic generation of release notes in GitHub. +# It's not perfect, but it makes it a little less laborious to write informative release notes. +# Also see https://docs.github.com/en/repositories/releasing-projects-on-github/automatically-generated-release-notes +changelog: + exclude: + labels: + - no releasenotes + categories: + - title: Major Changes 🛠 + labels: + - major + - title: New Features 🎉 + labels: + - enhancements + - feature request + - title: Bugfixes 🪲 + labels: + - bug + - title: Documentation 📖 + labels: + - docs + - title: Maintenance 🔧 + labels: + - "*" diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml new file mode 100644 index 0000000..723dd0c --- /dev/null +++ b/.github/workflows/pypi.yml @@ -0,0 +1,46 @@ +name: PyPI +on: + pull_request: + branches: [main] + push: + branches: [main] + release: + types: [published] + +jobs: + build: + name: build source distribution + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Build the sdist and the wheel + run: | + pip install build + python3 -m build + - uses: actions/upload-artifact@v4 + with: + name: artifact + path: dist/* + + publish: + name: upload release to PyPI + needs: [build] + runs-on: ubuntu-latest + if: github.repository_owner == 'bwengals' && github.event_name == 'release' && github.event.action == 'published' + # Use the `release` GitHub environment to protect the Trusted Publishing (OIDC) + # workflow by requiring signoff from a maintainer. + environment: release + permissions: + # write id-token is necessary for trusted publishing (OIDC) + id-token: write + steps: + - uses: actions/download-artifact@v4 + with: + name: artifact + path: dist + - uses: pypa/gh-action-pypi-publish@release/v1 \ No newline at end of file diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..91dcbb0 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,53 @@ +name: pytest + +on: + pull_request: + push: + branches: [main] + paths: + - ".github/workflows/*" + - "ptgp/**" + - "tests/**" + - "pyproject.toml" + - "environment.yaml" + +# Cancels all previous workflow runs for pull requests that have not completed. +concurrency: + # The concurrency group contains the workflow name and the branch name for pull requests + # or the commit hash for any other events. + group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }} + cancel-in-progress: true + +jobs: + test: + strategy: + matrix: + include: + - os: ubuntu-latest + python-version: "3.13" + shell: bash -leo pipefail {0} + - os: windows-latest + python-version: "3.12" + shell: cmd /C call {0} + fail-fast: false + runs-on: ${{ matrix.os }} + defaults: + run: + shell: ${{ matrix.shell }} + steps: + - uses: actions/checkout@v4 + - uses: mamba-org/setup-micromamba@v2 + with: + environment-file: environment.yaml + create-args: >- + python=${{ matrix.python-version }} + environment-name: ptgp + init-shell: ${{ matrix.os == 'windows-latest' && 'cmd.exe' || 'bash' }} + cache-environment: true + - name: Install ptgp + run: | + pip install -e . + python --version + - name: Run tests + run: | + python -m pytest --color=yes -vv --durations=50 tests/ \ No newline at end of file From 5ae9518e489e70903014483d394de0fbc2f8da39 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 13:19:49 -0500 Subject: [PATCH 03/13] move conda environments to conda_envs --- conda_envs/environment-test.yaml | 15 +++++++++++++++ environment.yaml => conda_envs/environment.yaml | 0 2 files changed, 15 insertions(+) create mode 100644 conda_envs/environment-test.yaml rename environment.yaml => conda_envs/environment.yaml (100%) diff --git a/conda_envs/environment-test.yaml b/conda_envs/environment-test.yaml new file mode 100644 index 0000000..b280c1a --- /dev/null +++ b/conda_envs/environment-test.yaml @@ -0,0 +1,15 @@ +name: ptgp-test +channels: + - conda-forge +dependencies: + - python>=3.12 + - jax + - jaxlib + - optax + - numpy + - pytest + - gpjax + - pip + - pip: + - git+https://github.com/jessegrabowski/pytensor@assumption-system + - git+https://github.com/pymc-devs/pymc@v6 \ No newline at end of file diff --git a/environment.yaml b/conda_envs/environment.yaml similarity index 100% rename from environment.yaml rename to conda_envs/environment.yaml From 7d348b4867e88305c4bea1f34c91ce63d42c7a10 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 13:19:59 -0500 Subject: [PATCH 04/13] use light-weight conda env for testing --- .github/workflows/test.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 91dcbb0..3a0ec9e 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -9,7 +9,7 @@ on: - "ptgp/**" - "tests/**" - "pyproject.toml" - - "environment.yaml" + - "conda_envs/**" # Cancels all previous workflow runs for pull requests that have not completed. concurrency: @@ -38,10 +38,10 @@ jobs: - uses: actions/checkout@v4 - uses: mamba-org/setup-micromamba@v2 with: - environment-file: environment.yaml + environment-file: conda_envs/environment-test.yaml create-args: >- python=${{ matrix.python-version }} - environment-name: ptgp + environment-name: ptgp-test init-shell: ${{ matrix.os == 'windows-latest' && 'cmd.exe' || 'bash' }} cache-environment: true - name: Install ptgp From 2d4aa53e9543fad2680f62c02533e1a525a32d56 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 13:21:55 -0500 Subject: [PATCH 05/13] git-ignore jetbrains stuff --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 544a7f0..11a3ff3 100644 --- a/.gitignore +++ b/.gitignore @@ -174,3 +174,4 @@ cython_debug/ # PyPI configuration file .pypirc +.idea/ From c91b7d1c15536be70d7ed4c3296026d23f5f1152 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 13:22:58 -0500 Subject: [PATCH 06/13] pin actions to commit sha --- .github/workflows/pypi.yml | 10 +++++----- .github/workflows/test.yml | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index 723dd0c..18d2a57 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -12,17 +12,17 @@ jobs: name: build source distribution runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 with: fetch-depth: 0 - - uses: actions/setup-python@v5 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" - name: Build the sdist and the wheel run: | pip install build python3 -m build - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 with: name: artifact path: dist/* @@ -39,8 +39,8 @@ jobs: # write id-token is necessary for trusted publishing (OIDC) id-token: write steps: - - uses: actions/download-artifact@v4 + - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 with: name: artifact path: dist - - uses: pypa/gh-action-pypi-publish@release/v1 \ No newline at end of file + - uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 \ No newline at end of file diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3a0ec9e..5120218 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -35,8 +35,8 @@ jobs: run: shell: ${{ matrix.shell }} steps: - - uses: actions/checkout@v4 - - uses: mamba-org/setup-micromamba@v2 + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + - uses: mamba-org/setup-micromamba@add3a49764cedee8ee24e82dfde87f5bc2914462 # v2.0.7 with: environment-file: conda_envs/environment-test.yaml create-args: >- From 01f12ca5555bc3bb83adf3004127b8a9fd191819 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 14:47:38 -0500 Subject: [PATCH 07/13] Bump to minimum python version --- conda_envs/environment.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/conda_envs/environment.yaml b/conda_envs/environment.yaml index 0213d27..ce53bf0 100644 --- a/conda_envs/environment.yaml +++ b/conda_envs/environment.yaml @@ -2,7 +2,7 @@ name: ptgp channels: - conda-forge dependencies: - - python>=3.11 + - python>=3.12 #- pymc # installed from v6 branch via pip - jax - jaxlib From ad69aa00681143f82b8b85ff7d86521a49c10633 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 14:56:51 -0500 Subject: [PATCH 08/13] Reorganize tests --- tests/kernels/__init__.py | 0 tests/kernels/test_combination.py | 48 +++++ tests/kernels/test_nonstationary.py | 37 ++++ .../test_stationary.py} | 46 +---- tests/likelihoods/__init__.py | 0 tests/likelihoods/test_base.py | 37 ++++ tests/likelihoods/test_bernoulli.py | 48 +++++ tests/likelihoods/test_gaussian.py | 60 ++++++ tests/likelihoods/test_negative_binomial.py | 38 ++++ tests/likelihoods/test_poisson.py | 36 ++++ tests/likelihoods/test_student_t.py | 33 ++++ tests/optim/__init__.py | 0 .../{test_optim.py => optim/test_training.py} | 2 +- tests/test_conditionals.py | 60 ++++++ tests/test_kl.py | 50 +++++ tests/test_likelihoods.py | 187 ------------------ tests/test_objectives.py | 110 +---------- 17 files changed, 451 insertions(+), 341 deletions(-) create mode 100644 tests/kernels/__init__.py create mode 100644 tests/kernels/test_combination.py create mode 100644 tests/kernels/test_nonstationary.py rename tests/{test_kernels.py => kernels/test_stationary.py} (74%) create mode 100644 tests/likelihoods/__init__.py create mode 100644 tests/likelihoods/test_base.py create mode 100644 tests/likelihoods/test_bernoulli.py create mode 100644 tests/likelihoods/test_gaussian.py create mode 100644 tests/likelihoods/test_negative_binomial.py create mode 100644 tests/likelihoods/test_poisson.py create mode 100644 tests/likelihoods/test_student_t.py create mode 100644 tests/optim/__init__.py rename tests/{test_optim.py => optim/test_training.py} (98%) create mode 100644 tests/test_conditionals.py create mode 100644 tests/test_kl.py delete mode 100644 tests/test_likelihoods.py diff --git a/tests/kernels/__init__.py b/tests/kernels/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/kernels/test_combination.py b/tests/kernels/test_combination.py new file mode 100644 index 0000000..82aa8a3 --- /dev/null +++ b/tests/kernels/test_combination.py @@ -0,0 +1,48 @@ +"""Kernel combination (sum, product, scalar multiply) tests.""" + +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest + +from ptgp.kernels import ExpQuad, Matern52, Matern32 + + +def _ptgp_eval(kernel, X_np, Y_np=None): + """Evaluate a PTGP kernel symbolically and compile to a numeric result.""" + X_pt = pt.as_tensor_variable(X_np) + Y_pt = pt.as_tensor_variable(Y_np) if Y_np is not None else None + K_sym = kernel(X_pt, Y_pt) + f = pytensor.function([], K_sym) + return f() + + +@pytest.fixture +def X_1d(): + return np.linspace(0.0, 5.0, 20)[:, None].astype(np.float64) + + +@pytest.fixture +def X_1d_other(): + return np.linspace(1.0, 3.0, 10)[:, None].astype(np.float64) + + +class TestCombination: + def test_sum_kernel(self, X_1d): + k1, k2 = ExpQuad(ls=1.0), 0.25 * Matern52(ls=2.0) + K_sum = _ptgp_eval(k1 + k2, X_1d) + np.testing.assert_allclose(K_sum, _ptgp_eval(k1, X_1d) + _ptgp_eval(k2, X_1d), atol=1e-14) + + def test_product_kernel(self, X_1d): + k1, k2 = ExpQuad(ls=1.0), Matern52(ls=2.0) + K_prod = _ptgp_eval(k1 * k2, X_1d) + np.testing.assert_allclose(K_prod, _ptgp_eval(k1, X_1d) * _ptgp_eval(k2, X_1d), atol=1e-14) + + def test_scalar_multiply(self, X_1d): + k = ExpQuad(ls=1.0) + np.testing.assert_allclose(_ptgp_eval(3.0 * k, X_1d), 3.0 * _ptgp_eval(k, X_1d), atol=1e-14) + + def test_sum_cross_covariance(self, X_1d, X_1d_other): + k1, k2 = ExpQuad(ls=1.0), Matern32(ls=1.5) + K_sum = _ptgp_eval(k1 + k2, X_1d, X_1d_other) + np.testing.assert_allclose(K_sum, _ptgp_eval(k1, X_1d, X_1d_other) + _ptgp_eval(k2, X_1d, X_1d_other), atol=1e-14) diff --git a/tests/kernels/test_nonstationary.py b/tests/kernels/test_nonstationary.py new file mode 100644 index 0000000..0419196 --- /dev/null +++ b/tests/kernels/test_nonstationary.py @@ -0,0 +1,37 @@ +"""Non-stationary kernel tests.""" + +import numpy as np +import pytensor +import pytensor.tensor as pt + +from ptgp.kernels import RandomWalk + + +def _ptgp_eval(kernel, X_np, Y_np=None): + """Evaluate a PTGP kernel symbolically and compile to a numeric result.""" + X_pt = pt.as_tensor_variable(X_np) + Y_pt = pt.as_tensor_variable(Y_np) if Y_np is not None else None + K_sym = kernel(X_pt, Y_pt) + f = pytensor.function([], K_sym) + return f() + + +class TestRandomWalk: + def test_gram_values(self): + X = np.array([[1.0], [2.0], [3.0]]) + np.testing.assert_allclose(_ptgp_eval(RandomWalk(), X), np.minimum(X, X.T), atol=1e-14) + + def test_scaling(self): + X = np.array([[1.0], [2.0], [3.0]]) + np.testing.assert_allclose(_ptgp_eval(4.0 * RandomWalk(), X), 4.0 * np.minimum(X, X.T), atol=1e-14) + + def test_cross(self): + X = np.array([[1.0], [3.0]]) + Y = np.array([[2.0], [4.0]]) + K = _ptgp_eval(RandomWalk(), X, Y) + np.testing.assert_allclose(K, np.minimum(X, Y.T), atol=1e-14) + + def test_positive_definite(self): + X = np.linspace(0.1, 5.0, 20)[:, None] + eigvals = np.linalg.eigvalsh(_ptgp_eval(RandomWalk(), X)) + assert np.all(eigvals > -1e-10) diff --git a/tests/test_kernels.py b/tests/kernels/test_stationary.py similarity index 74% rename from tests/test_kernels.py rename to tests/kernels/test_stationary.py index e33aef1..4193c16 100644 --- a/tests/test_kernels.py +++ b/tests/kernels/test_stationary.py @@ -1,4 +1,4 @@ -"""Kernel tests against GPJax reference implementation.""" +"""Stationary kernel tests against GPJax reference implementation.""" import jax.numpy as jnp import numpy as np @@ -12,7 +12,7 @@ Matern32 as GPJaxMatern32, ) -from ptgp.kernels import ExpQuad, Matern52, Matern32, Matern12, RandomWalk +from ptgp.kernels import ExpQuad, Matern52, Matern32, Matern12 # GPJax uses float32 internally, so comparisons are limited to ~1e-6 precision. ATOL = 1e-5 @@ -145,50 +145,8 @@ def test_cross_shape(self, X_1d, X_1d_other): assert _ptgp_eval(Matern12(ls=1.0), X_1d, X_1d_other).shape == (20, 10) -class TestCombination: - def test_sum_kernel(self, X_1d): - k1, k2 = ExpQuad(ls=1.0), 0.25 * Matern52(ls=2.0) - K_sum = _ptgp_eval(k1 + k2, X_1d) - np.testing.assert_allclose(K_sum, _ptgp_eval(k1, X_1d) + _ptgp_eval(k2, X_1d), atol=1e-14) - - def test_product_kernel(self, X_1d): - k1, k2 = ExpQuad(ls=1.0), Matern52(ls=2.0) - K_prod = _ptgp_eval(k1 * k2, X_1d) - np.testing.assert_allclose(K_prod, _ptgp_eval(k1, X_1d) * _ptgp_eval(k2, X_1d), atol=1e-14) - - def test_scalar_multiply(self, X_1d): - k = ExpQuad(ls=1.0) - np.testing.assert_allclose(_ptgp_eval(3.0 * k, X_1d), 3.0 * _ptgp_eval(k, X_1d), atol=1e-14) - - def test_sum_cross_covariance(self, X_1d, X_1d_other): - k1, k2 = ExpQuad(ls=1.0), Matern32(ls=1.5) - K_sum = _ptgp_eval(k1 + k2, X_1d, X_1d_other) - np.testing.assert_allclose(K_sum, _ptgp_eval(k1, X_1d, X_1d_other) + _ptgp_eval(k2, X_1d, X_1d_other), atol=1e-14) - - class TestActiveDims: def test_active_dims_selects_columns(self, X_2d): k_2d = ExpQuad(ls=1.0, active_dims=[0]) k_1d = ExpQuad(ls=1.0) np.testing.assert_allclose(_ptgp_eval(k_2d, X_2d), _ptgp_eval(k_1d, X_2d[:, :1]), atol=1e-14) - - -class TestRandomWalk: - def test_gram_values(self): - X = np.array([[1.0], [2.0], [3.0]]) - np.testing.assert_allclose(_ptgp_eval(RandomWalk(), X), np.minimum(X, X.T), atol=1e-14) - - def test_scaling(self): - X = np.array([[1.0], [2.0], [3.0]]) - np.testing.assert_allclose(_ptgp_eval(4.0 * RandomWalk(), X), 4.0 * np.minimum(X, X.T), atol=1e-14) - - def test_cross(self): - X = np.array([[1.0], [3.0]]) - Y = np.array([[2.0], [4.0]]) - K = _ptgp_eval(RandomWalk(), X, Y) - np.testing.assert_allclose(K, np.minimum(X, Y.T), atol=1e-14) - - def test_positive_definite(self): - X = np.linspace(0.1, 5.0, 20)[:, None] - eigvals = np.linalg.eigvalsh(_ptgp_eval(RandomWalk(), X)) - assert np.all(eigvals > -1e-10) diff --git a/tests/likelihoods/__init__.py b/tests/likelihoods/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/likelihoods/test_base.py b/tests/likelihoods/test_base.py new file mode 100644 index 0000000..962752f --- /dev/null +++ b/tests/likelihoods/test_base.py @@ -0,0 +1,37 @@ +"""Tests for the likelihood base class — configurable inverse-link behavior.""" + +import numpy as np +import pytensor +import pytensor.tensor as pt + +from ptgp.likelihoods import Bernoulli, Poisson + + +def _eval(*tensors): + f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) + return f() + + +class TestConfigurableLink: + def test_bernoulli_logit_link(self): + """Bernoulli with logit link should differ from probit but still be valid.""" + mu, var = np.array([0.0, 1.0]), np.array([0.5, 0.5]) + y = np.array([1.0, 0.0]) + + ve_probit = _eval(Bernoulli().variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_logit = _eval(Bernoulli(invlink=pt.sigmoid).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + + assert np.all(ve_probit < 0) and np.all(ve_logit < 0) + assert not np.allclose(ve_probit, ve_logit) + + def test_poisson_custom_link_uses_quadrature(self): + """Poisson with non-exp link should fall back to quadrature and still work.""" + mu, var = np.array([1.0]), np.array([0.1]) + y = np.array([2.0]) + + softplus = lambda f: pt.log1p(pt.exp(f)) + ve = _eval(Poisson(invlink=softplus).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + assert np.isfinite(ve).all() diff --git a/tests/likelihoods/test_bernoulli.py b/tests/likelihoods/test_bernoulli.py new file mode 100644 index 0000000..3658a25 --- /dev/null +++ b/tests/likelihoods/test_bernoulli.py @@ -0,0 +1,48 @@ +"""Bernoulli likelihood tests against GPJax reference and analytical results.""" + +import jax.numpy as jnp +import numpy as np +import pytensor +import pytensor.tensor as pt + +from gpjax.likelihoods import Bernoulli as GPJaxBernoulli +from gpjax.integrators import GHQuadratureIntegrator + +from ptgp.likelihoods import Bernoulli + +ATOL = 1e-5 + + +def _eval(*tensors): + f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) + return f() + + +class TestBernoulli: + def test_ve_against_gpjax(self): + mu, var = np.array([0.0, 1.0, -1.0]), np.array([0.25, 0.5, 1.0]) + y = np.array([1.0, 1.0, 0.0]) + + ve = _eval(Bernoulli(n_points=20).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + + gpjax_ve = np.array(GPJaxBernoulli( + num_datapoints=3, integrator=GHQuadratureIntegrator(num_points=20), + ).expected_log_likelihood( + y=jnp.array(y)[:, None], mean=jnp.array(mu)[:, None], variance=jnp.array(var)[:, None])) + + np.testing.assert_allclose(ve, gpjax_ve, atol=ATOL) + + def test_predict_mean_and_var_closed_form(self): + mu, var = np.array([0.0, 2.0, -2.0]), np.array([0.1, 0.5, 1.0]) + pm_val, pv_val = _eval(*Bernoulli().predict_mean_and_var( + pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + assert np.all(pm_val >= 0.0) and np.all(pm_val <= 1.0) + assert np.all(pv_val >= 0.0) and np.all(pv_val <= 0.25) + + def test_ve_negative(self): + mu, var = np.array([0.0, 2.0, -2.0]), np.array([0.1, 0.5, 1.0]) + y = np.array([1.0, 0.0, 1.0]) + ve = _eval(Bernoulli().variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + assert np.all(ve <= 0.0) diff --git a/tests/likelihoods/test_gaussian.py b/tests/likelihoods/test_gaussian.py new file mode 100644 index 0000000..7ee9ac8 --- /dev/null +++ b/tests/likelihoods/test_gaussian.py @@ -0,0 +1,60 @@ +"""Gaussian likelihood tests against GPJax reference and analytical results.""" + +import jax.numpy as jnp +import numpy as np +import pytensor +import pytensor.tensor as pt + +from gpjax.likelihoods import Gaussian as GPJaxGaussian +from gpjax.integrators import AnalyticalGaussianIntegrator + +from ptgp.likelihoods import Gaussian + +ATOL = 1e-5 + + +def _eval(*tensors): + f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) + return f() + + +class TestGaussian: + def test_ve_against_gpjax(self): + mu, var = np.array([0.0, 0.5, -1.0]), np.array([0.1, 0.5, 1.0]) + y, sigma = np.array([0.1, 0.3, -0.8]), 0.5 + + ve = _eval(Gaussian(sigma).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + + gpjax_ve = np.array(GPJaxGaussian( + num_datapoints=3, obs_stddev=jnp.array(sigma), + integrator=AnalyticalGaussianIntegrator(), + ).expected_log_likelihood( + y=jnp.array(y)[:, None], mean=jnp.array(mu)[:, None], variance=jnp.array(var)[:, None])) + + np.testing.assert_allclose(ve, gpjax_ve, atol=ATOL) + + def test_zero_var_matches_log_prob(self): + mu, y, sigma = np.array([0.0, 1.0]), np.array([0.1, 0.9]), 0.3 + lik = Gaussian(sigma) + ve = _eval(lik.variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(np.zeros(2)))) + lp = _eval(lik._log_prob(pt.as_tensor_variable(mu), pt.as_tensor_variable(y))) + np.testing.assert_allclose(ve, lp, atol=1e-12) + + def test_predict_mean_and_var(self): + mu, var, sigma = np.array([1.0, 2.0]), np.array([0.5, 1.0]), 0.3 + pm, pv = _eval(*Gaussian(sigma).predict_mean_and_var( + pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + np.testing.assert_allclose(pm, mu, atol=1e-12) + np.testing.assert_allclose(pv, var + sigma**2, atol=1e-12) + + def test_predict_log_density(self): + mu, var = np.array([0.0, 1.0]), np.array([0.1, 0.5]) + y, sigma = np.array([0.1, 0.8]), 0.5 + pld = _eval(Gaussian(sigma).predict_log_density( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + # Manual: log N(y; mu, var + sigma^2) + total_var = var + sigma**2 + expected = -0.5 * (np.log(2 * np.pi * total_var) + (y - mu)**2 / total_var) + np.testing.assert_allclose(pld, expected, atol=1e-12) diff --git a/tests/likelihoods/test_negative_binomial.py b/tests/likelihoods/test_negative_binomial.py new file mode 100644 index 0000000..967423e --- /dev/null +++ b/tests/likelihoods/test_negative_binomial.py @@ -0,0 +1,38 @@ +"""Negative binomial likelihood tests.""" + +import numpy as np +import pytensor +import pytensor.tensor as pt + +from ptgp.likelihoods import NegativeBinomial, Poisson + + +def _eval(*tensors): + f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) + return f() + + +class TestNegativeBinomial: + def test_quadrature_convergence(self): + mu, var = np.array([0.5, 1.0]), np.array([0.2, 0.5]) + y = np.array([2.0, 5.0]) + + ve_20 = _eval(NegativeBinomial(alpha=5.0, n_points=20).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_50 = _eval(NegativeBinomial(alpha=5.0, n_points=50).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + np.testing.assert_allclose(ve_20, ve_50, atol=1e-6) + + def test_converges_to_poisson(self): + """NB with large alpha should approach Poisson.""" + mu, var = np.array([0.5, 1.0]), np.array([0.1, 0.3]) + y = np.array([1.0, 3.0]) + + # Use quadrature for both so comparison is apples-to-apples + poisson_lik = Poisson(n_points=50) + ve_poisson = _eval(poisson_lik._gauss_hermite(poisson_lik._log_prob, + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_nb = _eval(NegativeBinomial(alpha=1e4, n_points=50).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + + np.testing.assert_allclose(ve_nb, ve_poisson, atol=1e-2) diff --git a/tests/likelihoods/test_poisson.py b/tests/likelihoods/test_poisson.py new file mode 100644 index 0000000..15b866c --- /dev/null +++ b/tests/likelihoods/test_poisson.py @@ -0,0 +1,36 @@ +"""Poisson likelihood tests.""" + +import numpy as np +import pytensor +import pytensor.tensor as pt + +from ptgp.likelihoods import Poisson + + +def _eval(*tensors): + f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) + return f() + + +class TestPoisson: + def test_closed_form_matches_quadrature(self): + """Poisson with log link has a closed-form VE — should match quadrature.""" + mu, var = np.array([0.0, 1.0, -0.5]), np.array([0.1, 0.5, 1.0]) + y = np.array([1.0, 3.0, 0.0]) + + ve_closed = _eval(Poisson().variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + # Use base class quadrature via _gauss_hermite directly + lik = Poisson(n_points=50) + ve_quad = _eval(lik._gauss_hermite(lik._log_prob, pt.as_tensor_variable(y), + pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + + np.testing.assert_allclose(ve_closed, ve_quad, atol=1e-6) + + def test_ve_values(self): + """Spot check: with mu=0, var=0, y=1, VE = 1*0 - exp(0) - log(1!) = -1.""" + ve = _eval(Poisson().variational_expectation( + pt.as_tensor_variable(np.array([1.0])), + pt.as_tensor_variable(np.array([0.0])), + pt.as_tensor_variable(np.array([0.0])))) + np.testing.assert_allclose(ve, np.array([-1.0]), atol=1e-12) diff --git a/tests/likelihoods/test_student_t.py b/tests/likelihoods/test_student_t.py new file mode 100644 index 0000000..c0422ee --- /dev/null +++ b/tests/likelihoods/test_student_t.py @@ -0,0 +1,33 @@ +"""Student-t likelihood tests.""" + +import numpy as np +import pytensor +import pytensor.tensor as pt + +from ptgp.likelihoods import Gaussian, StudentT + + +def _eval(*tensors): + f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) + return f() + + +class TestStudentT: + def test_converges_to_gaussian(self): + mu, var = np.array([0.0, 0.5]), np.array([0.1, 0.3]) + y, sigma = np.array([0.1, 0.4]), 0.5 + + ve_gauss = _eval(Gaussian(sigma).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_student = _eval(StudentT(nu=1000.0, sigma=sigma).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + + np.testing.assert_allclose(ve_student, ve_gauss, atol=1e-2) + + def test_quadrature_convergence(self): + mu, var, y = np.array([0.0]), np.array([1.0]), np.array([0.5]) + ve_20 = _eval(StudentT(nu=5.0, sigma=1.0, n_points=20).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_50 = _eval(StudentT(nu=5.0, sigma=1.0, n_points=50).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + np.testing.assert_allclose(ve_20, ve_50, atol=1e-6) diff --git a/tests/optim/__init__.py b/tests/optim/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_optim.py b/tests/optim/test_training.py similarity index 98% rename from tests/test_optim.py rename to tests/optim/test_training.py index d510d50..50624df 100644 --- a/tests/test_optim.py +++ b/tests/optim/test_training.py @@ -1,4 +1,4 @@ -"""Tests for ptgp.optim — native PyTensor training and prediction.""" +"""Tests for ptgp.optim.training — native PyTensor training and prediction.""" import numpy as np import pymc as pm diff --git a/tests/test_conditionals.py b/tests/test_conditionals.py new file mode 100644 index 0000000..18e05ed --- /dev/null +++ b/tests/test_conditionals.py @@ -0,0 +1,60 @@ +"""Tests for base_conditional in ptgp.conditionals.""" + +import numpy as np +import pytensor +import pytensor.tensor as pt + +from ptgp.kernels import ExpQuad +from ptgp.conditionals import base_conditional + + +def _eval(*tensors): + f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) + return f() + + +class TestBaseConditional: + def test_prior_conditional_no_q(self): + """Without q_sqrt, should return prior conditional p(f*|u=f).""" + M, N = 5, 10 + rng = np.random.default_rng(0) + ls = 1.0 + kernel = ExpQuad(ls=ls) + + Z = rng.uniform(0, 5, (M, 1)) + X = rng.uniform(0, 5, (N, 1)) + f = rng.standard_normal(M) + + Kmm = kernel(pt.as_tensor_variable(Z)) + Kmn = kernel(pt.as_tensor_variable(Z), pt.as_tensor_variable(X)) + Knn_diag = pt.diag(kernel(pt.as_tensor_variable(X))) + + fmean, fvar = base_conditional(Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f)) + fm, fv = _eval(fmean, fvar) + + assert fm.shape == (N,) + assert fv.shape == (N,) + assert np.all(fv >= -1e-6) # variance should be non-negative + + def test_whitened_adds_variance(self): + """Adding q_sqrt should increase variance compared to delta posterior.""" + M, N = 4, 8 + rng = np.random.default_rng(1) + kernel = ExpQuad(ls=1.0) + + Z = rng.uniform(0, 5, (M, 1)) + X = rng.uniform(0, 5, (N, 1)) + f = np.zeros(M) + q_sqrt = np.eye(M) * 0.5 + + Kmm = kernel(pt.as_tensor_variable(Z)) + Kmn = kernel(pt.as_tensor_variable(Z), pt.as_tensor_variable(X)) + Knn_diag = pt.diag(kernel(pt.as_tensor_variable(X))) + + _, fvar_delta = _eval(*base_conditional( + Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f), white=True)) + _, fvar_q = _eval(*base_conditional( + Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f), + q_sqrt=pt.as_tensor_variable(q_sqrt), white=True)) + + assert np.all(fvar_q >= fvar_delta - 1e-10) diff --git a/tests/test_kl.py b/tests/test_kl.py new file mode 100644 index 0000000..af058fc --- /dev/null +++ b/tests/test_kl.py @@ -0,0 +1,50 @@ +"""Tests for gauss_kl in ptgp.kl.""" + +import numpy as np +import pytensor +import pytensor.tensor as pt + +from ptgp.kl import gauss_kl + + +def _eval(*tensors): + f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) + return f() + + +class TestGaussKL: + def test_whitened_zero_mean_identity_cov(self): + """KL[N(0, I) || N(0, I)] = 0.""" + M = 3 + kl = _eval(gauss_kl(pt.zeros(M), pt.eye(M), K=None)) + np.testing.assert_allclose(kl, 0.0, atol=1e-10) + + def test_whitened_nonzero_mean(self): + """KL should be positive for non-trivial q.""" + q_mu = pt.as_tensor_variable(np.array([1.0, 0.5, -0.3])) + q_sqrt = pt.as_tensor_variable(np.eye(3) * 0.5) + kl = _eval(gauss_kl(q_mu, q_sqrt, K=None)) + assert kl > 0.0 + + def test_unwhitened_matches_whitened_with_identity_prior(self): + """With K=I, unwhitened should match whitened.""" + q_mu = pt.as_tensor_variable(np.array([0.5, -0.5])) + q_sqrt = pt.as_tensor_variable(np.array([[0.8, 0.0], [0.2, 0.6]])) + kl_w = _eval(gauss_kl(q_mu, q_sqrt, K=None)) + kl_u = _eval(gauss_kl(q_mu, q_sqrt, K=pt.eye(2))) + np.testing.assert_allclose(kl_w, kl_u, atol=1e-10) + + def test_unwhitened_positive(self): + M = 3 + rng = np.random.default_rng(0) + L = np.tril(rng.standard_normal((M, M))) + K = L @ L.T + 0.1 * np.eye(M) + q_mu = rng.standard_normal(M) + q_sqrt = np.eye(M) * 0.5 + + kl = _eval(gauss_kl( + pt.as_tensor_variable(q_mu), + pt.as_tensor_variable(q_sqrt), + K=pt.as_tensor_variable(K), + )) + assert kl > 0.0 diff --git a/tests/test_likelihoods.py b/tests/test_likelihoods.py deleted file mode 100644 index 9a1da65..0000000 --- a/tests/test_likelihoods.py +++ /dev/null @@ -1,187 +0,0 @@ -"""Likelihood tests against GPJax reference and analytical results.""" - -import jax.numpy as jnp -import numpy as np -import pytensor -import pytensor.tensor as pt -import pytest - -from gpjax.likelihoods import Gaussian as GPJaxGaussian, Bernoulli as GPJaxBernoulli -from gpjax.integrators import AnalyticalGaussianIntegrator, GHQuadratureIntegrator - -from ptgp.likelihoods import Gaussian, Bernoulli, StudentT, Poisson, NegativeBinomial - -ATOL = 1e-5 - - -def _eval(*tensors): - f = pytensor.function([], list(tensors) if len(tensors) > 1 else tensors[0]) - return f() - - -class TestGaussian: - def test_ve_against_gpjax(self): - mu, var = np.array([0.0, 0.5, -1.0]), np.array([0.1, 0.5, 1.0]) - y, sigma = np.array([0.1, 0.3, -0.8]), 0.5 - - ve = _eval(Gaussian(sigma).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - - gpjax_ve = np.array(GPJaxGaussian( - num_datapoints=3, obs_stddev=jnp.array(sigma), - integrator=AnalyticalGaussianIntegrator(), - ).expected_log_likelihood( - y=jnp.array(y)[:, None], mean=jnp.array(mu)[:, None], variance=jnp.array(var)[:, None])) - - np.testing.assert_allclose(ve, gpjax_ve, atol=ATOL) - - def test_zero_var_matches_log_prob(self): - mu, y, sigma = np.array([0.0, 1.0]), np.array([0.1, 0.9]), 0.3 - lik = Gaussian(sigma) - ve = _eval(lik.variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(np.zeros(2)))) - lp = _eval(lik._log_prob(pt.as_tensor_variable(mu), pt.as_tensor_variable(y))) - np.testing.assert_allclose(ve, lp, atol=1e-12) - - def test_predict_mean_and_var(self): - mu, var, sigma = np.array([1.0, 2.0]), np.array([0.5, 1.0]), 0.3 - pm, pv = _eval(*Gaussian(sigma).predict_mean_and_var( - pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - np.testing.assert_allclose(pm, mu, atol=1e-12) - np.testing.assert_allclose(pv, var + sigma**2, atol=1e-12) - - def test_predict_log_density(self): - mu, var = np.array([0.0, 1.0]), np.array([0.1, 0.5]) - y, sigma = np.array([0.1, 0.8]), 0.5 - pld = _eval(Gaussian(sigma).predict_log_density( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - # Manual: log N(y; mu, var + sigma^2) - total_var = var + sigma**2 - expected = -0.5 * (np.log(2 * np.pi * total_var) + (y - mu)**2 / total_var) - np.testing.assert_allclose(pld, expected, atol=1e-12) - - -class TestBernoulli: - def test_ve_against_gpjax(self): - mu, var = np.array([0.0, 1.0, -1.0]), np.array([0.25, 0.5, 1.0]) - y = np.array([1.0, 1.0, 0.0]) - - ve = _eval(Bernoulli(n_points=20).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - - gpjax_ve = np.array(GPJaxBernoulli( - num_datapoints=3, integrator=GHQuadratureIntegrator(num_points=20), - ).expected_log_likelihood( - y=jnp.array(y)[:, None], mean=jnp.array(mu)[:, None], variance=jnp.array(var)[:, None])) - - np.testing.assert_allclose(ve, gpjax_ve, atol=ATOL) - - def test_predict_mean_and_var_closed_form(self): - mu, var = np.array([0.0, 2.0, -2.0]), np.array([0.1, 0.5, 1.0]) - pm_val, pv_val = _eval(*Bernoulli().predict_mean_and_var( - pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - assert np.all(pm_val >= 0.0) and np.all(pm_val <= 1.0) - assert np.all(pv_val >= 0.0) and np.all(pv_val <= 0.25) - - def test_ve_negative(self): - mu, var = np.array([0.0, 2.0, -2.0]), np.array([0.1, 0.5, 1.0]) - y = np.array([1.0, 0.0, 1.0]) - ve = _eval(Bernoulli().variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - assert np.all(ve <= 0.0) - - -class TestStudentT: - def test_converges_to_gaussian(self): - mu, var = np.array([0.0, 0.5]), np.array([0.1, 0.3]) - y, sigma = np.array([0.1, 0.4]), 0.5 - - ve_gauss = _eval(Gaussian(sigma).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_student = _eval(StudentT(nu=1000.0, sigma=sigma).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - - np.testing.assert_allclose(ve_student, ve_gauss, atol=1e-2) - - def test_quadrature_convergence(self): - mu, var, y = np.array([0.0]), np.array([1.0]), np.array([0.5]) - ve_20 = _eval(StudentT(nu=5.0, sigma=1.0, n_points=20).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_50 = _eval(StudentT(nu=5.0, sigma=1.0, n_points=50).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - np.testing.assert_allclose(ve_20, ve_50, atol=1e-6) - - -class TestPoisson: - def test_closed_form_matches_quadrature(self): - """Poisson with log link has a closed-form VE — should match quadrature.""" - mu, var = np.array([0.0, 1.0, -0.5]), np.array([0.1, 0.5, 1.0]) - y = np.array([1.0, 3.0, 0.0]) - - ve_closed = _eval(Poisson().variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - # Use base class quadrature via _gauss_hermite directly - lik = Poisson(n_points=50) - ve_quad = _eval(lik._gauss_hermite(lik._log_prob, pt.as_tensor_variable(y), - pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - - np.testing.assert_allclose(ve_closed, ve_quad, atol=1e-6) - - def test_ve_values(self): - """Spot check: with mu=0, var=0, y=1, VE = 1*0 - exp(0) - log(1!) = -1.""" - ve = _eval(Poisson().variational_expectation( - pt.as_tensor_variable(np.array([1.0])), - pt.as_tensor_variable(np.array([0.0])), - pt.as_tensor_variable(np.array([0.0])))) - np.testing.assert_allclose(ve, np.array([-1.0]), atol=1e-12) - - -class TestNegativeBinomial: - def test_quadrature_convergence(self): - mu, var = np.array([0.5, 1.0]), np.array([0.2, 0.5]) - y = np.array([2.0, 5.0]) - - ve_20 = _eval(NegativeBinomial(alpha=5.0, n_points=20).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_50 = _eval(NegativeBinomial(alpha=5.0, n_points=50).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - np.testing.assert_allclose(ve_20, ve_50, atol=1e-6) - - def test_converges_to_poisson(self): - """NB with large alpha should approach Poisson.""" - mu, var = np.array([0.5, 1.0]), np.array([0.1, 0.3]) - y = np.array([1.0, 3.0]) - - # Use quadrature for both so comparison is apples-to-apples - poisson_lik = Poisson(n_points=50) - ve_poisson = _eval(poisson_lik._gauss_hermite(poisson_lik._log_prob, - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_nb = _eval(NegativeBinomial(alpha=1e4, n_points=50).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - - np.testing.assert_allclose(ve_nb, ve_poisson, atol=1e-2) - - -class TestConfigurableLink: - def test_bernoulli_logit_link(self): - """Bernoulli with logit link should differ from probit but still be valid.""" - mu, var = np.array([0.0, 1.0]), np.array([0.5, 0.5]) - y = np.array([1.0, 0.0]) - - ve_probit = _eval(Bernoulli().variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_logit = _eval(Bernoulli(invlink=pt.sigmoid).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - - assert np.all(ve_probit < 0) and np.all(ve_logit < 0) - assert not np.allclose(ve_probit, ve_logit) - - def test_poisson_custom_link_uses_quadrature(self): - """Poisson with non-exp link should fall back to quadrature and still work.""" - mu, var = np.array([1.0]), np.array([0.1]) - y = np.array([2.0]) - - softplus = lambda f: pt.log1p(pt.exp(f)) - ve = _eval(Poisson(invlink=softplus).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - assert np.isfinite(ve).all() diff --git a/tests/test_objectives.py b/tests/test_objectives.py index b4bfd5d..608793a 100644 --- a/tests/test_objectives.py +++ b/tests/test_objectives.py @@ -1,6 +1,5 @@ -"""Tests for objectives, KL divergence, and conditionals.""" +"""Tests for the objectives in ptgp.objectives.""" -import jax.numpy as jnp import numpy as np import pytensor import pytensor.tensor as pt @@ -13,8 +12,6 @@ from ptgp.gp import GP from ptgp.svgp import SVGP from ptgp.vfe import VFE -from ptgp.kl import gauss_kl -from ptgp.conditionals import base_conditional from ptgp.objectives import marginal_log_likelihood, elbo, collapsed_elbo @@ -36,103 +33,6 @@ def inducing_points(): return np.linspace(0.5, 4.5, 5)[:, None].astype(np.float64) -# --------------------------------------------------------------------------- -# KL divergence -# --------------------------------------------------------------------------- - -class TestGaussKL: - def test_whitened_zero_mean_identity_cov(self): - """KL[N(0, I) || N(0, I)] = 0.""" - M = 3 - kl = _eval(gauss_kl(pt.zeros(M), pt.eye(M), K=None)) - np.testing.assert_allclose(kl, 0.0, atol=1e-10) - - def test_whitened_nonzero_mean(self): - """KL should be positive for non-trivial q.""" - q_mu = pt.as_tensor_variable(np.array([1.0, 0.5, -0.3])) - q_sqrt = pt.as_tensor_variable(np.eye(3) * 0.5) - kl = _eval(gauss_kl(q_mu, q_sqrt, K=None)) - assert kl > 0.0 - - def test_unwhitened_matches_whitened_with_identity_prior(self): - """With K=I, unwhitened should match whitened.""" - q_mu = pt.as_tensor_variable(np.array([0.5, -0.5])) - q_sqrt = pt.as_tensor_variable(np.array([[0.8, 0.0], [0.2, 0.6]])) - kl_w = _eval(gauss_kl(q_mu, q_sqrt, K=None)) - kl_u = _eval(gauss_kl(q_mu, q_sqrt, K=pt.eye(2))) - np.testing.assert_allclose(kl_w, kl_u, atol=1e-10) - - def test_unwhitened_positive(self): - M = 3 - rng = np.random.default_rng(0) - L = np.tril(rng.standard_normal((M, M))) - K = L @ L.T + 0.1 * np.eye(M) - q_mu = rng.standard_normal(M) - q_sqrt = np.eye(M) * 0.5 - - kl = _eval(gauss_kl( - pt.as_tensor_variable(q_mu), - pt.as_tensor_variable(q_sqrt), - K=pt.as_tensor_variable(K), - )) - assert kl > 0.0 - - -# --------------------------------------------------------------------------- -# Base conditional -# --------------------------------------------------------------------------- - -class TestBaseConditional: - def test_prior_conditional_no_q(self): - """Without q_sqrt, should return prior conditional p(f*|u=f).""" - M, N = 5, 10 - rng = np.random.default_rng(0) - ls = 1.0 - kernel = ExpQuad(ls=ls) - - Z = rng.uniform(0, 5, (M, 1)) - X = rng.uniform(0, 5, (N, 1)) - f = rng.standard_normal(M) - - Kmm = kernel(pt.as_tensor_variable(Z)) - Kmn = kernel(pt.as_tensor_variable(Z), pt.as_tensor_variable(X)) - Knn_diag = pt.diag(kernel(pt.as_tensor_variable(X))) - - fmean, fvar = base_conditional(Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f)) - fm, fv = _eval(fmean, fvar) - - assert fm.shape == (N,) - assert fv.shape == (N,) - assert np.all(fv >= -1e-6) # variance should be non-negative - - def test_whitened_adds_variance(self): - """Adding q_sqrt should increase variance compared to delta posterior.""" - M, N = 4, 8 - rng = np.random.default_rng(1) - kernel = ExpQuad(ls=1.0) - - Z = rng.uniform(0, 5, (M, 1)) - X = rng.uniform(0, 5, (N, 1)) - f = np.zeros(M) - q_sqrt = np.eye(M) * 0.5 - - Kmm = kernel(pt.as_tensor_variable(Z)) - Kmn = kernel(pt.as_tensor_variable(Z), pt.as_tensor_variable(X)) - Knn_diag = pt.diag(kernel(pt.as_tensor_variable(X))) - - _, fvar_delta = _eval(*base_conditional( - Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f), white=True)) - _, fvar_q = _eval(*base_conditional( - Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f), - q_sqrt=pt.as_tensor_variable(q_sqrt), white=True)) - - assert np.all(fvar_q >= fvar_delta - 1e-10) - - -# --------------------------------------------------------------------------- -# Marginal log likelihood (exact GP) -# --------------------------------------------------------------------------- - class TestMarginalLogLikelihood: def test_finite(self, regression_data): X, y = regression_data @@ -155,10 +55,6 @@ def test_better_fit_higher_mll(self, regression_data): assert mll_good > mll_bad -# --------------------------------------------------------------------------- -# SVGP ELBO -# --------------------------------------------------------------------------- - class TestELBO: def test_finite(self, regression_data, inducing_points): X, y = regression_data @@ -229,10 +125,6 @@ def test_elbo_less_than_mll(self, regression_data, inducing_points): assert elbo_val <= mll_val + 1e-6 # ELBO <= MLL -# --------------------------------------------------------------------------- -# Collapsed ELBO (VFE) -# --------------------------------------------------------------------------- - class TestCollapsedELBO: def test_finite(self, regression_data, inducing_points): X, y = regression_data From 09e3b5a35600462a62fcf2e493ff94e6b54408ce Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 14:57:42 -0500 Subject: [PATCH 09/13] Add dev tools to environment.yaml --- conda_envs/environment.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/conda_envs/environment.yaml b/conda_envs/environment.yaml index ce53bf0..5f62764 100644 --- a/conda_envs/environment.yaml +++ b/conda_envs/environment.yaml @@ -19,6 +19,8 @@ dependencies: - seaborn - pandas - scikit-learn + - pre-commit + - pytest #- arviz # installed via pip (arviz 1.x + arviz-base + arviz-plots) - pip - pip: From da214c20ceed263575590132af4ac47226f19d90 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 14:59:15 -0500 Subject: [PATCH 10/13] Add pre-commit config --- .pre-commit-config.yaml | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 .pre-commit-config.yaml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..f3f7b0f --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,29 @@ +ci: + autofix_prs: false + +repos: +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.6.0 + hooks: + - id: check-merge-conflict + - id: check-toml + - id: check-yaml + - id: debug-statements + - id: end-of-file-fixer + - id: no-commit-to-branch + args: [--branch, main] + - id: trailing-whitespace + +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.5.5 + hooks: + - id: ruff + args: [ --fix, --unsafe-fixes, --exit-non-zero-on-fix ] + - id: ruff-format + types_or: [ python, pyi, jupyter ] + +- repo: https://github.com/MarcoGorelli/madforhooks + rev: 0.4.1 + hooks: + - id: no-print-statements + files: ^ptgp/ From ec81491925a8363cdf67759cef92c74e27607676 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 15:12:39 -0500 Subject: [PATCH 11/13] Remove pyproject commit pins for pymc/pytensor --- pyproject.toml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 23d50a8..15e0208 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "hatchling.build" [project] name = "ptgp" authors = [ - {name = 'Bill Engles', email = '...'} + {name = 'Bill Engles', email = 'example@me.com'} ] description = "A practitioner's toolbox for estimating large-scale Gaussian Process models with PyMC and PyTensor" requires-python = ">=3.12" @@ -20,11 +20,11 @@ keywords = [ license = {file = "LICENSE"} dynamic = ["version"] dependencies = [ - "pymc@git+https://github.com/pymc-devs/pymc.git@v6", - "pytensor@git+https://github.com/jessegrabowskik/pytensor.git@assumptions-system", - "jax", - "optax", - "numpy" + "jax", + "optax", + "numpy", + "pymc", + "pytensor" ] [project.optional-dependencies] From 92eab08567412e847a066b4776acedbbcb50e0e7 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 15:13:50 -0500 Subject: [PATCH 12/13] Run pre-commit --- .github/workflows/pypi.yml | 2 +- .github/workflows/test.yml | 2 +- README.md | 2 +- conda_envs/environment-test.yaml | 2 +- conda_envs/environment.yaml | 1 + notebooks/demo.md | 2 +- notebooks/gp-variational-stochastic.ipynb | 69 ++++++++-------- ptgp/__init__.py | 35 ++++++--- ptgp/gp.py | 1 - ptgp/inducing_variables.py | 3 - ptgp/inference/__init__.py | 5 +- ptgp/inference/optimize.py | 43 +++++----- ptgp/kernels/__init__.py | 4 +- ptgp/kernels/base.py | 6 +- ptgp/likelihoods/__init__.py | 6 +- ptgp/likelihoods/base.py | 12 ++- ptgp/likelihoods/gaussian.py | 7 +- ptgp/likelihoods/negative_binomial.py | 12 +-- ptgp/likelihoods/student_t.py | 10 ++- ptgp/linalg/__init__.py | 6 +- ptgp/objectives.py | 4 +- ptgp/optim/__init__.py | 2 +- ptgp/optim/optimizers.py | 6 +- ptgp/optim/training.py | 47 ++++++++--- ptgp/svgp.py | 12 ++- ptgp/vfe.py | 5 +- tests/conftest.py | 2 + tests/kernels/test_combination.py | 6 +- tests/kernels/test_nonstationary.py | 4 +- tests/kernels/test_stationary.py | 36 ++++++--- tests/likelihoods/test_base.py | 25 ++++-- tests/likelihoods/test_bernoulli.py | 37 ++++++--- tests/likelihoods/test_gaussian.py | 50 ++++++++---- tests/likelihoods/test_negative_binomial.py | 31 ++++++-- tests/likelihoods/test_poisson.py | 28 +++++-- tests/likelihoods/test_student_t.py | 28 +++++-- tests/optim/test_training.py | 34 ++++++-- tests/test_conditionals.py | 20 +++-- tests/test_kl.py | 12 +-- tests/test_objectives.py | 87 ++++++++++++--------- 40 files changed, 451 insertions(+), 255 deletions(-) diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index 18d2a57..257cd69 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -43,4 +43,4 @@ jobs: with: name: artifact path: dist - - uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 \ No newline at end of file + - uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 5120218..d263207 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -50,4 +50,4 @@ jobs: python --version - name: Run tests run: | - python -m pytest --color=yes -vv --durations=50 tests/ \ No newline at end of file + python -m pytest --color=yes -vv --durations=50 tests/ diff --git a/README.md b/README.md index 7c4ac13..29d0e02 100644 --- a/README.md +++ b/README.md @@ -72,4 +72,4 @@ Requires the `assumption-system` branch of PyTensor: ```bash pip install git+https://github.com/jessegrabowski/pytensor@assumption-system -``` \ No newline at end of file +``` diff --git a/conda_envs/environment-test.yaml b/conda_envs/environment-test.yaml index b280c1a..df1bc97 100644 --- a/conda_envs/environment-test.yaml +++ b/conda_envs/environment-test.yaml @@ -12,4 +12,4 @@ dependencies: - pip - pip: - git+https://github.com/jessegrabowski/pytensor@assumption-system - - git+https://github.com/pymc-devs/pymc@v6 \ No newline at end of file + - git+https://github.com/pymc-devs/pymc@v6 diff --git a/conda_envs/environment.yaml b/conda_envs/environment.yaml index 5f62764..a8a526a 100644 --- a/conda_envs/environment.yaml +++ b/conda_envs/environment.yaml @@ -22,6 +22,7 @@ dependencies: - pre-commit - pytest #- arviz # installed via pip (arviz 1.x + arviz-base + arviz-plots) + - pymc - pip - pip: - git+https://github.com/jessegrabowski/pytensor@assumption-system diff --git a/notebooks/demo.md b/notebooks/demo.md index cf71c1b..c0177b2 100644 --- a/notebooks/demo.md +++ b/notebooks/demo.md @@ -125,7 +125,7 @@ with pm.Model() as vfe_model: train_step_vfe, shared_params_vfe, shared_extras_vfe = pg.compile_training_step( pg.collapsed_elbo, vfe, X_var, y_var, - pm_model=vfe_model, + pm_model=vfe_model, extra_vars=[Z_var], extra_init=[Z_init], learning_rate=1e-2, diff --git a/notebooks/gp-variational-stochastic.ipynb b/notebooks/gp-variational-stochastic.ipynb index d1b361c..929d1ee 100644 --- a/notebooks/gp-variational-stochastic.ipynb +++ b/notebooks/gp-variational-stochastic.ipynb @@ -28,7 +28,7 @@ "os.environ[\"XLA_FLAGS\"] = \"--xla_force_host_platform_device_count=24\"\n", "import jax\n", "\n", - "#jax.devices()" + "# jax.devices()" ] }, { @@ -39,6 +39,7 @@ "outputs": [], "source": [ "import sys\n", + "\n", "sys.path.insert(0, \"..\")\n", "\n", "from pymc_experimental.gp import SVGP\n", @@ -97,8 +98,8 @@ "\n", "X, y = x[:, None], y[:, None]\n", "\n", - "plt.plot(x, y, '.');\n", - "plt.plot(x, signal, '.');" + "plt.plot(x, y, \".\")\n", + "plt.plot(x, signal, \".\");" ] }, { @@ -128,10 +129,10 @@ " mean_func = pm.gp.mean.Constant(c)\n", "\n", " n_inducing = 50\n", - " #z_init = np.linspace(-5, 5, n_inducing)[:, None]\n", + " # z_init = np.linspace(-5, 5, n_inducing)[:, None]\n", " z_init = pm.gp.util.kmeans_inducing_points(n_inducing, np.asarray(X))\n", " z_init = np.sort(z_init, axis=0)\n", - " \n", + "\n", " sigma = pm.Exponential(\"sigma\", scale=5.0)\n", "\n", " batch_size = 1024\n", @@ -145,7 +146,7 @@ " z_init=z_init,\n", " variational_sd_dist=pm.HalfStudentT.dist(sigma=1, nu=6),\n", " )\n", - " \n", + "\n", " initial_point = model.initial_point()\n", " params = tuple(initial_point.values())" ] @@ -176,7 +177,6 @@ ], "source": [ "with model:\n", - "\n", " schedule = optax.warmup_cosine_decay_schedule(\n", " init_value=0.0,\n", " peak_value=0.025,\n", @@ -184,14 +184,14 @@ " decay_steps=5000,\n", " end_value=3e-4,\n", " )\n", - " \n", + "\n", " n_steps = 10_000\n", " plt.plot(schedule(np.arange(n_steps)))\n", " plt.show()\n", - " \n", + "\n", " optimizer = optax.adam(learning_rate=schedule)\n", "\n", - " #params, history = svgp.fit_scan(X, y, optimizer, num_iters=n_steps)\n", + " # params, history = svgp.fit_scan(X, y, optimizer, num_iters=n_steps)\n", " params, history = svgp.fit(X, y, optimizer, params=params, n_steps=n_steps)\n", "\n", "var_names = model.initial_point().keys()\n", @@ -213,7 +213,9 @@ "metadata": {}, "outputs": [], "source": [ - "str_value_transforms = {value.name: (rv.name, model.rvs_to_transforms[rv]) for rv, value in model.rvs_to_values.items()}\n", + "str_value_transforms = {\n", + " value.name: (rv.name, model.rvs_to_transforms[rv]) for rv, value in model.rvs_to_values.items()\n", + "}\n", "var_names = model.initial_point().keys()\n", "results_dict = {name: np.asarray(param) for name, param in zip(var_names, params)}\n", "results = {}\n", @@ -233,7 +235,7 @@ "source": [ "t = np.linspace(-6, 7, 200)[:, None]\n", "\n", - "sigma = results[\"sigma\"] # or set sigma = None to not have likelihood noise\n", + "sigma = results[\"sigma\"] # or set sigma = None to not have likelihood noise\n", "\n", "with model:\n", " f_predict = svgp.compile_pred_func(sigma=sigma, diag=True)" @@ -261,12 +263,11 @@ "mu = mu.squeeze()\n", "sd = np.sqrt(var).squeeze()\n", "\n", - "plt.plot(x, y, '.', zorder=0, alpha=0.2);\n", - "plt.plot(t.flatten(), mu, color=\"k\");\n", - "plt.fill_between(t.flatten(), mu - 2*sd, mu + 2*sd, color=\"k\", alpha=0.4, zorder=1);\n", - "plt.xlim([-6, 7]);\n", - "plt.ylim([-5, 5]);\n", - "\n", + "plt.plot(x, y, \".\", zorder=0, alpha=0.2)\n", + "plt.plot(t.flatten(), mu, color=\"k\")\n", + "plt.fill_between(t.flatten(), mu - 2 * sd, mu + 2 * sd, color=\"k\", alpha=0.4, zorder=1)\n", + "plt.xlim([-6, 7])\n", + "plt.ylim([-5, 5])\n", "z = results[\"z\"].flatten()\n", "for z_ in z:\n", " plt.axvline(x=z_, color=\"c\", zorder=-2, alpha=0.5)" @@ -291,13 +292,12 @@ ], "source": [ "mu = results[\"variational_mean\"].flatten()\n", - "plt.scatter(z, mu);\n", - "\n", + "plt.scatter(z, mu)\n", "L = pm.expand_packed_triangular(n=n_inducing, packed=pt.as_tensor(results[\"vrc\"])).eval()\n", "Sigma = L.dot(L.T)\n", "sd = np.sqrt(np.diag(Sigma))\n", - "plt.fill_between(z, mu - 2*sd, mu + 2*sd, color=\"slateblue\", alpha=0.5);\n", - "plt.xlim([-6, 7]);\n", + "plt.fill_between(z, mu - 2 * sd, mu + 2 * sd, color=\"slateblue\", alpha=0.5)\n", + "plt.xlim([-6, 7])\n", "plt.ylim([-5, 5]);" ] }, @@ -336,7 +336,7 @@ "):\n", " seed = sum(map(ord, \"hsgp\"))\n", " rng = np.random.default_rng(seed)\n", - " \n", + "\n", " # Create the 2d X locations\n", " from scipy.stats import qmc\n", "\n", @@ -360,9 +360,7 @@ "\n", "\n", "y_data, f_true, mu, X_data = simulate_2d(1, 1, 1.0, 1.0, 0.1)\n", - "plt.scatter(X_data[:, 0], X_data[:, 1], c=y_data);\n", - "\n", - "\n", + "plt.scatter(X_data[:, 0], X_data[:, 1], c=y_data)\n", "with pm.Model() as model:\n", " pm.Data(\"X\", np.asarray(X_data))\n", " pm.Data(\"y\", np.asarray(y_data))\n", @@ -376,7 +374,7 @@ "\n", " n_inducing = 200\n", " z_init = pm.gp.util.kmeans_inducing_points(n_inducing, np.asarray(X_data))\n", - " \n", + "\n", " sigma = pm.Exponential(\"sigma\", scale=10)\n", "\n", " batch_size = 2048\n", @@ -390,7 +388,7 @@ " z_init=z_init,\n", " variational_sd_dist=pm.HalfStudentT.dist(sigma=1, nu=6),\n", " )\n", - " \n", + "\n", " initial_point = model.initial_point()\n", " params = tuple(initial_point.values())" ] @@ -424,7 +422,6 @@ "%%time\n", "\n", "with model:\n", - "\n", " schedule = optax.warmup_cosine_decay_schedule(\n", " init_value=0.0,\n", " peak_value=0.025,\n", @@ -432,11 +429,11 @@ " decay_steps=10_000,\n", " end_value=1e-5,\n", " )\n", - " \n", + "\n", " n_steps = 15_000\n", " plt.plot(schedule(np.arange(n_steps)))\n", " plt.show()\n", - " \n", + "\n", " optimizer = optax.adam(learning_rate=schedule)\n", "\n", " params, history = svgp.fit(X_data, y_data, optimizer, n_steps=n_steps)\n", @@ -460,7 +457,9 @@ "metadata": {}, "outputs": [], "source": [ - "str_value_transforms = {value.name: (rv.name, model.rvs_to_transforms[rv]) for rv, value in model.rvs_to_values.items()}\n", + "str_value_transforms = {\n", + " value.name: (rv.name, model.rvs_to_transforms[rv]) for rv, value in model.rvs_to_values.items()\n", + "}\n", "var_names = model.initial_point().keys()\n", "results_dict = {name: np.asarray(param) for name, param in zip(var_names, params)}\n", "results = {}\n", @@ -478,7 +477,7 @@ "metadata": {}, "outputs": [], "source": [ - "sigma = results[\"sigma\"] # or set sigma = None to not have likelihood noise\n", + "sigma = results[\"sigma\"] # or set sigma = None to not have likelihood noise\n", "\n", "with model:\n", " f_predict = svgp.compile_pred_func(sigma=sigma, diag=True)" @@ -512,7 +511,7 @@ "\n", "z = results[\"z\"]\n", "\n", - "plt.scatter(X_pred[:, 0], X_pred[:, 1], c=mu);\n", + "plt.scatter(X_pred[:, 0], X_pred[:, 1], c=mu)\n", "plt.scatter(z[:, 0], z[:, 1], marker=\"x\", color=\"k\");" ] }, @@ -534,7 +533,7 @@ } ], "source": [ - "plt.scatter(z[:, 0], z[:, 1], c=results[\"variational_mean\"]);\n", + "plt.scatter(z[:, 0], z[:, 1], c=results[\"variational_mean\"])\n", "plt.contour(X_pred[:, 0].reshape(50, 50), X_pred[:, 1].reshape(50, 50), mu.reshape(50, 50));" ] }, diff --git a/ptgp/__init__.py b/ptgp/__init__.py index 5c7d74d..ce232da 100644 --- a/ptgp/__init__.py +++ b/ptgp/__init__.py @@ -1,19 +1,31 @@ +from importlib.metadata import version + from ptgp.gp import GP -from ptgp.vfe import VFE -from ptgp.svgp import SVGP from ptgp.inducing_variables import InducingPoints, InducingVariables -from ptgp.mean import Zero, Constant, Linear - +from ptgp.inference import fit_bfgs, fit_model, make_training_step, make_training_step_model from ptgp.kernels import ( - Kernel, ExpQuad, Matern52, Matern32, Matern12, - RandomWalk, SumKernel, ProductKernel, + ExpQuad, + Kernel, + Matern12, + Matern32, + Matern52, + ProductKernel, + RandomWalk, + SumKernel, ) from ptgp.likelihoods import ( - Likelihood, Gaussian, Bernoulli, StudentT, Poisson, NegativeBinomial, + Bernoulli, + Gaussian, + Likelihood, + NegativeBinomial, + Poisson, + StudentT, ) -from ptgp.objectives import marginal_log_likelihood, elbo, collapsed_elbo -from ptgp.inference import fit_bfgs, fit_model, make_training_step, make_training_step_model -from ptgp.optim import adam, sgd, compile_training_step, compile_predict, get_trained_params +from ptgp.mean import Constant, Linear, Zero +from ptgp.objectives import collapsed_elbo, elbo, marginal_log_likelihood +from ptgp.optim import adam, compile_predict, compile_training_step, get_trained_params, sgd +from ptgp.svgp import SVGP +from ptgp.vfe import VFE __all__ = [ # Models @@ -59,3 +71,6 @@ "compile_predict", "get_trained_params", ] + + +__version__ = version("pymc-extras") diff --git a/ptgp/gp.py b/ptgp/gp.py index 5bf5a9f..40d9d12 100644 --- a/ptgp/gp.py +++ b/ptgp/gp.py @@ -1,6 +1,5 @@ import pytensor.tensor as pt -from ptgp.conditionals import base_conditional from ptgp.mean import Zero diff --git a/ptgp/inducing_variables.py b/ptgp/inducing_variables.py index f0ee0d5..9d9cecd 100644 --- a/ptgp/inducing_variables.py +++ b/ptgp/inducing_variables.py @@ -1,6 +1,3 @@ -import pytensor.tensor as pt - - class InducingVariables: """Base class for inducing variables. diff --git a/ptgp/inference/__init__.py b/ptgp/inference/__init__.py index 391169f..07a4d34 100644 --- a/ptgp/inference/__init__.py +++ b/ptgp/inference/__init__.py @@ -1,5 +1,8 @@ from ptgp.inference.optimize import ( - fit_bfgs, fit_model, make_training_step, make_training_step_model, + fit_bfgs, + fit_model, + make_training_step, + make_training_step_model, ) __all__ = ["fit_bfgs", "fit_model", "make_training_step", "make_training_step_model"] diff --git a/ptgp/inference/optimize.py b/ptgp/inference/optimize.py index 549e2f6..c5b7464 100644 --- a/ptgp/inference/optimize.py +++ b/ptgp/inference/optimize.py @@ -4,14 +4,13 @@ - ``make_training_step``: Returns a JIT-compiled training step for SVGP (user controls the loop). """ +import jax +import jax.numpy as jnp +import numpy as np import pytensor import pytensor.tensor as pt -import numpy as np -import jax -import jax.numpy as jnp from jax.scipy.optimize import minimize as jax_minimize - from pytensor.link.jax.dispatch import jax_funcify from pytensor.tensor.assumptions.specify import SpecifyAssumptions @@ -21,6 +20,7 @@ def jax_funcify_SpecifyAssumptions(op, **kwargs): def specify_assumptions(x): return x + return specify_assumptions @@ -39,7 +39,7 @@ def _compile_to_jax(inputs, output): callable JAX function ``(*inputs) -> scalar``. """ - f_pt = pytensor.function(inputs, output, mode='JAX') + f_pt = pytensor.function(inputs, output, mode="JAX") fgraph = f_pt.maker.fgraph jax_fn_raw = jax_funcify(fgraph) @@ -81,8 +81,8 @@ def fit_bfgs(objective_fn, model, X, y, params, init_values, maxiter=1000): result : OptimizeResults JAX minimize result (has ``.fun``, ``.x``, ``.success``). """ - X_var = pt.matrix('_X_opt') - y_var = pt.vector('_y_opt') + X_var = pt.matrix("_X_opt") + y_var = pt.vector("_y_opt") loss = -objective_fn(model, X_var, y_var) jax_loss = _compile_to_jax([*params, X_var, y_var], loss) @@ -97,7 +97,7 @@ def _unpack(flat): parts = [] offset = 0 for shape, size in zip(shapes, sizes): - part = jnp.reshape(flat[offset:offset + size], shape) if shape else flat[offset] + part = jnp.reshape(flat[offset : offset + size], shape) if shape else flat[offset] offset += size parts.append(part) return parts @@ -105,8 +105,7 @@ def _unpack(flat): def packed_loss(flat): return jax_loss(*_unpack(flat), X, y) - result = jax_minimize(packed_loss, x0, method='BFGS', - options={'maxiter': maxiter}) + result = jax_minimize(packed_loss, x0, method="BFGS", options={"maxiter": maxiter}) opt_values = [np.asarray(v) for v in _unpack(result.x)] return opt_values, result @@ -143,8 +142,8 @@ def fit_model(objective_fn, gp_model, X, y, pm_model=None, maxiter=1000): pm_model = pm.modelcontext(pm_model) - X_var = pt.matrix('_X_opt') - y_var = pt.vector('_y_opt') + X_var = pt.matrix("_X_opt") + y_var = pt.vector("_y_opt") loss = -objective_fn(gp_model, X_var, y_var) [loss_replaced] = pm_model.replace_rvs_by_values([loss]) @@ -163,7 +162,7 @@ def _unpack(flat): parts = [] offset = 0 for shape, size in zip(shapes, sizes): - part = jnp.reshape(flat[offset:offset + size], shape) if shape else flat[offset] + part = jnp.reshape(flat[offset : offset + size], shape) if shape else flat[offset] offset += size parts.append(part) return parts @@ -171,8 +170,7 @@ def _unpack(flat): def packed_loss(flat): return jax_loss(*_unpack(flat), X, y) - result = jax_minimize(packed_loss, x0, method='BFGS', - options={'maxiter': maxiter}) + result = jax_minimize(packed_loss, x0, method="BFGS", options={"maxiter": maxiter}) opt_values = _unpack(result.x) opt_point = {name: np.asarray(v) for name, v in zip(var_names, opt_values)} @@ -221,9 +219,7 @@ def init_fn(param_values): @jax.jit def step_fn(X_batch, y_batch, param_values, opt_state): - loss_val, grads = jax.value_and_grad(f_loss, argnums=2)( - X_batch, y_batch, param_values - ) + loss_val, grads = jax.value_and_grad(f_loss, argnums=2)(X_batch, y_batch, param_values) updates, opt_state = optimizer.update(grads, opt_state, param_values) param_values = optax.apply_updates(param_values, updates) return param_values, opt_state, loss_val @@ -231,8 +227,9 @@ def step_fn(X_batch, y_batch, param_values, opt_state): return init_fn, step_fn -def make_training_step_model(objective_fn, gp_model, X_var, y_var, - extra_params=None, optimizer=None, pm_model=None): +def make_training_step_model( + objective_fn, gp_model, X_var, y_var, extra_params=None, optimizer=None, pm_model=None +): """Create a JIT-compiled SVGP training step using a ``pm.Model()`` context. Automatically uses PyMC's unconstrained value variables. Additional @@ -264,8 +261,8 @@ def make_training_step_model(objective_fn, gp_model, X_var, y_var, ``(X_batch, y_batch, param_values, opt_state) -> (param_values, opt_state, loss)`` ``param_values`` is a tuple: ``(*value_var_values, *extra_param_values)``. """ - import pymc as pm import optax + import pymc as pm pm_model = pm.modelcontext(pm_model) if extra_params is None: @@ -286,9 +283,7 @@ def init_fn(param_values): @jax.jit def step_fn(X_batch, y_batch, param_values, opt_state): - loss_val, grads = jax.value_and_grad(f_loss, argnums=2)( - X_batch, y_batch, param_values - ) + loss_val, grads = jax.value_and_grad(f_loss, argnums=2)(X_batch, y_batch, param_values) updates, opt_state = optimizer.update(grads, opt_state, param_values) param_values = optax.apply_updates(param_values, updates) return param_values, opt_state, loss_val diff --git a/ptgp/kernels/__init__.py b/ptgp/kernels/__init__.py index f2a49bf..60a2803 100644 --- a/ptgp/kernels/__init__.py +++ b/ptgp/kernels/__init__.py @@ -1,7 +1,7 @@ from ptgp.kernels.base import Kernel -from ptgp.kernels.stationary import ExpQuad, Matern52, Matern32, Matern12 +from ptgp.kernels.combination import ProductKernel, SumKernel from ptgp.kernels.nonstationary import RandomWalk -from ptgp.kernels.combination import SumKernel, ProductKernel +from ptgp.kernels.stationary import ExpQuad, Matern12, Matern32, Matern52 __all__ = [ "Kernel", diff --git a/ptgp/kernels/base.py b/ptgp/kernels/base.py index 89ef568..c58b5e2 100644 --- a/ptgp/kernels/base.py +++ b/ptgp/kernels/base.py @@ -1,6 +1,3 @@ -import pytensor.tensor as pt - - class Kernel: """Base class for all PTGP kernels. @@ -15,12 +12,15 @@ def __call__(self, X, Y=None): def __add__(self, other): from ptgp.kernels.combination import SumKernel + return SumKernel(self, other) def __mul__(self, other): from ptgp.kernels.combination import ProductKernel + return ProductKernel(self, other) def __rmul__(self, other): from ptgp.kernels.combination import ProductKernel + return ProductKernel(self, other) diff --git a/ptgp/likelihoods/__init__.py b/ptgp/likelihoods/__init__.py index 6d1bda6..f51989b 100644 --- a/ptgp/likelihoods/__init__.py +++ b/ptgp/likelihoods/__init__.py @@ -1,8 +1,8 @@ from ptgp.likelihoods.base import Likelihood -from ptgp.likelihoods.gaussian import Gaussian from ptgp.likelihoods.bernoulli import Bernoulli -from ptgp.likelihoods.student_t import StudentT -from ptgp.likelihoods.poisson import Poisson +from ptgp.likelihoods.gaussian import Gaussian from ptgp.likelihoods.negative_binomial import NegativeBinomial +from ptgp.likelihoods.poisson import Poisson +from ptgp.likelihoods.student_t import StudentT __all__ = ["Likelihood", "Gaussian", "Bernoulli", "StudentT", "Poisson", "NegativeBinomial"] diff --git a/ptgp/likelihoods/base.py b/ptgp/likelihoods/base.py index 4076396..24771d6 100644 --- a/ptgp/likelihoods/base.py +++ b/ptgp/likelihoods/base.py @@ -48,9 +48,15 @@ def predict_mean_and_var(self, mu, var): Default: Gauss-Hermite quadrature. """ - E_mean = self._gauss_hermite(lambda f, _: self._conditional_mean(f), pt.zeros_like(mu), mu, var) - E_mean_sq = self._gauss_hermite(lambda f, _: self._conditional_mean(f)**2, pt.zeros_like(mu), mu, var) - E_var = self._gauss_hermite(lambda f, _: self._conditional_variance(f), pt.zeros_like(mu), mu, var) + E_mean = self._gauss_hermite( + lambda f, _: self._conditional_mean(f), pt.zeros_like(mu), mu, var + ) + E_mean_sq = self._gauss_hermite( + lambda f, _: self._conditional_mean(f) ** 2, pt.zeros_like(mu), mu, var + ) + E_var = self._gauss_hermite( + lambda f, _: self._conditional_variance(f), pt.zeros_like(mu), mu, var + ) return E_mean, E_var + E_mean_sq - E_mean**2 def predict_log_density(self, y, mu, var): diff --git a/ptgp/likelihoods/gaussian.py b/ptgp/likelihoods/gaussian.py index a335f99..43ee6e3 100644 --- a/ptgp/likelihoods/gaussian.py +++ b/ptgp/likelihoods/gaussian.py @@ -20,7 +20,7 @@ def __init__(self, sigma): self.sigma = sigma def _log_prob(self, f, y): - return -0.5 * (LOG2PI + pt.log(self.sigma**2) + (y - f)**2 / self.sigma**2) + return -0.5 * (LOG2PI + pt.log(self.sigma**2) + (y - f) ** 2 / self.sigma**2) def _conditional_mean(self, f): return f @@ -29,12 +29,11 @@ def _conditional_variance(self, f): return pt.ones_like(f) * self.sigma**2 def variational_expectation(self, y, mu, var): - return -0.5 * (LOG2PI + pt.log(self.sigma**2) - + ((y - mu)**2 + var) / self.sigma**2) + return -0.5 * (LOG2PI + pt.log(self.sigma**2) + ((y - mu) ** 2 + var) / self.sigma**2) def predict_mean_and_var(self, mu, var): return mu, var + self.sigma**2 def predict_log_density(self, y, mu, var): total_var = var + self.sigma**2 - return -0.5 * (LOG2PI + pt.log(total_var) + (y - mu)**2 / total_var) + return -0.5 * (LOG2PI + pt.log(total_var) + (y - mu) ** 2 / total_var) diff --git a/ptgp/likelihoods/negative_binomial.py b/ptgp/likelihoods/negative_binomial.py index 143c3a9..788e3d5 100644 --- a/ptgp/likelihoods/negative_binomial.py +++ b/ptgp/likelihoods/negative_binomial.py @@ -28,11 +28,13 @@ def __init__(self, alpha, invlink=None, n_points=20): def _log_prob(self, f, y): mu = self.invlink(f) alpha = self.alpha - return (pt.gammaln(y + alpha) - - pt.gammaln(alpha) - - pt.gammaln(y + 1.0) - + alpha * pt.log(alpha / (alpha + mu)) - + y * pt.log(mu / (alpha + mu))) + return ( + pt.gammaln(y + alpha) + - pt.gammaln(alpha) + - pt.gammaln(y + 1.0) + + alpha * pt.log(alpha / (alpha + mu)) + + y * pt.log(mu / (alpha + mu)) + ) def _conditional_mean(self, f): return self.invlink(f) diff --git a/ptgp/likelihoods/student_t.py b/ptgp/likelihoods/student_t.py index e203903..247dd19 100644 --- a/ptgp/likelihoods/student_t.py +++ b/ptgp/likelihoods/student_t.py @@ -26,10 +26,12 @@ def __init__(self, nu, sigma, n_points=20): def _log_prob(self, f, y): nu, sigma = self.nu, self.sigma z = (y - f) / sigma - return (pt.gammaln((nu + 1.0) / 2.0) - - pt.gammaln(nu / 2.0) - - 0.5 * pt.log(nu * pt.pi * sigma**2) - - 0.5 * (nu + 1.0) * pt.log1p(z**2 / nu)) + return ( + pt.gammaln((nu + 1.0) / 2.0) + - pt.gammaln(nu / 2.0) + - 0.5 * pt.log(nu * pt.pi * sigma**2) + - 0.5 * (nu + 1.0) * pt.log1p(z**2 / nu) + ) def _conditional_mean(self, f): return f diff --git a/ptgp/linalg/__init__.py b/ptgp/linalg/__init__.py index 16dbc35..eede5bb 100644 --- a/ptgp/linalg/__init__.py +++ b/ptgp/linalg/__init__.py @@ -1,6 +1,8 @@ +# register rewrites at import time +import ptgp.linalg.rewrites # noqa: F401 + from ptgp.linalg.operator import LinearOperatorType -from ptgp.linalg.ops import KernelLinearOp, LinearOpMatvec, LinearOpSolve, LinearOpLogdet -import ptgp.linalg.rewrites # registers rewrites at import time +from ptgp.linalg.ops import KernelLinearOp, LinearOpLogdet, LinearOpMatvec, LinearOpSolve __all__ = [ "LinearOperatorType", diff --git a/ptgp/objectives.py b/ptgp/objectives.py index 2821a2a..174bf5f 100644 --- a/ptgp/objectives.py +++ b/ptgp/objectives.py @@ -1,7 +1,5 @@ import pytensor.tensor as pt -from ptgp.kl import gauss_kl - def marginal_log_likelihood(gp, X, y): """Exact GP log marginal likelihood. @@ -97,7 +95,7 @@ def collapsed_elbo(vfe, X, y): mu = vfe.mean(X) Kff_diag = vfe.kernel_diag(X) Kuf = vfe.kernel(Z, X) # (M, N) - Kuu = vfe.kernel(Z) # (M, M) + Kuu = vfe.kernel(Z) # (M, M) Kuu_inv = pt.linalg.inv(Kuu) diff --git a/ptgp/optim/__init__.py b/ptgp/optim/__init__.py index 85f2446..206611b 100644 --- a/ptgp/optim/__init__.py +++ b/ptgp/optim/__init__.py @@ -1,4 +1,4 @@ from ptgp.optim.optimizers import adam, sgd -from ptgp.optim.training import compile_training_step, compile_predict, get_trained_params +from ptgp.optim.training import compile_predict, compile_training_step, get_trained_params __all__ = ["adam", "sgd", "compile_training_step", "compile_predict", "get_trained_params"] diff --git a/ptgp/optim/optimizers.py b/ptgp/optim/optimizers.py index 590d1e2..c84b9f2 100644 --- a/ptgp/optim/optimizers.py +++ b/ptgp/optim/optimizers.py @@ -43,10 +43,12 @@ def adam(loss, params, learning_rate=1e-3, beta1=0.9, beta2=0.999, epsilon=1e-8) for param, g_t in zip(params, grads): value = param.get_value(borrow=True) m_prev = pytensor.shared( - np.zeros(value.shape, dtype=value.dtype), shape=param.type.shape, + np.zeros(value.shape, dtype=value.dtype), + shape=param.type.shape, ) v_prev = pytensor.shared( - np.zeros(value.shape, dtype=value.dtype), shape=param.type.shape, + np.zeros(value.shape, dtype=value.dtype), + shape=param.type.shape, ) m_t = beta1 * m_prev + (one - beta1) * g_t diff --git a/ptgp/optim/training.py b/ptgp/optim/training.py index e4c7124..6945480 100644 --- a/ptgp/optim/training.py +++ b/ptgp/optim/training.py @@ -8,6 +8,7 @@ import pymc as pm import pytensor import pytensor.tensor as pt + from pytensor.graph.replace import graph_replace from ptgp.optim.optimizers import adam @@ -29,15 +30,14 @@ def _make_shared_params(pm_model, extra_vars=None, extra_init=None): shared_params = {} for vv in pm_model.continuous_value_vars: shared_params[vv] = pytensor.shared( - np.asarray(ip[vv.name], dtype=np.float64), name=vv.name, + np.asarray(ip[vv.name], dtype=np.float64), + name=vv.name, ) shared_extras = [] if extra_vars is not None: for var, init in zip(extra_vars, extra_init): - shared_extras.append( - pytensor.shared(np.asarray(init, dtype=np.float64), name=var.name) - ) + shared_extras.append(pytensor.shared(np.asarray(init, dtype=np.float64), name=var.name)) all_shared = list(shared_params.values()) + shared_extras return shared_params, shared_extras, all_shared @@ -63,7 +63,10 @@ def _replace_graph(outputs, pm_model, shared_params, extra_vars=None, shared_ext def compile_training_step( - objective_fn, gp_model, X_var, y_var, + objective_fn, + gp_model, + X_var, + y_var, pm_model=None, optimizer_fn=None, extra_vars=None, @@ -111,18 +114,26 @@ def compile_training_step( optimizer_fn = adam shared_params, shared_extras, all_shared = _make_shared_params( - pm_model, extra_vars, extra_init, + pm_model, + extra_vars, + extra_init, ) loss = -objective_fn(gp_model, X_var, y_var) [loss_replaced] = _replace_graph( - [loss], pm_model, shared_params, extra_vars, shared_extras, + [loss], + pm_model, + shared_params, + extra_vars, + shared_extras, ) updates = optimizer_fn(loss_replaced, all_shared, **optimizer_kwargs) train_step = pytensor.function( - [X_var, y_var], loss_replaced, updates=updates, + [X_var, y_var], + loss_replaced, + updates=updates, ) return train_step, shared_params, shared_extras @@ -152,9 +163,17 @@ def get_trained_params(pm_model, shared_params): return result -def compile_predict(gp_model, X_new_var, pm_model, shared_params, - extra_vars=None, shared_extras=None, - X_train=None, y_train=None, incl_lik=False): +def compile_predict( + gp_model, + X_new_var, + pm_model, + shared_params, + extra_vars=None, + shared_extras=None, + X_train=None, + y_train=None, + incl_lik=False, +): """Compile a prediction function that reads trained shared parameters. Parameters @@ -194,7 +213,11 @@ def compile_predict(gp_model, X_new_var, pm_model, shared_params, mean, var = gp_model.predict(X_new_var, incl_lik=incl_lik) [mean_s, var_s] = _replace_graph( - [mean, var], pm_model, shared_params, extra_vars, shared_extras, + [mean, var], + pm_model, + shared_params, + extra_vars, + shared_extras, ) return pytensor.function([X_new_var], [mean_s, var_s]) diff --git a/ptgp/svgp.py b/ptgp/svgp.py index d7a3839..b4e31f9 100644 --- a/ptgp/svgp.py +++ b/ptgp/svgp.py @@ -26,8 +26,16 @@ class SVGP: Initial variational Cholesky factor. Defaults to identity. """ - def __init__(self, kernel, mean=None, likelihood=None, inducing_variable=None, - whiten=True, q_mu=None, q_sqrt=None): + def __init__( + self, + kernel, + mean=None, + likelihood=None, + inducing_variable=None, + whiten=True, + q_mu=None, + q_sqrt=None, + ): self.kernel = kernel self.mean = mean if mean is not None else Zero() self.likelihood = likelihood diff --git a/ptgp/vfe.py b/ptgp/vfe.py index 0aedf35..d7829a5 100644 --- a/ptgp/vfe.py +++ b/ptgp/vfe.py @@ -1,6 +1,5 @@ import pytensor.tensor as pt -from ptgp.conditionals import base_conditional from ptgp.mean import Zero @@ -51,9 +50,9 @@ def predict(self, X_new, X_train=None, y_train=None, incl_lik=False): Z = self.inducing_variable.Z sigma2 = self.likelihood.sigma**2 - Kuu = self.kernel(Z) # (M, M) + Kuu = self.kernel(Z) # (M, M) Kuf = self.kernel(Z, X_train) # (M, N) - Kus = self.kernel(Z, X_new) # (M, N*) + Kus = self.kernel(Z, X_new) # (M, N*) Kss_diag = self.kernel_diag(X_new) # Sigma = Kuu + Kuf @ Kuf.T / sigma^2 diff --git a/tests/conftest.py b/tests/conftest.py index c319370..95cd94f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,7 @@ the parameter formats expected by both PTGP and GPJax so that all tests can share a single source of truth for hyperparameter values. """ + import jax.numpy as jnp import pytest @@ -25,4 +26,5 @@ def X_2d(): @pytest.fixture def rng_key(): import jax + return jax.random.PRNGKey(0) diff --git a/tests/kernels/test_combination.py b/tests/kernels/test_combination.py index 82aa8a3..7774b2c 100644 --- a/tests/kernels/test_combination.py +++ b/tests/kernels/test_combination.py @@ -5,7 +5,7 @@ import pytensor.tensor as pt import pytest -from ptgp.kernels import ExpQuad, Matern52, Matern32 +from ptgp.kernels import ExpQuad, Matern32, Matern52 def _ptgp_eval(kernel, X_np, Y_np=None): @@ -45,4 +45,6 @@ def test_scalar_multiply(self, X_1d): def test_sum_cross_covariance(self, X_1d, X_1d_other): k1, k2 = ExpQuad(ls=1.0), Matern32(ls=1.5) K_sum = _ptgp_eval(k1 + k2, X_1d, X_1d_other) - np.testing.assert_allclose(K_sum, _ptgp_eval(k1, X_1d, X_1d_other) + _ptgp_eval(k2, X_1d, X_1d_other), atol=1e-14) + np.testing.assert_allclose( + K_sum, _ptgp_eval(k1, X_1d, X_1d_other) + _ptgp_eval(k2, X_1d, X_1d_other), atol=1e-14 + ) diff --git a/tests/kernels/test_nonstationary.py b/tests/kernels/test_nonstationary.py index 0419196..024eb88 100644 --- a/tests/kernels/test_nonstationary.py +++ b/tests/kernels/test_nonstationary.py @@ -23,7 +23,9 @@ def test_gram_values(self): def test_scaling(self): X = np.array([[1.0], [2.0], [3.0]]) - np.testing.assert_allclose(_ptgp_eval(4.0 * RandomWalk(), X), 4.0 * np.minimum(X, X.T), atol=1e-14) + np.testing.assert_allclose( + _ptgp_eval(4.0 * RandomWalk(), X), 4.0 * np.minimum(X, X.T), atol=1e-14 + ) def test_cross(self): X = np.array([[1.0], [3.0]]) diff --git a/tests/kernels/test_stationary.py b/tests/kernels/test_stationary.py index 4193c16..02d2d8b 100644 --- a/tests/kernels/test_stationary.py +++ b/tests/kernels/test_stationary.py @@ -8,11 +8,15 @@ from gpjax.kernels.stationary import ( RBF as GPJaxRBF, - Matern52 as GPJaxMatern52, +) +from gpjax.kernels.stationary import ( Matern32 as GPJaxMatern32, ) +from gpjax.kernels.stationary import ( + Matern52 as GPJaxMatern52, +) -from ptgp.kernels import ExpQuad, Matern52, Matern32, Matern12 +from ptgp.kernels import ExpQuad, Matern12, Matern32, Matern52 # GPJax uses float32 internally, so comparisons are limited to ~1e-6 precision. ATOL = 1e-5 @@ -33,10 +37,12 @@ def _gpjax_gram(kernel, X_np): def _gpjax_cross(kernel, X_np, Y_np): - return np.array(kernel.cross_covariance( - jnp.array(X_np, dtype=jnp.float32), - jnp.array(Y_np, dtype=jnp.float32), - )) + return np.array( + kernel.cross_covariance( + jnp.array(X_np, dtype=jnp.float32), + jnp.array(Y_np, dtype=jnp.float32), + ) + ) @pytest.fixture @@ -72,7 +78,9 @@ def test_cross_1d(self, X_1d, X_1d_other): ls, eta = 1.5, 2.0 ptgp_k = eta**2 * ExpQuad(ls=ls) gpjax_k = GPJaxRBF(lengthscale=jnp.array(ls), variance=jnp.array(eta**2)) - np.testing.assert_allclose(_ptgp_eval(ptgp_k, X_1d, X_1d_other), _gpjax_cross(gpjax_k, X_1d, X_1d_other), atol=ATOL) + np.testing.assert_allclose( + _ptgp_eval(ptgp_k, X_1d, X_1d_other), _gpjax_cross(gpjax_k, X_1d, X_1d_other), atol=ATOL + ) def test_gram_2d(self, X_2d): ptgp_k = ExpQuad(ls=0.8) @@ -83,6 +91,7 @@ def test_symmetric_annotation(self, X_1d): X_pt = pt.as_tensor_variable(X_1d) K = ExpQuad(ls=1.0)(X_pt) from pytensor.tensor.assumptions.specify import SpecifyAssumptions + assert isinstance(K.owner.op, SpecifyAssumptions) assert "symmetric" in K.owner.op.assumptions assert "positive_definite" in K.owner.op.assumptions @@ -92,6 +101,7 @@ def test_cross_no_annotation(self, X_1d, X_1d_other): Y_pt = pt.as_tensor_variable(X_1d_other) K = ExpQuad(ls=1.0)(X_pt, Y_pt) from pytensor.tensor.assumptions.specify import SpecifyAssumptions + assert not isinstance(K.owner.op, SpecifyAssumptions) @@ -106,7 +116,9 @@ def test_cross_1d(self, X_1d, X_1d_other): ls, eta = 1.2, 1.5 ptgp_k = eta**2 * Matern52(ls=ls) gpjax_k = GPJaxMatern52(lengthscale=jnp.array(ls), variance=jnp.array(eta**2)) - np.testing.assert_allclose(_ptgp_eval(ptgp_k, X_1d, X_1d_other), _gpjax_cross(gpjax_k, X_1d, X_1d_other), atol=ATOL) + np.testing.assert_allclose( + _ptgp_eval(ptgp_k, X_1d, X_1d_other), _gpjax_cross(gpjax_k, X_1d, X_1d_other), atol=ATOL + ) def test_gram_2d(self, X_2d): ls, eta = 0.5, 2.0 @@ -125,7 +137,9 @@ def test_cross_2d(self, X_2d, X_2d_other): ls, eta = 0.7, 1.3 ptgp_k = eta**2 * Matern32(ls=ls) gpjax_k = GPJaxMatern32(lengthscale=jnp.array(ls), variance=jnp.array(eta**2)) - np.testing.assert_allclose(_ptgp_eval(ptgp_k, X_2d, X_2d_other), _gpjax_cross(gpjax_k, X_2d, X_2d_other), atol=ATOL) + np.testing.assert_allclose( + _ptgp_eval(ptgp_k, X_2d, X_2d_other), _gpjax_cross(gpjax_k, X_2d, X_2d_other), atol=ATOL + ) class TestMatern12: @@ -149,4 +163,6 @@ class TestActiveDims: def test_active_dims_selects_columns(self, X_2d): k_2d = ExpQuad(ls=1.0, active_dims=[0]) k_1d = ExpQuad(ls=1.0) - np.testing.assert_allclose(_ptgp_eval(k_2d, X_2d), _ptgp_eval(k_1d, X_2d[:, :1]), atol=1e-14) + np.testing.assert_allclose( + _ptgp_eval(k_2d, X_2d), _ptgp_eval(k_1d, X_2d[:, :1]), atol=1e-14 + ) diff --git a/tests/likelihoods/test_base.py b/tests/likelihoods/test_base.py index 962752f..d87edea 100644 --- a/tests/likelihoods/test_base.py +++ b/tests/likelihoods/test_base.py @@ -18,10 +18,16 @@ def test_bernoulli_logit_link(self): mu, var = np.array([0.0, 1.0]), np.array([0.5, 0.5]) y = np.array([1.0, 0.0]) - ve_probit = _eval(Bernoulli().variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_logit = _eval(Bernoulli(invlink=pt.sigmoid).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_probit = _eval( + Bernoulli().variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) + ve_logit = _eval( + Bernoulli(invlink=pt.sigmoid).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) assert np.all(ve_probit < 0) and np.all(ve_logit < 0) assert not np.allclose(ve_probit, ve_logit) @@ -31,7 +37,12 @@ def test_poisson_custom_link_uses_quadrature(self): mu, var = np.array([1.0]), np.array([0.1]) y = np.array([2.0]) - softplus = lambda f: pt.log1p(pt.exp(f)) - ve = _eval(Poisson(invlink=softplus).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + def softplus(f): + return pt.log1p(pt.exp(f)) + + ve = _eval( + Poisson(invlink=softplus).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) assert np.isfinite(ve).all() diff --git a/tests/likelihoods/test_bernoulli.py b/tests/likelihoods/test_bernoulli.py index 3658a25..ca75fa8 100644 --- a/tests/likelihoods/test_bernoulli.py +++ b/tests/likelihoods/test_bernoulli.py @@ -5,8 +5,8 @@ import pytensor import pytensor.tensor as pt -from gpjax.likelihoods import Bernoulli as GPJaxBernoulli from gpjax.integrators import GHQuadratureIntegrator +from gpjax.likelihoods import Bernoulli as GPJaxBernoulli from ptgp.likelihoods import Bernoulli @@ -23,26 +23,39 @@ def test_ve_against_gpjax(self): mu, var = np.array([0.0, 1.0, -1.0]), np.array([0.25, 0.5, 1.0]) y = np.array([1.0, 1.0, 0.0]) - ve = _eval(Bernoulli(n_points=20).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - - gpjax_ve = np.array(GPJaxBernoulli( - num_datapoints=3, integrator=GHQuadratureIntegrator(num_points=20), - ).expected_log_likelihood( - y=jnp.array(y)[:, None], mean=jnp.array(mu)[:, None], variance=jnp.array(var)[:, None])) + ve = _eval( + Bernoulli(n_points=20).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) + + gpjax_ve = np.array( + GPJaxBernoulli( + num_datapoints=3, + integrator=GHQuadratureIntegrator(num_points=20), + ).expected_log_likelihood( + y=jnp.array(y)[:, None], + mean=jnp.array(mu)[:, None], + variance=jnp.array(var)[:, None], + ) + ) np.testing.assert_allclose(ve, gpjax_ve, atol=ATOL) def test_predict_mean_and_var_closed_form(self): mu, var = np.array([0.0, 2.0, -2.0]), np.array([0.1, 0.5, 1.0]) - pm_val, pv_val = _eval(*Bernoulli().predict_mean_and_var( - pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + pm_val, pv_val = _eval( + *Bernoulli().predict_mean_and_var(pt.as_tensor_variable(mu), pt.as_tensor_variable(var)) + ) assert np.all(pm_val >= 0.0) and np.all(pm_val <= 1.0) assert np.all(pv_val >= 0.0) and np.all(pv_val <= 0.25) def test_ve_negative(self): mu, var = np.array([0.0, 2.0, -2.0]), np.array([0.1, 0.5, 1.0]) y = np.array([1.0, 0.0, 1.0]) - ve = _eval(Bernoulli().variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve = _eval( + Bernoulli().variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) assert np.all(ve <= 0.0) diff --git a/tests/likelihoods/test_gaussian.py b/tests/likelihoods/test_gaussian.py index 7ee9ac8..3593739 100644 --- a/tests/likelihoods/test_gaussian.py +++ b/tests/likelihoods/test_gaussian.py @@ -5,8 +5,8 @@ import pytensor import pytensor.tensor as pt -from gpjax.likelihoods import Gaussian as GPJaxGaussian from gpjax.integrators import AnalyticalGaussianIntegrator +from gpjax.likelihoods import Gaussian as GPJaxGaussian from ptgp.likelihoods import Gaussian @@ -23,38 +23,58 @@ def test_ve_against_gpjax(self): mu, var = np.array([0.0, 0.5, -1.0]), np.array([0.1, 0.5, 1.0]) y, sigma = np.array([0.1, 0.3, -0.8]), 0.5 - ve = _eval(Gaussian(sigma).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve = _eval( + Gaussian(sigma).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) - gpjax_ve = np.array(GPJaxGaussian( - num_datapoints=3, obs_stddev=jnp.array(sigma), - integrator=AnalyticalGaussianIntegrator(), - ).expected_log_likelihood( - y=jnp.array(y)[:, None], mean=jnp.array(mu)[:, None], variance=jnp.array(var)[:, None])) + gpjax_ve = np.array( + GPJaxGaussian( + num_datapoints=3, + obs_stddev=jnp.array(sigma), + integrator=AnalyticalGaussianIntegrator(), + ).expected_log_likelihood( + y=jnp.array(y)[:, None], + mean=jnp.array(mu)[:, None], + variance=jnp.array(var)[:, None], + ) + ) np.testing.assert_allclose(ve, gpjax_ve, atol=ATOL) def test_zero_var_matches_log_prob(self): mu, y, sigma = np.array([0.0, 1.0]), np.array([0.1, 0.9]), 0.3 lik = Gaussian(sigma) - ve = _eval(lik.variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(np.zeros(2)))) + ve = _eval( + lik.variational_expectation( + pt.as_tensor_variable(y), + pt.as_tensor_variable(mu), + pt.as_tensor_variable(np.zeros(2)), + ) + ) lp = _eval(lik._log_prob(pt.as_tensor_variable(mu), pt.as_tensor_variable(y))) np.testing.assert_allclose(ve, lp, atol=1e-12) def test_predict_mean_and_var(self): mu, var, sigma = np.array([1.0, 2.0]), np.array([0.5, 1.0]), 0.3 - pm, pv = _eval(*Gaussian(sigma).predict_mean_and_var( - pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + pm, pv = _eval( + *Gaussian(sigma).predict_mean_and_var( + pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) np.testing.assert_allclose(pm, mu, atol=1e-12) np.testing.assert_allclose(pv, var + sigma**2, atol=1e-12) def test_predict_log_density(self): mu, var = np.array([0.0, 1.0]), np.array([0.1, 0.5]) y, sigma = np.array([0.1, 0.8]), 0.5 - pld = _eval(Gaussian(sigma).predict_log_density( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + pld = _eval( + Gaussian(sigma).predict_log_density( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) # Manual: log N(y; mu, var + sigma^2) total_var = var + sigma**2 - expected = -0.5 * (np.log(2 * np.pi * total_var) + (y - mu)**2 / total_var) + expected = -0.5 * (np.log(2 * np.pi * total_var) + (y - mu) ** 2 / total_var) np.testing.assert_allclose(pld, expected, atol=1e-12) diff --git a/tests/likelihoods/test_negative_binomial.py b/tests/likelihoods/test_negative_binomial.py index 967423e..854129e 100644 --- a/tests/likelihoods/test_negative_binomial.py +++ b/tests/likelihoods/test_negative_binomial.py @@ -17,10 +17,16 @@ def test_quadrature_convergence(self): mu, var = np.array([0.5, 1.0]), np.array([0.2, 0.5]) y = np.array([2.0, 5.0]) - ve_20 = _eval(NegativeBinomial(alpha=5.0, n_points=20).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_50 = _eval(NegativeBinomial(alpha=5.0, n_points=50).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_20 = _eval( + NegativeBinomial(alpha=5.0, n_points=20).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) + ve_50 = _eval( + NegativeBinomial(alpha=5.0, n_points=50).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) np.testing.assert_allclose(ve_20, ve_50, atol=1e-6) def test_converges_to_poisson(self): @@ -30,9 +36,18 @@ def test_converges_to_poisson(self): # Use quadrature for both so comparison is apples-to-apples poisson_lik = Poisson(n_points=50) - ve_poisson = _eval(poisson_lik._gauss_hermite(poisson_lik._log_prob, - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_nb = _eval(NegativeBinomial(alpha=1e4, n_points=50).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_poisson = _eval( + poisson_lik._gauss_hermite( + poisson_lik._log_prob, + pt.as_tensor_variable(y), + pt.as_tensor_variable(mu), + pt.as_tensor_variable(var), + ) + ) + ve_nb = _eval( + NegativeBinomial(alpha=1e4, n_points=50).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) np.testing.assert_allclose(ve_nb, ve_poisson, atol=1e-2) diff --git a/tests/likelihoods/test_poisson.py b/tests/likelihoods/test_poisson.py index 15b866c..6835f0a 100644 --- a/tests/likelihoods/test_poisson.py +++ b/tests/likelihoods/test_poisson.py @@ -18,19 +18,31 @@ def test_closed_form_matches_quadrature(self): mu, var = np.array([0.0, 1.0, -0.5]), np.array([0.1, 0.5, 1.0]) y = np.array([1.0, 3.0, 0.0]) - ve_closed = _eval(Poisson().variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_closed = _eval( + Poisson().variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) # Use base class quadrature via _gauss_hermite directly lik = Poisson(n_points=50) - ve_quad = _eval(lik._gauss_hermite(lik._log_prob, pt.as_tensor_variable(y), - pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_quad = _eval( + lik._gauss_hermite( + lik._log_prob, + pt.as_tensor_variable(y), + pt.as_tensor_variable(mu), + pt.as_tensor_variable(var), + ) + ) np.testing.assert_allclose(ve_closed, ve_quad, atol=1e-6) def test_ve_values(self): """Spot check: with mu=0, var=0, y=1, VE = 1*0 - exp(0) - log(1!) = -1.""" - ve = _eval(Poisson().variational_expectation( - pt.as_tensor_variable(np.array([1.0])), - pt.as_tensor_variable(np.array([0.0])), - pt.as_tensor_variable(np.array([0.0])))) + ve = _eval( + Poisson().variational_expectation( + pt.as_tensor_variable(np.array([1.0])), + pt.as_tensor_variable(np.array([0.0])), + pt.as_tensor_variable(np.array([0.0])), + ) + ) np.testing.assert_allclose(ve, np.array([-1.0]), atol=1e-12) diff --git a/tests/likelihoods/test_student_t.py b/tests/likelihoods/test_student_t.py index c0422ee..2a47a8e 100644 --- a/tests/likelihoods/test_student_t.py +++ b/tests/likelihoods/test_student_t.py @@ -17,17 +17,29 @@ def test_converges_to_gaussian(self): mu, var = np.array([0.0, 0.5]), np.array([0.1, 0.3]) y, sigma = np.array([0.1, 0.4]), 0.5 - ve_gauss = _eval(Gaussian(sigma).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_student = _eval(StudentT(nu=1000.0, sigma=sigma).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_gauss = _eval( + Gaussian(sigma).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) + ve_student = _eval( + StudentT(nu=1000.0, sigma=sigma).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) np.testing.assert_allclose(ve_student, ve_gauss, atol=1e-2) def test_quadrature_convergence(self): mu, var, y = np.array([0.0]), np.array([1.0]), np.array([0.5]) - ve_20 = _eval(StudentT(nu=5.0, sigma=1.0, n_points=20).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) - ve_50 = _eval(StudentT(nu=5.0, sigma=1.0, n_points=50).variational_expectation( - pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var))) + ve_20 = _eval( + StudentT(nu=5.0, sigma=1.0, n_points=20).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) + ve_50 = _eval( + StudentT(nu=5.0, sigma=1.0, n_points=50).variational_expectation( + pt.as_tensor_variable(y), pt.as_tensor_variable(mu), pt.as_tensor_variable(var) + ) + ) np.testing.assert_allclose(ve_20, ve_50, atol=1e-6) diff --git a/tests/optim/test_training.py b/tests/optim/test_training.py index 50624df..74b16a9 100644 --- a/tests/optim/test_training.py +++ b/tests/optim/test_training.py @@ -40,8 +40,12 @@ def test_compile_training_step_gp(gp_data): y_var = pt.vector("y") train_step, shared_params, shared_extras = pg.compile_training_step( - pg.marginal_log_likelihood, gp, X_var, y_var, - pm_model=model, learning_rate=1e-2, + pg.marginal_log_likelihood, + gp, + X_var, + y_var, + pm_model=model, + learning_rate=1e-2, ) losses = [] @@ -68,8 +72,12 @@ def test_compile_predict_gp(gp_data): y_var = pt.vector("y") train_step, shared_params, shared_extras = pg.compile_training_step( - pg.marginal_log_likelihood, gp, X_var, y_var, - pm_model=model, learning_rate=1e-2, + pg.marginal_log_likelihood, + gp, + X_var, + y_var, + pm_model=model, + learning_rate=1e-2, ) for _ in range(200): @@ -77,8 +85,12 @@ def test_compile_predict_gp(gp_data): X_new_var = pt.matrix("X_new") predict_fn = pg.compile_predict( - gp, X_new_var, model, shared_params, - X_train=X, y_train=y, + gp, + X_new_var, + model, + shared_params, + X_train=X, + y_train=y, ) X_test = np.linspace(0, 5, 20)[:, None] @@ -120,7 +132,10 @@ def test_compile_training_step_svgp(svgp_data): y_var = pt.vector("y") train_step, shared_params, shared_extras = pg.compile_training_step( - pg.elbo, svgp, X_var, y_var, + pg.elbo, + svgp, + X_var, + y_var, pm_model=model, extra_vars=[q_mu_var, q_sqrt_var], extra_init=[np.zeros(M), np.eye(M)], @@ -151,7 +166,10 @@ def test_sgd_optimizer(gp_data): y_var = pt.vector("y") train_step, shared_params, shared_extras = pg.compile_training_step( - pg.marginal_log_likelihood, gp, X_var, y_var, + pg.marginal_log_likelihood, + gp, + X_var, + y_var, pm_model=model, optimizer_fn=pg.sgd, learning_rate=1e-3, diff --git a/tests/test_conditionals.py b/tests/test_conditionals.py index 18e05ed..bf1f4a2 100644 --- a/tests/test_conditionals.py +++ b/tests/test_conditionals.py @@ -4,8 +4,8 @@ import pytensor import pytensor.tensor as pt -from ptgp.kernels import ExpQuad from ptgp.conditionals import base_conditional +from ptgp.kernels import ExpQuad def _eval(*tensors): @@ -51,10 +51,18 @@ def test_whitened_adds_variance(self): Kmn = kernel(pt.as_tensor_variable(Z), pt.as_tensor_variable(X)) Knn_diag = pt.diag(kernel(pt.as_tensor_variable(X))) - _, fvar_delta = _eval(*base_conditional( - Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f), white=True)) - _, fvar_q = _eval(*base_conditional( - Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f), - q_sqrt=pt.as_tensor_variable(q_sqrt), white=True)) + _, fvar_delta = _eval( + *base_conditional(Kmn, Kmm, Knn_diag, pt.as_tensor_variable(f), white=True) + ) + _, fvar_q = _eval( + *base_conditional( + Kmn, + Kmm, + Knn_diag, + pt.as_tensor_variable(f), + q_sqrt=pt.as_tensor_variable(q_sqrt), + white=True, + ) + ) assert np.all(fvar_q >= fvar_delta - 1e-10) diff --git a/tests/test_kl.py b/tests/test_kl.py index af058fc..e14d412 100644 --- a/tests/test_kl.py +++ b/tests/test_kl.py @@ -42,9 +42,11 @@ def test_unwhitened_positive(self): q_mu = rng.standard_normal(M) q_sqrt = np.eye(M) * 0.5 - kl = _eval(gauss_kl( - pt.as_tensor_variable(q_mu), - pt.as_tensor_variable(q_sqrt), - K=pt.as_tensor_variable(K), - )) + kl = _eval( + gauss_kl( + pt.as_tensor_variable(q_mu), + pt.as_tensor_variable(q_sqrt), + K=pt.as_tensor_variable(K), + ) + ) assert kl > 0.0 diff --git a/tests/test_objectives.py b/tests/test_objectives.py index 608793a..6dd1b1c 100644 --- a/tests/test_objectives.py +++ b/tests/test_objectives.py @@ -5,14 +5,14 @@ import pytensor.tensor as pt import pytest +from ptgp.gp import GP +from ptgp.inducing_variables import InducingPoints from ptgp.kernels import ExpQuad from ptgp.likelihoods import Gaussian from ptgp.mean import Zero -from ptgp.inducing_variables import InducingPoints -from ptgp.gp import GP +from ptgp.objectives import collapsed_elbo, elbo, marginal_log_likelihood from ptgp.svgp import SVGP from ptgp.vfe import VFE -from ptgp.objectives import marginal_log_likelihood, elbo, collapsed_elbo def _eval(*tensors): @@ -37,8 +37,7 @@ class TestMarginalLogLikelihood: def test_finite(self, regression_data): X, y = regression_data gp = GP(kernel=ExpQuad(ls=1.0), mean=Zero(), likelihood=Gaussian(sigma=0.1)) - mll = _eval(marginal_log_likelihood(gp, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + mll = _eval(marginal_log_likelihood(gp, pt.as_tensor_variable(X), pt.as_tensor_variable(y))) assert np.isfinite(mll) def test_better_fit_higher_mll(self, regression_data): @@ -47,10 +46,12 @@ def test_better_fit_higher_mll(self, regression_data): gp_good = GP(kernel=ExpQuad(ls=1.0), mean=Zero(), likelihood=Gaussian(sigma=0.1)) gp_bad = GP(kernel=ExpQuad(ls=0.01), mean=Zero(), likelihood=Gaussian(sigma=10.0)) - mll_good = _eval(marginal_log_likelihood(gp_good, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) - mll_bad = _eval(marginal_log_likelihood(gp_bad, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + mll_good = _eval( + marginal_log_likelihood(gp_good, pt.as_tensor_variable(X), pt.as_tensor_variable(y)) + ) + mll_bad = _eval( + marginal_log_likelihood(gp_bad, pt.as_tensor_variable(X), pt.as_tensor_variable(y)) + ) assert mll_good > mll_bad @@ -59,22 +60,24 @@ class TestELBO: def test_finite(self, regression_data, inducing_points): X, y = regression_data svgp = SVGP( - kernel=ExpQuad(ls=1.0), mean=Zero(), likelihood=Gaussian(sigma=0.1), + kernel=ExpQuad(ls=1.0), + mean=Zero(), + likelihood=Gaussian(sigma=0.1), inducing_variable=InducingPoints(pt.as_tensor_variable(inducing_points)), ) - elbo_val = _eval(elbo(svgp, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + elbo_val = _eval(elbo(svgp, pt.as_tensor_variable(X), pt.as_tensor_variable(y))) assert np.isfinite(elbo_val) def test_unwhitened_finite(self, regression_data, inducing_points): X, y = regression_data svgp = SVGP( - kernel=ExpQuad(ls=1.0), mean=Zero(), likelihood=Gaussian(sigma=0.1), + kernel=ExpQuad(ls=1.0), + mean=Zero(), + likelihood=Gaussian(sigma=0.1), inducing_variable=InducingPoints(pt.as_tensor_variable(inducing_points)), whiten=False, ) - elbo_val = _eval(elbo(svgp, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + elbo_val = _eval(elbo(svgp, pt.as_tensor_variable(X), pt.as_tensor_variable(y))) assert np.isfinite(elbo_val) def test_whitened_and_unwhitened_agree_at_prior(self, regression_data, inducing_points): @@ -86,22 +89,27 @@ def test_whitened_and_unwhitened_agree_at_prior(self, regression_data, inducing_ # Whitened: q_mu=0, q_sqrt=I is the prior q(v)=N(0,I) svgp_w = SVGP( - kernel=kernel, mean=Zero(), likelihood=Gaussian(sigma=0.1), - inducing_variable=InducingPoints(Z), whiten=True, + kernel=kernel, + mean=Zero(), + likelihood=Gaussian(sigma=0.1), + inducing_variable=InducingPoints(Z), + whiten=True, ) - elbo_w = _eval(elbo(svgp_w, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + elbo_w = _eval(elbo(svgp_w, pt.as_tensor_variable(X), pt.as_tensor_variable(y))) # Unwhitened: q_mu=0, q_sqrt=Luu is the prior q(u)=N(0, Kuu) Kuu = _eval(kernel(Z)) Luu = np.linalg.cholesky(Kuu) svgp_u = SVGP( - kernel=kernel, mean=Zero(), likelihood=Gaussian(sigma=0.1), - inducing_variable=InducingPoints(Z), whiten=False, - q_mu=pt.zeros(5), q_sqrt=pt.as_tensor_variable(Luu), + kernel=kernel, + mean=Zero(), + likelihood=Gaussian(sigma=0.1), + inducing_variable=InducingPoints(Z), + whiten=False, + q_mu=pt.zeros(5), + q_sqrt=pt.as_tensor_variable(Luu), ) - elbo_u = _eval(elbo(svgp_u, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + elbo_u = _eval(elbo(svgp_u, pt.as_tensor_variable(X), pt.as_tensor_variable(y))) np.testing.assert_allclose(elbo_w, elbo_u, atol=1e-6) @@ -112,15 +120,17 @@ def test_elbo_less_than_mll(self, regression_data, inducing_points): kernel = ExpQuad(ls=ls) gp = GP(kernel=kernel, mean=Zero(), likelihood=Gaussian(sigma=sigma)) - mll_val = _eval(marginal_log_likelihood(gp, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + mll_val = _eval( + marginal_log_likelihood(gp, pt.as_tensor_variable(X), pt.as_tensor_variable(y)) + ) svgp = SVGP( - kernel=kernel, mean=Zero(), likelihood=Gaussian(sigma=sigma), + kernel=kernel, + mean=Zero(), + likelihood=Gaussian(sigma=sigma), inducing_variable=InducingPoints(pt.as_tensor_variable(inducing_points)), ) - elbo_val = _eval(elbo(svgp, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + elbo_val = _eval(elbo(svgp, pt.as_tensor_variable(X), pt.as_tensor_variable(y))) assert elbo_val <= mll_val + 1e-6 # ELBO <= MLL @@ -129,11 +139,12 @@ class TestCollapsedELBO: def test_finite(self, regression_data, inducing_points): X, y = regression_data vfe_model = VFE( - kernel=ExpQuad(ls=1.0), mean=Zero(), likelihood=Gaussian(sigma=0.1), + kernel=ExpQuad(ls=1.0), + mean=Zero(), + likelihood=Gaussian(sigma=0.1), inducing_variable=InducingPoints(pt.as_tensor_variable(inducing_points)), ) - celbo = _eval(collapsed_elbo(vfe_model, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + celbo = _eval(collapsed_elbo(vfe_model, pt.as_tensor_variable(X), pt.as_tensor_variable(y))) assert np.isfinite(celbo) def test_collapsed_elbo_less_than_mll(self, regression_data, inducing_points): @@ -143,14 +154,16 @@ def test_collapsed_elbo_less_than_mll(self, regression_data, inducing_points): kernel = ExpQuad(ls=ls) gp = GP(kernel=kernel, mean=Zero(), likelihood=Gaussian(sigma=sigma)) - mll_val = _eval(marginal_log_likelihood(gp, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + mll_val = _eval( + marginal_log_likelihood(gp, pt.as_tensor_variable(X), pt.as_tensor_variable(y)) + ) vfe_model = VFE( - kernel=kernel, mean=Zero(), likelihood=Gaussian(sigma=sigma), + kernel=kernel, + mean=Zero(), + likelihood=Gaussian(sigma=sigma), inducing_variable=InducingPoints(pt.as_tensor_variable(inducing_points)), ) - celbo = _eval(collapsed_elbo(vfe_model, - pt.as_tensor_variable(X), pt.as_tensor_variable(y))) + celbo = _eval(collapsed_elbo(vfe_model, pt.as_tensor_variable(X), pt.as_tensor_variable(y))) assert celbo <= mll_val + 1e-6 # collapsed ELBO <= MLL From 3225b704b6648854681f6c99ee4442794f8ce6eb Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Wed, 15 Apr 2026 15:26:03 -0500 Subject: [PATCH 13/13] fix version --- .gitignore | 4 ++++ ptgp/__init__.py | 7 ++++--- pyproject.toml | 2 +- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index 11a3ff3..435e586 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,7 @@ +# Don't check in hatch-vcs generated version file +ptgp/_version.py + + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/ptgp/__init__.py b/ptgp/__init__.py index ce232da..801ecc9 100644 --- a/ptgp/__init__.py +++ b/ptgp/__init__.py @@ -1,5 +1,3 @@ -from importlib.metadata import version - from ptgp.gp import GP from ptgp.inducing_variables import InducingPoints, InducingVariables from ptgp.inference import fit_bfgs, fit_model, make_training_step, make_training_step_model @@ -73,4 +71,7 @@ ] -__version__ = version("pymc-extras") +try: + from ptgp._version import __version__ +except ImportError: + __version__ = "0.0.0+unknown" diff --git a/pyproject.toml b/pyproject.toml index 15e0208..981c93e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,7 @@ Issues = "https://github.com/bwengals/ptgp/issues" source = 'vcs' [tool.hatch.build.hooks.vcs] -version-file = "_version.py" +version-file = "ptgp/_version.py" [tool.hatch.build.targets.sdist] exclude = [