def join_masks( orig_mask_in, new_mask_in, order='bilinear', operation='or', outfile=None, thresh=0.5, ): """ Reproject and combine a new mask Parameters: ----------- orig_mask_in : string or SpectralCube The original mask. new_mask_in : string or SpectralCube The new mask Keywords: --------- order : string Order of interpolation. Passed to spectral cube. operation : string method to combine the masks 'or' or 'and' outfile : string Filename where the mask will be written. The mask is also returned. thresh : float Floating point value above which the mask is considered true. Relevant because of interpolation. Default 0.5 . """ # TBD - check for two dimensional case # &%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&% # Read the data # &%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&% if type(orig_mask_in) is str: orig_mask = SpectralCube.read(orig_mask_in) elif type(orig_mask_in) is SpectralCube: orig_mask = orig_mask_in else: logging.error('Unrecognized input type for orig_mask_in') raise NotImplementedError # Enable large operations orig_mask.allow_huge_operations = True if type(new_mask_in) is str: new_mask = SpectralCube.read(new_mask_in) elif type(new_mask_in) is SpectralCube: new_mask = new_mask_in else: logging.error('Unrecognized input type for new_mask_in') raise NotImplementedError # Enable large operations new_mask.allow_huge_operations = True # &%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&% # Reproject the new mask onto the original WCS # &%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&% if order == 'fast_nearest_neighbor': logger.warn('Joining masks with nearest neighbor coordinate lookup') # Grab WCS out of template mask and map to other mask x, y, _ = new_mask.wcs.wcs_world2pix(*(orig_mask.world[0, :, :][::-1]), 0) _, _, z = new_mask.wcs.wcs_world2pix(*(orig_mask.world[:, 0, 0][::-1]), 0) x = np.rint(x).astype(np.int) y = np.rint(y).astype(np.int) z = np.rint(z).astype(np.int) new_mask_data = np.array(new_mask.filled_data[:].value > thresh, dtype=np.bool) # Create new mask new_mask_vals = np.zeros(orig_mask.shape, dtype=np.bool) # Find all values that are in bounds inbounds = reduce((lambda A, B: np.logical_and(A, B)), [(z[:, np.newaxis, np.newaxis] >= 0), (z[:, np.newaxis, np.newaxis] < new_mask.shape[0]), (y[np.newaxis, :, :] >= 0), (y[np.newaxis, :, :] < new_mask.shape[1]), (x[np.newaxis, :, :] >= 0), (x[np.newaxis, :, :] < new_mask.shape[2])]) # inbounds = np.logical_and.reduce( # Look 'em up in the new_mask new_mask_vals[inbounds] = new_mask_data[ (z[:, np.newaxis, np.newaxis] * np.ones(inbounds.shape, dtype=np.int))[inbounds], (y[np.newaxis, :, :] * np.ones(inbounds.shape, dtype=np.int))[inbounds], (x[np.newaxis, :, :] * np.ones(inbounds.shape, dtype=np.int))[inbounds]] else: logger.warn('Joining masks with reprojection') new_mask = new_mask.reproject(orig_mask.header, order=order) new_mask = new_mask.spectral_interpolate(orig_mask.spectral_axis) new_mask_vals = np.array(new_mask.filled_data[:].value > thresh, dtype=np.bool) # &%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&% # Combine # &%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&% if operation.strip().lower() == 'or': mask = np.logical_or( np.array(orig_mask.filled_data[:].value > thresh, dtype=np.bool), new_mask_vals) elif operation.strip().lower() == 'and': mask = np.logical_and( np.array(orig_mask.filled_data[:].value > thresh, dtype=np.bool), new_mask_vals) elif operation.strip().lower() == 'sum': mask = (orig_mask.filled_data[:].value) + new_mask_vals else: logging.error('Unrecognized operation. Not "and" or "or".') raise NotImplementedError # &%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&% # Write to disk, return output, etc. # &%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&%&% mask = SpectralCube(mask.astype(np.int), wcs=orig_mask.wcs, header=orig_mask.header, meta={ 'BUNIT': ' ', 'BTYPE': 'Mask' }) # Write to disk, if desired if outfile is not None: header = mask.header header['DATAMAX'] = 1 header['DATAMIN'] = 0 hdu = fits.PrimaryHDU(np.array(mask.filled_data[:], dtype=np.uint8), header=header) hdu.writeto(outfile, overwrite=overwrite) # mask.write(outfile, overwrite=overwrite) return (mask)
def make_vfield_mask(cube, vfield, window, outfile=None, overwrite=True): """Make a mask that includes only pixels within +/- some velocity window of a provided velocity field, which can be two-d or one-d. Parameters: ----------- data : string or SpectralCube The original data cube. vfield : float or two-d array or string The velocity field to use as a template. If a string, it's read as a file. window : float or two-d array or string The velocity field to use as a template. If a string, it's read as a file. Keywords: --------- TBD """ # ------------------------------------------------- # Get spectral information from the cube # ------------------------------------------------- spaxis = cube.spectral_axis spunit = spaxis.unit spvalue = spaxis.value nz, ny, nx = cube.shape # ------------------------------------------------- # Convert and align the velocity field as needed # ------------------------------------------------- # Make sure the velocity is a quanity if type(vfield) != u.quantity.Quantity: # Guess matched units vfield = u.quantity.Quantity(vfield,spunit) # Just convert if it's a scalar if np.ndim(vfield.data) <= 1: vfield = vfield.to(spunit) vfield = u.quantity.Quantity(np.ones((ny, nx))*vfield.value, vfield.unit) else: # reproject if it's a projection if type(vfield) is Projection: vfield = convert_and_reproject(vfield, template=cube.header, unit=spunit) else: vfield = vfield.to(spunit) # Check sizes if (cube.shape[1] != vfield.shape[0]) or (cube.shape[2] != vfield.shape[1]): return(np.nan) # ------------------------------------------------- # Convert and align the window as needed # ------------------------------------------------- # Make sure the window is a quanity if type(window) != u.quantity.Quantity: # Guess matched units window = u.quantity.Quantity(window,spunit) # Just convert if it's a scalar if np.ndim(window.data) <= 1: window = window.to(spunit) window = u.quantity.Quantity(np.ones((ny, nx))*window.value, window.unit) else: # reproject if it's a projection if type(window) is Projection: window = convert_and_reproject(window, template=cube.header, unit=spunit) else: window = window.to(spunit) # Check sizes if (cube.shape[1] != window.shape[0]) or (cube.shape[2] != window.shape[1]): return(np.nan) # ------------------------------------------------- # Generate a velocity cube and a vfield cube # ------------------------------------------------- spaxis_cube = np.ones((ny, nx))[None,:,:] * spvalue[:,None,None] vfield_cube = (vfield.value)[None,:,:] * np.ones(nz)[:,None,None] window_cube = (window.value)[None,:,:] * np.ones(nz)[:,None,None] # ------------------------------------------------- # Make the mask # ------------------------------------------------- mask = mask_around_value( \ spaxis_cube, target=vfield_cube, delta=window_cube) # ------------------------------------------------- # Write to disk # ------------------------------------------------- mask = SpectralCube(mask.astype(np.int), wcs=cube.wcs, header=cube.header, meta={'BUNIT': ' ', 'BTYPE': 'Mask'}) # Write to disk, if desired if outfile is not None: header = mask.header header['DATAMAX'] = 1 header['DATAMIN'] = 0 hdu = fits.PrimaryHDU(np.array(mask.filled_data[:], dtype=np.uint8), header=header) hdu.writeto(outfile, overwrite=overwrite) return(mask)