Skip to main content
`Ctrl`+`K`

# Source code for astropy.coordinates.polarization

    from collections.abc import Generator
    from contextlib import contextmanager
    from typing import NamedTuple

    import numpy as np

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

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

    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

    @contextmanager
    def custom_stokes_symbol_mapping(
        mapping: dict[int, StokesSymbol], replace: bool = False
    ) -> Generator[None, None, 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(out.info, 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]

    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
                self.info = stokes.info
            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=None):
            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__:
                new.info = self.info

            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__()