aboutsummaryrefslogtreecommitdiffstats
path: root/astro/calc_psd.py
blob: ad7659e5e7a95ae856aa78833b42ac8cd9469ed5 (plain)
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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
#!/usr/bin/env python3
#
# Copyright (c) 2015-2017 Aaron LI
# MIT License
#

"""
Compute the radial (i.e., azimuthally averaged) power spectral density
(a.k.a. power spectrum) of a FITS image.

NOTE: The input image must be square.

Credit
------
* Radially averaged power spectrum of 2D real-valued matrix
  Evan Ruzanski
  'raPsd2d.m'
  https://www.mathworks.com/matlabcentral/fileexchange/23636-radially-averaged-power-spectrum-of-2d-real-valued-matrix
"""

import os
import argparse
from functools import lru_cache

import numpy as np
from astropy.io import fits

import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure

plt.style.use("ggplot")


class PSD:
    """
    Computes the 2D power spectral density and the azimuthally averaged power
    spectral density (i.e., 1D radial power spectrum).

    Parameters
    ----------
    image : 2D `~numpy.ndarray`
        Input image array
    pixel : (float, str), optional
        Specify the pixel size and its unit of the image.
        e.g., (0.33, "arcmin")
    step : float, optional
        By default, a logarithmic grid with the specified step ratio
        (default: 1.1) will be used to do the azimuthal averages.
        If specified a value <=1 or None, then a evenly pixel-by-pixel
        (along radial direction) is adopted.
    meanstd : bool, optional
        By default, the median and 16% and 84% percentiles (i.e., 68% error)
        will be calculated for each averaged annulus.  If this option is
        ``True`` then calculate the mean and standard deviation instead.
    """
    def __init__(self, image, pixel=(1.0, "pixel"), step=1.1,
                 meanstd=False, bunit=None):
        self.image = np.array(image, dtype=float)
        self.shape = self.image.shape
        if self.shape[0] != self.shape[1]:
            raise ValueError("input image is not square!")

        self.pixel = pixel
        if step is None or step <= 1:
            self.step = None
        else:
            self.step = step

        self.meanstd = meanstd
        self.bunit = bunit

    @property
    @lru_cache()
    def radii(self):
        """
        The radial (frequency) points where to calculate the powers.
        If ``self.step`` is ``None``, then the powers at every frequency
        point are calculated.  If ``self.step`` is specified, then a
        log-even grid is adopted, which can greatly save computation time
        for large images.
        """
        dim_half = (self.shape[0] + 1) // 2
        x = np.arange(dim_half)
        if self.step is None:
            return x
        else:
            xmax = x.max()
            x2 = list(x[x*(self.step-1) <= 1])
            v1 = x[len(x2)]
            while v1 < xmax:
                x2.append(v1)
                v1 *= self.step
            x2.append(xmax)
            return np.array(x2)

    @property
    @lru_cache()
    def frequencies(self):
        """
        The (spatial) frequencies w.r.t. the above radii.
        """
        radii = self.radii
        freqs = (1 / (self.shape[0] * self.pixel[0])) * radii
        return freqs

    def calc_psd2d(self):
        """
        Computes the 2D power spectral density of the given image.
        Note that the low frequency components are shifted to the center
        of the FFT'ed image.

        NOTE
        ----
        The zero-frequency component is shifted to position of index (0-based)
            (ceil((n-1) / 2), ceil((m-1) / 2)),
        where (n, m) are the number of rows and columns of the image/psd2d.

        Returns
        -------
        2D power spectral density, which has dimension of ${input_unit}^2.
        """
        print("Calculating 2D power spectral density ... ", end="", flush=True)
        rows, cols = self.shape
        # Compute the power spectral density (i.e., power spectrum)
        imgf = np.fft.fftshift(np.fft.fft2(self.image))
        # Normalization w.r.t. image size
        norm = rows * cols * self.pixel[0]**2
        self.psd2d = (np.abs(imgf) ** 2) / norm
        print("DONE", flush=True)
        return self.psd2d

    def calc_psd(self):
        """
        Azimuthally average the above 2D power spectral density to generate
        the 1D radial power spectral density.

        Returns
        -------
        frequencies : 1D float `~numpy.ndarray`
            Spatial frequencies, [{pixel_unit}^(-1)]
        psd1d : 1D float `~numpy.ndarray`
            The median or mean (``self.meanstd=True``) of the powers within
            each (radial) spatial frequency bin.
        psd1d_errl, psd1d_erru : 1D float `~numpy.ndarray`
            The lower and upper errors of the powers.  By default, they are
            determined from the 16% and 84% percentiles w.r.t. the median.
            If ``self.meanstd=True`` then they are the standard deviation.

        Attributes
        ----------
        psd1d
        psd1d_errl
        psd1d_erru
        """
        if not hasattr(self, "ps2d") or self.psd2d is None:
            self.calc_psd2d()

        print("Azimuthally averaging 2D power spectral density ... ",
              end="", flush=True)
        dim = self.shape[0]
        dim_half = (dim+1) // 2
        # NOTE:
        # The zero-frequency component is shifted to position of index
        # (0-based): (ceil((n-1) / 2), ceil((m-1) / 2))
        px = np.arange(dim_half-dim, dim_half)
        x, y = np.meshgrid(px, px)
        rho, phi = self.cart2pol(x, y)

        radii = self.radii
        nr = len(radii)
        if nr > 100:
            print("\n    ... %d data points, may take a while ... " % nr,
                  end="", flush=True)
        else:
            print(" %d data points ... " % nr, end="", flush=True)
        psd1d = np.zeros(nr)
        psd1d_errl = np.zeros(nr)  # lower error
        psd1d_erru = np.zeros(nr)  # upper error
        for i, r in enumerate(radii):
            if (i+1) % 100 == 0:
                percent = 100 * (i+1) / nr
                print("%.1f%% ... " % percent, end="", flush=True)
            ii, jj = (rho <= r).nonzero()
            rho[ii, jj] = np.inf
            data = self.psd2d[ii, jj]
            if self.meanstd:
                psd1d[i] = np.mean(data)
                std = np.std(data)
                psd1d_errl[i] = std
                psd1d_erru[i] = std
            else:
                median, q16, q84 = np.percentile(data, q=(50, 16, 84))
                psd1d[i] = median
                psd1d_errl[i] = median - q16
                psd1d_erru[i] = q84 - median
        print("DONE", flush=True)

        self.psd1d = psd1d
        self.psd1d_errl = psd1d_errl
        self.psd1d_erru = psd1d_erru
        return (self.frequencies, psd1d, psd1d_errl, psd1d_erru)

    @staticmethod
    def cart2pol(x, y):
        """
        Convert Cartesian coordinates to polar coordinates.
        """
        rho = np.sqrt(x**2 + y**2)
        phi = np.arctan2(y, x)
        return (rho, phi)

    @staticmethod
    def pol2cart(rho, phi):
        """
        Convert polar coordinates to Cartesian coordinates.
        """
        x = rho * np.cos(phi)
        y = rho * np.sin(phi)
        return (x, y)

    def save(self, outfile):
        data = np.column_stack((self.frequencies, self.psd1d,
                                self.psd1d_errl, self.psd1d_erru))
        header = [
            "pixel: %s [%s]" % self.pixel,
            "frequency: [%s^-1]" % self.pixel[1],
        ]
        if self.meanstd:
            header += [
                "psd1d: *mean* powers of radial spectral annuli",
                "psd1d_errl, psd1d_erru: *standard deviation* (lower, upper)",
            ]
        else:
            header += [
                "psd1d: *median* powers of radial spectral annuli",
                "psd1d_errl, psd1d_erru: 16% and 84% *percentiles*",
            ]
        header += ["", "frequency   psd1d   psd1d_errl   psd1d_erru"]
        np.savetxt(outfile, data, header="\n".join(header))
        print("Saved PSD data to: %s" % outfile)

    def plot(self, ax):
        """
        Make a plot of the 1D radial power spectrum.
        """
        freqs = self.frequencies
        xmin = freqs[1] / 1.2  # ignore the first 0
        xmax = freqs[-1] * 1.1
        ymin = np.min(self.psd1d) / 10.0
        ymax = np.max(self.psd1d[1:] + self.psd1d_erru[1:]) * 1.5

        if self.meanstd:
            label = "mean"
            labelerr = "standard deviation"
        else:
            label = "median"
            labelerr = "68% percentile range"
        if self.bunit:
            ylabel = r"Power [(%s)$^2$]" % self.bunit
        else:
            ylabel = "Power"

        yerr = np.row_stack((self.psd1d_errl, self.psd1d_erru))
        ax.errorbar(freqs, self.psd1d, yerr=yerr,
                    fmt="none", label=labelerr)
        ax.plot(freqs, self.psd1d, marker="o", label=label)
        ax.set(xscale="log", yscale="log",
               xlim=(xmin, xmax), ylim=(ymin, ymax),
               title="Radial (Azimuthally Averaged) Power Spectral Density",
               xlabel=r"k [%s$^{-1}$]" % self.pixel[1],
               ylabel=ylabel)
        ax.legend()

        if self.pixel[1] != "pixel":
            # Add an additional X axis for pixel-based frequencies
            ax2 = ax.twiny()
            ax2.set_xscale(ax.get_xscale())
            pix_ticks = np.logspace(-4, 0, num=5)  # [pixel^-1]
            ax2.set_xticks(pix_ticks)
            ax2.set_xticklabels([r"10$^{%d}$" % ep
                                 for ep in np.log10(pix_ticks)])
            x1_min, x1_max = ax.get_xlim()
            x2_min, x2_max = x1_min*self.pixel[0], x1_max*self.pixel[0]
            ax2.set_xlim(x2_min, x2_max)
            ax2.set_xlabel(r"k [pixel$^{-1}$] (1 pixel = %.2f %s)" %
                           self.pixel)
            ax2.grid(False)
            # Raise title position to avoid overlapping
            ax.title.set_position([0.5, 1.1])
            return (ax, ax2)
        else:
            return ax


def open_image(infile):
    """
    Open the slice image and return its header and 2D image data.

    NOTE
    ----
    The input slice image may have following dimensions:
    * NAXIS=2: [Y, X]
    * NAXIS=3: [FREQ=1, Y, X]
    * NAXIS=4: [STOKES=1, FREQ=1, Y, X]

    NOTE
    ----
    Only open slice image that has only ONE frequency and ONE Stokes
    parameter.

    Returns
    -------
    header : `~astropy.io.fits.Header`
    image : 2D `~numpy.ndarray`
        The 2D [Y, X] image part of the slice image.
    """
    with fits.open(infile) as f:
        header = f[0].header
        data = f[0].data
    if data.ndim == 2:
        # NAXIS=2: [Y, X]
        image = data
    elif data.ndim == 3 and data.shape[0] == 1:
        # NAXIS=3: [FREQ=1, Y, X]
        image = data[0, :, :]
    elif data.ndim == 4 and data.shape[0] == 1 and data.shape[1] == 1:
        # NAXIS=4: [STOKES=1, FREQ=1, Y, X]
        image = data[0, 0, :, :]
    else:
        raise ValueError("Slice '{0}' has invalid dimensions: {1}".format(
            infile, data.shape))
    return (header, image)


def main():
    parser = argparse.ArgumentParser(
            description="Calculate radial power spectral density")
    parser.add_argument("-C", "--clobber", dest="clobber", action="store_true",
                        help="overwrite the output files if already exist")
    parser.add_argument("-s", "--step", dest="step", type=float, default=1.1,
                        help="step ratio (>1) between 2 consecutive radial " +
                        "frequency points, i.e., a logarithmic grid used. " +
                        "if specified a value <=1, then the power at every " +
                        "radial frequency point will be calculated, " +
                        "i.e., using a even grid, which may be very slow " +
                        "for very large images!")
    parser.add_argument("-p", "--pixelsize", dest="pixelsize", type=float,
                        help="image spatial pixel size [arcsec] " +
                        "(will try to obtain from FITS header)")
    parser.add_argument("-m", "--mean-std", dest="meanstd",
                        action="store_true",
                        help="calculate the mean and standard deviation " +
                        "for each averaged annulus instead of the median " +
                        "16%% and 84%% percentiles (i.e., 68%% error)")
    parser.add_argument("-P", "--plot", dest="plot", action="store_true",
                        help="plot the PSD and save as a PNG image")
    parser.add_argument("-i", "--infile", dest="infile", nargs="+",
                        help="input FITS image(s); if multiple images " +
                        "are provided, they are added first.")
    parser.add_argument("-o", "--outfile", dest="outfile", required=True,
                        help="output TXT file to save the PSD data")
    args = parser.parse_args()

    if args.plot:
        plotfile = os.path.splitext(args.outfile)[0] + ".png"

    # Check output files whether already exists
    if (not args.clobber) and os.path.exists(args.outfile):
        raise OSError("outfile '%s' already exists" % args.outfile)
    if args.plot:
        if (not args.clobber) and os.path.exists(plotfile):
            raise OSError("output plot file '%s' already exists" % plotfile)

    header, image = open_image(args.infile[0])
    bunit = header.get("BUNIT", "???")
    print("Read image from: %s" % args.infile[0])
    print("Image size: %dx%d" % tuple(reversed(image.shape)))
    print("Data unit: %s" % bunit)

    if args.pixelsize:
        pixel = (args.pixelsize/60, "arcmin")  # [arcsec]->[arcmin]
    else:
        try:
            pixel = (header["PixSize"]/60, "arcmin")  # [arcsec]->[arcmin]
        except KeyError:
            try:
                pixel = (abs(header["CDELT1"])*60, "arcmin")  # [deg]->[arcmin]
            except KeyError:
                pixel = (1.0, "pixel")
    print("Image pixel size: %.2f [%s]" % pixel)

    for fn in args.infile[1:]:
        print("Adding additional image: %s" % fn)
        header2, image2 = open_image(fn)
        bunit2 = header2.get("BUNIT", "???")
        if bunit2 == bunit:
            image += image2
        else:
            raise ValueError("image has different unit: %s" % bunit2)

    psd = PSD(image=image, pixel=pixel, step=args.step,
              meanstd=args.meanstd, bunit=bunit)
    psd.calc_psd()
    psd.save(args.outfile)

    if args.plot:
        # Make and save a plot
        fig = Figure(figsize=(8, 8))
        FigureCanvas(fig)
        ax = fig.add_subplot(1, 1, 1)
        psd.plot(ax=ax)
        fig.tight_layout()
        fig.savefig(plotfile, format="png", dpi=150)
        print("Plotted PSD and saved to image: %s" % plotfile)


if __name__ == "__main__":
    main()