source: branches/alma/python/asapfitter.py @ 1701

Last change on this file since 1701 was 1701, checked in by WataruKawasaki, 14 years ago

New Development: Yes

JIRA Issue: Yes (CAS-1800 + CAS-1807)

Ready to Release: Yes

Interface Changes: Yes

What Interface Changed: added new methods to scantable and fitter.

Test Programs:

Put in Release Notes: No

Module(s): sdfit, sdflag

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