"""Fit data to sine/cosine, circle, ellipse or Gaussian functions."""
#############################################################################
# Copyright (c) 2023 by Paul Scherrer Institute, Switzerland
# All rights reserved.
# Authors: Zheqiao Geng
#############################################################################
'''
#########################################################################
Here collects routines for fitting algorithms
Implemented:
- fit_sincos : fit the sine or cosine function
- fit_circle : fit the 2D points to a circle function
- fit_ellipse : fit the 2D points to an ellipse function
- fit_Gaussian : fit a 1D Guassian function
#########################################################################
'''
import numpy as np
from scipy.optimize import curve_fit
[docs]
def fit_sincos(X_rad, Y, target = 'cos'):
'''
Fit the sine or cosine function
- ``y = A * cos(x + phi) + c = (A*cos(phi)) * cos(x) - (A*sin(phi)) * sin(x) + c``
- ``y = A * sin(x + phi) + c = (A*cos(phi)) * sin(x) + (A*sin(phi)) * cos(x) + c``
Parameters:
X_rad: numpy array, phase array in radian
Y: numpy array, value of the sine or cosine function
target: 'sin' or 'cos', determine which function to fit to
Returns:
status: boolean, success (True) or fail (False)
A: float, amplitude of the function
phi_rad: float, phase of the function, rad
c: float, offset of the function
'''
# check the input
if (not X_rad.shape == Y.shape) or (X_rad.shape[0] < 3):
return False, None, None, None
# make the least-square fitting
if target == 'cos':
P = np.linalg.lstsq(np.vstack((np.cos(X_rad), -np.sin(X_rad), np.ones(Y.shape))).T, Y, rcond = None)
else:
P = np.linalg.lstsq(np.vstack((np.sin(X_rad), np.cos(X_rad), np.ones(Y.shape))).T, Y, rcond = None)
# calculate the result
A = np.sqrt(P[0][0]**2 + P[0][1]**2) # P[0][0] = A*cos(phi); P[0][1] = A*sin(phi)
phi_rad = np.arctan2(P[0][1], P[0][0])
c = P[0][2]
return True, A, phi_rad, c
[docs]
def fit_circle(X, Y):
'''
Fit a circle. Given ``x``, ``y`` data points and find ``x0``, ``y0`` and ``r`` following
``(x - x0)^2 + (y - y0)^2 = r^2``.
Parameters:
X: numpy array, x-coordinate of the points
Y: numpy array, y-coordinate of the points
Returns:
status: boolean, success (True) or fail (False)
x0, y0: float, center coordinate of the circle
r: float, radius of the circle
'''
# check the input
if (not X.shape == Y.shape) or (X.shape[0] < 3):
return False, None, None, None
# make the fit again with least square
P = np.linalg.lstsq(np.vstack((2*X, 2*Y, np.ones(Y.shape))).T, X**2+Y**2, rcond = None)
# calculate the result
x0 = P[0][0]
y0 = P[0][1]
r = np.sqrt(P[0][2] + x0**2 + y0**2)
return True, x0, y0, r
[docs]
def fit_ellipse(X, Y):
'''
Fit an ellipse. Get the general ellipse function and its characteristics, including
semi-major axis ``a``, semi-minor axis ``b``, center coordinates ``(x0, y0)`` and rotation
angle ``sita`` (the angle from the positive horizontal axis to the ellipse's major axis).
The general ellipse equation is ``A*X^2 + B*X*Y + C*Y^2 + D*X + E*Y + F = 0``.
When making the fitting, we divide the equation by ``C`` to normalize the coefficient of
``Y^2`` to 1 and move it to the right side of the fitting equation. The ellipse is derived
with the following steps:
1. define a canonical ellipse:
- ``X1 = a * cos(phi)``
- ``Y1 = b * sin(phi)``
where ``phi`` is a phase vector covering from 0 to 2*pi
2. rotate the canonical ellipse:
- ``X2 = X1 * cos(sita) - Y1 * sin(sita)``
- ``Y2 = X1 * sin(sita) + Y1 * cos(sita)``
3. add offset to the rotated ellipse:
- ``X = X2 + x0``
- ``Y = Y2 + y0``
See the webpage: https://en.wikipedia.org/wiki/Ellipse.
Parameters:
X: numpy array, x-coordinate of the points
Y: numpy array, y-coordinate of the points
Returns:
status: boolean, success (True) or fail (False)
Coef: list, coefficiets derived from the least-square fitting
a: float, semi-major
b: float, semi-minor
x0, y0: float, center of the ellipse
sita: float, angle of the ellipse (see above), rad
'''
# check the input
if (not X.shape == Y.shape) or (X.shape[0] < 3):
return (False,) + (None,)*6
# make the fit with least square
P = np.linalg.lstsq(np.vstack((X**2, X*Y, X, Y, np.ones(Y.shape))).T, -Y**2, rcond = None)
# calculate the characteristics
A = P[0][0]
B = P[0][1]
C = 1.0
D = P[0][2]
E = P[0][3]
F = P[0][4]
a = -np.sqrt(2*(A*E**2 + C*D**2 - B*D*E + (B**2 - 4*A*C)*F)*((A+C) + np.sqrt((A-C)**2 + B**2))) / (B**2 - 4*A*C)
b = -np.sqrt(2*(A*E**2 + C*D**2 - B*D*E + (B**2 - 4*A*C)*F)*((A+C) - np.sqrt((A-C)**2 + B**2))) / (B**2 - 4*A*C)
x0 = (2*C*D - B*E) / (B**2 - 4*A*C)
y0 = (2*A*E - B*D) / (B**2 - 4*A*C)
if B != 0.0:
sita = np.arctan((C - A - np.sqrt((A-C)**2 + B**2)) / B)
else:
if A <= C:
sita = 0
else:
sita = np.pi/2
return True, [A,B,C,D,E,F], a, b, x0, y0, sita
[docs]
def fit_Gaussian(X, Y):
'''
Fit a Guassian distribution function.
See the webpage: https://pythonguides.com/scipy-normal-distribution/.
Parameters:
X: numpy array, x-coordinate of the points
Y: numpy array, y-coordinate of the points
Returns:
status: boolean, success (True) or fail (False)
a: float, magnitude scale of the un-normalized distribution
mu: float, mean value
sigma: float, standard deviation
'''
# check the input
if (not X.shape == Y.shape) or (X.shape[0] < 3):
return (False,) + (None,)*3
# define the Gaussian (normal) distribution
def norm_dist(x, a, mu, var):
y = a * np.exp(-0.5 / var * (x - mu)**2) # since X,Y are not normalized, "a" can be arbitrary value
return y
# get the initial guess of the parameters
mu_est = sum(X * Y) / sum(Y)
var_est = sum((X - mu_est)**2 * Y) / sum(Y)
# make the fit
P, cov = curve_fit(norm_dist, X, Y, p0 = [1.0, mu_est, var_est])
a, mu, var = P[0], P[1], P[2]
return True, a, mu, np.sqrt(var)