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

Last change on this file since 1627 was 1627, checked in by Takeshi Nakazato, 15 years ago

New Development: No

JIRA Issue: No

Ready to Release: Yes

Interface Changes: No

What Interface Changed: Please list interface changes

Test Programs: List test programs

Put in Release Notes: No

Module(s): Module Names change impacts.

Description: Describe your changes here...

Bug fix.

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