Arctan2 & alternatives to pm.math.switch?

Questions version agnostic
1

Hi,
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 https://en.wikipedia.org/wiki/Atan2
    '''
    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…?

2

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(pt.gt(x, 0),
              pt.arctan(y/x),
              pt.switch(pt.gt(y, 0),
                        pt.pi/2 - pt.arctan(x/y),
                        ...))
1 Like