Hi all,

I am relatively new to PyMC and am still trying to understand, how everything works. For my current project I need to estimate a diffusion equation numerically as desribed here.

After trying different things in PyTensor and Aesara, I realized that I would need to write my own Op to run a root-finding algorithm, which can find the first 4 roots of the equation (23) in the cited work.

My implementation is:

```
import pymc as pm
import pytensor.tensor as at
from pytensor import *
from pytensor.graph.op import Op
from pytensor.graph.basic import Apply
from scipy.optimize import fsolve
def root_finder_loop_perform(D, Sf, Sb, thickness):
# Define an Aesara expression for the function to be rooted
def f(x, D, Sf, Sb, thickness):
return np.tan(x*thickness*1e-7)+(D*(Sb + Sf)*x)/(Sf*Sb-D**2*x**2)
# Set up an array to store the roots for each value of a
root_list = np.zeros((4))
# Loop over the values of a and find the roots
# Set the initial guess for the root
x = 100.
b = 0
# Loop over the first four roots
while b < len(root_list):
# Find the root using scipy.optimize.root
sol = fsolve(f, x, args=(D,Sf,Sb, thickness))#, fprime=jac)
# Check if the root is unique and add it to the root list if it is
if np.round(sol,0) not in np.round(root_list,0) and np.round(sol,0) > 0:
root_list[b] = sol
b += 1
else:
x = x*1.1
return np.array(root_list)
class RootFinder(Op):
def __init__(self, thickness):
self.thickness = thickness
def make_node(self, D, Sf, Sb):
outputs = [at.vector(dtype='float64')]
return Apply(self, [D, Sf, Sb], outputs)
def perform(self, node, inputs, outputs_storage):
D, Sf, Sb = inputs
root_list = root_finder_loop_perform(D, Sf, Sb, self.thickness)
outputs_storage[0][0] = root_list
def grad(self, inputs, grads):
D, Sf, Sb = inputs
x_grad = self(D, Sf, Sb)
D_grad_list = (Sf+Sb)*x_grad*(Sf*Sb*x_grad-D**2*x_grad**3)/(Sf*Sb-D**2*x_grad**2)**2
Sf_grad_list = -D*x_grad*(Sb**2+D**2 * x_grad**2)/(Sf*Sb-D**2*x_grad**2)**2
Sb_grad_list = -D*x_grad*(Sf**2+D**2 * x_grad**2)/(Sf*Sb-D**2*x_grad**2)**2
# Compute the gradients with respect to D, Sf, and Sb
return at.dot(D_grad_list,grads[0]), at.dot(Sf_grad_list,grads[0]), at.dot(Sb_grad_list,grads[0])
```

Where D, Sb and Sf are random variables and thickness is a constant.

In the main part of the code I then use:

```
rootfinder = RootFinder(thickness=thickness)
beta = rootfinder(D, Sf, Sb)
```

While this does run with step = pm.NUTS, it is very slow (several s/it) and I am worried that there is an issue with the implementation of the grad() method. I also donāt fully understand which gradients are stored in grads[0]. Any help would be appreciated!

Thank you alreadyā¦