source: tags/casa3.2.0asap/python/asapfitter.py

Last change on this file was 2047, checked in by WataruKawasaki, 13 years ago

New Development: Yes

JIRA Issue: Yes CAS-2847

Ready for Test: Yes

Interface Changes: No

What Interface Changed:

Test Programs:

Put in Release Notes: Yes

Module(s): scantable

Description: added {auto_}sinusoid_baseline() for sinusoidal baseline fitting. also minor bug fixes for asapfitter.


  • Property svn:eol-style set to native
  • Property svn:keywords set to Author Date Id Revision
File size: 25.1 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        indicationg 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        self._p.hold()
546        self._p.clear()
547        self._p.set_panels()
548        self._p.palette(0)
549        tlab = 'Spectrum'
550        xlab = 'Abcissa'
551        ylab = 'Ordinate'
552        from numpy import ma,logical_not,logical_and,array
553        m = self.mask
554        if self.data:
555            tlab = self.data._getsourcename(self._fittedrow)
556            xlab = self.data._getabcissalabel(self._fittedrow)
557            m =  logical_and(self.mask,
558                             array(self.data._getmask(self._fittedrow),
559                                   copy=False))
560
561            ylab = self.data._get_ordinate_label()
562
563        colours = ["#777777","#dddddd","red","orange","purple","green","magenta", "cyan"]
564        nomask=True
565        for i in range(len(m)):
566            nomask = nomask and m[i]
567        label0='Masked Region'
568        label1='Spectrum'
569        if ( nomask ):
570            label0=label1
571        else:
572            y = ma.masked_array( self.y, mask = m )
573            self._p.palette(1,colours)
574            self._p.set_line( label = label1 )
575            self._p.plot( self.x, y )
576        self._p.palette(0,colours)
577        self._p.set_line(label=label0)
578        y = ma.masked_array(self.y,mask=logical_not(m))
579        self._p.plot(self.x, y)
580        if residual:
581            self._p.palette(7)
582            self._p.set_line(label='Residual')
583            y = ma.masked_array(self.get_residual(),
584                                  mask=logical_not(m))
585            self._p.plot(self.x, y)
586        self._p.palette(2)
587        if components is not None:
588            cs = components
589            if isinstance(components,int): cs = [components]
590            if plotparms:
591                self._p.text(0.15,0.15,str(self.get_parameters()['formatted']),size=8)
592            n = len(self.components)
593            self._p.palette(3)
594            for c in cs:
595                if 0 <= c < n:
596                    lab = self.fitfuncs[c]+str(c)
597                    self._p.set_line(label=lab)
598                    y = ma.masked_array(self.fitter.evaluate(c),
599                                          mask=logical_not(m))
600
601                    self._p.plot(self.x, y)
602                elif c == -1:
603                    self._p.palette(2)
604                    self._p.set_line(label="Total Fit")
605                    y = ma.masked_array(self.fitter.getfit(),
606                                          mask=logical_not(m))
607                    self._p.plot(self.x, y)
608        else:
609            self._p.palette(2)
610            self._p.set_line(label='Fit')
611            y = ma.masked_array(self.fitter.getfit(),
612                                  mask=logical_not(m))
613            self._p.plot(self.x, y)
614        xlim=[min(self.x),max(self.x)]
615        self._p.axes.set_xlim(xlim)
616        self._p.set_axes('xlabel',xlab)
617        self._p.set_axes('ylabel',ylab)
618        self._p.set_axes('title',tlab)
619        self._p.release()
620        if (not rcParams['plotter.gui']):
621            self._p.save(filename)
622
623    @asaplog_post_dec
624    def auto_fit(self, insitu=None, plot=False):
625        """
626        Return a scan where the function is applied to all rows for
627        all Beams/IFs/Pols.
628
629        """
630        from asap import scantable
631        if not isinstance(self.data, scantable) :
632            msg = "Data is not a scantable"
633            raise TypeError(msg)
634        if insitu is None: insitu = rcParams['insitu']
635        if not insitu:
636            scan = self.data.copy()
637        else:
638            scan = self.data
639        rows = xrange(scan.nrow())
640        # Save parameters of baseline fits as a class attribute.
641        # NOTICE: This does not reflect changes in scantable!
642        if len(rows) > 0: self.blpars=[]
643        asaplog.push("Fitting:")
644        for r in rows:
645            out = " Scan[%d] Beam[%d] IF[%d] Pol[%d] Cycle[%d]" % (scan.getscan(r),
646                                                                   scan.getbeam(r),
647                                                                   scan.getif(r),
648                                                                   scan.getpol(r),
649                                                                   scan.getcycle(r))
650            asaplog.push(out, False)
651            self.x = scan._getabcissa(r)
652            self.y = scan._getspectrum(r)
653            self.mask = mask_and(self.mask, scan._getmask(r))
654            self.data = None
655            self.fit()
656            x = self.get_parameters()
657            fpar = self.get_parameters()
658            if plot:
659                self.plot(residual=True)
660                x = raw_input("Accept fit ([y]/n): ")
661                if x.upper() == 'N':
662                    self.blpars.append(None)
663                    continue
664            scan._setspectrum(self.fitter.getresidual(), r)
665            self.blpars.append(fpar)
666        if plot:
667            self._p.unmap()
668            self._p = None
669        return scan
Note: See TracBrowser for help on using the repository browser.