A general method to reversing transformations in custom step?

Hello, I am working on a custom step function that uses some random variables that I can access via point[]. My question is what is the best way of transforming them back into the variable space given some transform marker.

For example: sigma = pm.Gamma('sigma', alpha=0.7, beta=2) will be transformed using the log transform and then sampled in the log space. To be able to use the value of sigma from the last step within my custom step I can do:

sigma = np.exp(point[self.sigma.transformed.name])

Another variable’s prior is lamb = pm.Beta('lamb', alpha=1.3, beta=2.0), with the transform marker: lamb_logodds__. I can do the backward transform within the custom step with

import scipy as sp

lamb = sp.special.expit(point[self.lamb.transformed.name])

I have other random variables that have other transform markers (_lowerbound__).

Now I wonder if there is a best way of doing backtransforms that is also general - I have come across this part of the pymc3 documentation that discusses transformations. But I am not sure if I can use that in my case.
Or if it is best to write a function that figures out all the necessary transforms itself, kind of like this:

def undo_transform(x):

    if transform_marker = '_log__':
         value = np.exp(x)
    elif transform_marker = '_logodds__':
         value =  sp.special.expit(x)
    # etc.
    return value

Ideally I would be able to use something like this inside the step

lamb = self.lamb.transform.backward(point[self.lamb.transformed.name])

and that way figure out what transform is required automatically. I feel like this should be possible (according to the documentation above), but I have not been able to get it to work so far.

1 Like