diff --git a/ndcube/coords.py b/ndcube/coords.py index 350a5155c..33dfa8b80 100644 --- a/ndcube/coords.py +++ b/ndcube/coords.py @@ -5,8 +5,23 @@ def __init__(self, ndcube): if not isinstance(ndcube, NDCubeABC): raise TypeError("ndcube must be an instance of NDCubeABC.") self._ndcube = ndcube + # Create aliases to coordinate methods already on NDCube. self.axis_world_coords = self._ndcube.axis_world_coords self.axis_world_coords_values = self._ndcube.axis_world_coords_values + # Exposure properties to get values of coordinate types based on + # physical types present in NDCube.wcs and NDCube.extra_coords. + self._celestial_words = ["pos"] + self._time_words = ["time"] + self._spectral_words = ["em"] + self._stokes_words = ["stokes"] + if _words_in_physical_types(self._celestial_words) + self.celestial = self._celestial + if _words_in_physical_types(self._time_words) + self.time = self._time + if _words_in_physical_types(self._spectral_words) + self.spectral = self._spectral + if _words_in_physical_types(self._stokes_words) + self.stokes = self._stokes @property def wcs(self): @@ -28,8 +43,8 @@ def global_coords(self): def array_axis_physical_types(self): return self._ndcube.array_axis_physical_types - @property - def celestial(self): + @property + def _celestial(self): """ Returns celestial coordinates for all array indices in relevant axes. Celestial world axis physical type name must contain 'pos.'. @@ -42,10 +57,10 @@ def celestial(self): ------- : `astropy.coordinates.SkyCoord` """ - return self._get_coords_by_word("pos") + return self._get_coords_by_word(self._celestial_words) @property - def time(self): + def _time(self): """ Returns time coordinates for all array indices in relevant axes. Time world axis physical type name must contain 'time'. @@ -58,10 +73,10 @@ def time(self): ------- : `astropy.time.Time` """ - return self._get_coords_by_word("time") + return self._get_coords_by_word(self._time_words) @property - def spectral(self): + def _spectral(self): """ Returns spectral coordinates from WCS. Spectral world axis physical type name must contain 'em.'. @@ -74,10 +89,10 @@ def spectral(self): ------- : `astropy.coordinates.SpectralCoord` or `astropy.units.Quantity` """ - return self._get_coords_by_word("em") + return self._get_coords_by_word(self._spectral_words) @property - def stokes(self): + def _stokes(self): """ Returns stokes polarization for all array indices in relevant axes. Stokes world axis physical type name must contain 'stokes'. @@ -90,39 +105,52 @@ def stokes(self): ------- : """ - return self._get_coords_by_word("stokes") + return self._get_coords_by_word(self._stokes_words) def _get_coords_by_word(self, coord_words): - """ - Returns coordinates from a WCS corresponding to a world axis physical type. - """ - if isinstance(coord_words, str): - coord_words = [coord_words] - # Check if words are in physical types of main WCS and extra coords. - wcs_phys_types = [] - ec_phys_types = [] - for wcs, physical_types in zip([self.wcs, self.extra_coords.wcs], [wcs_phys_types, ec_phys_types]): - if wcs is not None: - for name in wcs.world_axis_physical_types: - words = set(name.replace(":", ".").split(".")) - if any(word in words for word in coord_words): - physical_types.append(name) - # Determine type of WCS to use based on whether the words were - # found in WCS and/or extra coords. - if len(wcs_phys_types) > 0: - if len(ec_phys_types) > 0: - wcs = self.combined_wcs - physical_types = wcs_phys_types + ec_phys_types - else: - wcs = self.wcs - physical_types = wcs_phys_types - elif len(ec_phys_types) > 0: - wcs = self.extra_coords - physical_types = ec_phys_types - else: - return None - # Calculate coordinates. - coords = self.axis_world_coords(*physical_types, wcs=wcs) - if len(coords) == 1: - coords = coords[0] - return coords + """ + Returns coordinates from a WCS corresponding to a world axis physical type. + """ + if isinstance(coord_words, str): + coord_words = [coord_words] + # Check if words are in physical types of main WCS and extra coords. + wcs_phys_types, ec_phys_types = _get_physical_types_from_words(coord_words, [self.wcs, self.extra_coords.wcs]) + # Determine type of WCS to use based on whether the words were + # found in WCS and/or extra coords. + if len(wcs_phys_types) > 0: + if len(ec_phys_types) > 0: + wcs = self.combined_wcs + physical_types = wcs_phys_types + ec_phys_types + else: + wcs = self.wcs + physical_types = wcs_phys_types + elif len(ec_phys_types) > 0: + wcs = self.extra_coords + physical_types = ec_phys_types + else: + return None + # Calculate coordinates. + coords = self.axis_world_coords(*physical_types, wcs=wcs) + if len(coords) == 1: + coords = coords[0] + return coords + + +def _get_physical_types_from_words(coord_words, wcses): + found_phys_types = [] + for i, wcs in enumerate(wcses): + found_phys_types.append([]) + if wcs is not None: + for phys_type in wcs.world_axis_physical_types: + words = set(phys_type.replace(":", ".").split(".")) + if any(word in words for word in coord_words): + found_phys_types[i].append(name) + return found_phys_types + + +def _words_in_physical_types(coord_words, wcses): + found_phys_types = _get_physical_types_from_words(coord_words, wcses) + if any(len(fpt) > 0 for fpt in found_phys_types): + return True + else: + return False