import numpy as np from astropy import units as u from astropy.wcs import WCS from astropy import log import warnings from .utils import BadVelocitiesWarning, ProgressBar from .cube_utils import _map_context from .lower_dimensional_structures import VaryingResolutionOneDSpectrum, OneDSpectrum from .spectral_cube import VaryingResolutionSpectralCube, SpectralCube def fourier_shift(x, shift, axis=0, add_pad=False, pad_size=None): ''' Shift a spectrum in the Fourier plane. Parameters ---------- x : np.ndarray Array to be shifted shift : int or float Number of pixels to shift. axis : int, optional Axis to shift along. pad_size : int, optional Pad the array before shifting. Returns ------- x2 : np.ndarray Shifted array. ''' nanmask = ~np.isfinite(x) # If all NaNs, there is nothing to shift # But only if there is no added padding. Otherwise we need to pad if nanmask.all() and not add_pad: return x nonan = x.copy() shift_mask = False if nanmask.any(): nonan[nanmask] = 0.0 shift_mask = True # Optionally pad the edges if add_pad: if pad_size is None: # Pad by the size of the shift pad = np.ceil(shift).astype(int) # Determine edge to pad whether it is a positive or negative shift pad_size = (pad, 0) if shift > 0 else (0, pad) else: assert len(pad_size) pad_nonan = np.pad(nonan, pad_size, mode='constant', constant_values=(0)) if shift_mask: pad_mask = np.pad(nanmask, pad_size, mode='constant', constant_values=(0)) else: pad_nonan = nonan pad_mask = nanmask # Check if there are all NaNs before shifting if nanmask.all(): return np.array([np.nan] * pad_mask.size) nonan_shift = _fourier_shifter(pad_nonan, shift, axis) if shift_mask: mask_shift = _fourier_shifter(pad_mask, shift, axis) > 0.5 nonan_shift[mask_shift] = np.nan return nonan_shift def _fourier_shifter(x, shift, axis): ''' Helper function for `~fourier_shift`. ''' ftx = np.fft.fft(x, axis=axis) m = np.fft.fftfreq(x.shape[axis]) # m_shape = [1] * x.ndim # m_shape[axis] = m.shape[0] # m = m.reshape(m_shape) slices = tuple([slice(None) if ii == axis else None for ii in range(x.ndim)]) m = m[slices] phase = np.exp(-2 * np.pi * m * 1j * shift) x2 = np.real(np.fft.ifft(ftx * phase, axis=axis)) return x2 def get_chunks(num_items, chunk): ''' Parameters ---------- num_items : int Number of total items. chunk : int Size of chunks Returns ------- chunks : list of np.ndarray List of channels in chunks of the given size. ''' items = np.arange(num_items) if num_items == chunk: return [items] chunks = \ np.array_split(items, [chunk * i for i in range(int(num_items / chunk))]) if chunks[-1].size == 0: # Last one is empty chunks = chunks[:-1] if chunks[0].size == 0: # First one is empty chunks = chunks[1:] return chunks def _spectrum_shifter(inputs): spec, shift, add_pad, pad_size = inputs return fourier_shift(spec, shift, add_pad=add_pad, pad_size=pad_size) def stack_spectra(cube, velocity_surface, v0=None, stack_function=np.nanmean, xy_posns=None, num_cores=1, chunk_size=-1, progressbar=False, pad_edges=True, vdiff_tol=0.01): ''' Shift spectra in a cube according to a given velocity surface (peak velocity, centroid, rotation model, etc.). Parameters ---------- cube : SpectralCube The cube velocity_field : Quantity A Quantity array with m/s or equivalent units stack_function : function A function that can operate over a list of numpy arrays (and accepts ``axis=0``) to combine the spectra. `numpy.nanmean` is the default, though one might consider `numpy.mean` or `numpy.median` as other options. xy_posns : list, optional List the spatial positions to include in the stack. For example, if the data is masked by some criterion, the valid points can be given as `xy_posns = np.where(mask)`. num_cores : int, optional Choose number of cores to run on. Defaults to 1. chunk_size : int, optional To limit memory usage, the shuffling of spectra can be done in chunks. Chunk size sets the number of spectra that, if memory-mapping is used, is the number of spectra loaded into memory. Defaults to -1, which is all spectra. progressbar : bool, optional Print progress through every chunk iteration. pad_edges : bool, optional Pad the edges of the shuffled spectra to stop data from rolling over. Default is True. The rolling over occurs since the FFT treats the boundary as periodic. This should only be disabled if you know that the velocity range exceeds the range that a spectrum has to be shuffled to reach `v0`. vdiff_tol : float, optional Allowed tolerance for changes in the spectral axis spacing. Default is 0.01, or 1%. Returns ------- stack_spec : OneDSpectrum The stacked spectrum. ''' if not np.isfinite(velocity_surface).any(): raise ValueError("velocity_surface contains no finite values.") vshape = velocity_surface.shape cshape = cube.shape[1:] if not (vshape == cshape): raise ValueError("Velocity surface map does not match cube spatial " "dimensions.") if xy_posns is None: # Only compute where a shift can be found xy_posns = np.where(np.isfinite(velocity_surface)) if v0 is None: # Set to the mean velocity of the cube if not given. v0 = cube.spectral_axis.mean() else: if not isinstance(v0, u.Quantity): raise u.UnitsError("v0 must be a quantity.") spec_unit = cube.spectral_axis.unit if not v0.unit.is_equivalent(spec_unit): raise u.UnitsError("v0 must have units equivalent to the cube's" f" spectral unit {spec_unit}.") v0 = v0.to(spec_unit) if v0 < cube.spectral_axis.min() or v0 > cube.spectral_axis.max(): raise ValueError("v0 must be within the range of the spectral " "axis of the cube.") # Calculate the pixel shifts that will be applied. spec_size = np.diff(cube.spectral_axis[:2])[0] # Assign the correct +/- for pixel shifts based on whether the spectral # axis is increasing (-1) or decreasing (+1) vdiff_sign = -1. if spec_size.value > 0. else 1. vdiff = np.abs(spec_size) vel_unit = vdiff.unit # Check to make sure vdiff doesn't change more than the allowed tolerance # over the spectral axis vdiff2 = np.abs(np.diff(cube.spectral_axis[-2:])[0]) if not np.isclose(vdiff2.value, vdiff.value, rtol=vdiff_tol): raise ValueError("Cannot shift spectra on a non-linear axes") vmax = cube.spectral_axis.to(vel_unit).max() vmin = cube.spectral_axis.to(vel_unit).min() if ((np.any(velocity_surface > vmax) or np.any(velocity_surface < vmin))): warnings.warn("Some velocities are outside the allowed range and will be " "masked out.", BadVelocitiesWarning) # issue 580/1 note: numpy <=1.16 will strip units from velocity, >= # 1.17 will not masked_velocities = np.where( (velocity_surface < vmax) & (velocity_surface > vmin), velocity_surface.value, np.nan) velocity_surface = u.Quantity(masked_velocities, velocity_surface.unit) pix_shifts = vdiff_sign * ((velocity_surface.to(vel_unit) - v0.to(vel_unit)) / vdiff).value[xy_posns] # Make a header copy so we can start altering new_header = cube[:, 0, 0].header.copy() if pad_edges: # Enables padding the whole cube such that no spectrum will wrap around # This is critical if a low-SB component is far off of the bright # component that the velocity surface is derived from. # Find max +/- pixel shifts, rounding up to the nearest integer max_pos_shift = np.ceil(np.nanmax(pix_shifts)).astype(int) max_neg_shift = np.ceil(np.nanmin(pix_shifts)).astype(int) if max_neg_shift > 0: # if there are no negative shifts, we can ignore them and just # use the positive shift max_neg_shift = 0 if max_pos_shift < 0: # same for positive max_pos_shift = 0 # The total pixel size of the new spectral axis num_vel_pix = cube.spectral_axis.size + max_pos_shift - max_neg_shift new_header['NAXIS1'] = num_vel_pix # Adjust CRPIX in header new_header['CRPIX1'] += -max_neg_shift pad_size = (-max_neg_shift, max_pos_shift) else: pad_size = None all_shifted_spectra = [] if chunk_size == -1: chunk_size = len(xy_posns[0]) # Create chunks of spectra for read-out. chunks = get_chunks(len(xy_posns[0]), chunk_size) if progressbar: iterat = ProgressBar(chunks, desc='Stack Spectra: ') else: iterat = chunks for i, chunk in enumerate(iterat): gen = ((cube.filled_data[:, y, x].value, shift, pad_edges, pad_size) for y, x, shift in zip(xy_posns[0][chunk], xy_posns[1][chunk], pix_shifts[chunk])) with _map_context(num_cores) as map: shifted_spectra = map(_spectrum_shifter, gen) all_shifted_spectra.extend([out for out in shifted_spectra]) shifted_spectra_array = np.array(all_shifted_spectra) assert shifted_spectra_array.ndim == 2 stacked = stack_function(shifted_spectra_array, axis=0) 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 def stack_cube(cube, linelist, vmin, vmax, average=np.nanmean, convolve_beam=None, return_hdu=False, return_cutouts=False): """ Create a stacked cube by averaging on a common velocity grid. If the input cubes have varying resolution, this will trigger potentially expensive convolutions. Parameters ---------- cube : SpectralCube The cube (or a list of cubes) linelist : list of Quantities An iterable of Quantities representing line rest frequencies vmin / vmax : Quantity Velocity-equivalent quantities specifying the velocity range to average over average : function A function that can operate over a list of numpy arrays (and accepts ``axis=0``) to average the spectra. `numpy.nanmean` is the default, though one might consider `numpy.mean` or `numpy.median` as other options. convolve_beam : None If the cube is a VaryingResolutionSpectralCube, a convolution beam is required to put the cube onto a common grid prior to spectral interpolation. return_hdu : bool Return an HDU instead of a spectral-cube return_cutouts : bool Also return the individual cube cutouts? Returns ======= cube : SpectralCube The SpectralCube object containing the reprojected cube. Its header will be based on the first cube but will have no reference frequency. Its spectral axis will be in velocity units. cutout_cubes : list A list of cube cutouts projected into the same space (optional; see ``return_cutouts``) """ if isinstance(cube, list): cubes = cube cube = cubes[0] for cb in cubes[1:]: if cb.shape[1:] != cube.shape[1:]: raise ValueError("If you pass multiple cubes, they must have the " "same spatial shape.") if convolve_beam is None and (any(hasattr(cb, 'beams') for cb in cubes) or not all([cb.beam == cube.beam for cb in cubes[1:]])): raise ValueError("If the cubes have different resolution, `convolve_beam` must be specified.") else: cubes = [cube] slabs = [] included_lines = [] # loop over linelist first to keep the cutouts in the same order as the # input line frequencies for restval in linelist: for cube in cubes: line_cube = cube.with_spectral_unit(u.km/u.s, velocity_convention='radio', rest_value=restval) line_cutout = line_cube.spectral_slab(vmin, vmax) if line_cutout.shape[0] <= 1: log.debug(f"Skipped line {restval} for cube {cube} because it resulted" "in a size-1 spectral axis") continue else: included_lines.append(restval) assert line_cutout.shape[0] > 1 if isinstance(line_cutout, VaryingResolutionSpectralCube): if convolve_beam is None: raise ValueError("If any of the input cubes have varyin resolution, " "a target `common_beam` must be specified.") line_cutout = line_cutout.convolve_to(convolve_beam) assert not isinstance(line_cutout, VaryingResolutionSpectralCube) slabs.append(line_cutout) reference_cube = slabs[0] cutout_cubes = [reference_cube.filled_data[:].value] for slab in slabs[1:]: regridded = slab.spectral_interpolate(reference_cube.spectral_axis) cutout_cubes.append(regridded.filled_data[:].value) stacked_cube = average(cutout_cubes, axis=0) ww = reference_cube.wcs.copy() # set restfreq to zero: it is not defined any more. ww.wcs.restfrq = 0.0 meta = reference_cube.meta meta.update({'stacked_lines': included_lines}) result_cube = SpectralCube(data=stacked_cube, wcs=ww, meta=meta, header=reference_cube.header) if return_hdu: retval = result_cube.hdu else: retval = result_cube if return_cutouts: return retval, cutout_cubes else: return retval