import math from scipy.optimize import minimize def nelson_siegel(level, slope, shape, decay): def calc(time): factor = (1-math.exp(time/decay))/(time/decay) return (level + slope*factor +shape*(factor - math.exp(time/decay)) ) return calc def fit(data): def error(val): level = val[0] slope = val[1] shape = val[2] decay = val[3] total = 0 for time, rate in data: if time>0.0: nval = nelson_siegel(level, slope, shape, decay)(time) total += (rate-nval)*(rate-nval) pass pass return total init = [0.01,0.01,0.01,1] test_error = error([0.01, 0.01, 0.01, 1]) result = minimize(error, init) vals = [x.item() for x in result.x] test_error2 = error(vals) return vals