base_1basketw.py 5.01 KB
title = "BasketWeavingMaps"
tip = "Corrects baseline effects by fitting polynomials in both scanning directions iteratively"
onein = False

import numpy as np

from guidata.qt.QtGui import QMessageBox
from guidata.dataset.datatypes import DataSet
from guidata.dataset.dataitems import (IntItem, StringItem, ChoiceItem, FloatItem, BoolItem)
from guiqwt.config import _

class NOD3_App:

    def __init__(self, parent):
        self.parent = parent
        self.parent.activateWindow()

    def Error(self, msg):
        QMessageBox.critical(self.parent.parent(), title,
                              _(u"Error:")+"\n%s" % str(msg))

    def compute_app(self, **args):
        class FuncParam(DataSet):
            order = IntItem('Order', default=7, max=13, min=1)
            autocal = BoolItem('Autocal', default=False)
        name = title.replace(" ", "")
        if args == {}:
           #param = FuncParam(_("Press"), "Apply a n-order polynomial fit in scanning direction")
           param = FuncParam(_(title), "Apply a n-order polynomial fit in both scanning directions")
        else:
           param = self.parent.ScriptParameter(name, args)

        # if no parameter needed set param to None. activate next line
        #param = None
        self.parent.compute_11(name, lambda m, p: self.function(m, p), param, onein) 

    def weight(self, data, sigma):
        #mask = -np.isnan(x)
        sigma2 = 2*sigma*sigma
        w = 0.0*data
        rows, cols = w.shape
        for row in range(rows):
            for col in range(cols):
                val = data[row][col]**2
                w[row][col] = np.exp(-val/sigma2) 
        return w
        #return lambda x: np.exp(-x*x/sigma2)

    def autocal(self, data1, data2, dummy):
        mask = (data1 != dummy)
        d1 = list(np.ravel(data1[mask]))
        mask = (data2 != dummy)
        d2 = list(np.ravel(data2[mask]))
        d1.sort()
        d2.sort()
        nl = max(10, len(d1)/10)
        a1, a0 = np.polyfit(d2[-nl:], d1[-nl:], 1)
        rows, cols = data2.shape
        for row in range(rows):
            for col in range(cols):
                if data2[row, col] != dummy:
                   data2[row, col] = a0 + a1*data2[row, col]
        return data1, data2

    def function(self, ms, p):
        dummy = 0.0
        lon = 0
        lat = 0
        lon_data = []
        lat_data = []
        for m in ms:
            data, w = self.parent.nan_check(m.data, dummy, weight=True)
            if m.header['SCANDIR'] in ('ALON', 'XLON', 'ULON', 'GLON', 'RA', 'HA'):
               if lon_data == []:
                  lon_data = data
                  wlon = w
               else:
                  lon_data += data
                  wlon += w
               lon += 1
            elif m.header['SCANDIR'] in ('ALAT', 'XLAT', 'ULAT', 'GLAT', 'DEC'):
               if lat_data == []:
                  lat_data = data
                  wlat = w
               else:
                  lat_data += data
                  wlat += w
               lat += 1
            else:
               self.Error(str("sorry, SCANDIR=%s not defined" % m.header['SCANDIR']))
               return [], p
        if lon == 0 or lat == 0:
           if lon == 0: sd = "Longitude"
           else: sd = "Latitude"
           self.Error(str("sorry, missing maps scanning direction %s" % sd))
           return [], p
        lon_data, w1 = self.parent.nan_check(lon_data/wlon, dummy, weight=True)
        lat_data, w2 = self.parent.nan_check(lat_data/wlat, dummy, weight=True)
        if p.autocal: lon_data, lat_data = self.autocal(lon_data, lat_data, dummy)
        w = w1 + w2
        rows, cols = lon_data.shape
        pfit_col = np.zeros((rows, cols))
        pfit_row = np.zeros((rows, cols))
        for i in range(10):
            diff = (lon_data - lat_data) / w
            aver = (lon_data + lat_data) / w
            x = np.arange(cols)
            y = np.arange(rows)
            rms = 5.0*np.std(diff) 
            for row in range(rows):
                mask = (lon_data[row] != dummy) & (abs(diff[row]) < rms)
                #mask = (lon_data[row] != dummy) & (lat_data[row] != dummy)
                if len(x[mask]) > p.order:
                   poly = np.poly1d(np.polyfit(x[mask], diff[row][mask], p.order))
                   pfit_row[row] = poly(x)
            lon_data -= pfit_row
            for col in range(cols):
                mask = (lat_data[:,col] != dummy) & (abs(diff[:,col]) < rms)
                #mask = (lon_data[:,col] != dummy) & (lat_data[:,col] != dummy)
                if len(y[mask]) > p.order:
                   poly = np.poly1d(np.polyfit(y[mask], diff[:,col][mask], p.order))
                   pfit_col[:,col] = poly(y)
            lat_data += pfit_col
        for m in ms:
            if m.header['SCANDIR'] in ('ALON', 'XLON', 'ULON', 'GLON', 'RA', 'HA'):
               m.data = lon_data - pfit_row
            else:
               m.data = lat_data + pfit_col
        #m.data = aver + pfit_col - pfit_row
        #if not m.header['SCANDIR'] in ("ALON", "ALAT"): m.header.__delitem__('SCANDIR')
        return ms, p