/
curve_fit.py
103 lines (80 loc) · 2.75 KB
/
curve_fit.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
########################################################################
#
# Classes for curve-fitting data.
#
########################################################################
#
# This file is part of FinanceOps:
#
# https://github.com/Hvass-Labs/FinanceOps
#
# Published under the MIT License. See the file LICENSE for details.
#
# Copyright 2018 by Magnus Erik Hvass Pedersen
#
########################################################################
from scipy.optimize import curve_fit
########################################################################
class CurveFit:
"""
Base-class for curve-fitting.
"""
def __init__(self, x=None, y=None):
"""
Pass numpy-arrays as the x and y args for fitting.
:param x: Optional numpy array with input-values.
:param y: Optional numpy array with output-values.
"""
if x is not None and y is not None:
self.fit(x=x, y=y)
def _f(self, x, *args, **kwargs):
"""Function to be fitted. Override this!"""
raise NotImplementedError()
def predict(self, x):
"""
Use the fitted function to predict new output-values.
Call fit() before calling predict().
:param x: Numpy array with input-values.
:return: Predicted output-values.
"""
return self._f(x, *self.params)
def fit(self, x, y):
"""
Fit the function parameters to the given data.
Call this before predict().
:param x: Numpy array with input-values.
:param y: Numpy array with output-values.
:return: Nothing.
"""
self.params, self.covar = curve_fit(self._f, x, y)
class CurveFitLinear(CurveFit):
"""
Linear curve-fitting: y = a * x + b
First call fit() then predict().
"""
def __init__(self, *args, **kwargs):
"""
Pass numpy-arrays as the x and y args for fitting.
:param x: Optional numpy array with input-values.
:param y: Optional numpy array with output-values.
"""
CurveFit.__init__(self, *args, **kwargs)
def _f(self, x, a, b):
"""Linear function to be fitted."""
return a * x + b
class CurveFitReciprocal(CurveFit):
"""
Reciprocal curve-fitting: y = a / x + b
First call fit() then predict().
"""
def __init__(self, *args, **kwargs):
"""
Pass numpy-arrays as the x and y args for fitting.
:param x: Optional numpy array with input-values.
:param y: Optional numpy array with output-values.
"""
CurveFit.__init__(self, *args, **kwargs)
def _f(self, x, a, b):
"""Reciprocal function to be fitted."""
return a / x + b
########################################################################