source: tags/Release2.1.1b/python/asapfitter.py @ 1271

Last change on this file since 1271 was 1271, checked in by mar637, 18 years ago

fix for fitting data not from scantable. maks was not set correctly

  • Property svn:eol-style set to native
  • Property svn:keywords set to Author Date Id Revision
File size: 20.6 KB
Line 
1import _asap
2from asap import rcParams
3from asap import print_log
4from asap import NUM
5
6class fitter:
7    """
8    The fitting class for ASAP.
9    """
10
11    def __init__(self):
12        """
13        Create a fitter object. No state is set.
14        """
15        self.fitter = _asap.fitter()
16        self.x = None
17        self.y = None
18        self.mask = None
19        self.fitfunc = None
20        self.fitfuncs = None
21        self.fitted = False
22        self.data = None
23        self.components = 0
24        self._fittedrow = 0
25        self._p = None
26        self._selection = None
27
28    def set_data(self, xdat, ydat, mask=None):
29        """
30        Set the absissa and ordinate for the fit. Also set the mask
31        indicationg valid points.
32        This can be used for data vectors retrieved from a scantable.
33        For scantable fitting use 'fitter.set_scan(scan, mask)'.
34        Parameters:
35            xdat:    the abcissa values
36            ydat:    the ordinate values
37            mask:    an optional mask
38
39        """
40        self.fitted = False
41        self.x = xdat
42        self.y = ydat
43        if mask == None:
44            self.mask = NUM.ones(len(xdat))
45        else:
46            self.mask = mask
47        return
48
49    def set_scan(self, thescan=None, mask=None):
50        """
51        Set the 'data' (a scantable) of the fitter.
52        Parameters:
53            thescan:     a scantable
54            mask:        a msk retireved from the scantable
55        """
56        if not thescan:
57            msg = "Please give a correct scan"
58            if rcParams['verbose']:
59                print msg
60                return
61            else:
62                raise TypeError(msg)
63        self.fitted = False
64        self.data = thescan
65        self.mask = None
66        if mask is None:
67            self.mask = NUM.ones(self.data.nchan())
68        else:
69            self.mask = mask
70        return
71
72    def set_function(self, **kwargs):
73        """
74        Set the function to be fit.
75        Parameters:
76            poly:    use a polynomial of the order given
77            gauss:   fit the number of gaussian specified
78        Example:
79            fitter.set_function(gauss=2) # will fit two gaussians
80            fitter.set_function(poly=3)  # will fit a 3rd order polynomial
81        """
82        #default poly order 0
83        n=0
84        if kwargs.has_key('poly'):
85            self.fitfunc = 'poly'
86            n = kwargs.get('poly')
87            self.components = [n]
88        elif kwargs.has_key('gauss'):
89            n = kwargs.get('gauss')
90            self.fitfunc = 'gauss'
91            self.fitfuncs = [ 'gauss' for i in range(n) ]
92            self.components = [ 3 for i in range(n) ]
93        else:
94            msg = "Invalid function type."
95            if rcParams['verbose']:
96                print msg
97                return
98            else:
99                raise TypeError(msg)
100
101        self.fitter.setexpression(self.fitfunc,n)
102        self.fitted = False
103        return
104
105    def fit(self, row=0, estimate=False):
106        """
107        Execute the actual fitting process. All the state has to be set.
108        Parameters:
109            row:        specify the row in the scantable
110            estimate:   auto-compute an initial parameter set (default False)
111                        This can be used to compute estimates even if fit was
112                        called before.
113        Example:
114            s = scantable('myscan.asap')
115            s.set_cursor(thepol=1)        # select second pol
116            f = fitter()
117            f.set_scan(s)
118            f.set_function(poly=0)
119            f.fit(row=0)                  # fit first row
120        """
121        if ((self.x is None or self.y is None) and self.data is None) \
122               or self.fitfunc is None:
123            msg = "Fitter not yet initialised. Please set data & fit function"
124            if rcParams['verbose']:
125                print msg
126                return
127            else:
128                raise RuntimeError(msg)
129
130        else:
131            if self.data is not None:
132                self.x = self.data._getabcissa(row)
133                self.y = self.data._getspectrum(row)
134                from asap import asaplog
135                asaplog.push("Fitting:")
136                i = row
137                out = "Scan[%d] Beam[%d] IF[%d] Pol[%d] Cycle[%d]" % (self.data.getscan(i),self.data.getbeam(i),self.data.getif(i),self.data.getpol(i), self.data.getcycle(i))
138                asaplog.push(out,False)
139        self.fitter.setdata(self.x, self.y, self.mask)
140        if self.fitfunc == 'gauss':
141            ps = self.fitter.getparameters()
142            if len(ps) == 0 or estimate:
143                self.fitter.estimate()
144        try:
145            fxdpar = list(self.fitter.getfixedparameters())
146            if len(fxdpar) and fxdpar.count(0) == 0:
147                 raise RuntimeError,"No point fitting, if all parameters are fixed."
148            converged = self.fitter.fit()
149            if not converged:
150                raise RuntimeError,"Fit didn't converge."
151        except RuntimeError, msg:
152            if rcParams['verbose']:
153                print msg
154            else:
155                raise
156        self._fittedrow = row
157        self.fitted = True
158        print_log()
159        return
160
161    def store_fit(self, filename=None):
162        """
163        Save the fit parameters.
164        Parameters:
165            filename:    if specified save as an ASCII file, if None (default)
166                         store it in the scnatable
167        """
168        if self.fitted and self.data is not None:
169            pars = list(self.fitter.getparameters())
170            fixed = list(self.fitter.getfixedparameters())
171            from asap.asapfit import asapfit
172            fit = asapfit()
173            fit.setparameters(pars)
174            fit.setfixedparameters(fixed)
175            fit.setfunctions(self.fitfuncs)
176            fit.setcomponents(self.components)
177            fit.setframeinfo(self.data._getcoordinfo())
178            if filename is not None:
179                import os
180                filename = os.path.expandvars(os.path.expanduser(filename))
181                if os.path.exists(filename):
182                    raise IOError("File '%s' exists." % filename)
183                fit.save(filename)
184            else:
185                self.data._addfit(fit,self._fittedrow)
186
187    #def set_parameters(self, params, fixed=None, component=None):
188    def set_parameters(self,*args,**kwargs):
189        """
190        Set the parameters to be fitted.
191        Parameters:
192              params:    a vector of parameters
193              fixed:     a vector of which parameters are to be held fixed
194                         (default is none)
195              component: in case of multiple gaussians, the index of the
196                         component
197        """
198        component = None
199        fixed = None
200        params = None
201
202        if len(args) and isinstance(args[0],dict):
203            kwargs = args[0]
204        if kwargs.has_key("fixed"): fixed = kwargs["fixed"]
205        if kwargs.has_key("params"): params = kwargs["params"]
206        if len(args) == 2 and isinstance(args[1], int):
207            component = args[1]
208        if self.fitfunc is None:
209            msg = "Please specify a fitting function first."
210            if rcParams['verbose']:
211                print msg
212                return
213            else:
214                raise RuntimeError(msg)
215        if self.fitfunc == "gauss" and component is not None:
216            if not self.fitted and sum(self.fitter.getparameters()) == 0:
217                pars = list(NUM.zeros(len(self.components)*3))
218                fxd = list(NUM.zeros(len(pars)))
219            else:
220                pars = list(self.fitter.getparameters())
221                fxd = list(self.fitter.getfixedparameters())
222            i = 3*component
223            pars[i:i+3] = params
224            fxd[i:i+3] = fixed
225            params = pars
226            fixed = fxd
227        self.fitter.setparameters(params)
228        if fixed is not None:
229            self.fitter.setfixedparameters(fixed)
230        print_log()
231        return
232
233    def set_gauss_parameters(self, peak, centre, fwhm,
234                             peakfixed=0, centerfixed=0,
235                             fwhmfixed=0,
236                             component=0):
237        """
238        Set the Parameters of a 'Gaussian' component, set with set_function.
239        Parameters:
240            peak, centre, fwhm:  The gaussian parameters
241            peakfixed,
242            centerfixed,
243            fwhmfixed:           Optional parameters to indicate if
244                                 the paramters should be held fixed during
245                                 the fitting process. The default is to keep
246                                 all parameters flexible.
247            component:           The number of the component (Default is the
248                                 component 0)
249        """
250        if self.fitfunc != "gauss":
251            msg = "Function only operates on Gaussian components."
252            if rcParams['verbose']:
253                print msg
254                return
255            else:
256                raise ValueError(msg)
257        if 0 <= component < len(self.components):
258            d = {'params':[peak, centre, fwhm],
259                 'fixed':[peakfixed, centerfixed, fwhmfixed]}
260            self.set_parameters(d, component)
261        else:
262            msg = "Please select a valid  component."
263            if rcParams['verbose']:
264                print msg
265                return
266            else:
267                raise ValueError(msg)
268
269    def get_area(self, component=None):
270        """
271        Return the area under the fitted gaussian component.
272        Parameters:
273              component:   the gaussian component selection,
274                           default (None) is the sum of all components
275        Note:
276              This will only work for gaussian fits.
277        """
278        if not self.fitted: return
279        if self.fitfunc == "gauss":
280            pars = list(self.fitter.getparameters())
281            from math import log,pi,sqrt
282            fac = sqrt(pi/log(16.0))
283            areas = []
284            for i in range(len(self.components)):
285                j = i*3
286                cpars = pars[j:j+3]
287                areas.append(fac * cpars[0] * cpars[2])
288        else:
289            return None
290        if component is not None:
291            return areas[component]
292        else:
293            return sum(areas)
294
295    def get_errors(self, component=None):
296        """
297        Return the errors in the parameters.
298        Parameters:
299            component:    get the errors for the specified component
300                          only, default is all components
301        """
302        if not self.fitted:
303            msg = "Not yet fitted."
304            if rcParams['verbose']:
305                print msg
306                return
307            else:
308                raise RuntimeError(msg)
309        errs = list(self.fitter.geterrors())
310        cerrs = errs
311        if component is not None:
312            if self.fitfunc == "gauss":
313                i = 3*component
314                if i < len(errs):
315                    cerrs = errs[i:i+3]
316        return cerrs
317
318    def get_parameters(self, component=None, errors=False):
319        """
320        Return the fit paramters.
321        Parameters:
322             component:    get the parameters for the specified component
323                           only, default is all components
324        """
325        if not self.fitted:
326            msg = "Not yet fitted."
327            if rcParams['verbose']:
328                print msg
329                return
330            else:
331                raise RuntimeError(msg)
332        pars = list(self.fitter.getparameters())
333        fixed = list(self.fitter.getfixedparameters())
334        errs = list(self.fitter.geterrors())
335        area = []
336        if component is not None:
337            if self.fitfunc == "gauss":
338                i = 3*component
339                cpars = pars[i:i+3]
340                cfixed = fixed[i:i+3]
341                cerrs = errs[i:i+3]
342                a = self.get_area(component)
343                area = [a for i in range(3)]
344            else:
345                cpars = pars
346                cfixed = fixed
347                cerrs = errs
348        else:
349            cpars = pars
350            cfixed = fixed
351            cerrs = errs
352            if self.fitfunc == "gauss":
353                for c in range(len(self.components)):
354                  a = self.get_area(c)
355                  area += [a for i in range(3)]
356        fpars = self._format_pars(cpars, cfixed, errors and cerrs, area)
357        if rcParams['verbose']:
358            print fpars
359        return {'params':cpars, 'fixed':cfixed, 'formatted': fpars,
360                'errors':cerrs}
361
362    def _format_pars(self, pars, fixed, errors, area):
363        out = ''
364        if self.fitfunc == 'poly':
365            c = 0
366            for i in range(len(pars)):
367                fix = ""
368                if len(fixed) and fixed[i]: fix = "(fixed)"
369                if errors :
370                    out += '  p%d%s= %3.6f (%1.6f),' % (c,fix,pars[i], errors[i])
371                else:
372                    out += '  p%d%s= %3.6f,' % (c,fix,pars[i])
373                c+=1
374            out = out[:-1]  # remove trailing ','
375        elif self.fitfunc == 'gauss':
376            i = 0
377            c = 0
378            aunit = ''
379            ounit = ''
380            if self.data:
381                aunit = self.data.get_unit()
382                ounit = self.data.get_fluxunit()
383            while i < len(pars):
384                if len(area):
385                    out += '  %2d: peak = %3.3f %s , centre = %3.3f %s, FWHM = %3.3f %s\n      area = %3.3f %s %s\n' % (c,pars[i],ounit,pars[i+1],aunit,pars[i+2],aunit, area[i],ounit,aunit)
386                else:
387                    out += '  %2d: peak = %3.3f %s , centre = %3.3f %s, FWHM = %3.3f %s\n' % (c,pars[i],ounit,pars[i+1],aunit,pars[i+2],aunit,ounit,aunit)
388                c+=1
389                i+=3
390        return out
391
392    def get_estimate(self):
393        """
394        Return the parameter estimates (for non-linear functions).
395        """
396        pars = self.fitter.getestimate()
397        fixed = self.fitter.getfixedparameters()
398        if rcParams['verbose']:
399            print self._format_pars(pars,fixed,None)
400        return pars
401
402    def get_residual(self):
403        """
404        Return the residual of the fit.
405        """
406        if not self.fitted:
407            msg = "Not yet fitted."
408            if rcParams['verbose']:
409                print msg
410                return
411            else:
412                raise RuntimeError(msg)
413        return self.fitter.getresidual()
414
415    def get_chi2(self):
416        """
417        Return chi^2.
418        """
419        if not self.fitted:
420            msg = "Not yet fitted."
421            if rcParams['verbose']:
422                print msg
423                return
424            else:
425                raise RuntimeError(msg)
426        ch2 = self.fitter.getchi2()
427        if rcParams['verbose']:
428            print 'Chi^2 = %3.3f' % (ch2)
429        return ch2
430
431    def get_fit(self):
432        """
433        Return the fitted ordinate values.
434        """
435        if not self.fitted:
436            msg = "Not yet fitted."
437            if rcParams['verbose']:
438                print msg
439                return
440            else:
441                raise RuntimeError(msg)
442        return self.fitter.getfit()
443
444    def commit(self):
445        """
446        Return a new scan where the fits have been commited (subtracted)
447        """
448        if not self.fitted:
449            msg = "Not yet fitted."
450            if rcParams['verbose']:
451                print msg
452                return
453            else:
454                raise RuntimeError(msg)
455        from asap import scantable
456        if not isinstance(self.data, scantable):
457            msg = "Not a scantable"
458            if rcParams['verbose']:
459                print msg
460                return
461            else:
462                raise TypeError(msg)
463        scan = self.data.copy()
464        scan._setspectrum(self.fitter.getresidual())
465        print_log()
466        return scan
467
468    def plot(self, residual=False, components=None, plotparms=False, filename=None):
469        """
470        Plot the last fit.
471        Parameters:
472            residual:    an optional parameter indicating if the residual
473                         should be plotted (default 'False')
474            components:  a list of components to plot, e.g [0,1],
475                         -1 plots the total fit. Default is to only
476                         plot the total fit.
477            plotparms:   Inidicates if the parameter values should be present
478                         on the plot
479        """
480        if not self.fitted:
481            return
482        if not self._p or self._p.is_dead:
483            if rcParams['plotter.gui']:
484                from asap.asaplotgui import asaplotgui as asaplot
485            else:
486                from asap.asaplot import asaplot
487            self._p = asaplot()
488        self._p.hold()
489        self._p.clear()
490        self._p.set_panels()
491        self._p.palette(0)
492        tlab = 'Spectrum'
493        xlab = 'Abcissa'
494        ylab = 'Ordinate'
495        from matplotlib.numerix import ma,logical_not,logical_and,array
496        m = self.mask
497        if self.data:
498            tlab = self.data._getsourcename(self._fittedrow)
499            xlab = self.data._getabcissalabel(self._fittedrow)
500            m =  logical_and(self.mask,
501                             array(self.data._getmask(self._fittedrow)),
502                                   copy=False)
503                             
504            ylab = self.data._get_ordinate_label()
505
506        colours = ["#777777","#dddddd","red","orange","purple","green","magenta", "cyan"]
507        self._p.palette(0,colours)
508        self._p.set_line(label='Spectrum')
509        y = ma.masked_array(self.y,mask=logical_not(m))
510        self._p.plot(self.x, y)
511        if residual:
512            self._p.palette(1)
513            self._p.set_line(label='Residual')
514            y = ma.masked_array(self.get_residual(),
515                                  mask=logical_not(m))
516            self._p.plot(self.x, y)
517        self._p.palette(2)
518        if components is not None:
519            cs = components
520            if isinstance(components,int): cs = [components]
521            if plotparms:
522                self._p.text(0.15,0.15,str(self.get_parameters()['formatted']),size=8)
523            n = len(self.components)
524            self._p.palette(3)
525            for c in cs:
526                if 0 <= c < n:
527                    lab = self.fitfuncs[c]+str(c)
528                    self._p.set_line(label=lab)
529                    y = ma.masked_array(self.fitter.evaluate(c),
530                                          mask=logical_not(m))
531
532                    self._p.plot(self.x, y)
533                elif c == -1:
534                    self._p.palette(2)
535                    self._p.set_line(label="Total Fit")
536                    y = ma.masked_array(self.fitter.getfit(),
537                                          mask=logical_not(m))
538                    self._p.plot(self.x, y)
539        else:
540            self._p.palette(2)
541            self._p.set_line(label='Fit')
542            y = ma.masked_array(self.fitter.getfit(),
543                                  mask=logical_not(m))
544            self._p.plot(self.x, y)
545        xlim=[min(self.x),max(self.x)]
546        self._p.axes.set_xlim(xlim)
547        self._p.set_axes('xlabel',xlab)
548        self._p.set_axes('ylabel',ylab)
549        self._p.set_axes('title',tlab)
550        self._p.release()
551        if (not rcParams['plotter.gui']):
552            self._p.save(filename)
553        print_log()
554
555    def auto_fit(self, insitu=None, plot=False):
556        """
557        Return a scan where the function is applied to all rows for
558        all Beams/IFs/Pols.
559
560        """
561        from asap import scantable
562        if not isinstance(self.data, scantable) :
563            msg = "Data is not a scantable"
564            if rcParams['verbose']:
565                print msg
566                return
567            else:
568                raise TypeError(msg)
569        if insitu is None: insitu = rcParams['insitu']
570        if not insitu:
571            scan = self.data.copy()
572        else:
573            scan = self.data
574        rows = xrange(scan.nrow())
575        from asap import asaplog
576        asaplog.push("Fitting:")
577        for r in rows:
578            out = " Scan[%d] Beam[%d] IF[%d] Pol[%d] Cycle[%d]" % (scan.getscan(r),scan.getbeam(r),scan.getif(r),scan.getpol(r), scan.getcycle(r))
579            asaplog.push(out, False)
580            self.x = scan._getabcissa(r)
581            self.y = scan._getspectrum(r)
582            self.data = None
583            self.fit()
584            x = self.get_parameters()
585            if plot:
586                self.plot(residual=True)
587                x = raw_input("Accept fit ([y]/n): ")
588                if x.upper() == 'N':
589                    continue
590            scan._setspectrum(self.fitter.getresidual(), r)
591        if plot:
592            self._p.unmap()
593            self._p = None
594        print_log()
595        return scan
596
Note: See TracBrowser for help on using the repository browser.