Skip to content

Commit

Permalink
fixed the tests
Browse files Browse the repository at this point in the history
  • Loading branch information
keflavich committed Dec 6, 2018
1 parent 546685f commit 6778b4c
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 22 deletions.
18 changes: 12 additions & 6 deletions spectral_cube/analysis_utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@

from .utils import BadVelocitiesWarning
from .cube_utils import _map_context
from .lower_dimensional_structures import OneDSpectrum
from .lower_dimensional_structures import VaryingResolutionOneDSpectrum, OneDSpectrum
from .spectral_cube import VaryingResolutionSpectralCube


Expand Down Expand Up @@ -302,11 +302,17 @@ def stack_spectra(cube, velocity_surface, v0=None,

stacked = stack_function(shifted_spectra_array, axis=0)

stack_spec = \
OneDSpectrum(stacked, unit=cube.unit, wcs=WCS(new_header),
header=new_header,
meta=cube.meta, spectral_unit=vel_unit,
beams=cube.beams if hasattr(cube, "beams") else None)
if hasattr(cube, 'beams'):
stack_spec = VaryingResolutionOneDSpectrum(stacked, unit=cube.unit,
wcs=WCS(new_header),
header=new_header,
meta=cube.meta,
spectral_unit=vel_unit,
beams=cube.beams)
else:
stack_spec = OneDSpectrum(stacked, unit=cube.unit, wcs=WCS(new_header),
header=new_header, meta=cube.meta,
spectral_unit=vel_unit, beam=cube.beam)

return stack_spec

Expand Down
13 changes: 6 additions & 7 deletions spectral_cube/lower_dimensional_structures.py
Original file line number Diff line number Diff line change
Expand Up @@ -701,9 +701,6 @@ def __new__(cls, value, unit=None, dtype=None, copy=True, wcs=None,
)

if beam is not None:
if beams is not None:
raise ValueError("Both 'beam' and 'beams' specified when "
"creating 1dspec")
self.beam = beam
self.meta['beam'] = beam

Expand Down Expand Up @@ -1063,7 +1060,10 @@ def _new_spectrum_with(self, data=None, wcs=None, mask=None, meta=None,

if beams is None:
if hasattr(self, 'beams'):
beams = self.beams
kwargs['beams'] = beams
elif beam is None:
if hasattr(self, 'beam'):
kwargs['beam'] = beam

if beams is not None:
cls = VaryingResolutionOneDSpectrum
Expand All @@ -1074,7 +1074,7 @@ def _new_spectrum_with(self, data=None, wcs=None, mask=None, meta=None,
spectrum = cls(value=data, wcs=wcs, mask=mask, meta=meta, unit=unit,
fill_value=fill_value, header=header or self._header,
wcs_tolerance=wcs_tolerance or self._wcs_tolerance,
beams=beams, beam=beam, **kwargs)
**kwargs)
spectrum._spectral_unit = spectral_unit

return spectrum
Expand All @@ -1095,5 +1095,4 @@ class VaryingResolutionOneDSpectrum(OneDSpectrum, MultiBeamMixinClass):
def __new__(cls, value, beams=None, **kwargs):
if beams is not None:
cls.beams = beams
return super(OneDSpectrum, cls).__new__(value, **kwargs)

return super(VaryingResolutionOneDSpectrum, cls).__new__(cls, value, **kwargs)
41 changes: 32 additions & 9 deletions spectral_cube/spectral_cube.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,16 +395,20 @@ def apply_numpy_function(self, function, fill=np.nan,
if set(axis) == set((1,2)):
new_wcs = self._wcs.sub([wcs.WCSSUB_SPECTRAL])
header = self._nowcs_header
if hasattr(self, 'beam'):
bmarg = {'beam': self.beam}
elif hasattr(self, 'beams'):
bmarg = {'beams': self.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,
beams=(self.beams
if hasattr(self,'beams')
else None),
**bmarg
)
else:
warnings.warn("Averaging over a spatial and a spectral "
Expand Down Expand Up @@ -544,15 +548,20 @@ def mean(self, axis=None, how='cube', **kwargs):
unit=self.unit, header=self._nowcs_header)
elif axis == (1,2):
newwcs = self._wcs.sub([wcs.WCSSUB_SPECTRAL])
if hasattr(self, 'beam'):
bmarg = {'beam': self.beam}
elif hasattr(self, 'beams'):
bmarg = {'beams': self.beams}
else:
bmarg = {}
return self._oned_spectrum(value=out,
wcs=newwcs,
copy=False,
unit=self.unit,
spectral_unit=self._spectral_unit,
beams=(self.beams
if hasattr(self, 'beams')
else None),
meta=self.meta)
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
Expand Down Expand Up @@ -1153,13 +1162,20 @@ def __getitem__(self, view):
newwcs = self._wcs.sub([a
for a in (1,2,3)
if a not in [x+1 for x in intslices]])
if hasattr(self, 'beam'):
bmarg = {'beam': self.beam}
elif hasattr(self, 'beams'):
bmarg = {'beams': self.beams}
else:
bmarg = {}
return self._oned_spectrum(value=self._data[view],
wcs=newwcs,
copy=False,
unit=self.unit,
spectral_unit=self._spectral_unit,
mask=self.mask[view],
meta=meta,
**bmarg
)

# only one element, so drop an axis
Expand Down Expand Up @@ -3351,14 +3367,21 @@ def __getitem__(self, view):
newwcs = self._wcs.sub([a
for a in (1,2,3)
if a not in [x+1 for x in intslices]])
if hasattr(self, 'beam'):
bmarg = {'beam': self.beam}
elif hasattr(self, 'beams'):
bmarg = {'beams': self.beams[specslice]}
else:
bmarg = {}
return self._oned_spectrum(value=self._data[view],
wcs=newwcs,
copy=False,
unit=self.unit,
spectral_unit=self._spectral_unit,
mask=self.mask[view],
beams=self.beams[specslice],
meta=meta)
meta=meta,
**bmarg
)

# only one element, so drop an axis
newwcs = wcs_utils.drop_axis(self._wcs, intslices[0])
Expand Down

0 comments on commit 6778b4c

Please sign in to comment.