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.