-
Notifications
You must be signed in to change notification settings - Fork 60
/
spectral_cube.py
4286 lines (3565 loc) · 169 KB
/
spectral_cube.py
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
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""
A class to represent a 3-d position-position-velocity spectral cube.
"""
from __future__ import print_function, absolute_import, division
import warnings
from functools import wraps
import operator
import re
import itertools
import copy
import tempfile
import textwrap
from pathlib import PosixPath
import six
from six.moves import zip, range
import dask.array as da
import astropy.wcs
from astropy import units as u
from astropy.io.fits import PrimaryHDU, BinTableHDU, Header, Card, HDUList
from astropy.utils.console import ProgressBar
from astropy import log
from astropy import wcs
from astropy import convolution
from astropy import stats
from astropy.constants import si
from astropy.io.registry import UnifiedReadWriteMethod
import numpy as np
from radio_beam import Beam, Beams
from . import cube_utils
from . import wcs_utils
from . import spectral_axis
from .masks import (LazyMask, LazyComparisonMask, BooleanArrayMask, MaskBase,
is_broadcastable_and_smaller)
from .ytcube import ytCube
from .lower_dimensional_structures import (Projection, Slice, OneDSpectrum,
LowerDimensionalObject,
VaryingResolutionOneDSpectrum
)
from .base_class import (BaseNDClass, SpectralAxisMixinClass,
DOPPLER_CONVENTIONS, SpatialCoordMixinClass,
MaskableArrayMixinClass, MultiBeamMixinClass,
HeaderMixinClass, BeamMixinClass,
)
from .utils import (cached, warn_slow, VarianceWarning, BeamWarning,
UnsupportedIterationStrategyWarning, WCSMismatchWarning,
NotImplementedWarning, SliceWarning, SmoothingWarning,
StokesWarning, ExperimentalImplementationWarning,
BeamAverageWarning, NonFiniteBeamsWarning, BeamWarning,
WCSCelestialError, BeamUnitsError)
from .spectral_axis import (determine_vconv_from_ctype, get_rest_value_from_wcs,
doppler_beta, doppler_gamma, doppler_z)
from .io.core import SpectralCubeRead, SpectralCubeWrite
from packaging.version import Version, parse
__all__ = ['BaseSpectralCube', 'SpectralCube', 'VaryingResolutionSpectralCube']
# apply_everywhere, world: do not have a valid cube to test on
__doctest_skip__ = ['BaseSpectralCube._apply_everywhere']
try:
from scipy import ndimage
scipyOK = True
except ImportError:
scipyOK = False
warnings.filterwarnings('ignore', category=wcs.FITSFixedWarning, append=True)
SIGMA2FWHM = 2. * np.sqrt(2. * np.log(2.))
# convenience structures to keep track of the reversed index
# conventions between WCS and numpy
np2wcs = {2: 0, 1: 1, 0: 2}
_NP_DOC = """
Ignores excluded mask elements.
Parameters
----------
axis : int (optional)
The axis to collapse, or None to perform a global aggregation
how : cube | slice | ray | auto
How to compute the aggregation. All strategies give the same
result, but certain strategies are more efficient depending
on data size and layout. Cube/slice/ray iterate over
decreasing subsets of the data, to conserve memory.
Default='auto'
""".replace('\n', '\n ')
def aggregation_docstring(func):
@wraps(func)
def wrapper(*args, **kwargs):
return func(*args, **kwargs)
wrapper.__doc__ += _NP_DOC
return wrapper
_PARALLEL_DOC = """
Other Parameters
----------------
parallel : bool
Use joblib to parallelize the operation.
If set to ``False``, will force the use of a single core without
using ``joblib``.
num_cores : int or None
The number of cores to use when applying this function in parallel
across the cube.
use_memmap : bool
If specified, a memory mapped temporary file on disk will be
written to rather than storing the intermediate spectra in memory.
"""
def parallel_docstring(func):
@wraps(func)
def wrapper(*args, **kwargs):
return func(*args, **kwargs)
line1 = wrapper.__doc__.split("\n")[1]
indentation = " "*(len(line1) - len(line1.lstrip()))
try:
wrapper.__doc__ += textwrap.indent(_PARALLEL_DOC, indentation)
except AttributeError:
# python2.7
wrapper.__doc__ = textwrap.dedent(wrapper.__doc__) + _PARALLEL_DOC
return wrapper
def _apply_spectral_function(arguments, outcube, function, **kwargs):
"""
Helper function to apply a function to a spectrum.
Needs to be declared toward the top of the code to allow pickling by
joblib.
"""
(spec, includemask, ii, jj) = arguments
if np.any(includemask):
outcube[:,jj,ii] = function(spec, **kwargs)
else:
outcube[:,jj,ii] = spec
def _apply_spatial_function(arguments, outcube, function, **kwargs):
"""
Helper function to apply a function to an image.
Needs to be declared toward the top of the code to allow pickling by
joblib.
"""
(img, includemask, ii) = arguments
if np.any(includemask):
outcube[ii, :, :] = function(img, **kwargs)
else:
outcube[ii, :, :] = img
class BaseSpectralCube(BaseNDClass, MaskableArrayMixinClass,
SpectralAxisMixinClass, SpatialCoordMixinClass,
HeaderMixinClass):
def __init__(self, data, wcs, mask=None, meta=None, fill_value=np.nan,
header=None, allow_huge_operations=False, wcs_tolerance=0.0):
# Deal with metadata first because it can affect data reading
self._meta = meta or {}
# must extract unit from data before stripping it
if 'BUNIT' in self._meta:
self._unit = cube_utils.convert_bunit(self._meta["BUNIT"])
elif hasattr(data, 'unit'):
self._unit = data.unit
else:
self._unit = None
# data must not be a quantity when stored in self._data
if hasattr(data, 'unit'):
# strip the unit so that it can be treated as cube metadata
data = data.value
# TODO: mask should be oriented? Or should we assume correctly oriented here?
self._data, self._wcs = cube_utils._orient(data, wcs)
self._wcs_tolerance = wcs_tolerance
self._spectral_axis = None
self._mask = mask # specifies which elements to Nan/blank/ignore
# object or array-like object, given that WCS needs
# to be consistent with data?
#assert mask._wcs == self._wcs
self._fill_value = fill_value
self._header = Header() if header is None else header
if not isinstance(self._header, Header):
raise TypeError("If a header is given, it must be a fits.Header")
# We don't pass the spectral unit via the initializer since the user
# should be using ``with_spectral_unit`` if they want to set it.
# However, we do want to keep track of what units the spectral axis
# should be returned in, otherwise astropy's WCS can change the units,
# e.g. km/s -> m/s.
# This can be overridden with Header below
self._spectral_unit = u.Unit(self._wcs.wcs.cunit[2])
# This operation is kind of expensive?
header_specaxnum = astropy.wcs.WCS(header).wcs.spec
header_specaxunit = spectral_axis.unit_from_header(self._header,
spectral_axis_number=header_specaxnum+1)
# Allow the original header spectral axis unit to override the default
# unit
if header_specaxunit is not None:
self._spectral_unit = header_specaxunit
self._spectral_scale = spectral_axis.wcs_unit_scale(self._spectral_unit)
self.allow_huge_operations = allow_huge_operations
self._cache = {}
@property
def _is_huge(self):
return cube_utils.is_huge(self)
@property
def _new_thing_with(self):
return self._new_cube_with
def _new_cube_with(self, data=None, wcs=None, mask=None, meta=None,
fill_value=None, spectral_unit=None, unit=None,
wcs_tolerance=None, **kwargs):
data = self._data if data is None else data
if unit is None and hasattr(data, 'unit'):
if data.unit != self.unit:
raise u.UnitsError("New data unit '{0}' does not"
" match cube unit '{1}'. You can"
" override this by specifying the"
" `unit` keyword."
.format(data.unit, self.unit))
unit = data.unit
elif unit is not None:
# convert string units to Units
if not isinstance(unit, u.Unit):
unit = u.Unit(unit)
if hasattr(data, 'unit'):
if u.Unit(unit) != data.unit:
raise u.UnitsError("The specified new cube unit '{0}' "
"does not match the input unit '{1}'."
.format(unit, data.unit))
elif self._unit is not None:
unit = self.unit
wcs = self._wcs if wcs is None else wcs
mask = self._mask if mask is None else mask
if meta is None:
meta = {}
meta.update(self._meta)
if unit is not None:
meta['BUNIT'] = unit.to_string(format='FITS')
fill_value = self._fill_value if fill_value is None else fill_value
spectral_unit = self._spectral_unit if spectral_unit is None else u.Unit(spectral_unit)
cube = self.__class__(data=data, wcs=wcs, mask=mask, meta=meta,
fill_value=fill_value, header=self._header,
allow_huge_operations=self.allow_huge_operations,
wcs_tolerance=wcs_tolerance or self._wcs_tolerance,
**kwargs)
cube._spectral_unit = spectral_unit
cube._spectral_scale = spectral_axis.wcs_unit_scale(spectral_unit)
return cube
read = UnifiedReadWriteMethod(SpectralCubeRead)
write = UnifiedReadWriteMethod(SpectralCubeWrite)
@property
def unit(self):
""" The flux unit """
if self._unit:
return self._unit
else:
return u.one
@property
def shape(self):
""" Length of cube along each axis """
return self._data.shape
@property
def size(self):
""" Number of elements in the cube """
return self._data.size
@property
def base(self):
""" The data type 'base' of the cube - useful for, e.g., joblib """
return self._data.base
def __len__(self):
return self.shape[0]
@property
def ndim(self):
""" Dimensionality of the data """
return self._data.ndim
def __repr__(self):
s = "{1} with shape={0}".format(self.shape, self.__class__.__name__)
if self.unit is u.one:
s += ":\n"
else:
s += " and unit={0}:\n".format(self.unit)
s += (" n_x: {0:6d} type_x: {1:8s} unit_x: {2:5s}"
" range: {3:12.6f}:{4:12.6f}\n".format(self.shape[2],
self.wcs.wcs.ctype[0],
self.wcs.wcs.cunit[0],
self.longitude_extrema[0],
self.longitude_extrema[1],))
s += (" n_y: {0:6d} type_y: {1:8s} unit_y: {2:5s}"
" range: {3:12.6f}:{4:12.6f}\n".format(self.shape[1],
self.wcs.wcs.ctype[1],
self.wcs.wcs.cunit[1],
self.latitude_extrema[0],
self.latitude_extrema[1],
))
s += (" n_s: {0:6d} type_s: {1:8s} unit_s: {2:5s}"
" range: {3:12.3f}:{4:12.3f}".format(self.shape[0],
self.wcs.wcs.ctype[2],
self._spectral_unit,
self.spectral_extrema[0],
self.spectral_extrema[1],
))
return s
@property
@cached
def spectral_extrema(self):
_spectral_min = self.spectral_axis.min()
_spectral_max = self.spectral_axis.max()
return u.Quantity((_spectral_min, _spectral_max))
def apply_numpy_function(self, function, fill=np.nan,
reduce=True, how='auto',
projection=False,
unit=None,
check_endian=False,
progressbar=False,
includemask=False,
**kwargs):
"""
Apply a numpy function to the cube
Parameters
----------
function : Numpy ufunc
A numpy ufunc to apply to the cube
fill : float
The fill value to use on the data
reduce : bool
reduce indicates whether this is a reduce-like operation,
that can be accumulated one slice at a time.
sum/max/min are like this. argmax/argmin/stddev are not
how : cube | slice | ray | auto
How to compute the moment. All strategies give the same
result, but certain strategies are more efficient depending
on data size and layout. Cube/slice/ray iterate over
decreasing subsets of the data, to conserve memory.
Default='auto'
projection : bool
Return a :class:`~spectral_cube.lower_dimensional_structures.Projection` if the resulting array is 2D or a
OneDProjection if the resulting array is 1D and the sum is over both
spatial axes?
unit : None or `astropy.units.Unit`
The unit to include for the output array. For example,
`SpectralCube.max` calls
``SpectralCube.apply_numpy_function(np.max, unit=self.unit)``,
inheriting the unit from the original cube.
However, for other numpy functions, e.g. `numpy.argmax`, the return
is an index and therefore unitless.
check_endian : bool
A flag to check the endianness of the data before applying the
function. This is only needed for optimized functions, e.g. those
in the `bottleneck <https://pypi.python.org/pypi/Bottleneck>`_ package.
progressbar : bool
Show a progressbar while iterating over the slices through the
cube?
kwargs : dict
Passed to the numpy function.
Returns
-------
result : :class:`~spectral_cube.lower_dimensional_structures.Projection` or `~astropy.units.Quantity` or float
The result depends on the value of ``axis``, ``projection``, and
``unit``. If ``axis`` is None, the return will be a scalar with or
without units. If axis is an integer, the return will be a
:class:`~spectral_cube.lower_dimensional_structures.Projection` if ``projection`` is set
"""
# leave axis in kwargs to avoid overriding numpy defaults, e.g. if the
# default is axis=-1, we don't want to force it to be axis=None by
# specifying that in the function definition
axis = kwargs.get('axis', None)
if how == 'auto':
strategy = cube_utils.iterator_strategy(self, axis)
else:
strategy = how
out = None
log.debug("applying numpy function {0} with strategy {1}"
.format(function, strategy))
if strategy == 'slice' and reduce:
out = self._reduce_slicewise(function, fill, check_endian,
includemask=includemask,
progressbar=progressbar, **kwargs)
elif how == 'ray':
out = self.apply_function(function, **kwargs)
elif how not in ['auto', 'cube']:
warnings.warn("Cannot use how=%s. Using how=cube" % how,
UnsupportedIterationStrategyWarning)
if out is None:
out = function(self._get_filled_data(fill=fill,
check_endian=check_endian),
**kwargs)
if axis is None:
# return is scalar
if unit is not None:
return u.Quantity(out, unit=unit)
else:
return out
elif projection and reduce:
meta = {'collapse_axis': axis}
meta.update(self._meta)
if hasattr(axis, '__len__') and len(axis) == 2:
# if operation is over two spatial dims
if set(axis) == set((1,2)):
new_wcs = self._wcs.sub([wcs.WCSSUB_SPECTRAL])
header = self._nowcs_header
if cube_utils._has_beam(self):
bmarg = {'beam': self.beam}
elif cube_utils._has_beams(self):
bmarg = {'beams': self.unmasked_beams}
else:
bmarg = {}
return self._oned_spectrum(value=out,
wcs=new_wcs,
copy=False,
unit=unit,
header=header,
meta=meta,
spectral_unit=self._spectral_unit,
**bmarg
)
else:
warnings.warn("Averaging over a spatial and a spectral "
"dimension cannot produce a Projection "
"quantity (no units or WCS are preserved).",
SliceWarning
)
return out
else:
new_wcs = wcs_utils.drop_axis(self._wcs, np2wcs[axis])
header = self._nowcs_header
return Projection(out, copy=False, wcs=new_wcs, meta=meta,
unit=unit, header=header)
else:
return out
def _reduce_slicewise(self, function, fill, check_endian,
includemask=False, progressbar=False, **kwargs):
"""
Compute a numpy aggregation by grabbing one slice at a time
"""
ax = kwargs.pop('axis', None)
full_reduce = ax is None
ax = ax or 0
if isinstance(ax, tuple):
assert len(ax) == 2 # we only work with cubes...
iterax = [x for x in range(3) if x not in ax][0]
else:
iterax = ax
log.debug("reducing slicewise with axis = {0}".format(ax))
if includemask:
planes = self._iter_mask_slices(iterax)
else:
planes = self._iter_slices(iterax, fill=fill, check_endian=check_endian)
result = next(planes)
if progressbar:
progressbar = ProgressBar(self.shape[iterax])
pbu = progressbar.update
else:
pbu = lambda: True
if isinstance(ax, tuple):
# have to make a result a list of itself, since we already "got"
# the first plane above
result = [function(result, axis=(0,1), **kwargs)]
for plane in planes:
# apply to axes 0 and 1, because we're fully reducing the plane
# to a number if we're applying over two axes
result.append(function(plane, axis=(0,1), **kwargs))
pbu()
result = np.array(result)
else:
for plane in planes:
# axis = 2 means we're stacking two planes, the previously
# computed one and the current one
result = function(np.dstack((result, plane)), axis=2, **kwargs)
pbu()
if full_reduce:
result = function(result)
return result
def get_mask_array(self):
"""
Convert the mask to a boolean numpy array
"""
return self._mask.include(data=self._data, wcs=self._wcs,
wcs_tolerance=self._wcs_tolerance)
def _naxes_dropped(self, view):
"""
Determine how many axes are being selected given a view.
(1,2) -> 2
None -> 3
1 -> 1
2 -> 1
"""
if hasattr(view,'__len__'):
return len(view)
elif view is None:
return 3
else:
return 1
@aggregation_docstring
@warn_slow
def sum(self, axis=None, how='auto', **kwargs):
"""
Return the sum of the cube, optionally over an axis.
"""
from .np_compat import allbadtonan
projection = self._naxes_dropped(axis) in (1,2)
return self.apply_numpy_function(allbadtonan(np.nansum), fill=np.nan,
how=how, axis=axis, unit=self.unit,
projection=projection, **kwargs)
@aggregation_docstring
@warn_slow
def mean(self, axis=None, how='cube', **kwargs):
"""
Return the mean of the cube, optionally over an axis.
"""
projection = self._naxes_dropped(axis) in (1,2)
if how == 'slice':
# two-pass approach: first total the # of points,
# then total the value of the points, then divide
# (a one-pass approach is possible but requires
# more sophisticated bookkeeping)
counts = self._count_nonzero_slicewise(axis=axis,
progressbar=kwargs.get('progressbar'))
ttl = self.apply_numpy_function(np.nansum, fill=np.nan, how=how,
axis=axis, unit=None,
projection=False, **kwargs)
out = ttl / counts
if projection:
if self._naxes_dropped(axis) == 1:
new_wcs = wcs_utils.drop_axis(self._wcs, np2wcs[axis])
meta = {'collapse_axis': axis}
meta.update(self._meta)
return Projection(out, copy=False, wcs=new_wcs,
meta=meta,
unit=self.unit, header=self._nowcs_header)
elif axis == (1,2):
newwcs = self._wcs.sub([wcs.WCSSUB_SPECTRAL])
if cube_utils._has_beam(self):
bmarg = {'beam': self.beam}
elif cube_utils._has_beams(self):
bmarg = {'beams': self.unmasked_beams}
else:
bmarg = {}
return self._oned_spectrum(value=out,
wcs=newwcs,
copy=False,
unit=self.unit,
spectral_unit=self._spectral_unit,
meta=self.meta,
**bmarg
)
else:
# this is a weird case, but even if projection is
# specified, we can't return a Quantity here because of WCS
# issues. `apply_numpy_function` already does this
# silently, which is unfortunate.
warnings.warn("Averaging over a spatial and a spectral "
"dimension cannot produce a Projection "
"quantity (no units or WCS are preserved).",
SliceWarning
)
return out
else:
return out
return self.apply_numpy_function(np.nanmean, fill=np.nan, how=how,
axis=axis, unit=self.unit,
projection=projection, **kwargs)
def _count_nonzero_slicewise(self, axis=None, progressbar=False):
"""
Count the number of finite pixels along an axis slicewise. This is a
helper function for the mean and std deviation slicewise iterators.
"""
counts = self.apply_numpy_function(np.sum, fill=np.nan,
how='slice', axis=axis,
unit=None,
projection=False,
progressbar=progressbar,
includemask=True)
return counts
@aggregation_docstring
@warn_slow
def std(self, axis=None, how='cube', ddof=0, **kwargs):
"""
Return the standard deviation of the cube, optionally over an axis.
Other Parameters
----------------
ddof : int
Means Delta Degrees of Freedom. The divisor used in calculations
is ``N - ddof``, where ``N`` represents the number of elements. By
default ``ddof`` is zero.
"""
projection = self._naxes_dropped(axis) in (1,2)
if how == 'slice':
if axis is None:
raise NotImplementedError("The overall standard deviation "
"cannot be computed in a slicewise "
"manner. Please use a "
"different strategy.")
if hasattr(axis, '__len__') and len(axis) == 2:
return self.apply_numpy_function(np.nanstd,
axis=axis,
how='slice',
projection=projection,
unit=self.unit,
**kwargs)
else:
counts = self._count_nonzero_slicewise(axis=axis)
ttl = self.apply_numpy_function(np.nansum, fill=np.nan, how='slice',
axis=axis, unit=None,
projection=False, **kwargs)
# Equivalent, but with more overhead:
# ttl = self.sum(axis=axis, how='slice').value
mean = ttl/counts
planes = self._iter_slices(axis, fill=np.nan, check_endian=False)
result = (next(planes)-mean)**2
for plane in planes:
result = np.nansum(np.dstack((result, (plane-mean)**2)), axis=2)
out = (result/(counts-ddof))**0.5
if projection:
new_wcs = wcs_utils.drop_axis(self._wcs, np2wcs[axis])
meta = {'collapse_axis': axis}
meta.update(self._meta)
return Projection(out, copy=False, wcs=new_wcs,
meta=meta,
unit=self.unit, header=self._nowcs_header)
else:
return out
# standard deviation cannot be computed as a trivial step-by-step
# process. There IS a one-pass algorithm for std dev, but it is not
# implemented, so we must force cube here. We could and should also
# implement raywise reduction
return self.apply_numpy_function(np.nanstd, fill=np.nan, how=how,
axis=axis, unit=self.unit,
projection=projection, **kwargs)
@aggregation_docstring
@warn_slow
def mad_std(self, axis=None, how='cube', **kwargs):
"""
Use astropy's mad_std to computer the standard deviation
"""
if int(astropy.__version__[0]) < 2:
raise NotImplementedError("mad_std requires astropy >= 2")
projection = self._naxes_dropped(axis) in (1,2)
if how == 'ray' and not hasattr(axis, '__len__'):
# no need for fill here; masked-out data are simply not included
return self.apply_numpy_function(stats.mad_std,
axis=axis,
how='ray',
unit=self.unit,
projection=projection,
ignore_nan=True,
)
elif how == 'slice' and hasattr(axis, '__len__') and len(axis) == 2:
return self.apply_numpy_function(stats.mad_std,
axis=axis,
how='slice',
projection=projection,
unit=self.unit,
fill=np.nan,
ignore_nan=True,
**kwargs)
elif how in ('ray', 'slice'):
raise NotImplementedError('Cannot run mad_std slicewise or raywise '
'unless the dimensionality is also reduced in the same direction.')
else:
return self.apply_numpy_function(stats.mad_std,
fill=np.nan,
axis=axis,
unit=self.unit,
ignore_nan=True,
how=how,
projection=projection, **kwargs)
@aggregation_docstring
@warn_slow
def max(self, axis=None, how='auto', **kwargs):
"""
Return the maximum data value of the cube, optionally over an axis.
"""
projection = self._naxes_dropped(axis) in (1,2)
return self.apply_numpy_function(np.nanmax, fill=np.nan, how=how,
axis=axis, unit=self.unit,
projection=projection, **kwargs)
@aggregation_docstring
@warn_slow
def min(self, axis=None, how='auto', **kwargs):
"""
Return the minimum data value of the cube, optionally over an axis.
"""
projection = self._naxes_dropped(axis) in (1,2)
return self.apply_numpy_function(np.nanmin, fill=np.nan, how=how,
axis=axis, unit=self.unit,
projection=projection, **kwargs)
@aggregation_docstring
@warn_slow
def argmax(self, axis=None, how='auto', **kwargs):
"""
Return the index of the maximum data value.
The return value is arbitrary if all pixels along ``axis`` are
excluded from the mask.
"""
return self.apply_numpy_function(np.nanargmax, fill=-np.inf,
reduce=False, projection=False,
how=how, axis=axis, **kwargs)
@aggregation_docstring
@warn_slow
def argmin(self, axis=None, how='auto', **kwargs):
"""
Return the index of the minimum data value.
The return value is arbitrary if all pixels along ``axis`` are
excluded from the mask
"""
return self.apply_numpy_function(np.nanargmin, fill=np.inf,
reduce=False, projection=False,
how=how, axis=axis, **kwargs)
def _argmaxmin_world(self, axis, method, **kwargs):
'''
Return the spatial or spectral index of the maximum or minimum value.
Use `argmax_world` and `argmin_world` directly.
'''
operation_name = '{}_world'.format(method)
if wcs_utils.is_pixel_axis_to_wcs_correlated(self.wcs, axis):
raise WCSCelestialError("{} requires the celestial axes"
" to be aligned along image axes."
.format(operation_name))
if method == 'argmin':
arg_pixel_plane = self.argmin(axis=axis, **kwargs)
elif method == 'argmax':
arg_pixel_plane = self.argmax(axis=axis, **kwargs)
else:
raise ValueError("`method` must be 'argmin' or 'argmax'")
# Convert to WCS coordinates.
out = cube_utils.world_take_along_axis(self, arg_pixel_plane, axis)
# Compute whether the mask has any valid data along `axis`
collapsed_mask = self.mask.include().any(axis=axis)
out[~collapsed_mask] = np.NaN
# Return a Projection.
new_wcs = wcs_utils.drop_axis(self._wcs, np2wcs[axis])
meta = {'collapse_axis': axis}
meta.update(self._meta)
return Projection(out, copy=False, wcs=new_wcs, meta=meta,
unit=out.unit, header=self._nowcs_header)
@warn_slow
def argmax_world(self, axis, **kwargs):
'''
Return the spatial or spectral index of the maximum value
along a line of sight.
Parameters
----------
axis : int
The axis to return the peak location along. e.g., `axis=0`
will return the value of the spectral axis at the peak value.
kwargs : dict
Passed to `~SpectralCube.argmax`.
'''
return self._argmaxmin_world(axis, 'argmax', **kwargs)
@warn_slow
def argmin_world(self, axis, **kwargs):
'''
Return the spatial or spectral index of the minimum value
along a line of sight.
Parameters
----------
axis : int
The axis to return the peak location along. e.g., `axis=0`
will return the value of the spectral axis at the peak value.
kwargs : dict
Passed to `~SpectralCube.argmin`.
'''
return self._argmaxmin_world(axis, 'argmin', **kwargs)
def chunked(self, chunksize=1000):
"""
Not Implemented.
Iterate over chunks of valid data
"""
raise NotImplementedError()
def _get_flat_shape(self, axis):
"""
Get the shape of the array after flattening along an axis
"""
iteraxes = [0, 1, 2]
iteraxes.remove(axis)
# x,y are defined as first,second dim to iterate over
# (not x,y in pixel space...)
nx = self.shape[iteraxes[0]]
ny = self.shape[iteraxes[1]]
return nx, ny
@warn_slow
def _apply_everywhere(self, function, *args, check_units=True):
"""
Return a new cube with ``function`` applied to all pixels
Private because this doesn't have an obvious and easy-to-use API
Parameters
----------
function : function
An operator that takes the data (self) and any number of additional
arguments
check_units : bool
When doing the initial test before running the full operation,
should units be included on the 'fake' test quantity? This is
specifically added as an option to enable using the subtraction and
addition operators without checking unit compatibility here because
they _already_ enforce unit compatibility.
Examples
--------
>>> newcube = cube.apply_everywhere(np.add, 0.5*u.Jy)
"""
try:
if check_units:
test_result = function(np.ones([1,1,1])*self.unit, *args)
new_unit = test_result.unit
else:
test_result = function(np.ones([1,1,1]), *args)
new_unit = self.unit
# First, check that function returns same # of dims?
assert test_result.ndim == 3,"Output is not 3-dimensional"
except Exception as ex:
raise AssertionError("Function could not be applied to a simple "
"cube. The error was: {0}".format(ex))
# We don't need to convert to a quantity here because the shape check
data_in = self._get_filled_data(fill=self._fill_value)
data = function(data_in, *args)
# strip the unit because data_in does not have a unit
# (we calculate the appropriate unit above and pass it on below)
if hasattr(data, 'unit'):
data = data.value
return self._new_cube_with(data=data, unit=new_unit)
@warn_slow
def _cube_on_cube_operation(self, function, cube, equivalencies=[], **kwargs):
"""
Apply an operation between two cubes. Inherits the metadata of the
left cube.
Parameters
----------
function : function
A function to apply to the cubes
cube : SpectralCube
Another cube to put into the function
equivalencies : list
A list of astropy equivalencies
kwargs : dict
Passed to np.testing.assert_almost_equal
"""
assert cube.shape == self.shape
if not self.unit.is_equivalent(cube.unit, equivalencies=equivalencies):
raise u.UnitsError("{0} is not equivalent to {1}"
.format(self.unit, cube.unit))
if not wcs_utils.check_equality(self.wcs, cube.wcs, warn_missing=True,
**kwargs):
warnings.warn("Cube WCSs do not match, but their shapes do",
WCSMismatchWarning)
try:
test_result = function(np.ones([1,1,1])*self.unit,
np.ones([1,1,1])*self.unit)
# First, check that function returns same # of dims?
assert test_result.shape == (1,1,1)
except Exception as ex:
raise AssertionError("Function {1} could not be applied to a "
"pair of simple "
"cube. The error was: {0}".format(ex,
function))
cube = cube.to(self.unit)
data = function(self._data, cube._data)
try:
# multiplication, division, etc. are valid inter-unit operations
unit = function(self.unit, cube.unit)
except TypeError:
# addition, subtraction are not
unit = self.unit
return self._new_cube_with(data=data, unit=unit)
def apply_function(self, function, axis=None, weights=None, unit=None,
projection=False, progressbar=False,
update_function=None, keep_shape=False, **kwargs):