zachopy/twod.py at master · TESScience/zachopy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
'''Tools for dealing with 1D arrays, particularly timeseries of images.'''
import numpy as np
import scipy.ndimage
import matplotlib.pyplot as plt
import displays.ds9

def scatter(cube, axis=0):
    shape = np.array(cube.shape)
    shape[axis] = 1

    med = np.median(cube, axis=axis)
    mad = np.median(np.abs(cube - med.reshape(shape)), axis=axis)
    return 1.48*mad


def stack(cube, axis=0, threshold=5.0):
    '''Combine a cube of images into one mean image, using a MAD noise estimator to reject outliers.'''
    shape = np.array(cube.shape)
    shape[axis] = 1

    med = np.median(cube, axis=axis).reshape(shape)
    noise = scatter(cube, axis=axis).reshape(shape)
    good = (np.abs(cube - med) < threshold*noise) | (noise == 0)
    mean = np.sum(good*cube, axis=axis)/np.sum(good, axis=axis)

    return mean, noise#.squeeze()

def interpolateOverBadPixels(image, bad, scale=2, visualize=False):
    '''Take an image and a bad pixel mask (=1 where bad, =0 otherwise) and interpolate over the bad pixels, using a Gaussian smoothing.'''
    smoothed = scipy.ndimage.filters.gaussian_filter(image*(bad == False), sigma=[scale,scale])
    weights = scipy.ndimage.filters.gaussian_filter(np.array(bad == False).astype(np.float), sigma=[scale,scale])
    smoothed /=weights
    corrected = image + 0.0
    corrected[bad.astype(np.bool)] = smoothed[bad.astype(np.bool)]
    if visualize:
        gs = plt.matplotlib.gridspec.GridSpec(1,4, wspace=0,hspace=0)
        orig = plt.subplot(gs[0])
        smoo = plt.subplot(gs[1], sharex=orig, sharey=orig)
        weig = plt.subplot(gs[2], sharex=orig, sharey=orig)
        corr = plt.subplot(gs[3], sharex=orig, sharey=orig)
        kw = dict(cmap='gray', interpolation='nearest')
        orig.imshow(np.log(image), **kw)
        smoo.imshow(np.log(smoothed), **kw)
        weig.imshow(bad, **kw)
        corr.imshow(np.log(corrected), **kw)
        plt.draw()
        a = raw_input('test?')
    return corrected


def polyInterpolate(image, bad, axis=0, order=2, visualize=True):
    '''Take an image and a bad pixel mask (=1 where bad, =0 otherwise), fit polynomials to the good data in one dimension, and return this polynomial smoothed version.'''
    n = image.shape[axis]
    smoothed = np.zeros_like(image)
    if visualize:
        plt.figure('interpolating in a 2D image')
        gs = plt.matplotlib.gridspec.GridSpec(1,1)
        ax = plt.subplot(gs[0,0])
        d = displays.ds9('poly')

    for i in range(image.shape[axis]):
        ydata = image.take(i, axis=axis)
        baddata = bad.take(i, axis=axis)
        xdata = np.arange(len(ydata))
        ok = baddata == 0
        med = np.median(ydata[ok])
        mad = np.median(np.abs(ydata[ok] - med))
        ok = ok*(np.abs(ydata - med) < 4*1.48*mad)


        fit = np.polynomial.polynomial.polyfit(xdata[ok], ydata[ok], order)
        poly = np.polynomial.polynomial.Polynomial(fit)

        if visualize:
            ax.cla()
            ax.plot(xdata, ydata, alpha=0.3, color='gray', marker='o')
            ax.plot(xdata[ok], ydata[ok], alpha=1, color='black', marker='o', linewidth=0)
            ax.plot(xdata, poly(xdata), alpha=0.3, color='red', linewidth=5)
            ax.set_ylim(0, np.percentile(image[bad == 0], 95))
            plt.draw()

        if axis == 0:
            smoothed[i,:] =  poly(xdata)
        else:
            smoothed[:,i] = poly(xdata)

    if visualize:
        d.update(smoothed)
    return smoothed
        #a = raw_input('?')
    '''smoothed = scipy.ndimage.filters.gaussian_filter(image*(bad == False), sigma=[scale,scale])
    weights = scipy.ndimage.filters.gaussian_filter(np.array(bad == False).astype(np.float), sigma=[scale,scale])
    smoothed /=weights
    corrected = image + 0.0
    corrected[bad.astype(np.bool)] = smoothed[bad.astype(np.bool)]
    if visualize:
        gs = plt.matplotlib.gridspec.GridSpec(1,4, wspace=0,hspace=0)
        orig = plt.subplot(gs[0])
        smoo = plt.subplot(gs[1], sharex=orig, sharey=orig)
        weig = plt.subplot(gs[2], sharex=orig, sharey=orig)
        corr = plt.subplot(gs[3], sharex=orig, sharey=orig)
        kw = dict(cmap='gray', interpolation='nearest')
        orig.imshow(np.log(image), **kw)
        smoo.imshow(np.log(smoothed), **kw)
        weig.imshow(bad, **kw)
        corr.imshow(np.log(corrected), **kw)
        plt.draw()
        a = raw_input('test?')
    return corrected'''

def estimateBackground(image, axis=-1):

    display = displays.ds9('background subtraction')
    display.one(image, clobber=True)

    roughSky1d = np.median(image, axis)
    shape = np.array(image.shape)
    shape[axis] = 1
    roughSkyImage = np.ones_like(image)*roughSky1d.reshape(shape)
    display.one(roughSkyImage)
    display.one(image - roughSkyImage)
    assert(False)