File: _moments.py

package info (click to toggle)
spectral-cube 0.4.0-2
  • links: PTS, VCS
  • area: main
  • in suites: stretch
  • size: 1,856 kB
  • ctags: 1,955
  • sloc: python: 13,201; makefile: 156; ansic: 88
file content (188 lines) | stat: -rw-r--r-- 4,754 bytes parent folder | download | duplicates (3)
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
from __future__ import print_function, absolute_import, division

import numpy as np

from .cube_utils import iterator_strategy
from .np_compat import allbadtonan

"""
Functions to compute moment maps in a variety of ways
"""


def _moment_shp(cube, axis):
    """
    Return the shape of the moment map

    Parameters
    -----------
    cube : SpectralCube
       The cube to collapse
    axis : int
       The axis to collapse along (numpy convention)

    Returns
    -------
    ny, nx
    """
    return cube.shape[:axis] + cube.shape[axis + 1:]


def _slice0(cube, axis):
    """
    0th moment along an axis, calculated slicewise

    Parameters
    ----------
    cube : SpectralCube
    axis : int

    Returns
    -------
    moment0 : array
    """
    shp = _moment_shp(cube, axis)
    result = np.zeros(shp)

    view = [slice(None)] * 3

    valid = np.zeros(shp, dtype=np.bool)
    for i in range(cube.shape[axis]):
        view[axis] = i
        plane = cube._get_filled_data(fill=np.nan, view=view)
        valid |= np.isfinite(plane)
        result += np.nan_to_num(plane) * cube._pix_size_slice(axis)
    result[~valid] = np.nan
    return result


def _slice1(cube, axis):
    """
    1st moment along an axis, calculated slicewise

    Parameters
    ----------
    cube : SpectralCube
    axis : int

    Returns
    -------
    moment1 : array
    """
    shp = _moment_shp(cube, axis)
    result = np.zeros(shp)

    view = [slice(None)] * 3
    pix_size = cube._pix_size_slice(axis)
    pix_cen = cube._pix_cen()[axis]
    weights = np.zeros(shp)

    for i in range(cube.shape[axis]):
        view[axis] = i
        plane = cube._get_filled_data(fill=0, view=view)
        result += (plane *
                   pix_cen[view] *
                   pix_size)
        weights += plane * pix_size
    return result / weights


def moment_slicewise(cube, order, axis):
    """
    Compute moments by accumulating the result 1 slice at a time
    """
    if order == 0:
        return _slice0(cube, axis)
    if order == 1:
        return _slice1(cube, axis)

    shp = _moment_shp(cube, axis)
    result = np.zeros(shp)

    view = [slice(None)] * 3
    pix_size = cube._pix_size_slice(axis)
    pix_cen = cube._pix_cen()[axis]
    weights = np.zeros(shp)

    # would be nice to get mom1 and momn in single pass over data
    # possible for mom2, not sure about general case
    mom1 = _slice1(cube, axis)

    for i in range(cube.shape[axis]):
        view[axis] = i
        plane = cube._get_filled_data(fill=0, view=view)
        result += (plane *
                   (pix_cen[view] - mom1) ** order *
                   pix_size)
        weights += plane * pix_size

    return (result / weights)


def moment_raywise(cube, order, axis):
    """
    Compute moments by accumulating the answer one ray at a time
    """
    shp = _moment_shp(cube, axis)
    out = np.zeros(shp) * np.nan

    pix_cen = cube._pix_cen()[axis]
    pix_size = cube._pix_size_slice(axis)

    for x, y, slc in cube._iter_rays(axis):
        # the intensity, i.e. the weights
        include = cube._mask.include(data=cube._data, wcs=cube._wcs, view=slc,
                                     wcs_tolerance=cube._wcs_tolerance)
        if not include.any():
            continue

        data = cube.flattened(slc).value * pix_size

        if order == 0:
            out[x, y] = data.sum()
            continue

        order1 = (data * pix_cen[slc][include]).sum() / data.sum()
        if order == 1:
            out[x, y] = order1
            continue

        ordern = (data * (pix_cen[slc][include] - order1) ** order).sum()
        ordern /= data.sum()

        out[x, y] = ordern
    return out

def moment_cubewise(cube, order, axis):
    """
    Compute the moments by working with the entire data at once
    """

    pix_cen = cube._pix_cen()[axis]
    data = cube._get_filled_data() * cube._pix_size_slice(axis)

    if order == 0:
        return allbadtonan(np.nansum)(data, axis=axis)

    if order == 1:
        return (np.nansum(data * pix_cen, axis=axis) /
                np.nansum(data, axis=axis))
    else:
        mom1 = moment_cubewise(cube, 1, axis)

        # insert an axis so it broadcasts properly
        shp = list(_moment_shp(cube, axis))
        shp.insert(axis, 1)
        mom1 = mom1.reshape(shp)

        return (np.nansum(data * (pix_cen - mom1) ** order, axis=axis) /
                np.nansum(data, axis=axis))


def moment_auto(cube, order, axis):
    """
    Build a moment map, choosing a strategy to balance speed and memory.
    """
    strategy = dict(cube=moment_cubewise, ray=moment_raywise,
                    slice=moment_slicewise)
    return strategy[iterator_strategy(cube, axis)](cube, order, axis)