Source code for astropy.coordinates.polarization

from __future__ import annotations

from contextlib import contextmanager
from typing import NamedTuple

import numpy as np

from astropy.utils import unbroadcast
from astropy.utils.compat import COPY_IF_NEEDED
from astropy.utils.data_info import MixinInfo
from astropy.utils.shapes import ShapedLikeNDArray

__all__ = ["StokesCoord", "custom_stokes_symbol_mapping", "StokesSymbol"]

[docs] class StokesSymbol(NamedTuple): """Symbol for a Stokes coordinate.""" symbol: str = "" description: str = ""
# This is table 29 in the FITS 4.0 paper FITS_STOKES_VALUE_SYMBOL_MAP = { 1: StokesSymbol("I", "Standard Stokes unpolarized"), 2: StokesSymbol("Q", "Standard Stokes linear"), 3: StokesSymbol("U", "Standard Stokes linear"), 4: StokesSymbol("V", "Standard Stokes circular"), -1: StokesSymbol("RR", "Right-right circular: <RR*>"), -2: StokesSymbol("LL", "Left-left circular: <LL*>"), -3: StokesSymbol("RL", "Right-left cross-circular: Re(<RL*>))"), -4: StokesSymbol("LR", "Left-right cross-circular: Re(<LR*>)=Im(<RL*>)"), -5: StokesSymbol("XX", "X parallel linear: <XX*>"), -6: StokesSymbol("YY", "Y parallel linear: <YY*>"), -7: StokesSymbol("XY", "XY cross linear: Re(<XY*>)"), -8: StokesSymbol("YX", "YX cross linear: Im(<XY*>)"), } STOKES_VALUE_SYMBOL_MAP = FITS_STOKES_VALUE_SYMBOL_MAP.copy() UNKNOWN_STOKES_VALUE = -99999
[docs] @contextmanager def custom_stokes_symbol_mapping( mapping: dict[int, StokesSymbol], replace: bool = False ) -> None: """ Add a custom set of mappings from values to Stokes symbols. Parameters ---------- mappings A list of dictionaries with custom mappings between values (integers) and `.StokesSymbol` classes. replace Replace all mappings with this one. """ global STOKES_VALUE_SYMBOL_MAP original_mapping = STOKES_VALUE_SYMBOL_MAP.copy() if not replace: STOKES_VALUE_SYMBOL_MAP = {**original_mapping, **mapping} else: STOKES_VALUE_SYMBOL_MAP = mapping yield STOKES_VALUE_SYMBOL_MAP = original_mapping
class StokesCoordInfo(MixinInfo): # The attributes containing actual information. _represent_as_dict_attrs = {"value"} # Since there is only one attribute, use a column with the name to represent it # (rather than as name.value) _represent_as_dict_primary_data = "value" # Attributes that should be presented as positional arguments to # the class initializer (which takes "stokes" as an argument, not "value"). _construct_from_dict_args = ("value",) @property def unit(self): return None @property def dtype(self): return self._parent._data.dtype @staticmethod def default_format(val): return f"{val.symbol}" def new_like(self, cols, length, metadata_conflicts="warn", name=None): """ Return a new StokesCoord instance which is consistent with the input ``cols`` and has ``length`` rows. This is intended for creating an empty column object whose elements can be set in-place for table operations like join or vstack. Parameters ---------- cols : list List of input columns length : int Length of the output column object metadata_conflicts : str ('warn'|'error'|'silent') How to handle metadata conflicts name : str Output column name Returns ------- col : `~astropy.coordinates.StokesCoord` (or subclass) Empty instance of this class consistent with ``cols`` """ # Get merged info attributes like shape, dtype, format, description, etc. attrs = self.merge_cols_attributes( cols, metadata_conflicts, name, ("meta", "format", "description") ) # Make an empty StokesCoord. shape = (length,) + attrs.pop("shape") data = np.zeros(shape=shape, dtype=attrs.pop("dtype")) # Get arguments needed to reconstruct class out = self._construct_from_dict({"value": data}) # Set remaining info attributes for attr, value in attrs.items(): setattr(, attr, value) return out def get_sortable_arrays(self): """ Return a list of arrays which can be lexically sorted to represent the order of the parent column. For StokesCoord this is just the underlying values. Returns ------- arrays : list of ndarray """ return [self._parent._data]
[docs] class StokesCoord(ShapedLikeNDArray): """ A representation of stokes coordinates with helpers for converting to profile names. Parameters ---------- stokes : array-like The numeric values representing stokes coordinates. """ info = StokesCoordInfo() def __init__(self, stokes, copy=False): if isinstance(stokes, type(self)): data = stokes._data.copy() if copy else stokes._data = else: data = np.asanyarray(stokes) if data.dtype.kind == "O": msg = "StokesCoord objects cannot be initialised with an object array." raise ValueError(msg) if data.dtype.kind == "U": data = self._from_symbols(data) else: data = data.copy() if copy and data is stokes else data self._data = data @property def shape(self): return self._data.shape @property def value(self): return self._data @property def dtype(self): return self._data.dtype def __array__(self, dtype=None, copy=COPY_IF_NEEDED): return self._data.astype(dtype, copy=copy) def _apply(self, method, *args, **kwargs): cls = type(self) if callable(method): new = cls(method(self._data, *args, **kwargs)) else: new = cls(getattr(self._data, method)(*args, **kwargs)) # Copy other 'info' attr only if it has actually been defined. # See PR #3898 for further explanation and justification, along # with Quantity.__array_finalize__ if "info" in self.__dict__: = return new @staticmethod def _from_symbols(symbols): """ Convert an array of symbols to an array of values """ values_array = np.full_like( symbols, UNKNOWN_STOKES_VALUE, dtype=int, subok=False ) for stokes_value, symbol in STOKES_VALUE_SYMBOL_MAP.items(): values_array[symbols == symbol.symbol] = stokes_value if (unknown_values := np.equal(values_array, UNKNOWN_STOKES_VALUE)).any(): raise ValueError( f"Unknown stokes symbols present in the input array: {np.unique(symbols[unknown_values])}" ) return values_array @property def symbol(self): """The coordinate represented as strings.""" known_symbols = tuple( ["?"] + [s.symbol for s in STOKES_VALUE_SYMBOL_MAP.values()] ) max_len = np.max([len(s) for s in known_symbols]) # Note we unbroadcast and re-broadcast here to prevent the new array # using more memory than the old one. unbroadcasted = np.round(unbroadcast(self.value)) symbolarr = np.full(unbroadcasted.shape, "?", dtype=f"<U{max_len}") for value, symbol in STOKES_VALUE_SYMBOL_MAP.items(): symbolarr[unbroadcasted == value] = symbol.symbol return np.broadcast_to(symbolarr, self.shape) def __setitem__(self, item, value): self._data[item] = type(self)(value)._data def __eq__(self, other): try: other = self.__class__(other) except Exception: return NotImplemented return self._data == other._data def __str__(self): arrstr = np.array2string(self.symbol, separator=", ", prefix=" ") return f"{type(self).__name__}({arrstr})" def __repr__(self): return self.__str__()