source: trunk/python/asapfitter.py @ 1859

Last change on this file since 1859 was 1859, checked in by Malte Marquarding, 14 years ago

Ticket #193: the rcParamsverbose? flag is only used in standard asap cli mode. Otherwise log messages are always send to the logger and one needs to call asaplog.disable()/asaplog.enable() to controls this. I have also added the function name as the log origin.

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