# Custom PyTensor Op to calculate an analytical gradient

Hello,

I would like to implement a custom PyTensor Op to give me out a value, EA, while taking an analytical gradient of this function and passing it to pymc to continue forward. The gradient of the function I want to use has been implemented in this paper (Equations A5 and A6)

Here:
\begin{align*} \frac{\partial E}{\partial M} &= \frac{1}{1-e\cos E(t)} \end{align*}

and

\begin{align*} \frac{\partial E}{\partial e} &= \frac{\sin E(t)}{1-e\cos E(t)} \end{align*}
The current function I already have is:

class KeplerSolverOp(Op):
__props__ = ()

def make_node(self, M, e):
M = as_tensor_variable(M)
e = as_tensor_variable(e)

return Apply(self, [M, e], [M.type(), e.type()])

def perform(self, node, inputs, output_storage):
M, e = inputs
EA = _calc_ecc_anom(M, e)  # Compute ecc-anomly
output_storage[0][0] = EA

M, e = inputs

EA = _calc_ecc_anom(M, e)
sea = pm.math.sin(EA)
cea = pm.math.cos(EA)
temp = 1 - e * cea
dEA_dM = 1.0 / temp
dEA_de = (sea * EA) / temp

return [dEA_dM, dEA_de]


The problem is that the output of “KeplerSolverOp” is a not an individual value that i’m expecting. As I was expecting to call it like:

kepler_solver_op = KeplerSolverOp()
eanom = kepler_solver_op(manom, ecc)


What could I be missing?

You defined an Op that has two outputs, M.type() and e.type(), so calling the Op returns two variables

Thank you for the fast response! If I want the output of my "KeplerSolverOp " to be “EA” from

EA = _calc_ecc_anom(M, e)


\begin{align*} \frac{\partial E}{\partial M} &= \frac{1}{1-e\cos E(t)} \end{align*}
and
\begin{align*} \frac{\partial E}{\partial e} &= \frac{\sin E(t)}{1-e\cos E(t)} \end{align*}
sent to PyMC to continue working then I should write the make_node like:

    def make_node(self, M, e):
M = as_tensor_variable(M)
e = as_tensor_variable(e)
EA = _calc_ecc_anom(M, e)
return Apply(self, [M, e], [EA])


When I do that I get:

Blockquote
Traceback (most recent call last):
File “/Users/jrob/Desktop/School/NU/Research/Jason/gj504/pymc_GJ504.py”, line 77, in
raoff, deoff = kepler.calc_orbit(epoch, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=epoch[0])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “/Users/jrob/Desktop/School/NU/Research/Jason/gj504/kepler.py”, line 227, in calc_orbit
f = pt.function([manom,ecc], KeplerSolverOp()(manom,ecc))
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “/Users/jrob/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/graph/op.py”, line 293, in call
node = self.make_node(*inputs, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “/Users/jrob/Desktop/School/NU/Research/Jason/gj504/kepler.py”, line 75, in make_node
return Apply(self, [M, e], [EA])
^^^^^^^^^^^^^^^^^^^^^^^^^
File “/Users/jrob/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/graph/basic.py”, line 158, in init
raise ValueError(
ValueError: All output variables passed to Apply must belong to it.

In make node you just specify the type of the input and output variables, your output is doing more than that (there shouldn’t be a need for a calling _calc_ecc_anom in it)

Type means an instance of tensor(shape=..., dtype=...) That’s all PyTensor needs to know about the output at the make_node method. Actual computation should be defined in the perform method

And your gradient must have only PyTensor pure operations, so you may need to call self if you need the output of the Op to define the gradient.

Perhaps this guided will clarify some of the things: Using a “black box” likelihood function — PyMC example gallery

I haven’t seen this yet. I’ll look through it and check back here when I implement it!

Alrighty, so I reorgianized my Custom Op’s to now look like this

class KeplerSolverOp(Op):
__props__ = ()

def make_node(self, M, e):
#setting inputs to function of M and e
M = at.as_tensor(M)
e = at.as_tensor_variable(e)

#setting input list for Apply
inputs = [M, e]
outputs = [M.type()]

return Apply(self, inputs, outputs)

def perform(self, node, inputs, outputs):
M, e = inputs
EA = _calc_ecc_anom(M, e)  # Compute ecc-anomly
outputs[0][0] = EA #may have to add something likept.value...

M, e = inputs

def make_node(self, M, e):
#setting inputs to function of M and e
M = at.as_tensor(M)
e = at.as_tensor_variable(e)

#setting input list for Apply
inputs = [M, e]

# There are two outputs with the same type as data,
# for the partial derivatives wrt to m, c
outputs = [M.type(), M.type()]

return Apply(self, inputs, outputs)

def perform(self, node, inputs, outputs):
M, e = inputs

EA = _calc_ecc_anom(M, e)
sea = pm.math.sin(EA)
cea = pm.math.cos(EA)
temp = 1 - e * cea
dEA_dM = 1.0 / temp
dEA_de = (sea * EA) / temp

outputs[0][0] = dEA_dM
outputs[1][0] = dEA_de



It seems to run, but I get a segfault. I don’t see a “core” to look at to even investigate it. I wonder what could be consuming so much RAM? (If that’s the true answer)

One thing I have tried is to run

import numpy as np
import pytensor.tensor as at
import pymc as pm
from pytensor.graph.op import Op, Apply
from pytensor import function
import pytensor as pt

def _calc_ecc_anom(manom, ecc):
"""
Computes the eccentric anomaly from the mean anomlay.
Code from Rob De Rosa's orbit solver (e < 0.95 use Newton, e >= 0.95 use Mikkola)

Args:
manom (float/np.array): mean anomaly, either a scalar or np.array of any shape
ecc (float/np.array): eccentricity, either a scalar or np.array of the same shape as manom

Return:
eanom (float/np.array): eccentric anomalies, same shape as manom

Written: Jason Wang, 2018
"""
# print(type(ecc))
alpha = (1.0 - ecc) / ((4.0 * ecc) + 0.5)
beta = (0.5 * manom) / ((4.0 * ecc) + 0.5)

aux = pm.math.sqrt(beta**2.0 + alpha**3.0)
z = pm.math.abs(beta + aux)**(1.0/3.0)

s0 = z - (alpha/z)
s1 = s0 - (0.078*(s0**5.0)) / (1.0 + ecc)
e0 = manom + (ecc * (3.0*s1 - 4.0*(s1**3.0)))

se0 = pm.math.sin(e0)
ce0 = pm.math.cos(e0)

f = e0-ecc*se0-manom

f1 = 1.0-ecc*ce0
f2 = ecc*se0
f3 = ecc*ce0
f4 = -f2
u1 = -f/f1
u2 = -f/(f1+0.5*f2*u1)
u3 = -f/(f1+0.5*f2*u2+(1.0/6.0)*f3*u2*u2)
u4 = -f/(f1+0.5*f2*u3+(1.0/6.0)*f3*u3*u3+(1.0/24.0)*f4*(u3**3.0))

return (e0 + u4)

class KeplerSolverOp(Op):
__props__ = ()

def make_node(self, M, e):
#setting inputs to function of M and e
M = at.as_tensor(M)
e = at.as_tensor_variable(e)

#setting input list for Apply
inputs = [M, e]
outputs = [M.type()]

return Apply(self, inputs, outputs)

def perform(self, node, inputs, outputs):
M, e = inputs
EA = _calc_ecc_anom(M, e)  # Compute ecc-anomly
outputs[0][0] = EA #may have to add something likept.value...

M, e = inputs

def make_node(self, M, e):
#setting inputs to function of M and e
M = at.as_tensor(M)
e = at.as_tensor_variable(e)

#setting input list for Apply
inputs = [M, e]

# There are two outputs with the same type as data,
# for the partial derivatives wrt to m, c
outputs = [M.type(), M.type()]

return Apply(self, inputs, outputs)

def perform(self, node, inputs, outputs):
M, e = inputs

EA = _calc_ecc_anom(M, e)
sea = pm.math.sin(EA)
cea = pm.math.cos(EA)
temp = 1 - e * cea
dEA_dM = 1.0 / temp
dEA_de = (sea * EA) / temp

outputs[0][0] = dEA_dM
outputs[1][0] = dEA_de

M = at.scalar("m")
e = at.scalar("e")

eval_out = out.eval({M: 2.96466117, e: 0.4781471212028443})
print(eval_out.eval())


and I got no errors and received “3.021801889806251” which is correct…

Also I have tried

out = keplergrad(M,e)



LogLikeGrad.0 [id A] <Scalar(float64, shape=())>