25 lines
731 B
Python
25 lines
731 B
Python
import numpy as np
|
|
from scipy.optimize import least_squares
|
|
|
|
def k_fit(y_dat):
|
|
"""
|
|
Fit linear and quadratic to K total
|
|
"""
|
|
def optimize_func(a, y, n):
|
|
result = []
|
|
i, j = 0, 0
|
|
for u in range(len(y)):
|
|
val = y[0] + a[0] * i + a[1] * j**2 - y[u]
|
|
i = i + 1
|
|
if u >= n: j = j + 1
|
|
result.append(val)
|
|
return result
|
|
mini = None
|
|
for n in range(len(y_dat)):
|
|
a_init = (0.0001, 0.0001)
|
|
result = least_squares(optimize_func, a_init, args=(y_dat, n))
|
|
if mini is None or result["cost"] < mini:
|
|
mini = result["cost"]
|
|
n_fit = n
|
|
a_fit = result["x"]
|
|
return (a_fit[0], a_fit[1], n_fit) |