source: trunk/python/asapfitter.py @ 2153

Last change on this file since 2153 was 2153, checked in by Kana Sugimoto, 13 years ago

New Development: No

JIRA Issue: No (a bug fix)

Ready for Test: Yes

Interface Changes: No

What Interface Changed: Please list interface changes

Test Programs:

Put in Release Notes: No

Module(s): asapfitter, sdfit

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