Defining Complicated function using Theano operations

I have a formula that is completely defined in numpy and scipy and I want to define it using theano ops for use in pymc3 and I want to be able to use HMC or NUTS.

The inputs are v \sin i which is a float, f which are frequencies from an associated Fourier transform, and F_\lambda which are flux densities (light energy per density per wavelength)

Here is the formula:

u_b = 2 \pi f v \sin i
K = \frac{j_1(u_b) }{u_b} - \frac{3 \cos (u_b)}{2 u_b^2} + \frac{3 \sin (u_b)}{2 u_b^3}
F_\lambda^* = F_\lambda \star K

This is essentially a broadening kernel used in astronomy. We want to implement this convolution by multiplying the kernel in the Fourier domain, so the python implementation looks like this

import numpy as np
from scipy.special import j1

def rotational_broaden(wavelengths, fluxes, vsini):
  # This is the frequency spacing with maximum information 
  dv = calculate_dv(wavelengths) 
  freq = np.fft.rfftfreq(fluxes.shape[-1], d=dv)
  flux_ff = np.fft.rfft(fluxes)
  # ignore the 0th frequency
  ub = 2. * np.pi * vsini * freq[1:]

  # Calculate the stellar broadening kernel (Gray 2008)
  sb = j1(ub) / ub - 3 * np.cos(ub) / (2 * ub ** 2) + 3. * np.sin(ub) / (2 * ub ** 3)
  flux_ff[:, 1:] *= sb
  flux_final = np.fft.irfft(flux_ff, n=fluxes.shape[-1])
  return flux_final

I am sort of able to incorporate this using theano, for instance there exists tt.j1. I’m not sure how to work with the Fourier transforms, though, the theano implementations don’t match the numpy versions. Also is the issue with the frequency spacing and only defining the kernel from freq[1:] to avoid f=0 in the denominator of the kernel.