Arbitrary Prior class transformations

[wd60622]

The Prior class looks for transformations in the name space of pytensor.tensor and pm.math in that order.

There isn't out of the box support for arbitrary transformations. However, either one of these modules can be patched to include a new transformation.

Here is an example of patching pm.math to include a new transformation:

import pymc as pm

def register_new_transformation(func, name: str) -> None:
    setattr(pm.math, name, func)

Then using in action:

import pymc as pm
import pytensor.tensor as pt

from pymc_marketing.prior import Prior

import matplotlib.pyplot as plt


def create_polynomial_transformation(*coefficients):
    def func(x):
        degree = len(coefficients)
        powers = pt.arange(degree)

        X = pt.pow(x, powers)
        return pt.dot(X, pt.as_tensor_variable(coefficients))

    return func

# f(x) = 3 + x - 2x^2 + 0.5x^3
my_polynomial = create_polynomial_transformation(3, 1, -2, 0.5)
register_new_transformation(my_polynomial, "my_polynomial")


distribution = Prior("Normal", transform="my_polynomial")

samples = distribution.sample_prior()

samples.to_dataframe().plot.scatter(x="var_raw", y="var")
plt.show()

The to_json and from_json methods will still work. The transform will just have to be registered before loaded from JSON.