/
lmfit_wrapper.py
188 lines (166 loc) · 6.22 KB
/
lmfit_wrapper.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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
#! /usr/bin/env python
# -*- coding: utf-8 -*-
""" A wrapper module to use lmfit as a backend for fitting a grism model.
So far, it works for single-row(-set), single transition only.
Maybe more sophisticated ways to do it can be implemented later.
Functions
---------
load_params()
Takes grism-generated DataFrame as input, loads values into an lmfit
Parameters() object.
build_model()
The mathematical function to be minimized.
This function is not necessary to call manually when fitting, but can be
practical to call for generating a plot of the model for a given set of
parameters.
fit_it
Takes as minimum a Parameters() set and a wavelength array x as arguments,
"""
import scipy as sp
import pandas as pd
try:
import lmfit as lf
except ImportError:
print 'Could not find working lmfit installation.'
raise
def load_params(indf):
df = indf.copy()
p = lf.Parameters()
p.clear()
df.Pos += df['Line center']
df.set_value('Contin', 'Identifier', 'x')
for comp in df.index:
print comp
if comp == 'Contin':
varval = df.ix[comp]['Ampl'] + .0001
p.add('Contin_Ampl', value=varval, min=0.,
vary=sp.invert(df.loc[comp]['Lock']))
continue
else:
for col in df.columns:
if col in ['Ampl', 'Sigma', 'Pos', 'Gamma',]:
#name = df.ix[comp]['Identifier']+'_'+col
name = comp + '_' + col
value = df.ix[comp][col]
if col == 'Pos':
varmin = value - 10
varmax = value + 10
vary = sp.invert(df.loc[comp]['Lock'][0])
elif col == 'Sigma':
varmin = 1e-8
varmax = 1e8
vary = sp.invert(df.loc[comp]['Lock'][1])
elif col == 'Ampl':
varmin = 0.
varmax = None
vary = sp.invert(df.loc[comp]['Lock'][2])
else:
varmin = None
varmax = None
p.add(name, value, min=varmin, max=varmax, vary=vary)
return p
def build_model(params, x, data=None, error=None):
""" Dynamically defines the model from a lmfit.Parameters() instance.
Then calculates the residuals by subtracting the data and model.
This is the function that will be minimized by minimize()
Returns either:
* A model generated by an input Parameters() object, with no attempt
being made at minimizing if either data nor errors are provided, OR
* An unweighted set of residuals (data - value) if no errors are
provided, OR
* A inverse-variances-weighted set of residuals if errors are provided.
"""
# Because parameters cannot be hierarcically grouped, we need to create a
# list of components and then cycle through that.
IDs_done = []
for par in params:
iden = par.split('_')[0]
kind = par.split('_')[1]
if not iden in IDs_done:
IDs_done.append(iden)
#print IDs_done
model = sp.zeros_like(x)
model += params['Contin_Ampl'].value
for comp in IDs_done:
if comp == 'Contin':
continue
else:
ampl = params[comp + '_Ampl'].value
pos = params[comp + '_Pos'].value
sigma = params[comp + '_Sigma'].value
line = sp.stats.norm.pdf(x, pos, sigma) \
* sigma * sp.sqrt(2 * sp.pi) * ampl
model += line
if data is None:
return model
elif error is None:
return (model - data)
else:
return sp.sqrt((model - data) ** 2. / error ** 2.)
def fit_it(params, args, method='nelder', kwargs=None):
""" Carries out the fit.
Parameters
----------
params : lmfit.Parameters() instance
Call load_params to generate.
args : tuple
Arguments to pass to the function to minimize. Must contain a
wavelength array as first argument, optional second and third argument
will be interpreted as data and errors, respectively.
arrays are optional.
kwargs : tuple
keyword arguments, will be passed directly to the lmfit.minimize()
function. See lmfit docs for options.
Returns
-------
result : lmfit.Minimizer() object
"""
# For testing:
x = args[0]
data = args[1]
errs = args[2]
earlymodel = build_model(params, x)
# Now: fitting.
print params
try:
result = lf.minimize(build_model, params, args=args, method=method)
result = lf.minimize(build_model, params, args=args, method=method)
result = lf.minimize(build_model, params, args=args, method='lestsq')
except:
result = lf.minimize(build_model, params, args=args, method=method)
result = lf.minimize(build_model, params, args=args, method='leastsq')
lf.report_errors(params)
# Now: moar plots
latemodel = build_model(result.params, x)
#import matplotlib.pyplot as plt
#plt.clf()
#plt.errorbar(x, data, yerr=errs, color='black', label='Data', lw=2.)
#plt.axhline(y=0., color='black')
#plt.plot(x, earlymodel, lw=1.6, label='Guess', color='green')
#plt.plot(x, latemodel, lw=1.6, label='Fit', color='orange')
#plt.legend(fancybox=True, shadow=True)
#plt.grid()
#plt.title('Plot of initial guess and LMfit best fit.')
#plt.xlabel(u'Wavelegth')
#plt.ylabel('Flux')
#plt.show()
return result
def params_to_grism(params, output_format='dict'):
""" Reads lmfit Params() object back into grism-readable formats.
Parameters
----------
output_format : str
Can either be 'dict' (default), which is read into the line profile
editor, or 'df', a pandas dataframe which can be inserted back into a
model dataframe by whatever means one wants.
"""
tempdf = pd.DataFrame()
for thekey in params:
compinfo = thekey.split('_')
compo = compinfo[0]
param = compinfo[1]
value = params[thekey].value
tempdf = tempdf.set_value(compo, param, value)
if output_format == 'dict':
tempdf = tempdf.T.to_dict()
return tempdf