Search code examples
python3dcurve-fittingsurfacemathematical-expressions

Fitting a coned semi torus surface to a scatter plot in python


So maybe this is more of a mathematical question but I'm a little stuck on this one. I have a scatter plot and am trying to fit a surface to it. Using two Gaussian's is working fairly well but its not perfect, and since I know the original data comes from a torus I'd like to force the fit to use an ellipse with an intensity.

What I have currently created the following fig: attempted fit

You can see it works kinda well but its not quite giving me that perfect Torus shape I desire. I now show you the fit functions I'm currently using and ones I've tried.

To generate the plot:

# Our function to fit is going to be a sum of two-dimensional Gaussians
def gaussian(x, y, x0, y0, xalpha, yalpha, A):
    return A * np.exp( -((x-x0)/xalpha)**2 -((y-y0)/yalpha)**2)

# This is the callable that is passed to curve_fit. M is a (2,N) array
# where N is the total number of data points in Z, which will be ravelled
# to one dimension.
def _gaussian(M, *args):
    x, y = M
    arr = np.zeros(x.shape)
    for i in range(len(args)//5):
       arr += gaussian(x, y, *args[i*5:i*5+5])
    return arr

And others I've tried:

def distancefromellips(x, y, h, a, k, b):
    xfact = (x-h)**2/a**2
    yfact = (x-k)**2/b**2
    return abs(xfact-yfact-1)

def distanceandgauss(x,y,h,a,k,b,x0, y0, xalpha, yalpha, A):
    return distancefromellips(x, y, h, a, k, b) * gaussian(x, y, x0, y0, xalpha, yalpha, A)

def _ellipseG(M, *args):
    x, y = M
    arr = np.zeros(x.shape)
    for i in range(len(args)//9):
       arr += distanceandgauss(x, y, *args[i*9:i*9+9])
    return arr

Let me know what you think. Maybe the solution is just more Gaussian's? TIA.


Solution

  • The solution is to use a function in polar coordinates. A "coned torus" is just a Gaussian over R with some variance on theta. A function such as:

    def polarf(x, y, a, b, c, d, A, B, C, lamb):
        r = np.sqrt(x**2 + y**2)
        #print(r)
        theta = np.arctan(y/x)
        #print(theta.min())
        return a * np.cos(theta*A) * np.exp(-(((r-b+B*np.cos(theta))**2)/c)) * (C*(np.exp(-lamb)*(lamb**r))/scipy.special.gamma(r+1)) #Standard best fit so far
    

    Will work well giving such an image:

    enter image description here