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 = 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), at.dot(Sf_grad_list,grads), at.dot(Sb_grad_list,grads)
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. Any help would be appreciated!
Thank you already…