Arctan2 & alternatives to pm.math.switch?

For an ABC model I’m trying to build I need the arctan2 output of a Tensor variable.
arctan is implemented in pm.math, but not arctan2, and so I tried implementing it myself using
a custom function with if and else statements. I then found out that Boolean operations are not allowed in pymc - and the alternative is to use pm.math.switch. as pointed out in this answer.

While trying to ‘convert’ arctan output to arctan2 output (using the if-else statements here), I’m not quite sure if the ‘switching’ in pm.math.switch would result in the same output as a set of order sensitive if-s and else-s. To clarify, here’s the standard Python code to convert from atan → atan2:

def custom_atan2(y,x):
    Based on
    if x>0:
        return np.arctan(y/x)
    elif y>0:
        return np.pi/2 - np.arctan(x/y)
    elif y<0:
        return -np.pi/2 -  np.arctan(x/y)
    elif x<0:
        return np.arctan(y/x) + np.pi
    elif x==0 and y==0:
        return np.nan

If I were to use pm.math.switch, the order-sensitivity of the return statements would be lost (e.g. if x is not >0, then I don’t want to switch to the condition where x<0 already). Does anyone have workarounds to this, or maybe I’m not able to see something very obvious…?

The signature of switch is switch(conditon, if, else), so you can achieve ifelse blocks by putting more switches in the second argument like this:

x = pt.dvector('x')
y = pt.dvector('y')

z = pt.switch(, 0),
              pt.switch(, 0),
                        pt.pi/2 - pt.arctan(x/y),
1 Like