Understanding pytensor's API vs pymc's API and which is more appropriate to use for a new project

[kratsg]

Following up from #1545 (reply in thread) .

The project: pyhs3. The README contains a hello-world example with the expect API using the HEP Statistics Serialization Standard (HS3) specification:

Hello World example

>>> import pyhs3
>>> import scipy
>>> import math
>>> workspace_data = {
...     "distributions": [
...         {
...             "name": "model",
...             "type": "gaussian_dist",
...             "x": "x",
...             "mean": "mu",
...             "sigma": "sigma",
...         }
...     ],
...     "parameter_points": [
...         {
...             "name": "default_values",
...             "parameters": [
...                 {"name": "x", "value": 0.0},
...                 {"name": "mu", "value": 0.0},
...                 {"name": "sigma", "value": 1.0},
...             ],
...         }
...     ],
...     "domains": [
...         {
...             "name": "default_domain",
...             "type": "product",
...             "axes": [
...                 {"name": "x", "min": -5.0, "max": 5.0},
...                 {"name": "mu", "min": -2.0, "max": 2.0},
...                 {"name": "sigma", "min": 0.1, "max": 3.0},
...             ],
...         }
...     ],
... }
>>> ws = pyhs3.Workspace(workspace_data)
>>> model = ws.model()
<BLANKLINE>
>>> print(model)
Model(
    mode: FAST_RUN
    parameters: 3 (x, mu, sigma)
    distributions: 1 (model)
    functions: 0 ()
)
>>> parameters = {par.name: par.value for par in model.parameterset}
>>> result = -2 * model.logpdf("model", **parameters)
>>> print(f"parameters: {parameters}")
parameters: {'x': 0.0, 'mu': 0.0, 'sigma': 1.0}
>>> print(f"nll: {result:.8f}")
nll: 1.83787707
>>> result_scipy = -2 * math.log(scipy.stats.norm.pdf(0, loc=0, scale=1))
>>> print(f"nll: {result_scipy:.8f}")
nll: 1.83787707

At the pytensor level, the gaussian_dist was re-implemented here :https://github.com/scipp-atlas/pyhs3/blob/0c3e46961d4ce7da2c12088e6ef4f210d16acffd/src/pyhs3/distributions.py#L118-L218 . It's not abundantly clear if this is duplication, or if one can use pymc instead.

Related issue: scipp-atlas/pyhs3#72

./cc @ricardoV94 .

[ricardoV94]

It sounds like you'll be reinventing less things if you build on top of PyMC (and reuse PyMC models) instead of directly on PyTensor.

Also for some related functionality: https://www.pymc.io/projects/extras/en/stable/generated/pymc_extras.prior.Prior.html

The tests may be more readable: https://github.com/pymc-devs/pymc-extras/blob/main/tests/test_prior.py

[kratsg]

Is there documentation on Prior? The linked documentation is slightly not clear to me how the expected usage is meant to be done. From what i can tell from the tests, Prior is sort of similar to what we'd be doing, but we would need to be able to provide a pdf and logpdf call (the latter is provided by model.logp(...) it seems.

Is there a way to use the compilation mode functionality of pytensor with pymc? It looks like it gets wrapped or hidden away.

[ricardoV94]

No more documentation on Prior, I didn't mean to say you should use it. Jjust shared to see if any bulbs would light up.

On PyMC/PyTensor. PyMC gives you the pytensor graph for the logp, you can then compile it into a PyTensor function using whichever mode you want (or differentiate, vectorize, ... whatever).

model_logp = model.logp()
model_logp_fn = pytensor.function([model_logp.value_vars, model_logp, mode="NUMBA")

For the pdf you could just exponentiate the logp, it's usually better to compute everything in log-scale anyway, to avoid numerical issues, and then exponentiate only at the end.

[ricardoV94]

@kratsg I wrote this snippet (by no means completely fleshed out), for another project with the aim of making it easier to work with PyTensor as a dev: https://gist.github.com/ricardoV94/a51785d4199b8c14ce0d337f2b437772#file-pytensor_jit-py

[ricardoV94]

Also note we may soon be opening a repo with common distribution functions (logpdf, logcdf, moments ... etc) written in PyTensor that other packages can use (and contribute to) to avoid duplicated work (or depending on a package like PyMC just to use this very restricted functionality): pymc-devs/distributions#1

[kratsg]

Ooh, this looks interesting. I have a slightly parallel but not fully overlapping effort -- where we have a way of serializing likelihoods and probability models to JSON -- and need to be able to load that into a built computational graph (e.g. pytensor) -- https://github.com/scipp-atlas/pyhs3 .

The slight difference is that I need to build the entire likelihood f(x | pars) without having evaluated the pars yet -- and typically x remains constant while the pars vary (some subset of all pars).

[ricardoV94]

Ooh, this looks interesting. I have a slightly parallel but not fully overlapping effort -- where we have a way of serializing likelihoods and probability models to JSON -- and need to be able to load that into a built computational graph (e.g. pytensor) -- https://github.com/scipp-atlas/pyhs3 .

The slight difference is that I need to build the entire likelihood f(x | pars) without having evaluated the pars yet -- and typically x remains constant while the pars vary (some subset of all pars).

That should be fine, we're probably just gonna keep the bare stuff there, ie the actual expressions for the densities, so in your case you could use them and pass x as constant and the other as parameters.