Commit 15f590c2 authored by LUSTIG Peter's avatar LUSTIG Peter

added fitmodels

parent d073b615
from astropy.modeling import Fittable1DModel, Parameter
from scipy.special import erf
import numpy as np
import matplotlib.pyplot as plt
from astropy.modeling.fitting import LevMarLSQFitter
class ErrorFunction(Fittable1DModel):
x0 = Parameter()
amplitude = Parameter()
offset = Parameter()
@staticmethod
def evaluate(x, x0, amplitude, offset):
return (amplitude * erf(x-x0)) + offset
@staticmethod
def fit_deriv(x, x0, amplitude, offset):
d_x0 = - 2. / np.sqrt(np.pi) * np.exp(-(x-x0)**2)
d_amplitude = erf(x-x0)
d_offset = np.ones_like(x)
return [d_x0, d_amplitude, d_offset]
class ModifiedErrorFunction(Fittable1DModel):
x0 = Parameter()
sigma = Parameter()
amplitude = Parameter()
offset = Parameter()
@staticmethod
def evaluate(x, x0, sigma, amplitude, offset):
return erf((x-x0) / (np.sqrt(2) * sigma)) / 2 * amplitude + offset
@staticmethod
def fit_deriv(x, x0, sigma, amplitude, offset):
d_x0 = - amplitude * np.exp(-(x-x0) ** 2 / (2 * sigma**2)) \
/ (np.sqrt(2 * np.pi) * sigma)
d_sigma = - amplitude * (x - x0) \
* np.exp(-(x-x0) ** 2 / (2 * sigma**2)) \
/ (np.sqrt(2 * np.pi) * sigma**2)
d_amplitude = .5 * erf(x - x0) / (np.sqrt(2) * sigma)
d_offset = np.ones_like(x)
return [d_x0, d_sigma, d_amplitude, d_offset]
class FluxBoosting(Fittable1DModel):
# Derivative ugly and therefore not implemented, is obtained numerically
Sd = Parameter()
sigma = Parameter()
@staticmethod
def evaluate(Si, Sd, sigma):
up = sigma * np.exp(-.5 * ((Si-Sd) / sigma) ** 2)
down = 1 + erf((Si - Sd)/(np.sqrt(2)*sigma))
return Si + np.sqrt(2. / np.pi) * up / down
class FluxRelativeBoosting(Fittable1DModel):
# sqrt(2/pi)*s*exp(-((x-g)/(sqrt(2)s))^2)/(1+erf((x-g)/(sqrt(2)*s)))
# Derivative ugly and therefore not implemented, is obtained numerically
Sd = Parameter()
sigma = Parameter()
@staticmethod
def evaluate(Si, Sd, sigma):
up = sigma * np.exp(-.5 * ((Si-Sd) / sigma) ** 2)
down = 1 + erf((Si - Sd)/(np.sqrt(2)*sigma))
return 1 + np.sqrt(2. / np.pi) * up / (down * Si)
def test_ErrorFunction():
# create test data
np.random.seed(0)
parameters = (3, .5, .5)
f = ErrorFunction(*parameters)
N = 100
x = np.linspace(-10, 10, N)
y = f(x) + np.random.randn(N)*.1
# Define model and fix parameters
mymodel = ErrorFunction(0, 0, 0)
# mymodel.offset.fixed = True
# mymodel.amplitude.fixed = True
# fit
fitter = LevMarLSQFitter()
res = fitter(mymodel, x, y)
# assert np.isclose(res.x0.value, 3., rtol=.2)
assert np.allclose([res.x0.value, res.amplitude.value, res.offset.value],
parameters, atol=.2)
'''
# %matplotlib tk
plt.plot(x, y)
plt.plot(x, res(x))
plt.show()
'''
def test_ModifiedErrorFunction():
np.random.seed(0)
parameters = (3, 4, .5, .5)
f = ModifiedErrorFunction(*parameters)
N = 100
x = np.linspace(-10+3, 10+3, N)
y = f(x) + np.random.randn(N)*.01
# Define model and fix parameters
mymodel = ModifiedErrorFunction(*parameters)
# mymodel.x0.fixed = True
# mymodel.x0.sigma = True
# mymodel.offset.fixed = True
# mymodel.amplitude.fixed = True
# fit
fitter = LevMarLSQFitter()
res = fitter(mymodel, x, y)
resparameters = [res.x0.value, res.sigma.value, res.amplitude.value,
res.offset.value]
print(resparameters)
# %matplotlib tk
# plt.plot(x, y, marker='o', linestyle=None)
# plt.plot(x, res(x))
assert np.allclose(resparameters, parameters, atol=.3)
if __name__ == '__main__':
test_ErrorFunction()
test_ModifiedErrorFunction()
# % matplotlib tk
f = FluxRelativeBoosting(1, 1)
x = np.linspace(.1, 10, 1000)
plt.plot(x, f(x))
ax = plt.gca()
ax.set_xscale('log')
ax.set_yscale('log')
# plt.show()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment