        Save the parameters of this data class to a file

        Parameters
        ----------
        filename : str
            the name of the file to write out the parameters to
        mode : str
            the mode to use when writing the parameters to file; i.e., 'w', 'a'


The global mininum wavenumber


The global maximum wavenumber


The number of data bins


Return the number of measurements


        The measurement `k` values, concatenated from each
        `PowerMeasurement` in `measurements`


        The measurement `mu` values, concatenated from each
        `PowerMeasurement` in `measurements`


        The measurement `ell` values, concatenated from each
        `PowerMeasurement` in `measurements`


        The measurement power values, concatenated from each `PowerMeasurement`
        in `self.measurement`


        The measurement `error` values, concatenated from each
        `PowerMeasurement` in `measurements`


Return `True` if the covariance matrix is diagonal.


        Return a list of length `size` that holds the slices needed
        to extract the `i-th` measurement from the flattened list
        of measurements.


    Flatten the input array using `Fortran` format
    (i.e., col #1 then col #2, etc), and remove
    any NaNs along the way


        The main data attribute, which stores the upper triangle + diagonal
        elements of the symmetric matrix


Dictionary of local attributes on this variable.


Rescale the inverse of the covariance matrix by this factor


        Given the matrix element (i, j) return the index of the corresponding
        element in `_data`


        Setup the class by veryifying the input and storing the
        upper triangle + diagonal elements of the input data (N*(N+1)/2 elements)

        This also initializes the GridIndexer object as the `_indexer` attribute


        Verify that the input data is symmetric and positive-semidefinite, as
        any covariance matrix must be


Return a deep copy of `self`


Save the covariance matrix as a pickle file


Return a pickled CovarianceMatrix


Save the covariance matrix to a plain text file


Load the covariance matrix from a plain text file


        Label-based indexing by coordinate name

        Notes
        -----
        Returns a new class instance if the resulting sliced data
        is symmetric, otherwise, returns the sliced data array

        Parameters
        ----------
        kwargs : dict
            key, value pairs where the keys represent dimension names and
            the values represent indexing objects, i.e., an integer, slice object,
            or array


        Interger-based indexing by coordinate name

        Notes
        -----
        Returns a new class instance if the resulting sliced data
        is symmetric, otherwise, returns the sliced data array

        Parameters
        ----------
        kwargs : dict
            key, value pairs where the keys represent dimension names and
            the values represent indexing objects, i.e., an integer, slice object,
            or array


A list of GridIndex objects for each axis


The dimension names


A list of dictionaries holding (dim, coord) pairs for each axis


The diagonal elements of the matrix


The inverse of the covariance matrix, returned as a 2D ndarray


Return the total 2D covariance matrix array


Return the normalized covariance matrix, i.e., the correlation matrix


        Plot the correlation matrix (normalized covariance matrix), optionally
        saving if `filename != None`


r"""
        Return the Gaussian prediction for the covariance between
        :math:`P(k,\mu)` wedges for a periodic box simulation, where the number
        density is constant.

        See eq. 17 of Grieb et al. 2015 arxiv:1509.04293

        Parameters
        ----------
        model : GalaxySpectrum, QuasarSpectrum
            the model instance used to evaluate the theoretical :math:`P(k,\mu)`
        k : array_like
            the array of wavenumbers (units of :math:`h/\mathrm{Mpc}`) where
            the covariance matrix will be evaluated
        mu_edges : array_like
            the edges of the :math:`\mu` bins
        nbar : float
            the constant number density in the box in units of
            :math:`(\mathrm{Mpc}/h)^{-3}` -- the shot noise contribution
            to the covariance is the inverse of this value
        volume : float
            the volume of the box, in units of :math:`(\mathrm{Mpc}/h)^3`
        Nmu : int, optional
            the number of mu bins to use when performing the multipole integration
            over :math:`\mu`

        Returns
        -------
        PkmuCovarianceMatrix :
            the covariance matrix object holding the Gaussian prediction for
            the covariance between the specified wedges

        Examples
        --------
        >>> import numpy
        >>> from pyRSD.rsd import GalaxySpectrum
        >>> from pyRSD.rsdfit.data import PkmuCovarianceMatrix
        >>> volume = 1380.0**3
        >>> nbar = 3e-4
        >>> model = GalaxySpectrum(params='boss_dr12_fidcosmo.ini')
        >>> k = numpy.arange(0., 0.4, 0.005) + 0.005/2
        >>> mu_edges = [0., 0.2, 0.4, 0.6, 0.8, 1.0]
        >>> C = PkmuCovarianceMatrix.periodic_gaussian_covariance(model, k, ells, nbar, volume)


Enumerate across the upper mu triangle of blocks


        Trim the k bounds of the covariance matrix, returning a new, trimmed
        PkmuCovarianceMatrix object

        Parameters
        ----------
        kmin : float or array_like
            the minimum k value; if an array is provided, it is interpreted as the values
            for each mu bin
        kmax : float or array_like
            the maximum k value; if an array is provided, it is interpreted as the values
            for each mu bin


        The unique center values of the `mu` bins, optionally
        returning all (non-unique) values


        The unique center values of the `k` bins, optionally
        returning all (non-unique) values


        Plot the square-root of the diagonal elements for a specfic mu value

        Parameters
        ---------------
        ax : Axes, optional
            Axes instance to plot to. If `None`, plot to the current axes
        mu : {float, int},
            If not `None`, only plot the diagonals for a single `mu` value, else
            plot for all `mu` values
        options : dict
            Dictionary of options with the following keywords:
                show_gaussian : bool, (`False`)
                    Overplot the Gaussian prediction, i.e.,
                    :math: 2/N_{modes} * [P(k) + Pshot(k)]**2
                norm_by_gaussian : bool, (`False`)
                    Normalize the diagonal elements by the Gaussian prediction
                norm_by_power : bool, (`False`)
                    Normalize the diagonal elements by the mean power,
                    P(k) + Pshot(k)
                subtract_gaussian : bool, (`False`)
                    Subtract out the Gaussian prediction


r"""
        Plot the off diagonal elements, specifically:

        :math: \sqrt(Cov(k, kbar) / P(k, mu) / P(kbar, mu))

        Parameters
        ---------------
        kbar : {float, int}
            The specific value of k to plot the covariances at
        ax : Axes, optional
            Axes instance to plot to. If `None`, plot to the current axes
        mu : {float, int},
            the first mu value
        mu_bar : {float, int},
            the second mu value
        options : dict
            Dictionary of options with the following keywords:
                show_diagonal : bool, (`True`)
                    If `True`, plot the diagonal value as a stem plot
                show_binned : bool, (`False`)
                    If `True`, plot a smoothed spline interpolation
                show_zero_line : bool, (`False`)
                    If `True`, plot a y=0 horizontal line


r"""
        Return the Gaussian prediction for the covariance between multipoles
        for a periodic box simulation, where the number density is constant.

        See eq. 16 of Grieb et al. 2015 arxiv:1509.04293

        Parameters
        ----------
        model : GalaxySpectrum, QuasarSpectrum
            the model instance used to evaluate the theoretical :math:`P(k,\mu)`
        k : array_like
            the array of wavenumbers (units of :math:`h/\mathrm{Mpc}`) where
            the covariance matrix will be evaluated
        ells : list of int
            the list of multipole numbers to compute the covariance of
        nbar : float
            the constant number density in the box in units of
            :math:`(\mathrm{Mpc}/h)^{-3}` -- the shot noise contribution
            to the covariance is the inverse of this value
        volume : float
            the volume of the box, in units of :math:`(\mathrm{Mpc}/h)^3`
        Nmu : int, optional
            the number of mu bins to use when performing the multipole integration
            over :math:`\mu`

        Returns
        -------
        PoleCovarianceMatrix :
            the covariance matrix object holding the Gaussian prediction for
            the covariance between the specified multipoles

        Examples
        --------
        >>> import numpy
        >>> from pyRSD.rsd import GalaxySpectrum
        >>> from pyRSD.rsdfit.data import PoleCovarianceMatrix
        >>> volume = 1380.0**3
        >>> nbar = 3e-4
        >>> model = GalaxySpectrum(params='boss_dr12_fidcosmo.ini')
        >>> k = numpy.arange(0., 0.4, 0.005) + 0.005/2
        >>> ells = [0,2,4]
        >>> C = PoleCovarianceMatrix.periodic_gaussian_covariance(model, k, ells, nbar, volume)


r"""
        Return the Gaussian prediction for the covariance between multipoles
        for a "cutsky" survey, i.e., a survey with a varying :math:`n(z)`
        distribution

        Parameters
        ----------
        model : GalaxySpectrum, QuasarSpectrum
            the model instance used to evaluate the theoretical :math:`P(k,\mu)`
        k : array_like
            the array of wavenumbers (units of :math:`h/\mathrm{Mpc}`) where
            the covariance matrix will be evaluated
        ells : list of int
            the list of multipole numbers to compute the covariance of
        nbar : callable
            a callable function returning the number density as a function of
            redshift
        fsky : float
            the sky fraction, used to normalized the effective volume calculation
        zmin : float
            the minimum redshift value to include when performing the effective
            volume calculation
        zmax : float
            the maximum redshift value to include when performing the effective
            volume calculation
        FKP_P0 : float, optional
            the FKP P0 value to use, where the FKP weights are :math:`1/(1+n(z)P_0)`;
            default is 1e4
        Nmu : int, optional
            the number of mu bins to use when performing the multipole integration
            over :math:`\mu`
        Nz : int, optional
            the number of redshift bins to use when performing the integral over
            redshift

        Returns
        -------
        PoleCovarianceMatrix :
            the covariance matrix object holding the Gaussian prediction for
            the covariance between the specified multipoles

        Examples
        --------
        >>> import numpy
        >>> from scipy.interpolate import InterpolatedUnivariateSpline as spline
        >>> from pyRSD.rsd import GalaxySpectrum
        >>> from pyRSD.rsdfit.data import PoleCovarianceMatrix
        >>> filename = 'nbar_DR12v5_CMASSLOWZ_North_om0p31_Pfkp10000.dat'
        >>> nbar = numpy.loadtxt(filename, skiprows=3)
        >>> nbar = spline(nbar[:,0], nbar[:,3])
        >>> fsky = 0.1436
        >>> model = GalaxySpectrum(params='boss_dr12_fidcosmo.ini')
        >>> k = numpy.arange(0., 0.4, 0.005) + 0.005/2
        >>> ells = [0,2,4]
        >>> zmin = 0.2
        >>> zmax = 0.5
        >>> C = PoleCovarianceMatrix.cutsky_gaussian_covariance(model, k, ells, nbar, fsky, zmin, zmax)


Enumerate across the upper triangle of blocks


        The unique center values of the `ell` bins, optionally
        returning all (non-unique) values


        Plot the square-root of the diagonal elements for a specfic ell value

        Parameters
        ---------------
        ax : Axes, optional
            Axes instance to plot to. If `None`, plot to the current axes
        ell : int
            multipole number for the first axis
        ell_bar : int
            multipole number for the second axis
        options : dict
            Dictionary of options with the following keywords:
                show_gaussian : bool, (`False`)
                    Overplot the Gaussian prediction, i.e.,
                    :math: 2/N_{modes} * [P(k) + Pshot(k)]**2
                norm_by_gaussian : bool, (`False`)
                    Normalize the diagonal elements by the Gaussian prediction
                norm_by_power : bool, (`False`)
                    Normalize the diagonal elements by the mean power,
                    P(k) + Pshot(k)
                subtract_gaussian : bool, (`False`)
                    Subtract out the Gaussian prediction


r"""
        Plot the off diagonal elements, specifically:

        :math: \sqrt(Cov(k, kbar) / P(k, ell) / P(kbar, ell_bar))

        Parameters
        ---------------
        kbar : {float, int}
            The specific value of k to plot the covariances at
        ax : Axes, optional
            Axes instance to plot to. If `None`, plot to the current axes
        ell : int
            multipole axis for the first axis
        ell_bar : int
            multipole axis for the second axis
        options : dict
            Dictionary of options with the following keywords:
                show_diagonal : bool, (`True`)
                    If `True`, plot the diagonal value as a stem plot
                show_binned : bool, (`False`)
                    If `True`, plot a smoothed spline interpolation
                show_zero_line : bool, (`False`)
                    If `True`, plot a y=0 horizontal line


    Plot the traces of the specified parameters using :mod:`plotly`.

    .. note::
        This is designed to be used in a notebook environment.


    Make a 2D plot of the joint posterior distributions of the specified
    parameters using kernel density estimation and :mod:`plotly`.

    .. note::
        This is designed to be used in a notebook environment.


    Make a 1D histogram plot of the posterior distributions of the specified
    parameter using :mod:`plotly`.

    .. note::
        This is designed to be used in a notebook environment.


Perform kernel density estimation on two MCMC traces.


Convert matplotlib color map to plotly list of RGB tuples.


    A workaround in :mod:`plotly` where latex labels are broken in notebooks.

    See also: https://github.com/plotly/plotly.py/issues/515


Plot a data vs theory comparison using :mod:`plotly`.


Broken off into separate method to simplify code debugging


        If user did not supply range thresholds that varied by azimuth, then
        populate all azimuth bins with constant thresholds before developing
        the range mask.


Produces mask for all useful rain data in volume


Produces mask for all useful drizzle data in volume


Consolidate and invert the bad masks, then add more good masks.


Applies rain mask to volume, and bins up by azimuth


Applies drizzle mask to volume, and bins up by azimuth


Applies FSC mask to volume, and bins up by azimuth


Load KDP-method azimuth-based blockage stats from file.


        Makes a two-panel plot of reflectivity and differential reflectivity
        blockage as functions of azimuth.


        Internal method that actually produces an individual panel in the
        two-panel plot.


        Main processing method. This loops through the list of radar volume
        files, iteratively calling the beam blockage correction routines
        supplied by the BeamBlockSingleVolume class, then compiling the data
        from each file into master arrays that are periodically saved to file,
        both as raw data and as statistics, as well as plotted up.


        Initialize data lists. Basic structure is a 1-D array of 1-D variable-
        length lists. The master array is the same size as the number azimuth
        analysis bins (e.g., BeamBlockSingleVolume.bin_width=1 means size 360).
        The secondary lists, which are unique to each az bin, can be of any
        length and that depends on the amount of data meeting the specified
        criteria in each bin. This will typically grow as the number of files
        processed increases.

        The save attribute basically flags whether the user will start from
        scratch with a list of radar volume files, or if initial arrays will be
        populated with blockage data from a previous run (e.g., on a different
        dataset from the same radar).


Method to write out a blockage statistics file.


        Checks first if stats file, then if not checks if argument is list or
        a single basestring (i.e., one file)


        Saves data to HDF5 file.

        Parameters
        ----------
        obj : pyblock.BeamBlockMultiVolume
            BeamBlockMultiVolume object to save to file
        filename : str
            Path and name of file to save data too. If this file exists it will
            be removed first.
        gzip : bool
            True - Save file will be gzipped, removing previous gzipped file if
                   it exists.

            False - Save file will not be compressed.


        Loads data from HDF5 file.

        Parameters
        ----------
        filename : str
            Full path and name of file that hosts BeamBlockMultiVolume data.
        gzip : bool
            True - File to be loaded is gzipped.

            False - File to be loaded is not gzipped


Isolate data within user-defined blocked regions.


        method = Only relevant to KdpMethodAnalysis plots, can be either
                 'standard', 'strict', or 'loose'
        save = Full path and filename to save image file to
        zdr_range = Range of ZDR corrections displayed
        zh_range = Range of ZH corrections displayed
        return_flag = Set to True to return figure and axes objects


        azimuths = list of azimuths to adjust
        offsets = list of offsets to change to, same size and order as azimuths
        var = change to 'DR', 'ZDR', or 'ZD' to adjust ZDR rather than ZH
        method = 'standard', 'strict', or 'loose'


        Performs weighted multiple linear regression to obtain coefficients for
        self-consistency relationship:
        Zestimate = b * log10(Kdp) + c * Zdr + a (cf. Eq. 1, Lang et al., 2009)
        This is done using median values for unblocked data, with Zh being
        binned up and Kdp and Zdr medians within each bin being determined.

        zmin = Minimum Zh value to consider in regression. We control for
               Kdp <= 0 but just in case it is still noisy in low Zh, you
               can use zmin to filter that out.

        Key output
        ----------
        self.beta_hat = list of coefficients for above eq. (order = [b, c, a])
        self.a, self.b, self.c = more direct storage for these coefficients


        Define a, b, and c in the self-consistency relationship:
        Z = a + b * log10(Kdp) + c * Zdr


        Plots a simple 2-D histogram for visualizing the results & performance
        of the regress_unblocked_data() method.

        Keywords
        --------
        save = Name of image file to save figure to.
        domain = 2-element tuple to narrow down the display range of histogram.
        zmin = Minimum Zh to consider in regression (if redone).


        Compute the I1 and I2 integrals from Eq. 2-3 in Lang et al. (2009).
        Loops thru each of the blocked azimuth bins. First it will
        integrate the I1 side of the self-consistency relationship. Then it
        integrates the I2 side of the relationship and checks a variety of
        possible Zh offsets to see which one best matches I1. The resolution
        of the offset steps is delk, and kmax*delk is the maximum possible Zh
        offset examined. The zh_adjustments attribute then holds all of the Zh
        offsets to apply at each azimuth bin. If it is 0 then there is no
        offset to apply at that azimuth. If it is equal to delk*(kmax-1)
        then the solution likely did not converge and the block is too strong.
        Note that blocks ~30 dBZ or more are extremely difficult to correct,
        and even 20+ dBZ corrections are likely associated with greater error
        than weaker blocks.

        Assumes <ZDR> should match inherent offset in unblocked data when
        Z=0 dBZ. This is used to develop the zdr_offset attribute.

        Key Outputs
        -----------
        zh_adjustments = Add to Zh at each azimuth bin to correct for blockage.
        zdr_offsets = Difference between <ZDR(Z=0)> at each blocked azimuth and
                      the same parameter in unblocked data. Apply to the ZDR
                      data in blocked regions to force it to match unblocked
                      behavior. This is an alternative to correcting ZDR via
                      the KDP method.


        save = Full path and filename to save image file to
        zdr_range = Range of ZDR corrections displayed
        zh_range = Range of ZH corrections displayed
        return_flag = Set to True to return figure and axes objects


        azimuths = list of azimuths to adjust
        offsets = list of offsets to change to, same size and order as azimuths
        var = change to 'DR', 'ZDR', or 'ZD' to adjust ZDR rather than ZH


        Populates the unblocked data dictionary by consolidating all data from
        unblocked azimuths. Performs a simple test to make sure there is no
        weirdness in terms of data structure.


Simple check to make sure everything is the same size


Determine median Zdr/Kdp as function of Z for unblocked regions


Weight by number of Kdp/Zdr observations in each Z bin


Loop thru all possible reflectivity offsets up to delk*kmax


Obtain expected values for Zdr/Kdp for each blocked azimuth bin


        Prep work for determining the corrections ...
        Figure out first the differences!


Performs standard, loose, and strict correction calculations


Choose the right variables to consider depending on zh_flag


        Prints out Zh & Zdr corrections for all three approaches:
        standard, strict, & loose (values in dBZ or dB)


        Prints out Zh & Zdr corrections for SelfConsistentAnalysis, along with
        I1 and matched I2 integrated values.


Simple plotting of Zh and Zdr corrections via matplotlib


Translates the KdpMethodAnalysis object to plottable arrays


Translates the SelfConsistentAnalysis object to plottable arrays


    bad_masks = list of bad data masks to apply to array.
    Returns a mask that will provide all good data.


Calculate 95% confidence interval for median


    Simple unweighted least-squares multiple linear regression using
    numpy.linalg module.


    Simple weighted least-squares multiple linear regression using
    statsmodel module.


    Use np.dot() to quickly evaluate a linear expression to recover the
    dependent variable from any number of independent variables.


Spectrum Autoencoder model for the SDSS main galaxy sample.

    This model has S=10 latents and a restframe resolution R=5881 for a maximum redshift of z_max=0.5.

    See Spender paper I for details:
        Melchior et al. (2022): arXiv:2211.07890


Spectrum Autoencoder super-resolution model for the SDSS main galaxy sample.

    This model has S=8 latents and a restframe resolution R=11762 for a maximum redshift of z_max=0.5.

    See Spender paper I (section 4.2) for details:
        Melchior et al. (2022): arXiv:2211.07890


Spectrum Autoencoder model for the SDSS main galaxy sample.

    This model has S=6 latents and a restframe resolution R=7000 for a maximum redshift of z_max=0.5.
    It has been trained with fidelity, similarity, and consistency losses to provide a redshift-invariant latent space.

    See spender papers for details:
        Liang at al. (2023): arXiv:2302.02496
        Melchior et al. (2022): arXiv:2211.07890


Spectrum Autoencoder model for the DESI EDR Bright Galaxy Survey sample.

    This model has S=6 latents and a restframe resolution R=9780 for a maximum redshift of z_max=0.8.
    It has been trained with fidelity, similarity, and consistency losses to provide a redshift-invariant latent space.

    See spender papers for details:
        Liang at al. (2023b): arXiv:2307.07664
        Liang at al. (2023a): arXiv:2302.02496
        Melchior et al. (2022): arXiv:2211.07890


Spectrum Autoencoder model for the DESI EDR Milky Way Survey sample.

    This model has S=6 latents and a restframe resolution R=7864 for a maximum redshift of z_max=0.005.
    It has been trained with fidelity, and similarity but not consistency loss.

    See spender papers for details:
        Liang at al. (2023b): arXiv:2307.07664
        Liang at al. (2023a): arXiv:2302.02496
        Melchior et al. (2022): arXiv:2211.07890


Normalizing flow model for the DESI EDR BGS spender latent space.

    See spender papers for details:
        Liang at al. (2023b): arXiv:2307.07664


Normalizing flow model for the DESI EDR MWS spender latent space.

    See spender papers for details:
        Liang at al. (2023b): arXiv:2307.07664


    Train the neural density estimator based on input data.
    Here we use the ``log(P)`` loss. This function is not used in the ``popsed`` project.
    Parameters
    ----------
    n_epochs: int.
        Number of epochs to train.
    display: bool.
        Whether to display the training loss.
    suffix: str.
        Suffix to add to the output file.


Builds MAF to describe p(x).
    Args:
        dim: Batch of xs, used to infer dimensionality and (optional) z-scoring.
        hidden_features: Number of hidden features.
        num_transforms: Number of transforms.
        embedding_net: Optional embedding network for y.
        kwargs: Additional arguments that are passed by the build function but are not
            relevant for maf and are therefore ignored.
    Returns:
        Neural network.


Forward method

        Parameters
        ----------
        x: `torch.tensor`

        Returns
        -------
        x': `torch.tensor`, same shape as `x`


Forward method

        Parameters
        ----------
        y: `torch.tensor`, shape (N, L)
            Batch of observed spectra

        Returns
        -------
        s: `torch.tensor`, shape (N, n_latent)
            Batch of latents that encode `spectra`


Gradient of the attention weights

        Factor to compute the importance of attention for Grad-FAM method.

        Requires a previous `loss.backward` call for any scalar loss function based on
        outputs of this class's `forward` method. This functionality is switched off
        during training.


Decode latents into restframe spectrum

        Parameter
        ---------
        s: `torch.tensor`, shape (N, S)
            Batch of latents

        Returns
        -------
        x: `torch.tensor`, shape (N, L)
            Batch of restframe spectra


Forward method

        Parameter
        ---------
        s: `torch.tensor`, shape (N, S)
            Batch of latents
        instrument: :class:`spender.Instrument`
            Instrument to generate spectrum for
        z: `torch.tensor`, shape (N, 1)
            Redshifts for each spectrum

        Returns
        -------
        y: `torch.tensor`, shape (N, L)
            Batch of spectra at redshift `z` as observed by `instrument`


Transformations from restframe to observed frame

        Parameter
        ---------
        x: `torch.tensor`, shape (N, S)
            Batch of restframe spectra
        instrument: :class:`spender.Instrument`
            Instrument to generate spectrum for
        z: `torch.tensor`, shape (N, 1)
            Redshifts for each spectrum
        return_valid: bool
            Whether the range of valid wavelength in the model is returned
            Whatever this flag, the model is set to 0 at invalid wavelengths.

        Returns
        -------
        y: `torch.tensor`, shape (N, L)
            Batch of spectra at redshift `z` as observed by `instrument`


Encode from observed spectrum to latents

        Parameters
        ----------
        y: `torch.tensor`, shape (N, L)
            Batch of observed spectra

        Returns
        -------
        s: `torch.tensor`, shape (N, n_latent)
            Batch of latents that encode `spectra`


Apply normalization to the reconstructed spectrum using a flat region
        in the restframe spectrum, as determined by the model.

        Inputs:
        =======
        spectrum: torch.tensor
            Unnormalized observed spectrum.
        weights: torch.tensor
            Inverse variance weights of the observed spectrum.
        restframe: torch.tensor
            Restframe spectrum from the model.
        reconstruction: torch.tensor
            Reconstructed spectrum from the model.

        Outputs:
        ========
        restframe: torch.tensor
            Restframe spectrum, normalized to the observed spectrum.
        reconstruction: torch.tensor
            Reconstructed spectrum, normalized to the observed spectrum.


Perform a forward pass through the model to create a latent, restframe, and reconstruction from an observed spectrum, instrument, and redshift. If inverse variance weights are passed, also normalizes the reconstruction to the observed spectrum.


Forward method

        Transforms observed spectra into their reconstruction for a given intrument and redshift. If weights are passed, also normalizes the reconstruction to the observed spectrum.

        Parameter
        ---------
        y: `torch.tensor`, shape (N, L)
            Batch of observed spectra
        instrument: :class:`spender.Instrument`
            Instrument to generate spectrum for
        z: `torch.tensor`, shape (N, 1)
            Redshifts for each spectrum. When given, `aux` is ignored.
        s: `torch.tensor`, shape (N, S)
            (optional) Batch of latents. When given, encoding is omitted and these latents are used instead.
        normalize: bool
            (optional) Whether to normalize the reconstruction match the amplitude of the observed spectrum.
            If this is set to False, it requires that the spectra have been flux-normalized.
        weights: `torch.tensor`, shape (N, L)
            (optional) Inverse variance weights for each spectrum, used for normalization computation.

        Returns
        -------
        y: `torch.tensor`, shape (N, L)
            Batch of spectra at redshift `z` as observed by `instrument`


Weighted MSE loss

        Parameter
        --------
        y: `torch.tensor`, shape (N, L)
            Batch of observed spectra
        w: `torch.tensor`, shape (N, L)
            Batch of weights for observed spectra
        instrument: :class:`spender.Instrument`
            Instrument to generate spectrum for
        z: `torch.tensor`, shape (N, 1)
            Redshifts for each spectrum. When given, `aux` is ignored.
        s: `torch.tensor`, shape (N, S)
            (optional) Batch of latents. When given, encoding is omitted and these
            latents are used instead.
        normalize: bool
            (optional) Whether to normalize the reconstruction match the amplitude of the observed spectrum.
            If this is set to False, it requires that the spectra have been flux-normalized.
        individual: bool
            Whether the loss is computed for each spectrum individually or aggregated

        Returns
        -------
        float or `torch.tensor`, shape (N,) of weighted MSE loss


Load models from file

    Parameter
    ---------
    filename: str
        Path to file which contains the torch state dictionary
    instrument: :class:`spender.Instrument`
        Instrument to generate spectrum for
    device: `torch.Device`
        Device to load model structure into

    Returns
    -------
    model: `torch.nn.Module`
        The default :class`SpectrumAutoencoder` model loaded from file
    loss: `torch.tensor`
        Training and validation loss for this model


List all batch files

        Parameters
        ----------
        dir: string
            Root directory for data storage
        which: ['train', 'valid', 'test'] or None
            Which subset of the spectra to return. If `None`, returns all of them.
        tag: string
            Name to specify which batch files to load

        Returns
        -------
        list of filepaths


Save batch into a pickled file

        Parameters
        ----------
        dir: string
            Root directory for data storage
        batch: `torch.tensor`, shape (N, L)
            Spectrum batch
        tag: string
            Name to specify which batch file name
        counter: int
            Set to add a batch counter to the filename

        Returns
        -------
        None


Save all spectra for given ids into batch files

        Parameters
        ----------
        dir: string
            Root directory for data storage
        ids: list of (survey, program, healpix, target)
            Identifier of healpix
        tag: string
            Name to specify the batch file name
        batch_size: int
            Number of spectra in each batch

        Returns
        -------
        None


Prepare spectra in DESI healpix for analysis

        This method creates an extended mask, using the original DESI mask and
        the skyline mask of the instrument. The weights for all masked regions is set to
        0, but the spectrum itself is not altered.

        The spectrum and weights are then cast into a fixed-format data vector, which
        standardizes the variable wavelengths of each observation.

        A normalization is computed as the median flux in the relatively flat region
        between restframe 5300 and 5850 A. The spectrum is divided by this factor, the weights
        are muliplied with the square of this factor.

        Parameter
        ---------
        filename: string
            Path to local file containing the spectrum
        target: string
            target class: BGS, LRG, ELG, QSO

        Returns
        -------
        spec: `torch.tensor`, shape (L, )
            Normalized spectrum
        w: `torch.tensor`, shape (L, )
            Inverse variance weights of normalized spectrum
        norm: `torch.tensor`, shape (1, )
            Flux normalization factor
        z: `torch.tensor`, shape (1, )
            Redshift (only returned when argument z=None)
        zerr: `torch.tensor`, shape (1, )
            Redshift error (only returned when argument z=None)


Select SURVEY, PROGRAM, HEALPIX from healpix look up table for healpix that
        match `selection_fct`. Look up table only includes information on
        TILEID, SURVEY, PROGRAM, PETAL_LOC, HEALPIX

        Parameters
        ----------
        dir: string
            Root directory for data storage
        fields: list of string
            Catalog field names to return
        selection_fct: callable
            Function to select matches from all items in the main table

        Returns
        -------
        fields: `torch.tensor`, shape (N, F)
            Tensor of fields for the selected healpix


Augment spectra for greater diversity

        Parameters
        ----------
        batch: `torch.tensor`, shape (N, L)
            Spectrum batch
        redshift: bool
            Modify redshift by up to 0.2 (keeping it within 0...0.5)
        noise: bool
            Whether to add noise to the spectrum (up to 0.2*max(spectrum))
        mask: bool
            Whether to block out a fraction (given by `ratio`) of the spectrum
        ratio: float
            Fraction of the spectrum that will be masked
        z_new: float
            Adopt this redshift for all spectra in the batch

        Returns
        -------
        spec: `torch.tensor`, shape (N, L)
            Altered spectrum
        w: `torch.tensor`, shape (N, L)
            Altered inverse variance weights of spectrum
        z: `torch.tensor`, shape (N, )
            Altered redshift


Save all spectra for given ids into batch files

        Parameters
        ----------
        dir: string
            Root directory for data storage
        ids: list of (plate, mjd, fiberid)
            Identifier of spectrum
        tag: string
            Name to specify the batch file name
        batch_size: int
            Number of spectra in each batch

        Returns
        -------
        None


Prepare spectrum for analysis

        This method creates an extended mask, using the original SDSS `and_mask` and
        the skyline mask of the instrument. The weights for all masked regions is set to
        0, but the spectrum itself is not altered.

        The spectrum and weights are then cast into a fixed-format data vector, which
        standardizes the variable wavelengths of each observation.

        A normalization is computed as the median flux in the relatively flat region
        between restframe 5300 and 5850 A. The spectrum is divided by this factor, the weights
        are muliplied with the square of this factor.

        Parameter
        ---------
        filename: string
            Path to local file containing the spectrum
        z: float or None
            Redshift of the spectrum
            If None, redshift and its error will be read from the file and returned

        Returns
        -------
        spec: `torch.tensor`, shape (L, )
            Normalized spectrum
        w: `torch.tensor`, shape (L, )
            Inverse variance weights of normalized spectrum
        norm: `torch.tensor`, shape (1, )
            Flux normalization factor
        z: `torch.tensor`, shape (1, )
            Redshift (only returned when argument z=None)
        zerr: `torch.tensor`, shape (1, )
            Redshift error (only returned when argument z=None)


Make a batch of spectra from their IDs

        Parameters
        ----------
        dir: string
            Root directory for data storage
        fields: list of (plate, mjd, fiberid, [z, z_err])
            List of object qualifiers from query()

        Returns
        -------
        spec: `torch.tensor`, shape (N, L)
            Normalized spectrum
        w: `torch.tensor`, shape (N, L)
            Inverse variance weights of normalized spectrum
        z: `torch.tensor`, shape (N, )
            Redshift from SDSS pipeline
        norm: `torch.tensor`, shape (N, )
            Normalization factor
        zerr: torch.tensor`, shape (N, )
            Redshift error from SDSS pipeline


Select fields from main specrum table for objects that match `selection_fct`

        NOTE: This function will download the file `specObj-dr16.fits` from the data
        archive. This file is *not* small...

        Parameters
        ----------
        dir: string
            Root directory for data storage
        fields: list of string
            Catalog field names to return
        selection_fct: callable
            Function to select matches from all items in the main table

        Returns
        -------
        fields: `torch.tensor`, shape (N, F)
            Tensor of fields for the selected objects


Pyuvsim output for interpolated polarized beam.


Run the test of comparing matvis and pyuvsim visibilities.


Default CST UVBeam.


CST Beam made unpolarized.


Convert angle cosines in ENU coordinates into azimuth and zenith angle.

    For a pointing vector in East-North-Up (ENU) coordinates vec{p}, the input
    arguments are ``enu_e = vec{p}.hat{e}`` and ``enu_n = vec{p}.hat{n}`, where
    ``hat{e}`` is a unit vector in ENU coordinates etc.

    For a drift-scan telescope pointing at the zenith, the ``hat{e}`` direction
    is aligned with the ``U`` direction (in the UVW plane), which means that we
    can identify the direction cosines ``l = enu_e`` and ``m = enu_n``.

    Azimuth is oriented East of North, i.e. Az(N) = 0 deg, Az(E) = +90 deg in
    the astropy convention, and North of East, i.e. Az(N) = +90 deg, and
    Az(E) = 0 deg in the UVBeam convention.

    Parameters
    ----------
    enu_e, enu_n : array_like
        Normalized angle cosine coordinates on the interval (-1, +1).

    orientation : str, optional
        Orientation convention used for the azimuth angle. The default is
        ``'astropy'``, which uses an East of North convention (Az(N) = 0 deg,
        Az(E) = +90 deg). Alternatively, the ``'uvbeam'`` convention uses
        North of East (Az(N) = +90 deg, Az(E) = 0 deg).

    periodic_azimuth : bool, optional
        if True, constrain az to be betwee 0 and 2 * pi
        This avoids the issue that arctan2 outputs angles between -pi and pi
        while most CST beam formats store azimuths between 0 and 2pi which leads
        interpolation domain mismatches.

    Returns
    -------
    az, za : array_like
        Corresponding azimuth and zenith angles (in radians).


3x3 transformation matrix to rotate ECI to ENU coordinates.

    Transformation matrix to project Earth-Centered Inertial (ECI) coordinates
    to local observer-centric East-North-Up (ENU) coordinates at a given time
    and location.

    The ECI coordinates are aligned with the celestial pole, i.e. for (x,y,z)
    (RA=0  deg, Dec=0  deg) = (1, 0, 0)
    (RA=90 deg, Dec=0  deg) = (0, 1, 0)
    (RA=0  deg, Dec=90 deg) = (0, 0, 1)

    Note: This is a stripped-down version of the ``eq2top_m`` function.

    Parameters
    ----------
    ha : float
        Hour angle, in radians, where HA = LST - RA.

    lat : float
        Latitude of the observer, in radians.

    Returns
    -------
    m : array_like
        3x3 array containing the rotation matrix for a given time.


3x3 transformation matrix to rotate ENU to ECI coordinates.

    3x3 transformation matrix to project local observer-centric East-North-Up
    (ENU) coordinates at a given time and location to Earth-Centered Inertial
    (ECI) coordinates.

    The ECI coordinates are aligned with the celestial pole, i.e. for (x,y,z)
    (RA=0  deg, Dec=0  deg) = (1, 0, 0)
    (RA=90 deg, Dec=0  deg) = (0, 1, 0)
    (RA=0  deg, Dec=90 deg) = (0, 0, 1)

    Parameters
    ----------
    ha : float
        Hour angle, in radians, where HA = LST - RA.

    lat : float
        Latitude of the observer, in radians.

    Returns
    -------
    m : array_like
        3x3 array containing the rotation matrix for a given time.


Convert alt/az coordinates as given by Astropy, into ENU coordinates.

    Astropy has Az oriented East of North, i.e. Az(N) = 0 deg, Az(E) = +90 deg.
    See: https://gssc.esa.int/navipedia/index.php/Transformations_between_ECEF_and_ENU_coordinates


Coordinate transform of source locations from equatorial to Cartesian.

    This converts RA and Dec angles to a Cartesian x,y,z unit vector in an
    Earth-Centered Inertial (ECI) coordinate system aligned with the celestial
    pole, i.e. for (x,y,z)
    (RA=0  deg, Dec=0  deg) = (1, 0, 0)
    (RA=90 deg, Dec=0  deg) = (0, 1, 0)
    (RA=0  deg, Dec=90 deg) = (0, 0, 1)

    The RA and Dec are assumed to be in a particular reference frame that may
    not match-up with standard frames like ICRS/J2000. To convert coordinates
    from a standard system into the relevant frame, see
    :func:`~equatorial_to_eci_coords`.

    Parameters
    ----------
    ra, dec : array_like
        1D arrays of source positions in equatorial coordinates (radians).

    Returns
    -------
    array_like
        Equatorial coordinates of sources, in Cartesian system. Shape=(3, NSRCS).


Convert RA and Dec coordinates into the ECI system used by  matvis.

    Convert RA and Dec coordinates into the ECI (Earth-Centered Inertial)
    system used by  matvis. This ECI system is aligned with the celestial pole,
    not the Earth's axis.

    To ensure that all corrections are properly taken into account, this
    function uses Astropy to find the Alt/Az positions of the coordinates at a
    specified reference time and location. These are then transformed back to
    ENU (East-North-Up) coordinates, and then to Cartesian ECI coordinates,
    using the inverses of the same transforms that are used to do the forward
    transforms when running `` matvis``. The ECI coordinates are then finally
    converted back into adjusted RA and Dec coordinates in the ECI system.

    While the time-dependent corrections are determined for the reference time,
    the adjusted coordinates are expected to yield a good approximation to the
    true Azimuth and Zenith angle of the sources at other times (but the same
    location) when passed into the `` matvis`` function via the ``crd_eq``
    array and then converted using the standard conversions within `` matvis``.

    The following example code shows how to set up the Astropy ``Time`` and
    ``EarthLocation`` objects that are required by this function::

        # HERA location
        location = EarthLocation.from_geodetic(lat=-30.7215,
                                               lon=21.4283,
                                               height=1073.)
        # Observation time
        obstime = Time('2018-08-31T04:02:30.11', format='isot', scale='utc')

    Parameters
    ----------
    ra, dec : array_like
        Input RA and Dec positions. The units and reference frame of these
        positions can be set using the ``unit`` and ``frame`` kwargs.

    obstime : astropy.Time object
        ``Time`` object specifying the time of the reference observation.

    location : astropy.EarthLocation
        ``EarthLocation`` object specifying the location of the reference
        observation.

    unit : str, optional
        Which units the input RA and Dec values are in, using names intelligible
        to ``astropy.SkyCoord``. Default: 'rad'.

    frame : str, optional
        Which frame that input RA and Dec positions are specified in. Any
        system recognized by ``astropy.SkyCoord`` can be used. Default: 'icrs'.

    Returns
    -------
    eci_ra, eci_dec : array_like
        Arrays of RA and Dec coordinates with respect to the ECI system used
        by  matvis.


    Construct an axis-angle rotation matrix R that carries vector r1 to r2.

    This function has been adopted from pyradiosky.

    Parameters
    ----------
    r1, r2 : ndarray, shape (3, N)
        Sets of 3D vectors.

    Returns
    -------
    R : ndarray, shape (3, 3, N)
        Rotation matrices such that R[..., i] @ r1[:, i] = r2[:, i].


    Return the unit vector theta_hat in Cartesian coords for spherical angles.

    Parameters
    ----------
    theta : array_like
        Colatitude angle(s), in radians.
    phi : array_like
        Longitude angle(s), in radians.

    Returns
    -------
    vec : ndarray, shape (3, ...)
        Cartesian components of theta_hat.


    Return the unit vector phi_hat in Cartesian coords for spherical angles.

    Parameters
    ----------
    theta : array_like
        Colatitude angle(s), in radians.
    phi : array_like
        Longitude angle(s), in radians.

    Returns
    -------
    vec : ndarray, shape (3, ...)
        Cartesian components of phi_hat.


    Run a basic simulation using ``matvis``.

    This wrapper handles the necessary coordinate conversions etc.

    Parameters
    ----------
    ants : dict
        Dictionary of antenna positions. The keys are the antenna names
        (integers) and the values are the Cartesian x,y,z positions of the
        antennas (in meters) relative to the array center.
    fluxes : array_like
        2D array with the flux of each source as a function of frequency, of
        shape (NSRCS, NFREQS).
    ra, dec : array_like
        Arrays of source RA and Dec positions in radians. RA goes from [0, 2 pi]
        and Dec from [-pi/2, +pi/2].
    freqs : array_like
        Frequency channels for the simulation, in Hz.
    times : astropy.Time instance
        Times of the observation (can be an array of times).
    beams : list of ``UVBeam``, ``AnalyticBeam`` or ``BeamInterface`` objects
        Beam objects to use for each antenna.
    telescope_loc
        An EarthLocation object representing the center of the array.
    polarized : bool, optional
        If True, use polarized beams and calculate all available linearly-
        polarized visibilities, e.g. V_nn, V_ne, V_en, V_ee.
        Default: False (only uses the 'ee' polarization).
    precision : int, optional
        Which precision setting to use for :func:`~matvis`. If set to ``1``,
        uses the (``np.float32``, ``np.complex64``) dtypes. If set to ``2``,
        uses the (``np.float64``, ``np.complex128``) dtypes.
    use_feed
        Either 'x' or 'y'. Only used if polarized is False.
    use_gpu : bool, optional
        Whether to use the GPU for simulation.
    beam_spline_opts : dict, optional
        Options to be passed to :meth:`pyuvdata.uvbeam.UVBeam.interp` as `spline_opts`.
    beam_idx
        An array of integers, of the same length as ``ants``. Each entry is for an
        antenna of the same index, and its value should be the index of the beam in
        the beam list that corresponds to the antenna.
    antpairs
        A list of antpairs (in the form of 2-tuples of integers) to actually
        calculate visibility for. If None, all feed-pairs are calculated.
    source_buffer : float, optional
        The fraction of the total number of sources to use when allocating memory
        for the sources above horizon. For large numbers of sources, a fraction of
        ~0.55 should be sufficient.

    Returns
    -------
    vis : array_like
        Complex array of shape (NFREQS, NTIMES, NBLS, NFEED, NFEED)
        if ``polarized == True``, or (NFREQS, NTIMES, NBLS) otherwise.


Run the script.


Add common profiling options to a function.


Run profiling of matvis with a HERA-like array.


Ceiling division for integers.

    From https://stackoverflow.com/a/17511341/1467820


Get a human-readable data size.

    From: https://stackoverflow.com/a/43690506/1467820


    Log the memory usage since last call.

    Parameters
    ----------
    highest_peak : int
        The highest peak memory usage in bytes.

    Returns
    -------
    int
        The new highest peak memory usage.


Debug logging of the value of an array.


Logging of progress.


    Calculate visibility from an input intensity map and beam model.

    Parameters
    ----------
    antpos : array_like
        Antenna position array. Shape=(NANT, 3).
    freq : float
        Frequency to evaluate the visibilities at [GHz].
    I_sky : array_like
        Intensity distribution of sources/pixels on the sky, assuming intensity
        (Stokes I) only. The Stokes I intensity will be split equally between
        the two linear polarization channels, resulting in a factor of 0.5 from
        the value inputted here. This is done even if only one polarization
        channel is simulated.
        Shape=(NSRCS,).
    beam_list : list of UVBeam, optional
        If specified, evaluate primary beam values directly using UVBeam
        objects instead of using pixelized beam maps. Only one of ``bm_cube`` and
        ``beam_list`` should be provided.Note that if `polarized` is True,
        these beams must be efield beams, and conversely if `polarized` is False they
        must be power beams with a single polarization (either XX or YY).
    antpairs : array_like, optional
        Either a 2D array, shape ``(Npairs, 2)``, or list of 2-tuples of ints, with
        the list of antenna-pairs to return as visibilities (all feed-pairs are always
        calculated). If None, all feed-pairs are returned.
    precision : int, optional
        Which precision level to use for floats and complex numbers.
        Allowed values:
        - 1: float32, complex64
        - 2: float64, complex128
    polarized : bool, optional
        Whether to simulate a full polarized response in terms of nn, ne, en,
        ee visibilities. See Eq. 6 of Kohn+ (arXiv:1802.04151) for notation.
        Default: False.
    beam_idx
        Optional length-NANT array specifying a beam index for each antenna.
        By default, either a single beam is assumed to apply to all antennas or
        each antenna gets its own beam.
    beam_spline_opts : dict, optional
        Dictionary of options to pass to the beam interpolation function.
    max_progress_reports : int, optional
        Maximum number of progress reports to print to the screen (if logging level
        allows). Default is 100.
    matprod_method : str, optional
        The method to use for the final matrix multiplication. Default is 'CPUMatMul',
        which simply uses `np.dot` over the two full matrices. Currently, the other
        option is `CPUVectorLoop`, which uses a loop over the antenna pairs,
        computing the sum over sources as a vector dot product.
        Whether to calculate visibilities for each antpair in antpairs as a vector
        dot-product instead of using a full matrix-matrix multiplication for all
        possible pairs. Default is False. Setting to True can be faster for large
        arrays where `antpairs` is small (possibly from high redundancy). You should
        run a performance test before using this.
    coord_method : str, optional
        The method to use to transform coordinates from the equatorial to horizontal
        frame. The default is to use Astropy coordinate transforms. A faster option,
        which is accurate to within 6 mas, is to use "CoordinateTransformERFA".
    max_memory : int, optional
        The maximum memory (in bytes) to use for the visibility calculation. This is
        not a hard-set limit, but rather a guideline for how much memory to use. If the
        expected memory usage is more than this, the calculation will be broken up into
        chunks. Default is 512 MB.
    min_chunks : int, optional
        The minimum number of chunks to break the source axis into. Default is 1.
    source_buffer : float, optional
        The fraction of the total sources (per chunk) to pre-allocate memory for.
        Default is 0.55, which allows for 10% variance around half the sources
        (since half should be below the horizon). If you have a particular sky model in
        which you expect more or less sources to appear above the horizon at any time,
        set this to a different value.
    coord_method_params
        Parameters particular to the coordinate rotation method of choice. For example,
        for the CoordinateRotationERFA (and GPU version of the same) method, there
        is the parameter ``update_bcrs_every``, which should be a time in seconds, for
        which larger values speed up the computation.

    Returns
    -------
    vis : array_like
        Simulated visibilities. If `polarized = True`, the output will have
        shape (NTIMES, NBLS, NFEED, NFEED), otherwise it will have
        shape (NTIMES, NBLS).


Evaluate the beam on the CPU.

        This function will either interpolate the beam to the given coordinates tx, ty,
        or evaluate the beam there if it is an analytic beam.

        Parameters
        ----------
        tx, ty
            Coordinates to evaluate the beam at, in sin-projection.


Compute the topocentric coordinates from the given time and telescope loc.


        Apply light deflection by a solar-system body.

        Parameters
        ----------
        p : np.ndarray (3, nsrc)
            direction from observer to source (unit vector)
        e : np.ndarray (3,)
            direction from body to observer (unit vector)
        em
            distance from body to observer (au)
        dlim
            deflection limiter (Note 4)


        Apply aberration to transform natural direction into proper direction.

        Parameters
        ----------
        pnat : (3, nsrc) array
            natural direction to the source (unit vector)
        v : (3,) vector
            observer barycentric velocity in units of c
        s
            distance between the Sun and the observer (au)
        bm1
            ``sqrt(1-|v|^2)``: reciprocal of Lorenz factor


Rotate the coordinates into the observed frame.


Perform the source-summing operation for a single time and chunk.

        Parameters
        ----------
        z
            Complex integrand. Shape=(Nfeed, Nant, Nax, Nsrc).
        out
            Output array, shaped as (Nfeed, Nfeed, Npairs).


Computes a.conj() @ b.T.


Allocate memory for the visibilities.

        The shape here is (nchunks, nant, nfeed, nant, nfeed), which is backwards
        from what you'd expect (nfeed,nant, nfeed, nant). This is because the
        fortran ordering is used in CUBLAS, which is the same as the transpose of the
        expected shape.


Perform the source-summing operation for a single time and chunk.


Sum the chunks into the output array.

        Here we need to also re-shape the visibility array into the output array.

        Parameters
        ----------
        out
            The output visibilities, with shape (Nfeed, Nfeed, Npairs).


Allocate memory for the visibilities.


Sum the chunks into the output array.


GPU implementation of the visibility simulator.


Obtain coordinates for doing map_coordinates interpolation from a UVBeam.


Evaluate the beam on the GPU.

        This function will either interpolate the beam to the given coordinates tx, ty,
        or evaluate the beam there if it is an analytic beam.

        Parameters
        ----------
        tx, ty
            Coordinates to evaluate the beam at, in sin-projection.


Perform the beam interpolation, choosing between CPU and GPU as necessary.


    Interpolate beam values from a regular az/za grid using GPU.

    Parameters
    ----------
    beam
        The beam values. The shape of this array should be
        ``(nbeam, nax, nfeed, nza, naz)``. This is the axis ordering returned by
        UVBeam.interp. This array can either be real or complex. Either way, the output
        is complex.
    daz, dza
        The grid sizes in azimuth and zenith-angle respectively.
    az, za
        The azimuth and zenith-angle values of the sources to which to interpolate.
        These should be  1D arrays. They are not treated as a "grid".

    Returns
    -------
    beam_at_src
        The beam interpolated at the sources. The shape of the array is
        ``(nbeam, nfeed, nax, nsrc)``. The array is always complex (at single or
        double precision, depending on the input).


Given a BeamInterface, prepare it for an un-polarized simulation.


Perform all the operations and checks on the input beams.

    Checks that the beam indices match the number of antennas, pre-interpolates to the
    given frequency, and checks that the beam type is appropriate for the given
    polarization

    Parameters
    ----------
    beam_idx
        Index of the beam to use for each antenna.
    beam_list
        List of unique beams.
    polarized
        Whether to use beam polarization
    nant
        Number of antennas
    freq
        Frequency to interpolate beam to.


Perform the beam interpolation.

        This method must return an array of shape ``(nbeam, nfeed, nax, nsrcs_up)``.


        Set the chunk of coordinates to rotate.

        This function sets the chunk of coordinates to rotate. It is used following the
        `rotate` method and returns a chunk of sources coordinates and fluxes for the sources
        above the horizon. If the sky model is polarized, it also rotates the frame of the
        coherency matrix to the alt/az frame. The chunk size is determined by the `chunk_size`
        parameter.


        Rotate the frame of the coherency matrix.

        This function rotates the coherency matrix from the equatorial frame to the
        alt/az frame. It is used in the `rotate` method of subclasses.


Allocate memory for the visibilities.

        The shape of the visibilities must have a first axis of length nchunks,
        but then can be arbitrary shaped after that, so long as it is consistently
        used throughout the class.


        Perform the source-summing operation for a single time and chunk.

        Parameters
        ----------
        z
            Complex integrand. Shape=(Nant, Nfeed, Nax, Nsrc).
        out
            Output array, shaped like the visibilities set in `allocate_vis`, but
            without the first chunk axis.


        Sum the chunks into the output array.

        Parameters
        ----------
        out
            The output visibilities, with shape (Nfeed, Nfeed, Npairs).


Validate input shapes and types.


Decorator to mark a method as the handler for a particular VCS.


Extract embeddings from the AstroClip model and save them to a file


Functionality for plotting retrieved galaxy images


Functionality for plotting retrieved galaxy spectra


Generate a catalog from a list of files.


Cross-match two tables.


Get images from files.


Download Galaxy Zoo 5 classifications.


Get the locations of the Legacy Survey image files.


    Pairs Galaxy Zoo classifications with DECaLS images in an Astropy table.

    Args:
        root_dir (str): Root directory of DECaLS images.
        survey_path (str): Path to Galaxy Zoo survey.

    Returns:
        Table: Table of paired classifications.


Functionality for plotting radar chart


Function to train and evaluate a simple feedforward neural network on a dataset.


Functionality to plot redshift scatter plots for different models.


Download the PROVABGS data from the web and save it to the specified directory.


Get the best fit model for each galaxy.


Cross-match the AstroCLIP and PROVABGS datasets.


Train a few-shot model using a simple neural network


Train a zero-shot model using KNN


Use model to get predictions


Initialize the weights of a module based on the depth of the model.


Safely evaluate an arithmetic expression.

    :param s: expression to evaluate; should only contain numbers, spaces, or the
        symbols `+, -, *, /, _, (, )`; exponential notation is supported
    :param max_len: maximum length string that will be evaluated; longer strings raise
        a `ValueError`


Use a default .env but tell the user how to create their own.


Forward pass through the model.
        The training step performs masking before preprocessing,
        thus samples should not be preprocessed again as in forward()


Randomly masks contiguous sections of an unbatched sequence,
        ensuring separation between chunks is at least chunk_width.


Override to add your own encoders.


Momentum update of the key encoder.


Batch shuffle, for making use of BatchNorm.

        *** Only support DistributedDataParallel (DDP) model. ***


Undo batch shuffle.

        *** Only support DistributedDataParallel (DDP) model. ***


        Input:
            im_q: a batch of query images
            im_k: a batch of key images
            queue: a queue from which to pick negative samples
        Output:
            logits, targets


        Distributed synchronization of the metric
        Warning: does not synchronize the deque!


    Returns:
        True if distributed training is enabled


    Returns:
        True if the current process is the main one.


This function disables printing when not in the main process


Enable distributed mode

    Args:
        set_cuda_current_device: If True, call torch.cuda.set_device() to set the
            current PyTorch CUDA device to the one matching the local rank.
        overwrite: If True, overwrites already set variables. Else fails.


    Creates a dataset with the specified parameters.

    Args:
        dataset_str: A dataset string description (e.g. ImageNet:split=TRAIN).
        transform: A transform to apply to images.
        target_transform: A transform to apply to targets.

    Returns:
        The created dataset.


    Creates a data loader with the specified parameters.

    Args:
        dataset: A dataset (third party, LaViDa or WebDataset).
        batch_size: The size of batches to generate.
        num_workers: The number of workers to use.
        shuffle: Whether to shuffle samples.
        seed: The random seed to use.
        sampler_type: Which sampler to use: EPOCH, INFINITE, SHARDED_INFINITE, SHARDED_INFINITE_NEW, DISTRIBUTED or None.
        sampler_size: The number of images per epoch (when applicable) or -1 for the entire dataset.
        sampler_advance: How many samples to skip (when applicable).
        drop_last: Whether the last non-full batch of data should be dropped.
        persistent_workers: maintain the workers Dataset instances alive after a dataset has been consumed once.
        collate_fn: Function that performs batch collation


Import a torch package from a given path


Process a single file in the dataset


        begin method. This function is executed before the event loop.
        The antenna pattern provieder is initialized here.


Derives features used by ML models from the `progress` variable.

      Args:
        name: Base variable name from which features are derived.
        dims: List of the output feature dimensions, e.g. ("day", "lon").
        progress: Progress variable values.

      Returns:
        Dictionary of xarray variables derived from the `progress` values. It
        includes the original `progress` variable along with its sin and cos
        transformations.

      Raises:
        ValueError if the number of feature dimensions is not equal to the number
          of data dimensions.


Constructs year and day progress features.


Extracts inputs and targets for prediction, from a Dataset with a time dim.

      The input period is assumed to be contiguous (specified by a duration), but
      the targets can be a list of arbitrary lead times.

      Examples:

        # Use 18 hours of data as inputs, and two specific lead times as targets:
        # 3 days and 5 days after the final input.
        extract_inputs_targets(
            dataset,
            input_duration='18h',
            target_lead_times=('3d', '5d')
        )

        # Use 1 day of data as input, and all lead times between 6 hours and
        # 24 hours inclusive as targets. Demonstrates a friendlier supported string
        # syntax.
        extract_inputs_targets(
            dataset,
            input_duration='1 day',
            target_lead_times=slice('6 hours', '24 hours')
        )

        # Just use a single target lead time of 3 days:
        extract_inputs_targets(
            dataset,
            input_duration='24h',
            target_lead_times='3d'
        )

      Args:
        dataset: An xa.Dataset with a 'time' dimension whose coordinates are
          timedeltas. It's assumed that the time coordinates have a fixed offset /
          time resolution, and that the input_duration and target_lead_times are
          multiples of this.

      Returns:
        inputs:
        targets:
          Two datasets with the same shape as the input dataset except that a
          selection has been made from the time axis, and the origin of the
          time coordinate will be shifted to refer to lead times relative to the
          final input timestep. So for inputs the times will end at lead time 0,
          for targets the time coordinates will refer to the lead times requested.


Returns the minimum duration for the target lead times.


Loads a TESS *spoc* FITS file and returns TIME, PDCSAP_FLUX


    import matplotlib.pyplot as plt
    plt.scatter(time, flux, s=3, color='black')
    plt.scatter(time, clipped, s=3, color='orange')
    plt.show()


    ``flatten`` removes low frequency trends in time-series data.

    Parameters
    ----------
    time : array-like
        Time values
    flux : array-like
        Flux values for every time point
    window_length : float
        The length of the filter window in units of ``time`` (usually days), or in
        cadences (for cadence-based sliders ``savgol`` and ``medfilt``).
    method : string, default: ``biweight``
        Detrending method. Rime-windowed sliders: ``median``, ``biweight``, ``hodges``,
        ``tau``, ``welsch``, ``huber``, ``huber_psi``, ``andrewsinewave``, ``mean``,
        ``hampel``, ``ramsay``, ``trim_mean``, ``hampelfilt``, ``winsorize``. Cadence
        based slider: ``medfilt``. Splines: ``hspline``, ``rspline`, ``pspline``.
        Locally weighted scatterplot smoothing: ``lowess``. Savitzky-Golay filter:
        ``savgol``. Gaussian processes: ``gp``. Cosine Filtering with Autocorrelation
        Minimization: ``cofiam``.  Cosine fitting: ``cosine``, Friedman's Super-Smoother:
        ``supersmoother``. Gaussian regressions: ``ridge``, ``lasso``, ``elasticnet``.
    break_tolerance : float, default: window_length/2
        If there are large gaps in time (larger than ``window_length``/2), flatten will
        split the flux into several sub-lightcurves and apply the filter to each
        individually. ``break_tolerance`` must be in the same unit as ``time`` (usually
        days). To disable this feature, set ``break_tolerance`` to 0. If the method is
        ``supersmoother`` and no ``break_tolerance`` is provided, it will be taken as
        `1` in units of ``time``.
    edge_cutoff : float, default: None
        Trends near edges are less robust. Depending on the data, it may be beneficial
        to remove edges. The ``edge_cutoff`` defines the length (in units of time) to be
        cut off each edge. Default: Zero. Cut off is maximally ``window_length``/2, as
        this fills the window completely. Applicable only for time-windowed sliders.
    cval : float or int
        Tuning parameter for the robust estimators. See documentation for defaults.
        Larger values for make the estimate more efficient but less robust. For the
        super-smoother, cval determines the bass enhancement (smoothness) and can be
        `None` or in the range 0 < ``cval`` < 10. For the ``savgol``, ``cval``
        determines the (integer) polynomial order (default: 2).
    proportiontocut : float, default: 0.1
        Fraction to cut off (or filled) of both tails of the distribution using methods
        ``trim_mean`` (or ``winsorize``)
    kernel : str, default: `squared_exp`
        Choice of `squared_exp` (squared exponential), `matern`, `periodic`,
        `periodic_auto`.
    kernel_size : float, default: 1
        The length scale of the Gaussian Process kernel.
    kernel_period : float
        The periodicity of the Gaussian Process kernel (in units of ``time``). Must be
        provided for the kernel `periodic`. Can not be specified for the
        `periodic_auto`, for which it is determined automatically using a Lomb-Scargle
        periodogram pre-search.
    robust : bool, default: False
        If `True`, the fitting process will be run iteratively. In each iteration,
        2-sigma outliers from the fitted trend will be  clipped until convergence.
        Supported by the Gaussian Process kernels `squared_exp` and `matern`, as well as
        `cosine` fitting.

    return_trend : bool, default: False
        If `True`, the method will return a tuple of two elements
        (``flattened_flux``, ``trend_flux``) where ``trend_flux`` is the removed trend.

    Returns
    -------
    flatten_flux : array-like
        Flattened flux.
    trend_flux : array-like
        Trend in the flux. Only returned if ``return_trend`` is `True`.


Robust location estimators


Robust estimators of location and scale, with breakdown points of 50%.
    Also referred to as: Tau measure of location by Yohai and Zamar
    Source: Yohai and Zamar JASA, vol 83 (1988), pp 406-413 and
            Maronna and Zamar Technometrics, vol 44 (2002), pp. 307-317


One-step M-estimator of location using Huber's psi


Values beyond (cval * mad {=median absolute deviation}) are replaced with the
    median. Source: Ronald K. Pearson, Moncef Gabbouj (p. 147), "Nonlinear Digital
    Filtering with Python: An Introduction" recommend cval=3


Mean of array after trimming `proportiontocut` from both tails.


Mean of array after winsorizing `proportiontocut` from both tails.


Hodges-Lehmann-Sen robust location estimator


Planetary transit duration assuming a central transit on a circular orbit

    Parameters
    ----------
    R_s : float
        Stellar radius (in solar radii)
    M_s : float
        Stellar mass (in solar masses)
    P : float
        Planetary period (in days)
    small_planet : bool, default: `False`
        Small planet assumption if `True` (zero planetary radius). Otherwise, a
        planetary radius of 2 Jupiter radii is assumed (maximum reasonable planet size)

    Returns
    -------
    t14 : float
        Planetary transit duration (in days)


Iterator for a single time-series segment using time-series window sliders


With the import, we can't use numba. Thus this here is separate and slower


Regressions with scikit-learn


Clips data in the current segment


Sliding time-windowed outlier clipper.

    Parameters
    ----------
    time : array-like
        Time values
    flux : array-like
        Flux values for every time point
    window_length : float
        The length of the filter window in units of ``time`` (usually days)
    low : float or int
        Lower bound factor of clipping. Default is 3.
    high : float or int
        Lower bound factor of clipping. Default is 3.
    method : ``mad`` (median absolute deviation; default) or ``std`` (standard deviation)
        Outliers more than ``low`` and ``high`` times the ``mad`` (or the ``std``) from
        the middle point are clipped
    center : ``median`` (default) or ``mean``
        Method to determine the middle point

    Returns
    -------

    clipped : array-like
        Input array with clipped elements replaced by ``NaN`` values.


Takes numpy arrays with masks and non-float values.
    Returns unmasked cleaned arrays.


Robust B-Spline regression with scikit-learn


        Function that changes working directory to the run path (dir with the executable),
        replaces keywords in the CICE and ROMS in-files (as specified by the self._params
        attribute object) and finally calls the self._run() function to start the METROMS run.

        Args:
            runoption (int)    : Either run to run with MPI, OPENMP, SERIAL or DRY run
                                 (choose an option from the Contants module)
            debugoption (int)  : Wether or not METROMS is run in debug mode or not,
                                 (choose an option from the Contants module)
            architecture (int) : What architecture to run METROMS on
                                 (choose an option from the Contants module)


Contains collection of preprocessors common for all models.


        This function will replace the keywords in the given
        keyword.in-file.


        This function will replace the keywords in the given
        cice-keyword.in-file.


        Function that executes the ROMS model itself using MPI. Depending on the specified architecture,
        the ROMS executable is run using different binaries/commands.

        Args:
            ncpus (int)        : Total number of CPUs to run with
            infile (str)       : Filename of the roms.in parameter/config file
            debugoption (int)  : Wether or not METROMS is run in debug mode or not
                                 (choose an option from the Contants module)
            architecture (int) : What architecture to run METROMS on
                                 (choose an option from the Contants module)
        Returns:
            result (int) : 0 if all good, some other value if error with run or user choices


        Execute the ROMS model itself using OpenMP.

        Returns:
            result (int) : 0 if all good, some other value if error with run


        Execute the ROMS model itself in serial mode.

        Returns:
            result (int) : 0 if all good, 1 if some error with run


        Function that calls other functions for running METROMS depending on the specified
        architecture and and runoption.

        Args:
            runoption (int)    : Either run to run with MPI, OPENMP, SERIAL or DRY run
                                 (choose an option from the Contants module)
            debugoption (int)  : Wether or not METROMS is run in debug mode or not
                                 (choose an option from the Contants module)
            architecture (int) : What architecture to run METROMS on
                                 (choose an option from the Contants module)
        Returns:
            result (int) : 0 if all good, some other value if error with run or user choices


Function checking if the specified model start time is found in the
        CICE restart file located in cice restart pointer file, if not, it checks
        if model start time is found in any other restart files in the restart
        directory. All restart files are assumed to contain only one time and
        if the model start time is not found, an IOError is raised.


Function checking for match between <date> and time in cice rst file.


        Method that computes the time array for the boundary file using supplied
        input time data together with the self.time_units attribute.

        Args:
            nc_file (Dataset) : netcDF4.Dataset with a time variable
            time_name (str)   : Name of the time variable in file
        Returns:
            bry_time (ndarray) : 1D-array with time data for boundary file


Method that generates a BryVar object for each variable necessary to specify
        in the boundary file (most variable names are on the form <name>_{W,S,E,N}_bry).


        Method that creates a BryVar object for each boundary wall (W, S, E, N)
        based on the supplied input BryVar template input parameters.

        Args:
            name (str)        : Template for name of variable (replace bry_replace)
            dtype (type)      : Datatype of variable
            dims (tuple)      : Tuple of strings for dims of variable (replace bry_replace)
            attrs (dict)      : Dict of <str: str> (variable attributes) (replace bry_replace)
            bry_replace (str) : String to replace in <name>, <dims>, <attrs>
        Returns:
            bry_vars (list) : List of 4 BryVar objects (one for each boundary)


        Method writing global attributes to the boundary file. Some attributes are written
        by default, but can be supplied with more (ro to overwrite defaults) in a dictionary.

        Args:
            attr (dict) : Dict <name: val> (empty dict by default)


Method writing variables (metadata only) based on the boundary variables created
        earlier to the boundary file.


Method that writes variable attributes (from the BryVar objects in self.bry_vars)
        to variable attributes in the boundary file.


Method that is a wrapper for all the required steps to generate a boundary file. Calls all
        necessary methods and relies on <self> being an instance o a subclass ith an write_bry_data method.


        Method that takes the witten boundary file that contains a time dimension spanning
        only the relevant dates (for which data is computed) and expands the file into yearly
        files with time entries for the entire year. If the original written boundary data has
        a time dimension spanning across one or more year transistions, the data will be written
        to one file per year, e.g. if start date is 20191225 and end date is 20190105, you will
        get two files, BRY_2018.nc and BRY_2019.nc. (method must be called after self.write_bry_data())

        TODO: (IMPORTANT) Needs to be updated to account for year-to-year transition. In that case,
        have this method generate one 365/366 file for each year contained in the bry date range.
        Make it general so it works for both forecasting and hindcasts.

        Args:
            filename (str) : If supplied, copy written bry file to this filename and
                             expand into yearly files. Defaults to existing bry file.


Method that checks if the boundary file contains all necessary variables.


        Method that implements a hack to use the limited ice data from TOPAZ5 to generate
        data for the CICE boundary file. Relevant available variables from TOPAZ are only
        aice (ice conc.), hice (ice thick.) anf hsnow (snow thick.). From these, we generate
        the CICE variables aicen, vicen and vsnon at the boundaries. Ice conc. ice distributed
        amongst the ice categories (according the pre-defined thickness categories (self.cat_lims))
        and so is ice volume and snow volume.

        Additionally, atmosphere (from relevant model at corresponding times) and ocean (from
        clm file) data is used to help specify data for the variablesTsfc, Tinz and Sinz. Tsfc
        is assumed to be equal to the atmosphere surface temp. Tinz is linearly interpolated
        from the ocean surface temp to the air temperature (if there is less than 1cm of snow
        (at the particular grid cell in question)) and to ocean_temp+0.2*|ocean_temp - atm_temp|
        (if there is more than 1cm of snow). Sinz is linearly interpolated from the ocean surface
        salinity to 0.2 of the ocean surface salinity.

        Time-independent variables Ice_layers, TLON and TLAT are also as well as Time itself are
        also set here. All other time-dependent variables that are needed in the boundary file
        (see self.load_bry_vars()) are set to zero everywhere because of a lack of a better choice.


        Method that implements a hack to use the limited ice data from TOPAZ5 to generate
        data for the CICE boundary file. Relevant available variables from TOPAZ are only
        aice (ice conc.), hice (ice thick.) anf hsnow (snow thick.). From these, we generate
        the CICE variables aicen, vicen and vsnon at the boundaries by distributing TOPAZ4 data
        amongst the ice categories (according the pre-defined thickness categories (self.cat_lims))
        and so is ice volume and snow volume. Tsfcn is set to -5, alvln and vlvln to 1, Tinz is
        vertically interpolated from -5 at surface to -2 at bottom and Sinz is set to 5 everywhere

        Time-independent variables Ice_layers, TLON and TLAT as well as Time itself are also set
        here. All other time-dependent variables that are needed in the boundary file (see
        self.load_bry_vars()) are set to zero everywhere because of a lack of a better choice.


    Ported from temperature_mush() subroutine in ice_therm_mushy.f90 in CICE 5.1.2 source code.

    Args:
        enthalpy (array-like)  : Enthalpy of ice (scalar or array of any dimension)
        saltiness (array-like) : Salinity of ice (scalar or array of any dimension)


    Calculates the vector/tensor windows needed in the pure Cl
    calculation

    Notes
    -----
    We calculate both straight from the wlms of the scalar window because
    calculating the spin-2 from the spherical harmonic transform of a
    spin-1 window will result in larger errors.

    If you don't want zero pixels to be masked out input a full sky mask.


    Notes
    -----
    Input wlm should be wlm of the scalar map and not the s=0 wlm which
    differ by a negative sign.


Factor used when converting scalar alms to spin wlms,
    sign*sqrt((l+s)! / (l-s)!)


    Constructs the pure E/B alms from polarization maps apodized by the
    scalar, vector, and tensor window


    Calculates the pure E/B alms from the input Q/U maps and the apodized
    window.


This is apodizing the input QU maps by multiplying by the complex
    conjugate of the input window.


Similar to healpy.map2alm except calculates pure alms


Extension to H.map2alm which can take in 4 maps (I,Q,U,V).


Similar to healpy.map2alm except calculates hybrid alms.


Computes the pure power spectrum of a set of 3 (I,Q,U) or 4 (I,Q,U,V)
    Healpix maps, or the cross-spectrum between two sets of maps if *map2* is
    given.

    Parameters
    ----------
    map1 : float, array-like shape (3, Npix) or (4, Npix)
        A sequence of 3 (or 4) arrays representing I, Q, U, (V) maps
    swindow : float, array-like shape (Npix,)
        A Healpix map containing the window function
    map2 : float, array-like shape (3, Npix) or (4, Npix)
        A sequence of 3 (or 4) arrays representing I, Q, U, (V) maps
    mask : float, array-like shape (Npix,)
        The mask for the maps. If None, generated from the scalar window
    nspec : None or int, optional
        The number of spectra to return. If None, returns all, otherwise
        returns cls[:nspec]
    lmax : int, scalar, optional
        Maximum l of the power spectrum (default: 3*nside-1)
    mmax : int, scalar, optional
        Maximum m of the alm (default: lmax)
    iter : int, scalar, optional
        Number of iteration (default: 3)
    alm : bool, scalar, optional
        If True, returns both cl and alm, otherwise only cl is returned
    datapath : None or str, optional
        If given, the directory where to find the weights data.

    Returns
    -------
    res : array or sequence of arrays
        If *alm* is False, returns cl or a list of cl's (TT, EE, BB, etc.)
        Otherwise, returns a tuple (cl, alm), where cl is as above and
        alm is the spherical harmonic transform or a list of almT, almE, almB
        for polarized input


Computes the hybrid power spectrum of a set of 3 (I,Q,U) or 4 (I,Q,U,V)
    Healpix maps, or the cross-spectrum between two sets of maps if *map2* is
    given. Hybrid power spectra are ones where E is pseudo and B is pure.

    Parameters
    ----------
    map1 : float, array-like shape (3, Npix) or (4, Npix)
        A sequence of 3 (or 4) arrays representing I, Q, U, (V) maps
    swindow : float, array-like shape (Npix,)
        A Healpix map containing the window function
    map2 : float, array-like shape (3, Npix) or (4, Npix)
        A sequence of 3 (or 4) arrays representing I, Q, U, (V) maps
    mask : float, array-like shape (Npix,)
        The mask for the maps. If None, generated from the scalar window
    nspec : None or int, optional
        The number of spectra to return. If None, returns all, otherwise
        returns cls[:nspec]
    lmax : int, scalar, optional
        Maximum l of the power spectrum (default: 3*nside-1)
    mmax : int, scalar, optional
        Maximum m of the alm (default: lmax)
    iter : int, scalar, optional
        Number of iteration (default: 3)
    alm : bool, scalar, optional
        If True, returns both cl and alm, otherwise only cl is returned
    datapath : None or str, optional
        If given, the directory where to find the weights data.

    Returns
    -------
    res : array or sequence of arrays
        If *alm* is False, returns cl or a list of cl's (TT, EE, BB, etc.)
        Otherwise, returns a tuple (cl, alm), where cl is as above and
        alm is the spherical harmonic transform or a list of almT, almE, almB,
        (almV)


Computes the power spectrum of an Healpix map, or the cross-spectrum
    between two maps if *map2* is given.
    No removal of monopole or dipole is performed. This is an extension of
    healpy.anafast for a possibility of 4 Stokes parameters

    Parameters
    ----------
    map1 : float, array-like shape (Npix,) or (3, Npix) or (4, Npix)
      Either an array representing a map, or a sequence of
      3 (I,Q,U), or 4 (I,Q,U,V) arrays
    map2 : float, array-like shape (Npix,) or (3, Npix) or (4, Npix)
      Either an array representing a map, or a sequence of
      3 (I,Q,U), or 4 (I,Q,U,V) arrays
    nspec : None or int, optional
      The number of spectra to return. If None, returns all, otherwise
      returns cls[:nspec]
    lmax : int, scalar, optional
      Maximum l of the power spectrum (default: 3*nside-1)
    mmax : int, scalar, optional
      Maximum m of the alm (default: lmax)
    iter : int, scalar, optional
      Number of iteration (default: 3)
    alm : bool, scalar, optional
      If True, returns both cl and alm, otherwise only cl is returned
    pol : bool, optional
      If True, assumes input maps are TQU. Output will be TEB cl's and
      correlations (input must be 1 or 3 maps).
      If False, maps are assumed to be described by spin 0 spherical harmonics
      (input can be any number of maps)
      If there is only one input map, it has no effect. Default: True.
    datapath : None or str, optional
      If given, the directory where to find the weights data.

    Returns
    -------
    res : array or sequence of arrays
      If *alm* is False, returns cl or a list of cl's (TT, EE, BB, TE, EB, TB
      for polarized input map)
      Otherwise, returns a tuple (cl, alm), where cl is as above and
      alm is the spherical harmonic transform or a list of almT, almE, almB
      for polarized input


This is an implementation of the PolSpice algorithm in Python.

    Parameters
    ----------
    map1 : str or list
        Input filename to load or list of 3 Healpix maps (I, Q, U)

    map2 : str or list, optional
        Input filename to load or list of 3 Healpix maps (I, Q, U)

    window : str or numpy.ndarray, optional
        Filename or array giving the weighted window to use.

    mask : str or numpy.ndarray, optional
        Filename or array giving the mask to use

    window2 : str or numpy.ndarray, optional
        Filename or array giving the weighted window to use for
        the second map. If not input, the same window is used as
        for the first map.

    mask2 : str or numpy.ndarray, optional
        Filename or array giving the mask to use for the second
        map. If not input, the same window is used as for the
        first map.

    apodizesigma : float
        Scale factor of the correlation function apodization (degrees)

    apodizetype : int
        Type of apodization. 0 is a Gaussian window. 1 is a cosine window

    thetamax : float
        maximum value of theta used in the integrals to calculate the
        power spectra from the correlation functions.

    decouple : bool
        whether to compute the decoupled correlation functions

    returnxi : bool
        whether or not to return the correlation functions as an additional
        output of the function

    remove_monopole : bool
        whether or not to remove the monopole from the maps before analyzing
        them

    remove_dipole : bool
        whether or not to remove the dipole from the maps before analyzing
        them

    Returns
    -------
    cls : numpy.ndarray
        An array containing the power spectra of the maps. TT, EE, BB, TE, EB, TB

    xi : numpy.ndarray, optional
        An array containing the real space correlation functions.

    Notes
    -----
    If neither a window nor a mask are input, the functionality will
    be similar to anafast.


Correct the amplitude of the correlation function due to the
    presence of the mask/window

    Parameters
    ----------
    xi : array
        The correlation function of the input maps

    xi_mask : array
        The correlation function of the mask/window

    Returns
    -------
    xi_output : array
        The correlation function of the input maps
        corrected for the effect of the presence of
        the mask.


Calculates an apodization function to apply to
    the correlation functions

    Parameters
    ----------
    mu: array
        cos(thetas), where thetas is the angles at which the
        correlation function is calculated

    fwhm_degrees: float
        The full width half max in degrees for the apodization

    thetamax: float
        The maximum value of theta in the correlation function

    typef: int
        Either 0 or 1. Decides between two different types of
        apodization

    Returns
    ------
    output: array
        The apodized function


Calculate the functions of Legendre polyonmials that are needed
    for the correlation function transformations and used in several
    other functions.

    Parameters
    ----------
    x : array
        An array of values of cos(theta) over which to calculate the
        Legengre polynomial functions

    ell : int
        The ell value for the Legendre polynomials

    s : array, optional
        The values of sin(theta)

    Plm12 : array, optional
        The values of P_{ell-1}^2 at each value of x. If not input,
        these will be calculated.

    Returns
    -------
    Pl0 : array
        Legendre polynomial of order ell evaluated at x

    Pl2 : array
        Associated legendre polynomial (m=2) of order ell evaluated at x

    Gp : array
        Needed to calculate the small d-matrix used for the polarization
        correlation function calculations

    Gm : array
        Needed to calculate the small d-matrix used for the polarization
        correlation function calculations

    Nl : array
        sqrt((ell-2)!/(ell+2)!)


Calculate the C_l from the real-space correlation functions

    Parameters
    ----------
    xi : array
        The correlation function

    w : array
        The weights used in the Gauss-Legendre quadrature

    x : array
        The x values used in the Gauss-Legendre quadrature

    decouple : bool
        Whether or not the input correlation function is the decoupled
        polarization correlation function

    Note
    ----
    This does the integrals by Gauss-Legendre quatrature, therefore the
    real-space correlation functions must be input at the correct x values
    along with the correct weights to use.


Calculate an amplitude correction of the TE cross-correlation
    due to the apodization of the correlation function.

    Parameters
    ----------
    apfunction : array
        The function dependent on angle used to apodize the correlation
        function.

    x : array
        cos(angle), the Legendre-Gauss quadrature points

    w : array
        the weights needed for the Legendre-Gauss quadrature

    Returns
    -------
    fact : array
        TE correction function due to the fact that the correlation function
        is apodized


This replaces the polarized elements of xi with versions that only
    depend on E or B polarization, but not both.

    Note
    ----
    This corresponds to Section 5 in Chon, Challinor, Prunet, Hivon, and
    Szapudi (2004)


Calculate xi_+ corelation function.

    Note
    ----
    This is Eqs 17 and 9.


    Notes
    -----
    This calculates the two integrals in Eq. 90


    Solve for the true Cls from the pseudo/pure Cls

    Parameters
    ----------
    cls_in : array-like (ncls, nell)
        Input Cls, either 6 (T, E, B) or 10 (T, E, B, V)

    Mscal : array-like (nell, nell)
        Mode mixing matrix used for TT and VV

    Mpol: array-like (3*nell, 3*nell)
        Mode mixing matrix for EE, BB, EB mixing

    Mcross: array-like (2*nell, 2*nell)
        Mode mixing matrix for TE, TB or EV, BV mixing

    Notes
    -----
    This gets the true Cls by essentially solving the equation Ax=b by
    least squares procedure.


Solve for estimates of true Cls given input pseudo/pure Cls and
    mode-mixing matrices or apply the mode-mixing matrices to full sky
    Cls to get estimates of te pure/pseudo Cls.

    Parameters
    ----------
    cls_in : array-like (ncls, nell)
        Input Cls, either 6 (T, E, B) or 10 (T, E, B, V)

    Mscal : array-like (nell, nell)
        Mode mixing matrix used for TT and VV

    Mpol: array-like (3*nell, 3*nell)
        Mode mixing matrix for EE, BB, EB mixing

    Mcross: array-like (2*nell, 2*nell)
        Mode mixing matrix for TE, TB or EV, BV mixing

    inverse: boolean, optional (default: True)
        Whether to solve for the full sky Cls given in pseudo/pure Cls or
        calculate the pseudo/pure given the full sky Cls

    Notes
    -----
    For solving the inverse equation, we calculate the pseudo-inverse of
    the matrix (np.linalg.pinv).


    Calculate the E/B-type alms from the vector and tensor windows.
    Constructs the vector and tensor windows from the scalar wlms and
    then masks the resulting maps. This masking results in non-zero B-mode
    type wlms.


    Calculates the mode-mixing matrices given an input full sky window in
    Healpix format

    Parameters
    ----------
    comm : mpi4py.MPI.Intracomm
        MPI communicator (normally COMM_WORLD)

    window : array-like
        Healpix scalar window

    cl_type : 'pseudo', 'pure', or 'hybrid'
        The type of Cl that we want the mode-mixing matrices

    lmax : int, scalar, optional
        Maximum l of the power spectrum. Default: 3*nside-1

    scal : bool, optional
        Whether to calculate the mode-mixing matrix for temperature

    pol : bool, optional
        Whether to calculate the mode-mixing matrix for polarization

    cross : bool, optional
        Whether to calculate the mode-mixing matrix for temp-pol

    Returns
    -------
    Mout : list
        A list of the mode-mixing matrices in order: Mscal, Mpol, Mcross
        Only requested matrices are calculated and returned.


Calculate the temperature/V polarization mode-mixing matrix


Calculate the polarization mode mixing matrix


Calculate the temp-pol mode-mixing matrix.


Auxiliary function used in 3j recursion


Auxiliary function used in 6j recursion


Normalization function used in Wigner 3j calculation

    Notes
    -----
    The m3 input does not do anything. It is here so that the 3j and 6j
    normalization functions have the same signature and can both be called
    from a generic function that does the recursion.


The multiplicative sign of the 3j vector given the inputs


Normalization function used in Wigner 6j calculation


The multiplicative sign of the 6j vector given the inputs


Calculates a family of Wigner 3j symbols. Given the input values, it
    determines the range of valid j3 and calculates the Wigner 3j symbol for
    all these j3.

    Parameters
    ----------
    tj1 : int
        2*j1

    tj2 : int
        2*j2

    tm1 : int
        2*m1

    tm2 : int
        2*m2

    verbose : bool
        Whether there is verbose output

    Returns
    -------
    WigVal : int or array-like
        Returns an array of the Wigner 3j values or 0 if there is no valid j3

    Notes
    -----
    Calculates the Wigner 3j symbols using recursion and normalization.
    Input values are rounded to the nearest integer.


Calculates a family a Wigner 6j symbols

    tj1 : int
        2*j1

    tj2 : int
        2*j2

    tk1 : int
        2*k1

    tk2 : int
        2*k2

    tk3 : int
        2*k3

    Returns
    -------
    WigVal : int or array-like
        Returns an array of the Wigner 6j values or 0 is there is no valid j3

    Notes
    -----
    Calculates the Wigner 6j symbols using the same algorithm as the 3j
    symbols, excepts uses a different recursion relation. Input values
    are rounded to the nearest integer.