Hello,
I am having trouble getting a correct fit using a piecewise function that is split up over three different intervals. I tried using if and elif statements but there was an issue with a part of the function being an array and the other part being a list... is there another way to do this without using if and elif statements???? Below is my code. Any help would be greatly appreciated! Thanks! - Sam
from numpy import sqrt, pi, exp, linspace, loadtxt, arctan, log
import math
from lmfit import Model, Minimizer
import scipy.special
import matplotlib.pyplot as plt
data = loadtxt('data.dat')
r = data[:,0]
v_tot = data[:,3]
def rotation1(r, r_d, M, r_c, rho_c):
conv = 3.0857758e21
Sigma_0 = (M/(2.0*pi*(r_d*r_d)))
G = 0.0000000678
R = r*conv
x = (R/(2.0*r_d))
Ve_sqrd_1 = (pi*G*Sigma_0)*(R**2/r_d)
Ve_sqrd_2 = (scipy.special.i0(x)*scipy.special.k0(x)-scipy.special.i1(x)*scipy.special.k1(x))
Ve_sqrd = Ve_sqrd_1*Ve_sqrd_2
Ve = sqrt(Ve_sqrd)
Ve_kms = Ve*0.00001
xh = R/r_c
c = 0.866
e = sqrt(1-(c**2)) #c is the axis ratio
for cookie in xh:
if cookie < 0.3:
xi = (sqrt((e**2-cookie**2)))
psi = sqrt(cookie**2+e**2)
Vh_1 = (pi*G*rho_c*(r_c**2)*c)
Vh_2 = (-2.0/psi)*(arctan(xh)+arctan((e**2)/(cookie+(c*psi))))
Vh_3 = (1.0/psi)*log((cookie**2+1.0)*(((xh+psi)**2)/((cookie+(c*psi))**2+e**4)))
Vh_4 = (2.0/xi)*(arctan(xi/cookie)-arctan((c*xi)/(cookie+e**2)))
Vh = Vh_1*(Vh_2+Vh_3+Vh_4)
Vh_s = sqrt(Vh)
Vh_kms = Vh_s*0.00001
elif 0.3 < cookie < 1.2:
dx = cookie - e
Vh = sqrt(G*rho_c*(r_c**2))*(0.79273 + 1.18206*dx - 0.763731*(dx**2) + 0.175637*(dx**3) + 0.131187*(dx**4) - 0.125809*(dx**5) - 0.00628458*(dx**6) + 0.084855456*(dx**7) - 0.0606543127*(dx**8))
Vh_kms = Vh*0.00001
elif cookie > 1.2:
xi_m = sqrt(cookie**2-e**2)
xi_p = sqrt(cookie**2+e**2)
Vh_1 = (pi*G*rho_c*(r_c**2)*c)
Vh_2 = (-2.0/xi_p)*(arctan(cookie)+arctan((e**2)/(cookie+(c*xi_p))))
Vh_3 = (1.0/xi_p)*log((cookie**2+1.0)*(((cookie+xi_p)**2)/((cookie+(c*xi_p))**2+e**4)))
Vh_4 = (2.0/xi_m)*(log(cookie+1.0)*((cookie+xi_m)/(cookie+e**2+(c*xi_m))))
Vh = Vh_1*(Vh_2+Vh_3+Vh_4)
Vh_s = sqrt(Vh)
Vh_kms = Vh_s*0.00001
model = (sqrt(Ve_kms**2+Vh_kms**2))
return model
rmod1 = Model(rotation1)
result1 = rmod1.fit(v_tot, r=r, r_d = 4.0e22, M = 9.0e43, r_c = 4.0e22, rho_c = 2.0e-24)
print(result1.fit_report())
plt.plot(r, v_tot, 'bo', label = 'data)
plt.plot(r, result1.init_fit, 'k--', label = 'initial fit')
plt.plot(r, result1.best_fit, 'r-', label = 'best fit')
legend = plt.legend(loc='lower right')
plt.title('Data Total Fit')
plt.xlabel('Radius (kpc)')
plt.ylabel('Velocity (km/s)')
plt.show()