Search code examples
pythonbayesianpymc3

pymc3 Heaviside step function


In pymc3, how can the Heaviside step function be applied as a deterministic transform?

Any step function will work, but I'm asking about the Heaviside step function specifically to disambiguate it from what the documentation refers to as a "step function", which is talking about sampling steps, which are irrelevant to what I'm looking for and makes finding relevant documentation difficult to locate, if it exists. Also I wasn't able to find somewhere that all supported transforms are documented.

I want 1 if x >= 0 else 0, but writing it literally like that won't work because presumably x >= 0 evaluates to an object, and so the if expression will immediately evaluate to 1, and then h will always be 1.

import pymc3 as pm

with pm.Model() as model:
    x = pm.Normal('x', mu=0, sigma=1)
    h = pm.Deterministic('h', 1 if x >= 0 else 0)  # ???

Solution

  • Try using theano.tensor.switch. That is,

    import pymc3 as pm
    import theano.tensor as tt
    
    with pm.Model() as model:
        x = pm.Normal('x', mu=0, sigma=1)
        h = pm.Deterministic('h', tt.switch(x < 0, 0, 1))