ptgp 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (118) hide show
  1. ptgp-0.1.0/.gitignore +190 -0
  2. ptgp-0.1.0/.pre-commit-config.yaml +30 -0
  3. ptgp-0.1.0/AGENTS.md +219 -0
  4. ptgp-0.1.0/LICENSE +21 -0
  5. ptgp-0.1.0/PKG-INFO +165 -0
  6. ptgp-0.1.0/README.md +122 -0
  7. ptgp-0.1.0/conda_envs/environment-test.yaml +32 -0
  8. ptgp-0.1.0/conda_envs/environment.yaml +27 -0
  9. ptgp-0.1.0/docs/agents/ptgp-vfe/SKILL.md +62 -0
  10. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/M_too_small.md +59 -0
  11. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/bad_priors.md +88 -0
  12. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/categorical_inducing_dim.md +66 -0
  13. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/eta_collapse.md +65 -0
  14. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/excess_fit_per_n_negative.md +62 -0
  15. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/inducing_collapse.md +69 -0
  16. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/inducing_layout_poor.md +66 -0
  17. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/kuu_ill_conditioned.md +71 -0
  18. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/large_grad_at_init.md +60 -0
  19. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/lbfgsb_abnormal.md +69 -0
  20. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/lengthscale_runaway.md +79 -0
  21. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/non_finite_at_init.md +74 -0
  22. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/sigma_collapse.md +67 -0
  23. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/sigma_inflation.md +68 -0
  24. ptgp-0.1.0/docs/agents/ptgp-vfe/pitfalls/slow_convergence.md +52 -0
  25. ptgp-0.1.0/docs/agents/ptgp-vfe/reference/api.md +130 -0
  26. ptgp-0.1.0/docs/agents/ptgp-vfe/reference/choosing_M.md +76 -0
  27. ptgp-0.1.0/docs/agents/ptgp-vfe/reference/discrete_inputs.md +118 -0
  28. ptgp-0.1.0/docs/agents/ptgp-vfe/reference/interpretation.md +71 -0
  29. ptgp-0.1.0/docs/agents/ptgp-vfe/reference/workflow.md +155 -0
  30. ptgp-0.1.0/docs/agents/ptgp-vfe/scripts/_common.py +68 -0
  31. ptgp-0.1.0/docs/agents/ptgp-vfe/scripts/check_inducing.py +159 -0
  32. ptgp-0.1.0/docs/agents/ptgp-vfe/scripts/detect_collapse.py +262 -0
  33. ptgp-0.1.0/docs/agents/ptgp-vfe/scripts/plot_history.py +83 -0
  34. ptgp-0.1.0/ptgp/__init__.py +62 -0
  35. ptgp-0.1.0/ptgp/_version.py +24 -0
  36. ptgp-0.1.0/ptgp/conditionals.py +72 -0
  37. ptgp-0.1.0/ptgp/gp/__init__.py +11 -0
  38. ptgp-0.1.0/ptgp/gp/svgp.py +299 -0
  39. ptgp-0.1.0/ptgp/gp/unapproximated.py +69 -0
  40. ptgp-0.1.0/ptgp/gp/vfe.py +92 -0
  41. ptgp-0.1.0/ptgp/idata.py +189 -0
  42. ptgp-0.1.0/ptgp/inducing.py +723 -0
  43. ptgp-0.1.0/ptgp/inducing_fourier.py +380 -0
  44. ptgp-0.1.0/ptgp/kernels/__init__.py +21 -0
  45. ptgp-0.1.0/ptgp/kernels/base.py +69 -0
  46. ptgp-0.1.0/ptgp/kernels/categorical.py +101 -0
  47. ptgp-0.1.0/ptgp/kernels/combination.py +66 -0
  48. ptgp-0.1.0/ptgp/kernels/nonstationary.py +166 -0
  49. ptgp-0.1.0/ptgp/kernels/stationary.py +116 -0
  50. ptgp-0.1.0/ptgp/kl.py +59 -0
  51. ptgp-0.1.0/ptgp/likelihoods/__init__.py +8 -0
  52. ptgp-0.1.0/ptgp/likelihoods/base.py +98 -0
  53. ptgp-0.1.0/ptgp/likelihoods/bernoulli.py +48 -0
  54. ptgp-0.1.0/ptgp/likelihoods/gaussian.py +66 -0
  55. ptgp-0.1.0/ptgp/likelihoods/negative_binomial.py +44 -0
  56. ptgp-0.1.0/ptgp/likelihoods/poisson.py +41 -0
  57. ptgp-0.1.0/ptgp/likelihoods/student_t.py +40 -0
  58. ptgp-0.1.0/ptgp/linalg/__init__.py +2 -0
  59. ptgp-0.1.0/ptgp/linalg/operator.py +4 -0
  60. ptgp-0.1.0/ptgp/linalg/ops.py +5 -0
  61. ptgp-0.1.0/ptgp/linalg/rewrites.py +4 -0
  62. ptgp-0.1.0/ptgp/mean.py +40 -0
  63. ptgp-0.1.0/ptgp/objectives.py +327 -0
  64. ptgp-0.1.0/ptgp/optim/__init__.py +30 -0
  65. ptgp-0.1.0/ptgp/optim/api.py +214 -0
  66. ptgp-0.1.0/ptgp/optim/optimizers.py +140 -0
  67. ptgp-0.1.0/ptgp/optim/schedules.py +75 -0
  68. ptgp-0.1.0/ptgp/optim/training.py +1251 -0
  69. ptgp-0.1.0/ptgp/rewrites.py +172 -0
  70. ptgp-0.1.0/ptgp/utils.py +412 -0
  71. ptgp-0.1.0/pyproject.toml +116 -0
  72. ptgp-0.1.0/scripts/inplace_audit.py +156 -0
  73. ptgp-0.1.0/scripts/install_claude_skills.py +86 -0
  74. ptgp-0.1.0/scripts/joint_graph_analysis.py +244 -0
  75. ptgp-0.1.0/scripts/mypy-failing.txt +0 -0
  76. ptgp-0.1.0/scripts/run_mypy.py +180 -0
  77. ptgp-0.1.0/tests/__init__.py +0 -0
  78. ptgp-0.1.0/tests/_fixtures/__init__.py +0 -0
  79. ptgp-0.1.0/tests/_fixtures/generate_svgp_baseline.py +52 -0
  80. ptgp-0.1.0/tests/_fixtures/svgp_points_unwhitened_baseline.pkl +0 -0
  81. ptgp-0.1.0/tests/_fixtures/vff_kuu_oracle.py +105 -0
  82. ptgp-0.1.0/tests/kernels/__init__.py +0 -0
  83. ptgp-0.1.0/tests/kernels/test_categorical.py +178 -0
  84. ptgp-0.1.0/tests/kernels/test_combination.py +58 -0
  85. ptgp-0.1.0/tests/kernels/test_diag.py +30 -0
  86. ptgp-0.1.0/tests/kernels/test_nonstationary.py +147 -0
  87. ptgp-0.1.0/tests/kernels/test_stationary.py +195 -0
  88. ptgp-0.1.0/tests/likelihoods/__init__.py +0 -0
  89. ptgp-0.1.0/tests/likelihoods/test_base.py +48 -0
  90. ptgp-0.1.0/tests/likelihoods/test_bernoulli.py +61 -0
  91. ptgp-0.1.0/tests/likelihoods/test_gaussian.py +80 -0
  92. ptgp-0.1.0/tests/likelihoods/test_negative_binomial.py +53 -0
  93. ptgp-0.1.0/tests/likelihoods/test_poisson.py +48 -0
  94. ptgp-0.1.0/tests/likelihoods/test_student_t.py +45 -0
  95. ptgp-0.1.0/tests/optim/__init__.py +0 -0
  96. ptgp-0.1.0/tests/optim/test_api.py +99 -0
  97. ptgp-0.1.0/tests/optim/test_frozen_vars.py +173 -0
  98. ptgp-0.1.0/tests/optim/test_schedules_and_groups.py +185 -0
  99. ptgp-0.1.0/tests/optim/test_scipy_objective.py +266 -0
  100. ptgp-0.1.0/tests/optim/test_training.py +352 -0
  101. ptgp-0.1.0/tests/optim/test_training_domain.py +188 -0
  102. ptgp-0.1.0/tests/test_conditionals.py +68 -0
  103. ptgp-0.1.0/tests/test_cubic_floor.py +121 -0
  104. ptgp-0.1.0/tests/test_heteroskedastic_predict.py +190 -0
  105. ptgp-0.1.0/tests/test_idata.py +207 -0
  106. ptgp-0.1.0/tests/test_inducing.py +282 -0
  107. ptgp-0.1.0/tests/test_inducing_fourier.py +409 -0
  108. ptgp-0.1.0/tests/test_kl.py +84 -0
  109. ptgp-0.1.0/tests/test_objectives.py +190 -0
  110. ptgp-0.1.0/tests/test_predict_samples.py +127 -0
  111. ptgp-0.1.0/tests/test_rewrites.py +249 -0
  112. ptgp-0.1.0/tests/test_skills_install.py +31 -0
  113. ptgp-0.1.0/tests/test_svgp.py +340 -0
  114. ptgp-0.1.0/tests/test_svgp_scipy_ref.py +379 -0
  115. ptgp-0.1.0/tests/test_utils.py +167 -0
  116. ptgp-0.1.0/tests/test_vfe.py +18 -0
  117. ptgp-0.1.0/tests/test_vff_e2e.py +67 -0
  118. ptgp-0.1.0/tests/test_vff_paths.py +158 -0
ptgp-0.1.0/.gitignore ADDED
@@ -0,0 +1,190 @@
1
+ # Don't check in hatch-vcs generated version file
2
+ ptgp/_version.py
3
+
4
+ # Local working notes
5
+ FUTURE.md
6
+
7
+ # Archived prototypes / old notebooks (not part of the package)
8
+ archive/
9
+
10
+ # Per-user Claude Code settings (personal permission caches, not project config)
11
+ .claude/settings.json
12
+ .claude/settings.local.json
13
+
14
+
15
+ # Byte-compiled / optimized / DLL files
16
+ __pycache__/
17
+ *.py[cod]
18
+ *$py.class
19
+
20
+ # C extensions
21
+ *.so
22
+
23
+ # Distribution / packaging
24
+ .Python
25
+ build/
26
+ comparison_libraries/
27
+ develop-eggs/
28
+ dist/
29
+ downloads/
30
+ eggs/
31
+ .eggs/
32
+ lib/
33
+ lib64/
34
+ parts/
35
+ sdist/
36
+ var/
37
+ wheels/
38
+ share/python-wheels/
39
+ *.egg-info/
40
+ .installed.cfg
41
+ *.egg
42
+ MANIFEST
43
+
44
+ # PyInstaller
45
+ # Usually these files are written by a python script from a template
46
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
47
+ *.manifest
48
+ *.spec
49
+
50
+ # Installer logs
51
+ pip-log.txt
52
+ pip-delete-this-directory.txt
53
+
54
+ # Unit test / coverage reports
55
+ htmlcov/
56
+ .tox/
57
+ .nox/
58
+ .coverage
59
+ .coverage.*
60
+ .cache
61
+ nosetests.xml
62
+ coverage.xml
63
+ *.cover
64
+ *.py,cover
65
+ .hypothesis/
66
+ .pytest_cache/
67
+ cover/
68
+
69
+ # Translations
70
+ *.mo
71
+ *.pot
72
+
73
+ # Django stuff:
74
+ *.log
75
+ local_settings.py
76
+ db.sqlite3
77
+ db.sqlite3-journal
78
+
79
+ # Flask stuff:
80
+ instance/
81
+ .webassets-cache
82
+
83
+ # Scrapy stuff:
84
+ .scrapy
85
+
86
+ # Sphinx documentation
87
+ docs/_build/
88
+
89
+ # PyBuilder
90
+ .pybuilder/
91
+ target/
92
+
93
+ # Jupyter Notebook
94
+ .ipynb_checkpoints
95
+
96
+ # IPython
97
+ profile_default/
98
+ ipython_config.py
99
+
100
+ # pyenv
101
+ # For a library or package, you might want to ignore these files since the code is
102
+ # intended to run in multiple environments; otherwise, check them in:
103
+ # .python-version
104
+
105
+ # pipenv
106
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
107
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
108
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
109
+ # install all needed dependencies.
110
+ #Pipfile.lock
111
+
112
+ # UV
113
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
114
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
115
+ # commonly ignored for libraries.
116
+ #uv.lock
117
+
118
+ # poetry
119
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
120
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
121
+ # commonly ignored for libraries.
122
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
123
+ #poetry.lock
124
+
125
+ # pdm
126
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
127
+ #pdm.lock
128
+ # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
129
+ # in version control.
130
+ # https://pdm.fming.dev/latest/usage/project/#working-with-version-control
131
+ .pdm.toml
132
+ .pdm-python
133
+ .pdm-build/
134
+
135
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
136
+ __pypackages__/
137
+
138
+ # Celery stuff
139
+ celerybeat-schedule
140
+ celerybeat.pid
141
+
142
+ # SageMath parsed files
143
+ *.sage.py
144
+
145
+ # Environments
146
+ .env
147
+ .venv
148
+ env/
149
+ venv/
150
+ ENV/
151
+ env.bak/
152
+ venv.bak/
153
+
154
+ # Spyder project settings
155
+ .spyderproject
156
+ .spyproject
157
+
158
+ # Rope project settings
159
+ .ropeproject
160
+
161
+ # mkdocs documentation
162
+ /site
163
+
164
+ # mypy
165
+ .mypy_cache/
166
+ .dmypy.json
167
+ dmypy.json
168
+
169
+ # Pyre type checker
170
+ .pyre/
171
+
172
+ # pytype static type analyzer
173
+ .pytype/
174
+
175
+ # Cython debug symbols
176
+ cython_debug/
177
+
178
+ # PyCharm
179
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
180
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
181
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
182
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
183
+ #.idea/
184
+
185
+ # Ruff stuff:
186
+ .ruff_cache/
187
+
188
+ # PyPI configuration file
189
+ .pypirc
190
+ .idea/
@@ -0,0 +1,30 @@
1
+ ci:
2
+ autofix_prs: false
3
+
4
+ repos:
5
+ - repo: https://github.com/pre-commit/pre-commit-hooks
6
+ rev: v4.6.0
7
+ hooks:
8
+ - id: check-merge-conflict
9
+ - id: check-toml
10
+ - id: check-yaml
11
+ - id: debug-statements
12
+ - id: end-of-file-fixer
13
+ - id: no-commit-to-branch
14
+ args: [--branch, main]
15
+ - id: trailing-whitespace
16
+
17
+ - repo: https://github.com/astral-sh/ruff-pre-commit
18
+ rev: v0.5.5
19
+ hooks:
20
+ - id: ruff
21
+ args: [ --fix, --unsafe-fixes, --exit-non-zero-on-fix ]
22
+ - id: ruff-format
23
+ types_or: [ python, pyi, jupyter ]
24
+
25
+ - repo: https://github.com/MarcoGorelli/madforhooks
26
+ rev: 0.4.1
27
+ hooks:
28
+ - id: no-print-statements
29
+ files: ^ptgp/
30
+ exclude: ^ptgp/_cli\.py$
ptgp-0.1.0/AGENTS.md ADDED
@@ -0,0 +1,219 @@
1
+ # AGENTS.md
2
+
3
+ Guidance for AI coding assistants (Claude Code, Codex, Cursor, Aider, …)
4
+ working in this repository. Follows the [AGENTS.md](https://agents.md/)
5
+ convention.
6
+
7
+ ## Mental model
8
+
9
+ PTGP is a Gaussian process library whose central bet is: **write naive
10
+ linear algebra, let PyTensor's rewrite system pick efficient algorithms
11
+ based on declared matrix structure**. Three pieces hold this up:
12
+
13
+ 1. **Naive linalg in source.** Always write `pt.linalg.inv(K)`,
14
+ `pt.linalg.slogdet(K)`, `pt.linalg.solve(K, b)` — never hand-roll a
15
+ Cholesky + solve_triangular chain. The rewrite system picks the
16
+ algorithm.
17
+ 2. **Assumption annotations.** When a kernel evaluates `K(X, X)`, wrap
18
+ the result with `pta.assume(K, symmetric=True, positive_definite=True)`
19
+ (where `pta` is `pytensor.assumptions`) so the rewrite system can prove
20
+ it's PSD and lower it to Cholesky. Cross-covariance `K(X, Y)` is left
21
+ **unannotated**. This happens automatically in `Kernel.__call__` —
22
+ don't duplicate it in kernel subclasses.
23
+ 3. **Cubic-op floor.** `tests/test_cubic_floor.py` asserts each model's
24
+ joint (loss + all grads) graph compiles to the *minimum* count of
25
+ O(N³) factorizations: 1 for `Unapproximated`, 2 for `VFE`, 1 for
26
+ `SVGP`. This test is the canary for the whole rewrite story — run it
27
+ after any change to `ptgp/rewrites.py`, `ptgp/objectives.py`, or kernel
28
+ evaluation. When it fails, `scripts/joint_graph_analysis.py` prints the
29
+ full per-model op breakdown for diagnosis.
30
+
31
+ `ptgp/rewrites.py` and `ptgp/linalg/rewrites.py` register assumption rules
32
+ and structural rewrites into PyTensor's global registries at import time
33
+ (both are imported for side-effects in `ptgp/__init__.py` —
34
+ **don't remove those imports**). `ptgp/rewrites.py` also monkey-patches
35
+ `pta.assume` to accept `positive=True` for scalar positivity.
36
+
37
+ **PyMC is a prior container only.** `pm.Model()` holds priors on
38
+ hyperparameters; `model.logp()` is added to the objective for MAP
39
+ training. **Never call `pm.sample()`, `pm.find_MAP()`, or any PyMC
40
+ inference routine** — PTGP's training paths in `ptgp/optim/` are the
41
+ inference layer.
42
+
43
+ ## Running things
44
+
45
+ All commands assume the project's `ptgp` conda env is active (Python
46
+ 3.14, PyTensor installed from `pymc-devs/pytensor@main`). The env path
47
+ varies by machine — if no `ptgp` env is found, ask the user where their
48
+ env lives rather than guessing.
49
+
50
+ ```bash
51
+ # tests (pyproject sets -v and --doctest-modules, so docstrings run too)
52
+ python -m pytest tests/
53
+ python -m pytest tests/test_vfe.py # single file
54
+ python -m pytest tests/test_vfe.py::test_name # single test
55
+ python -m pytest -k vfe # by keyword
56
+ python -m pytest tests/test_cubic_floor.py # the rewrite-floor canary
57
+
58
+ # lint / format (also enforced by pre-commit; line-length 100)
59
+ ruff check ptgp/ tests/
60
+ ruff format ptgp/ tests/
61
+
62
+ # diagnostic scripts — not in the test suite, run by hand when needed
63
+ python scripts/joint_graph_analysis.py # op-count tables per model, with/without ptgp rewrites
64
+ python scripts/inplace_audit.py # which linalg ops are inplace and why others aren't
65
+
66
+ # execute the demo notebook in place
67
+ python -m jupyter nbconvert --to notebook --execute --inplace \
68
+ notebooks/demo.ipynb --ExecutePreprocessor.kernel_name=ptgp
69
+
70
+ # install the bundled agent-skill docs into a Claude Code skill dir (optional)
71
+ python scripts/install_claude_skills.py --project . # ./.claude/skills/
72
+ python scripts/install_claude_skills.py --user # ~/.claude/skills/
73
+ ```
74
+
75
+ The reference usage example for all three models is
76
+ [`notebooks/demo.ipynb`](notebooks/demo.ipynb).
77
+
78
+ ## Where things live
79
+
80
+ - **Models** — `ptgp/gp/`. `Unapproximated`, `VFE`, `SVGP`. Each exposes
81
+ `predict_marginal` / `predict_joint`; `SVGP` adds `predict_f_samples`
82
+ and `prior_kl`. `VariationalParams` + `init_variational_params(M)` build
83
+ the symbolic `q_mu`, `q_sqrt`, plus the underlying trainable
84
+ `extra_vars` / `extra_init` needed by the compile helpers.
85
+ - **Objectives** — `ptgp/objectives.py`. Standalone functions, **not**
86
+ methods on models, each returning a named tuple whose first field is
87
+ the scalar loss: `marginal_log_likelihood` → `MLLTerms`, `elbo` →
88
+ `ELBOTerms`, `collapsed_elbo` → `CollapsedELBOTerms`,
89
+ `fitc_log_marginal_likelihood` → `FITCTerms`, plus `vfe_diagnostics`
90
+ and `dpp_regularizer`.
91
+ - **Kernels** — `ptgp/kernels/`. `Kernel` base class; subclasses
92
+ implement `_eval(X, Y)` and (optionally) `diag(X)`. Stationary
93
+ (`ExpQuad`, `Matern12/32/52`), nonstationary (`Gibbs`, `RandomWalk`,
94
+ `WarpedInput`), categorical (`Overlap`, `LowRankCategorical`),
95
+ combination (`SumKernel`, `ProductKernel`, via `+` / `*`). `active_dims`
96
+ selects input columns.
97
+ - **Likelihoods** — `ptgp/likelihoods/`. `Gaussian`, `Bernoulli`,
98
+ `Poisson`, `NegativeBinomial`, `StudentT`. Non-Gaussian variants
99
+ implement `variational_expectation` for SVGP.
100
+ - **Inducing variables** — `ptgp/inducing.py` (`Points`,
101
+ `random_subsample_init`, `kmeans_init`, `greedy_variance_init`),
102
+ `ptgp/inducing_fourier.py` (`FourierFeatures1D` — 1D Matern VFF with
103
+ `M = 2 * num_frequencies + 1`). Init routines return rich diagnostics
104
+ (`GreedyVarianceDiagnostics`, `KernelHealthDiagnostics`, …).
105
+ - **Optim** — `ptgp/optim/`. Two training paths and a prediction
106
+ compiler:
107
+ - `compile_training_step` — Adam/SGD via `pytensor.shared` for SVGP /
108
+ minibatch. Auto-discovers PyMC free RVs as trainable; pass
109
+ `extra_vars` (e.g. `q_mu`, `q_sqrt`, trainable `Z`), optional
110
+ `frozen_vars` for staged training, `param_groups` for per-group
111
+ learning rates. Returns `(train_step, shared_params, shared_extras)`.
112
+ - `compile_scipy_objective` — L-BFGS-B path. Pair with
113
+ `tracked_minimize` (one phase, per-iteration diagnostics) or
114
+ `minimize_staged_vfe` (multi-phase VFE with per-phase diagnostic
115
+ capture). `compile_scipy_diagnostics` builds a sibling callable that
116
+ returns `VFEDiagnostics` for the same parameter vector.
117
+ - `compile_predict` — reads the same `shared_params` produced by
118
+ training, so prediction picks up the trained values automatically
119
+ (no model reconstruction).
120
+ - **Conditionals / KL** — `ptgp/conditionals.py`, `ptgp/kl.py`. Whitened
121
+ vs unwhitened SVGP posterior pieces; `gauss_kl` / `gauss_kl_structured`.
122
+ - **Linalg** — `ptgp/linalg/`. `LinearOperator` abstraction (matvec /
123
+ solve / logdet ops) + structural rewrites. Powers the Fourier-features
124
+ path.
125
+ - **Utils** — `ptgp/utils.py`. `check_init` (NaN/inf + grad-norm sanity
126
+ check at initialisation — call this before training), `get_initial_params`
127
+ (prior-median / prior-draw / unconstrained-zero), `save_fit` / `load_fit`.
128
+ - **Rewrites** — `ptgp/rewrites.py` + `ptgp/linalg/rewrites.py`. The
129
+ cubic-floor enablers; see the mental-model section.
130
+
131
+ `comparison_libraries/` is **gitignored** — it's a local-only directory
132
+ where some users keep checked-out reference repos (GPJax, GPflow,
133
+ GPyTorch, CoLA, linear_operator, PyTensor, PyMC). GPJax is the primary
134
+ reference for correctness; GPflow has the cleanest VFE/SVGP architecture.
135
+ If `comparison_libraries/` isn't present on this machine, fetch from the
136
+ upstream repos when reference code is needed.
137
+
138
+ ## Conventions
139
+
140
+ Things that aren't self-evident from reading the code:
141
+
142
+ - **Action-first names.** `minimize_staged_vfe`, not
143
+ `staged_vfe_minimize`. `compile_scipy_objective`, not
144
+ `objective_scipy_compile`.
145
+ - **Kernel scale = `eta`, lengthscale = `ls`.** Kernels are written
146
+ `eta**2 * ExpQuad(ls=ls)` — `eta` is **always squared**. Multiple of
147
+ either get numeric suffixes: `eta1`, `ls2`.
148
+ - **Shape-annotate every symbolic var.** Pass `shape=(...)` whenever a
149
+ dimension is known at construction time. Use `None` only for axes that
150
+ genuinely vary across calls. Standard letters:
151
+ - `N` — data points (almost always `None`)
152
+ - `M` — inducing points (concrete int)
153
+ - `D` — input dimension (concrete int)
154
+ - `K` — number of outputs (concrete int)
155
+
156
+ ```python
157
+ X = pt.matrix("X", shape=(None, D))
158
+ Z = pt.matrix("Z", shape=(M, D))
159
+ q_mu = pt.vector("q_mu", shape=(M,))
160
+ ```
161
+
162
+ Shape info lets rewrites specialize on known relationships (e.g. gating
163
+ on `M < N`) and catches mismatches at compile time. Even single-axis
164
+ annotations (`shape=(None, 1)`) are worth it.
165
+ - **Internal symbolic var names start with `_`.** Anything named inside
166
+ library code (e.g. `pt.matrix("_X")` in `greedy_variance_init`) gets
167
+ the underscore prefix so it's distinguishable from user-named vars in
168
+ error messages and graph dumps.
169
+ - **Name vars in complex examples.** When symbolic vars become dict keys
170
+ (`param_groups`, multi-phase training), pass `name="..."` — unnamed
171
+ vars render as opaque addresses in error text via
172
+ `var.name or repr(var)`.
173
+ - **Error messages: two short sentences.** First names the problem,
174
+ second names the reason or fix. Example: `"Variables appear in both
175
+ extra_vars and frozen_vars: [...]. They cannot be both trainable and
176
+ frozen."`
177
+ - **Tests use float64; GPJax uses float32.** For cross-library
178
+ correctness comparisons, use `atol=1e-5`.
179
+ - **Notebooks** are stored as `.ipynb` and edited directly — no Jupytext
180
+ pairing or markdown mirror.
181
+
182
+ ## Stage and scope
183
+
184
+ Prototype. Keep code and tests brief; the goal is a working prototype,
185
+ not completeness. Don't add features, abstractions, or backwards-compat
186
+ shims beyond what the task requires.
187
+
188
+ ## Skill docs (`docs/agents/`)
189
+
190
+ The repo ships a backend-agnostic agent-skill doc at
191
+ `docs/agents/ptgp-vfe/` covering VFE training diagnosis (pitfalls,
192
+ escalation workflow, interpretation of `VFEDiagnostics` /
193
+ `GreedyVarianceDiagnostics`). `scripts/install_claude_skills.py` is a
194
+ Claude-specific convenience that symlinks it into a Claude Code skill
195
+ directory (`~/.claude/skills/` or a project-local `.claude/skills/`);
196
+ other AI tools can read `docs/agents/` directly.
197
+
198
+ **When you change any of the following, check whether the skill needs
199
+ updating:**
200
+
201
+ | File | Skill files to scan |
202
+ | --------------------- | -------------------------------------------------------------- |
203
+ | `ptgp/objectives.py` | `reference/api.md`, `reference/interpretation.md` |
204
+ | `ptgp/optim/training.py` | `reference/api.md`, `pitfalls/lbfgsb_abnormal.md`, `pitfalls/slow_convergence.md` |
205
+ | `ptgp/inducing.py` | `reference/interpretation.md`, `pitfalls/inducing_*.md` |
206
+ | `ptgp/utils.py` | `pitfalls/non_finite_at_init.md`, `pitfalls/large_grad_at_init.md` |
207
+
208
+ The skill's pitfall pages reference fields on `VFEDiagnostics`,
209
+ `CollapsedELBOTerms`, and `GreedyVarianceDiagnostics` by name — any
210
+ change to those namedtuples needs a sweep through the skill.
211
+
212
+ ## Commits
213
+
214
+ - One-sentence subject for small changes. Add a multi-line body for
215
+ multi-file refactors, non-obvious design decisions, or behavior changes
216
+ that need context.
217
+ - After finishing a feature or plan, **prompt the user** about staging
218
+ and committing. Don't auto-commit.
219
+ - Never add a `Co-Authored-By: …` trailer or any AI-attribution footer.
ptgp-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 Bill Engels and the PyMC Developers
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
ptgp-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,165 @@
1
+ Metadata-Version: 2.4
2
+ Name: ptgp
3
+ Version: 0.1.0
4
+ Summary: A practitioner's toolbox for estimating large-scale Gaussian Process models with PyMC and PyTensor
5
+ Project-URL: Repository, https://github.com/pymc-devs/ptgp.git
6
+ Project-URL: Issues, https://github.com/pymc-devs/ptgp/issues
7
+ Author-email: Bill Engels <w.j.engels@gmail.com>, PyMC Developers <pymc.devs@gmail.com>
8
+ License: MIT License
9
+
10
+ Copyright (c) 2025 Bill Engels and the PyMC Developers
11
+
12
+ Permission is hereby granted, free of charge, to any person obtaining a copy
13
+ of this software and associated documentation files (the "Software"), to deal
14
+ in the Software without restriction, including without limitation the rights
15
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
16
+ copies of the Software, and to permit persons to whom the Software is
17
+ furnished to do so, subject to the following conditions:
18
+
19
+ The above copyright notice and this permission notice shall be included in all
20
+ copies or substantial portions of the Software.
21
+
22
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
23
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
24
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
25
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
26
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
27
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
28
+ SOFTWARE.
29
+ License-File: LICENSE
30
+ Keywords: bayesian,gaussian processes,probabilistic programming,pymc,pytensor,variational inference
31
+ Requires-Python: >=3.12
32
+ Requires-Dist: jax
33
+ Requires-Dist: numpy
34
+ Requires-Dist: optax
35
+ Requires-Dist: pymc>=6.0.0
36
+ Requires-Dist: pytensor>=3.0.2
37
+ Provides-Extra: dev
38
+ Requires-Dist: mypy; extra == 'dev'
39
+ Requires-Dist: polars; extra == 'dev'
40
+ Requires-Dist: pre-commit; extra == 'dev'
41
+ Requires-Dist: pytest; extra == 'dev'
42
+ Description-Content-Type: text/markdown
43
+
44
+ # PTGP
45
+
46
+ A Gaussian process library for building GP models that solve real-world problems.
47
+
48
+ ## Who this is for
49
+
50
+ PTGP is for practitioners who need flexible, well-supported GP modeling. The goal of PTGP is to be fully batteries-included and ready to work on real-world problems:
51
+
52
+ - **Practical GP algorithms:** exact GP, VFE with collapsed bound, SVGP with minibatch training, VFF (Variational Fourier Features)
53
+ - **Full kernel library:** ExpQuad, Matern52/32/12, RandomWalk, Gibbs, WarpedInput, categorical kernels for multi-class or categorical input variables, composition via `+` and `*`, `active_dims` for dimension selection
54
+ - **Non-Gaussian likelihoods:** Bernoulli, Poisson, NegativeBinomial, StudentT
55
+ - **PyMC priors:** set priors on any hyperparameter; use PyMC distributions for mean functions and noise models; MAP training by default
56
+ - **Training tools:** L-BFGS-B and Adam optimizers, per-parameter learning rates, staged optimization, frozen variables, inducing point initialization strategies, diagnostic-guided workflows; more are being added, such as carefully monitored training to help diagnose issues early
57
+ - **Agent-readable docs:** `docs/agents/` ships LLM-readable guides for debugging training issues and folk wisdom (VFE training covered). See the [Working with AI coding assistants](#working-with-ai-coding-assistants) section below.
58
+ - **More coming:** see the [issues](https://github.com/pymc-devs/ptgp/issues)
59
+
60
+ Researchers benefit from the underlying design: PTGP is built on PyTensor's symbolic graph and rewrite system, so you write GP math directly (`pt.linalg.inv(K)`, `pt.linalg.slogdet(K)`) and the compiler chooses efficient algorithms based on declared matrix structure. This makes it straightforward to implement new GP approximations and create custom models, and will eventually allow matrix structure like Kronecker, Toeplitz, and sparse to be taken advantage of automatically.
61
+
62
+ ## Models
63
+
64
+ | Model | Scale | Best for |
65
+ |-------|-------|----------|
66
+ | `gp.Unapproximated` | N < ~2,000 | Exact inference, model comparison |
67
+ | `gp.VFE` | N < ~50,000 | Medium-scale data with inducing points |
68
+ | `gp.SVGP` | N up to ~500,000 | Large data, non-Gaussian likelihoods, minibatch training |
69
+ | `FourierFeatures1D` | 1D Matern kernels | Structured Kuu via Fourier basis; no inducing point placement |
70
+
71
+ ## Quick start
72
+
73
+ ```python
74
+ import numpy as np
75
+ import pymc as pm
76
+ import pytensor.tensor as pt
77
+ import ptgp as pg
78
+
79
+ X = np.random.randn(200, 1)
80
+ y = np.sin(X.ravel()) + 0.1 * np.random.randn(200)
81
+ Z_init = np.linspace(-2, 2, 20)[:, None]
82
+ Z_var = pt.matrix("Z", shape=(20, 1))
83
+
84
+ with pm.Model() as model:
85
+ ls = pm.InverseGamma("ls", alpha=2.0, beta=1.0)
86
+ eta = pm.Exponential("eta", lam=1.0)
87
+ kernel = eta**2 * pg.kernels.Matern52(input_dim=1, ls=ls)
88
+
89
+ svgp = pg.gp.SVGP(
90
+ kernel=kernel,
91
+ likelihood=pg.likelihoods.Gaussian(sigma=0.1),
92
+ inducing_variable=pg.inducing.Points(Z_var, Z_init=Z_init),
93
+ variational_params=pg.gp.init_variational_params(M=20),
94
+ )
95
+ fit = pg.fit(svgp, X, y, method="L-BFGS-B")
96
+
97
+ mean, var = pg.predict(svgp, np.linspace(-3, 3, 100)[:, None], fit)
98
+ ```
99
+
100
+ `pg.fit` picks a default objective from the gp type (`Unapproximated` → `marginal_log_likelihood`, `VFE` → `collapsed_elbo`, `SVGP` → `elbo`) and returns a `FitResult` that `pg.predict` consumes. For stochastic mini-batch training, staged VFE, or per-group learning rates, drop down to `pg.optim.compile_training_step` / `pg.optim.compile_scipy_objective` — see [`notebooks/demo.ipynb`](notebooks/demo.ipynb):
101
+
102
+ ```python
103
+ X_var = pt.matrix("X")
104
+ y_var = pt.vector("y")
105
+
106
+ step, shared_params, shared_extras = pg.optim.compile_training_step(
107
+ pg.objectives.elbo, svgp, X_var, y_var, model, learning_rate=1e-2
108
+ )
109
+
110
+ for i in range(500):
111
+ loss = step(X, y)
112
+
113
+ predict_fn = pg.optim.compile_predict(
114
+ svgp, pt.matrix("X_new"), model, shared_params, shared_extras=shared_extras
115
+ )
116
+ mean, var = predict_fn(np.linspace(-3, 3, 100)[:, None])
117
+ ```
118
+
119
+ Training uses MAP by default: the PyMC log-prior is added to the objective. Pass `include_prior=False` for pure ELBO. For exact GPs and VFE, use `compile_scipy_objective` with L-BFGS-B instead. See [`notebooks/demo.ipynb`](notebooks/demo.ipynb) for end-to-end examples covering all three models.
120
+
121
+ ## How it works
122
+
123
+ PTGP is built on PyTensor's symbolic graph. Kernels, likelihoods, and GP models return symbolic tensors with naive linear algebra like `pt.linalg.inv(K)` that PyTensor's rewrite system automatically lowers to efficient Cholesky-based code using declared matrix properties. All models compile their full forward+gradient step down to the minimum number of cubic factorizations.
124
+
125
+ PTGP tries to distill some of the approaches of existing GP libraries and make them more accessible, mainly [GPJax](https://github.com/JaxGaussianProcesses/GPJax), [GPflow](https://github.com/GPflow/GPflow), and [GPyTorch](https://github.com/cornellius-gp/gpytorch).
126
+
127
+ ## Working with AI coding assistants
128
+
129
+ PTGP is set up to work nicely with AI coding assistants:
130
+
131
+ - **[`AGENTS.md`](AGENTS.md)** — project-level instructions for AI coding assistants (architecture, conventions, where things live, how to run tests). Follows the [AGENTS.md](https://agents.md/) cross-tool convention used by Codex, Cursor, Aider, and others.
132
+ - **[`docs/agents/`](docs/agents/)** — backend-agnostic agent-skill docs covering folk wisdom and training-debug recipes. Currently includes [`ptgp-vfe`](docs/agents/ptgp-vfe/) (VFE diagnostic skill: pitfalls, escalation workflow, interpretation of `VFEDiagnostics` and `GreedyVarianceDiagnostics`).
133
+
134
+ ### Claude Code users
135
+
136
+ Claude Code reads `CLAUDE.md`, not `AGENTS.md`. Symlink so they stay in sync:
137
+
138
+ ```bash
139
+ ln -s AGENTS.md CLAUDE.md
140
+ ```
141
+
142
+ To install the VFE skill into a Claude Code skill directory (so Claude auto-discovers it when you mention VFE), run:
143
+
144
+ ```bash
145
+ python scripts/install_claude_skills.py --project . # ./.claude/skills/
146
+ python scripts/install_claude_skills.py --user # ~/.claude/skills/
147
+ ```
148
+
149
+ ## Install
150
+
151
+ ```bash
152
+ pip install git+https://github.com/pymc-devs/ptgp.git
153
+ ```
154
+
155
+ To hack on PTGP itself, clone and install in editable mode:
156
+
157
+ ```bash
158
+ git clone https://github.com/pymc-devs/ptgp.git
159
+ cd ptgp
160
+ pip install -e .
161
+ ```
162
+
163
+ ## Contributing
164
+
165
+ See the [issues](https://github.com/pymc-devs/ptgp/issues) for what's being worked on. Feel free to propose issues, feature requests, or use cases you've been hoping could be made easier. PRs always welcome.