Calculation of a Wigner 3j symbol

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

    tj2 : int
        2*j2

    tj3 : int
        2*j3

    tm1 : int
        2*m1

    tm2 : int
        2*m2

    tm3 : int
        2*m3

    verbose : bool
        Whether there is verbose output

    Returns
    -------
    wigner3j : float
        Wigner3j(j1,j2,j3,m1,m2,m3)

    Notes
    -----
    The input values are double the requested j, m values because they are
    integers or half-integers. Code will round to the nearest integer.


Calculation of a Wigner 6j Symbol

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

    tj2 : int
        2*j2

    tj3 : int
        2*j3

    tk1 : int
        2*k1

    tk2 : int
        2*k2

    tk3 : int
        2*k3

    verbose : bool
        Whether there is verbose output

    Returns
    -------
    wigner6j : float
        Wigner6j(j1,j2,j3;k1,k2,k3)

    Notes
    -----
    We input 2*j because j can be an integer or half-integer, so this
    combination must be an integer. Code will round to the nearest
    integer.


Calculation of the Wigner 9j symbol.

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

    tj2 : int
        2*j2

    tj3 : int
        2*j3

    tj4 : int
        2*j4

    tj5 : int
        2*j5

    tj6 : int
        2*j6

    tj7 : int
        2*j7

    tj8 : int
        2*j8

    tj9 : int
        2*j9

    Returns
    -------
    wigner9j : float
        Wigner9j(j1,j2,j3,j4,j5,j6,j7,j8,j9)

    Notes
    -----
    Calculation is done following the reduction to 6j symbols seen at
    http://en.wikipedia.org/wiki/9-j_symbol and symmetry of rows/columns are
    done to move the x in the sum to the upper right part of the three terms


Calculate the Clebsch-Gordon coefficient <j1 jm1 j2 m2 | J M>

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

    tj2: int
        2*j2

    tm1: int
        2*m1

    tm2: int
        2*m2

    tj: int
        2*J

    tm: int
        2*M

    Returns
    -------
    cgcoeff: float
        Clebsch-Gordon coefficient


r'''Calculate the Gaunt coefficient which is an integral over three
    spherical harmonics

    Parameters
    ----------
    l1 : int
        ell value of the first spherical harmonic

    l2 : int
        ell value of the second spherical harmonic

    l3 : int
        ell value of the third spherical harmonic

    m1 : int
        m value of the first spherical harmonic

    m2 : int
        m value of the second spherical harmonic

    m3 : int
        m value of the third spherical harmonic

    Returns
    -------
    gaunt_coeff : float
        Gaunt coefficient \int Y_{l1, m1} Y_{l2, m2} Y_{l3, m3}


Calculate the Racah V coefficient

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

    tj2 : int
        2*j2

    tj : int
        2*j

    tm1 : int
        2*m1

    tm2 : int
        2*m2

    tm : int
        2*m

    Returns
    -------
    racahv : float
        RacahV(j1,j2,j2,m1,m2,m)


Calculate the Racah W coefficient: W(abcd;ef)

    Parameters
    ----------
    ta : int
        2*a

    tb : int
        2*b

    tc : int
        2*c

    td : int
        2*d

    te : int
        2*e

    tf : int
        2*f

    Returns
    -------
    racahw: float
        W(abcd;ef)


    In the case of where Yjmax and Yjmin are 0 we must use the three-term
    recursion for the calculation of all the symbols. Every other term in
    the set is 0. Yj=0 for everything


Calculation of Wigner 3j/6j symbols if Y(jmax) = 0 and Y(jmin) !=0


Calculation of Wigner 3j/6j symbols if Y(jmin) == 0 and Y(jmax) !=0


    Algorithm to calculate Wigner 3j/6j symbols when Y(jmin) !=0 and
    Y(jmax) != 0


projquiver is a wrapper around :func:`matplotlib.Axes.quiver` to take
        into account the spherical projection.

        You can call this function as::

           projquiver([X, Y], U, V, [C], **kw)
           projquiver([theta, phi], mag, ang, [C], **kw)

        Parameters
        ----------

        Notes
        -----
        Other keywords are transmitted to :func:`matplotlib.Axes.quiver`

        See Also
        --------
        projplot, projscatter, projtext


    Get the COBE pixel closest to the input longitude/latitude

    Parameters
    ----------
    lon : float, array-like shape (Npts,)
        either longitude (ra) or SkyCoord object (or array of these)
    lat : float, array-like shape (Npts,), optional
        latitude (dec) or array of latitudes. if present, lon is longitude
    coord : string, optional
        coordinate system of the input longitude, latitude. Default is celestial coordinates
        (icrs). Ignored if the input coordinates are SkyCoord objects
    res : string or int, optional
        COBE pixelization resolution. 'F' (FIRAS) = 'D' (DMR) = 6, 'B' (DIRBE) = 9.
        Default is 6. Resolution <= 15

    Returns
    -------
    pix : int, array-like shape (Npts,)
        output pixel number or array of pixel numbers in COBE pixelization


Returns pixel number given unit vector pointing to center
    of pixel resolution of the cude

    Parameters
    ----------
    c : astropy.coordinates.SkyCoord or array-like
        sky coordinate system in cartesian coordinates

    res : int, optional
        resolution of output pixel numbers, default=6

    Returns
    -------
    pixel : int or array-like (Npts, )
        pixel numbers


Converts position into nface number (0-5) and x, y in range 0-1

    Parameters
    ----------
    c : array-like (Npts, ) of astropy.coordinates.SkyCoord objects
        sky position in cartesian coordinates


This function takes an n by 3 element vector containing the
    face, column, and row number (the latter two within the face) of
    a pixel and converts it into an n-element pixel array for a given
    resolution

    Parameters
    ----------
    fij : array-like (Npts,)
        array of face, column, row numbers

    res : int
        COBE pixelization resolution for output pixel numbers

    Returns
    -------
    pixel : array-like (Npts,)
        COBE pixel numbers


Convert COBE quad-cube pixel number to sky coordinates (lon/lat, ra/dec, etc.)

    Parameters
    ----------
    pixel : int, array-like shape (Npts,)
        pixel number of array of pixel numbers
    res : int
        resolution of COBE pixelization. 'F'='D'=6, 'B'=9
    coord : string, optional
        coordinate system of output coordinates. Either in Healpy or Astropy coordinate string.
        Default is celestial coordinate system

    Returns
    -------
    c : astropy.coordinates.SkyCoord or array-like (Npts,)
        sky coordinates of the input pixel numbers


This function takes a n-element pixel array
    and generates an nx3 element array containing the
    corresponding face, column, and row number (the latter
    two within the face)

    Parameters
    ----------
    pixel : float or array-like (Npts,)
        pixel number array

    res : int
        COBE pixelization resolution

    Returns
    -------
    fij : array-like (Npts, 3)
        face, column, row for each pixel


Based on polynomial fit found using fcfit.for. Taken
    from forward_cube.for

    Parameters
    ----------
    x: float or array-like (Npts,)
        database coordinate

    y: float or array-like (Npts,)
        databse coordinate

    Returns
    -------
    xi : float or array-like (Npts,)
        tangent plane coordinate

    eta : float or array-like (Npts,)
        tangent plane coordinate


Converts face number an xi, eta into a unit vector

    Parameters
    ----------
    nface : int, array-like shape (Npts,)
        sky-cube face number
    xi : float, array-like shape (Npts,)
        tangent plane coordinate
    eta : float, array-like shape (Npts,)
        tangent plane coordinate

    Returns
    -------
    c : array-like (Npts, 3)
        unit vector


Convert unit vectors to longitude/latitude

    Parameters
    ----------
    vector : array-like (Npts, 3)
        unit vectors

    Returns
    -------
    lon_lat: array-like (Npts, 2)
        longitude/latitude array


Routine to set up the bit tables for use in the pixelization subroutines (extracted and
    generalized from existing routines)


Check for ix, iy being over the edge of a face
    Returns correct face, ix, iy in mface, jx, jy


Calculates the resolution given the input number of pixels


Packs an unfolded skycube into 2x3 format (no wasted space).

    Notes
    -----
    The dimensions compared to the IDL routine are reversed because IDL
    is column-major and numpy is row-major


Uncompresses the packed skycybe created by sixpack


This function creates a raster image (sky cube or face) given a pixel list
    and data array. The data array can be either a vector or 2d-array. In the
    latter case, the data for each raster image can be stored in either
    the columns or rows. The procedure also returns the x and y raster coordinates
    of each pixel.


Generates a raster image


This function creates a data array given either a list of
    pixels or a set of x and y raster coordinates and a raster
    image (sky cube or face). The skycube can be in either unfolded
    or six pack format. This routine is the "complement" to pix2xy.
    The program assumes a right oriented, ecliptic coordinate
    input raster.


Converts units vectors to projection (screen) coordinates


    Compute the non-linear power spectrum using a symbolic approximation
    to the linear power spectrum.

    Args:
        :k (torch.Tensor): k values [h/Mpc] with shape (n_k)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 9),
            the 9 parameters are :
                As: 10^9 times the amplitude of the primordial P(k)
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                mnu: Sum of neutrino masses [eV / c^2]
                w0: Time independent part of the dark energy EoS
                wa: Time dependent part of the dark energy EoS
                a: The scale factor to evaluate P(k) at

    Returns:
        :pk_nl (torch.Tensor): computed nonlinear power spectrum for each k and theta in the batch


    The offset of the emulated non-linear power spectrum from euclidemulator2

    Args:
        :k_batch (torch.Tensor): k values [h/Mpc] with shape (n_k,1)

    Returns:
        :offset (torch.Tensor): The offset of the emulated non-linear log_10 P(k)


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit k / h Mpc-1
        :z (Union[float, np.ndarray]): Redshift
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): Amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): Energy per SFR of the galactic winds
        :A_AGN1 (Union[float, np.ndarray]): Energy per black-hole accretion rate of the kinetic black-hole feedback
        :A_SN2 (Union[float, np.ndarray]): Wind speed of the galactic winds
        :A_AGN2 (Union[float, np.ndarray]): Energy per unit black-hole accretion rate of the thermal model of the black-hole feedback

    Returns:
        :S (Union[float, np.ndarray)): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit h Mpc-1
        :z (Union[float, np.ndarray]): Redshift
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): Amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): Energy per unit SFR of the galactic winds
        :A_AGN1 (Union[float, np.ndarray]): Energy per unit blach-hole accretion rate, not used in the model, only here for consistency with other simulators.
        :A_SN2 (Union[float, np.ndarray]): Wind speed of the galactic winds
        :A_AGN2 (Union[float, np.ndarray]): Ejection speed/burstiness of the kinetic mode of the black-hole feedback

    Returns:
        :S (Union[float, np.ndarray)): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit h Mpc-1
        :z (Union[float, np.ndarray]): Redshift
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): Amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): Mass loading of the galactic winds
        :A_AGN1 (Union[float, np.ndarray]): Momentum flux of the QSO & jet-mode black-hole feedback
        :A_SN2 (Union[float, np.ndarray]): Wind speed of the galactic winds
        :A_AGN2 (Union[float, np.ndarray]): Jet speed of the jet-mode black-hole feedback

    Returns:
        :S (Union[float, np.ndarray)): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit h Mpc-1
        :z (Union[float, np.ndarray]): Redshift
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): Thermal energy injected in each SNII event
        :A_AGN1 (Union[float, np.ndarray]): Scaling of the black hole Bondi accretion rate
        :A_SN2 (Union[float, np.ndarray]): Metallicity dependence of the stellar feedback fraction per unit stellar mass
        :A_AGN2 (Union[float, np.ndarray]): Temperature jump of gas particles in AGN feedback events

    Returns:
        :S (Union[float, np.ndarray)): Baryonic suppression of the matter power spectrum


    Compute nonlinear P(k) for the cosmology of interest

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :z (Union[float, np.ndarray]): Redshift.
        :sigma8 (Union[float, np.ndarray]): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (Union[float, np.ndarray]): The z=0 total matter density parameter, Omega_m
        :Ob (Union[float, np.ndarray]): The z=0 baryonic density parameter, Omega_b
        :logMc (Union[float, np.ndarray]): (Log of) mass fraction of hot gas in haloes
        :logeta (Union[float, np.ndarray]): (Log of) extent of ejected gas
        :logbeta (Union[float, np.ndarray]): (Log of) mass fraction of hot gas in haloes
        :logM1 (Union[float, np.ndarray]): (Log of) characteristic halo mass scale for central galaxy
        :logMinn (Union[float, np.ndarray]): (Log of) property describing density profile of hot gas in haloes
        :logthetainn (Union[float, np.ndarray]): (Log of) property describing density profile of hot gas in haloes

    Returns:
        :S (np.ndarray): Ratio of baryonic to non-baryonic power spectrum at ks provided.


    Calculate the typical error of the prediction of `Astrid`.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :z (Union[float, np.ndarray]): Redshift.

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S of Astrid.


    Calculate the typical error of the prediction of `IllustrisTNG`.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :z (Union[float, np.ndarray]): Redshift.

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S of IllustrisTNG.


    Calculate the typical error of the prediction of `SIMBA`.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :z (Union[float, np.ndarray]): Redshift.

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S of SIMBA.


    Calculate the typical error of the prediction of `Swift_EAGLE`.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :z (Union[float, np.ndarray]): Redshift.

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S of Swift-EAGLE.


    Correction to the growth factor for a batch of parameters.

    Args:
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 9),
            the 9 parameters are :
                As: 10^9 times the amplitude of the primordial P(k)
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                mnu: Sum of neutrino masses [eV / c^2]
                w0: Time independent part of the dark energy EoS
                wa: Time dependent part of the dark energy EoS
                a: The scale factor to evaluate P(k) at

    Returns:
        :result (torch.Tensor): The correction to the growth factor for the given parameters


    Corrections to the present-day linear power spectrum for a batch of parameters.

    Args:
        :k_batch (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 9),
            the 9 parameters are :
                As: 10^9 times the amplitude of the primordial P(k)
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                mnu: Sum of neutrino masses [eV / c^2]
                w0: Time independent part of the dark energy EoS
                wa: Time dependent part of the dark energy EoS
                a: The scale factor to evaluate P(k) at

    Returns:
        :result (torch.Tensor): Corrections to the present-day linear power spectrum


    Compute the emulated logarithm of the ratio between the true linear power spectrum
    and the Eisenstein & Hu 1998 fit for LCDM modified from implementation in linear.py (Bartlett et al. 2023).

    Args:
        :k_batch (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 9),
            the 9 parameters are :
                As: 10^9 times the amplitude of the primordial P(k)
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                mnu: Sum of neutrino masses [eV / c^2]
                w0: Time independent part of the dark energy EoS
                wa: Time dependent part of the dark energy EoS
                a: The scale factor to evaluate P(k) at

    Returns:
        :logF (torch.Tensor): The emulated logarithm of the ratio between the true linear power spectrum


    Fiducial power spectrum given in Bartlett et al. 2023 for a batch of parameters.

    Args:
        :k (torch.Tensor): k values [h/Mpc] with shape (n_k)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 9),
            the 9 parameters are :
                As: 10^9 times the amplitude of the primordial P(k)
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                mnu: Sum of neutrino masses [eV / c^2]
                w0: Time independent part of the dark energy EoS
                wa: Time dependent part of the dark energy EoS
                a: The scale factor to evaluate P(k) at

    Returns:
        :Pk (torch.Tensor): computed fiducial power spectrum for each k and theta in the batch


    Emulate the non-linear scale of halofit using the approximation of
    Bartlett et al. 2024

    Args:
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 6),
            the 6 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                a: The scale factor

    Returns:
        :ksigma (torch.Tensor): The non-linear scale (h/Mpc)


    Emulate the effective slope of the linear power spectrum at the non-linear scale
    using the approximation of Bartlett et al. 2024

    Args:
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 6),
            the 6 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                a: The scale factor

    Returns:
        :neff (torch.Tensor): The effective slope of the linear power spectrum at the
            non-linear scale


    Emulate the curvature of the linear power spectrum at the non-linear scale
    using the approximation of Bartlett et al. 2024

    Args:
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 6),
            the 6 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                a: The scale factor

    Returns:
        :neff (torch.Tensor): The curvature of the linear power spectrum at the
            non-linear scale


    Compute the correction term to halofit as described in Bartlett et al. 2024

    Args:
        :k_batch (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 6),
            the 6 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                a: The scale factor to evaluate P(k) at
        :ksigma (torch.Tensor, default=None): The non-linear scale (h/Mpc). If this is not
            provided, then we used the fitting formula of Bartlett et al. 2024
        :neff (torch.Tensor, default=None): The effective slope of the linear power spectrum
            at the non-linear scale. If this is not provided, then we used the fitting
            formula of Bartlett et al. 2024
        :C (torch.Tensor, default=None): The curvature of the linear power spectrum
            at the non-linear scale. If this is not provided, then we used the fitting
            formula of Bartlett et al. 2024

    Returns:
        :A (torch.Tensor): The correction factor for halofit

    Returns:


    Compute the non-linear power spectrum using halofit and a symbolic approximation
    to the linear power spectrum. The model sr-halofit of Bartlett et al. 2024 is
    obtained with the default parameters.

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 6),
            the 6 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                a: The scale factor to evaluate P(k) at
        :emulator (str, default='fiducial'): Which linear P(k)emulator to use from
            Bartlett et al. 2023. 'fiducial' uses the fiducial one, and 'max_precision'
            uses the most precise one.
        :extrapolate (bool, default=False): If using the Bartlett et al. 2023 linear
            fit, then if true we extrapolate outside range tested in paper. Otherwise,
            we use the E&H with baryons fit for this regime
        :which_params (str, default='Bartlett'): Which halofit parameters to use.
            Currently available: 'Takahashi' uses those from Takahashi et al. 2012, or
            'Bartlett' uses those from Bartlett et al. 2024.
        :add_correction (bool, default=True): Whether to add the correction 1+A from
            Bartlett et al. 2024 or not.

    Returns:
        :pk_lin (torch.Tensor): The emulated non-linear P(k) [(Mpc/h)^3]


    Compute the growth factor D(a) normalised to a=1.

    Args:
        :Om (torch.Tensor): The z=0 total matter density parameter
        :acosmo (torch.Tensor): The scale factor.

    Returns:
        :Da (torch.Tensor): The normalised growth factor D(a).


    Compute the approximate growth factor D(a), normalised to a=1, based on the Eisenstein & Hu (1998).
    This function provides a much faster computation of the growth factor, especially when using GPUs.

    Args:
        :Om (torch.Tensor): The z=0 total matter density parameter
        :a (torch.Tensor): The scale factor.

    Returns:
        :Da (torch.Tensor): The normalised growth factor D(a).


    Compute the Eisentein & Hu 1998 zero-baryon approximation to P(k) at z=0

    Args:
        :kbatch (torch.Tensor): k values [h/Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 5),
            the 5 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
        :integral_norm (bool, default=True): Whether to compute the normalisation of the
            power spectrum using an integral over k

    Returns:
        :pk_eh (torch.Tensor): The Eisenstein & Hu 1998 zero-baryon P(k) [(Mpc/h)^3] for each
            k and theta in the batch


    Compute the emulated logarithm of the ratio between the true linear power spectrum
    and the Eisenstein & Hu 1998 fit for LCDM as given in Bartlett et al. 2023.

    Args:
        :k_batch (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 5),
            the 5 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
        :extrapolate (bool, default=False): If True, then extrapolates the Bartlett
            et al. 2023 fit outside range tested in paper. Otherwise, uses E&H with
            baryons for k < kmin and k > kmax
        :kmin (float, default=9.e-3): Minimum k value to use Bartlett et al. formula
            if extrapolate=False
        :kmax (float, default=9): Maximum k value to use Bartlett et al. formula
            if extrapolate=False

    Returns:
        :logF (torch.Tensor): The emulated logarithm of the ratio between the true linear power spectrum and the
            Eisenstein & Hu 1998 zero-baryon fit


    Compute the emulated logarithm of the ratio between the true linear
    power spectrum and the Eisenstein & Hu 1998 fit. Here we use the mosy precide
    exprssion given in Bartlett et al. 2023.

    Args:
        :k_batch (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 5),
            the 5 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
        :extrapolate (bool, default=False): If True, then extrapolates the Bartlett
            et al. 2023 fit outside range tested in paper. Otherwise, uses E&H with
            baryons for k < kmin and k > kmax
        :kmin (float, default=9.e-3): Minimum k value to use Bartlett et al. formula
            if extrapolate=False
        :kmax (float, default=9): Maximum k value to use Bartlett et al. formula
            if extrapolate=False

    Returns:
        :logF (torch.Tensor): The logarithm of the ratio between the linear P(k) and the
            Eisenstein & Hu 1998 zero-baryon fit


    Compute the emulated linear matter power spectrum using the fits of
    Eisenstein & Hu 1998 and Bartlett et al. 2023.

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 6),
            the 6 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
                a: The scale factor to evaluate P(k) at
        :emulator (str, default='fiducial'): Which emulator to use from Bartlett et al.
            2023. 'fiducial' uses the fiducial one, and 'max_precision' uses the
            most precise one.
        :extrapolate (bool, default=False): If True, then extrapolates the Bartlett
            et al. 2023 fit outside range tested in paper. Otherwise, uses E&H with
            baryons for k < kmin and k > kmax
        :kmin (float, default=9.e-3): Minimum k value to use Bartlett et al. formula
            if extrapolate=False
        :kmax (float, default=9): Maximum k value to use Bartlett et al. formula
            if extrapolate=False
        :use_approx_D (bool, default=True): Whether to use the approximate growth factor

    Returns:
        :pk_lin (torch.Tensor): The emulated linear P(k) [(Mpc/h)^3]


    Compute the emulated conversion sigma8 -> As as given in Bartlett et al. 2023

    Args:
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 5),
            the 5 parameters are :
                sigma8: Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
        :old_equation (bool, default=False): Whether to use the version of the sigma8
            emulator which appeared in v1 of the paper on arXiv (True) or the final
            published version (and v2 on arXiv).

    Returns:
        :As (ftorch.Tensor): 10^9 times the amplitude of the primordial P(k)


    Compute the emulated conversion As -> sigma8 as given in Bartlett et al. 2023

    Args:
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 5),
            the 5 parameters are :
                As: 10^9 times the amplitude of the primordial P(k)
                Om: The z=0 total matter density parameter
                Ob: The z=0 baryonic density parameter
                h: Hubble constant, H0, divided by 100 km/s/Mpc
                ns: Spectral tilt of primordial power spectrum
        :old_equation (bool, default=False): Whether to use the version of the sigma8
            emulator which appeared in v1 of the paper on arXiv (True) or the final
            published version (and v2 on arXiv).

    Returns:
        :sigma8 (torch.Tensor): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h


    Compute the Eisentein & Hu 1998 zero-baryon approximation to P(k) at z=0

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :use_colossus (bool, default=False): Whether to use the external package colossus
            to compute this term
        :integral_norm (bool, default=True): Whether to compute the normalisation of the
            power spectrum using an integral over k

    Returns:
        :pk_eh (np.ndarray): The Eisenstein & Hu 1998 zero-baryon P(k) [(Mpc/h)^3]


    Compute the Eisentein & Hu 1998 baryon approximation to P(k) at z=0

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum

    Returns:
        :pk_eh (np.ndarray): The Eisenstein & Hu 1998 baryon P(k) [(Mpc/h)^3]


    Compute the emulated logarithm of the ratio between the true linear
    power spectrum and the Eisenstein & Hu 1998 fit. Here we use the fiducial exprssion
    given in Bartlett et al. 2023.

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :extrapolate (bool, default=False): If True, then extrapolates the Bartlett
            et al. 2023 fit outside range tested in paper. Otherwise, uses E&H with
            baryons for k < kmin and k > kmax. Due to problems with the Colossus
            Eisenstein & Hu fit on large scales (typically a 2 percent offset), it is
            strongly recommended to use extrapolate=False.
        :kmin (float, default=9.e-3): Minimum k value to use Bartlett et al. formula
            if extrapolate=False
        :kmax (float, default=9): Maximum k value to use Bartlett et al. formula
            if extrapolate=False

    Returns:
        :logF (np.ndarray): The logarithm of the ratio between the linear P(k) and the
            Eisenstein & Hu 1998 zero-baryon fit


    Compute the emulated logarithm of the ratio between the true linear
    power spectrum and the Eisenstein & Hu 1998 fit. Here we use the mosy precide
    exprssion given in Bartlett et al. 2023.

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :extrapolate (bool, default=False): If True, then extrapolates the Bartlett
            et al. 2023 fit outside range tested in paper. Otherwise, uses E&H with
            baryons for k < kmin and k > kmax. Due to problems with the Colossus
            Eisenstein & Hu fit on large scales (typically a 2 percent offset), it is
            strongly recommended to use extrapolate=False.
        :kmin (float, default=9.e-3): Minimum k value to use Bartlett et al. formula
            if extrapolate=False
        :kmax (float, default=9): Maximum k value to use Bartlett et al. formula
            if extrapolate=False

    Returns:
        :logF (np.ndarray): The logarithm of the ratio between the linear P(k) and the
            Eisenstein & Hu 1998 zero-baryon fit


    Compute the emulated linear matter power spectrum using the fits of
    Eisenstein & Hu 1998 and Bartlett et al. 2023.

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :a (float, default=1): The scale factor to evaluate P(k) at
        :emulator (str, default='fiducial'): Which emulator to use from Bartlett et al.
            2023. 'fiducial' uses the fiducial one, and 'max_precision' uses the
            most precise one.
        :extrapolate (bool, default=False): If using the Bartlett et al. 2023 fit, then
            if true we extrapolate outside range tested in paper. Otherwise, uses E&H with
            baryons for k < kmin and k > kmax. Due to problems with the Colossus
            Eisenstein & Hu fit on large scales (typically a 2 percent offset), it is
            strongly recommended to use extrapolate=False.
        :kmin (float, default=9.e-3): Minimum k value to use Bartlett et al. formula
            if extrapolate=False
        :kmax (float, default=9): Maximum k value to use Bartlett et al. formula
            if extrapolate=False

    Returns:
        :pk_lin (np.ndarray): The emulated linear P(k) [(Mpc/h)^3]


    Compute the emulated conversion sigma8 -> As as given in Bartlett et al. 2023

    Args:
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :old_equation (bool, default=False): Whether to use the version of the sigma8
            emulator which appeared in v1 of the paper on arXiv (True) or the final
            published version (and v2 on arXiv).

    Returns:
        :As (float): 10^9 times the amplitude of the primordial P(k)


    Compute the emulated conversion As -> sigma8 as given in Bartlett et al. 2023

    Args:
        :As (float): 10^9 times the amplitude of the primordial P(k)
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :old_equation (bool, default=False): Whether to use the version of the sigma8
            emulator which appeared in v1 of the paper on arXiv (True) or the final
            published version (and v2 on arXiv).

    Returns:
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h


    Compute the non-linear power spectrum using a symbolic approximation
    to the linear power spectrum.

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :As (float): 10^9 times the amplitude of the primordial P(k)
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS
        :a (float): Scale factor to consider

    Returns:
        :pk_nl (np.ndarray): The emulated non-linear P(k) [(Mpc/h)^3]


    the offset of the emulated non-linear power spectrum from euclidemulator2

    Parameters:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]

    Returns:
        :offset (np.ndarray): The offset of the emulated non-linear power spectrum


    Emulate the non-linear scale of halofit using the approximation of
    Bartlett et al. 2024

    Args:
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :a (float): The scale factor

    Returns:
        :ksigma (float): The non-linear scale (h/Mpc)


    Emulate the effective slope of the linear power spectrum at the non-linear scale
    using the approximation of Bartlett et al. 2024

    Args:
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :a (float): The scale factor

    Returns:
        :neff (float): The effective slope of the linear power spectrum at the
            non-linear scale


    Emulate the curvature of the linear power spectrum at the non-linear scale
    using the approximation of Bartlett et al. 2024

    Args:
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :a (float): The scale factor

    Returns:
        :neff (float): The curvature of the linear power spectrum at the
            non-linear scale


    Compute the correction term to halofit as described in Bartlett et al. 2024

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :a (float): The scale factor
        :ksigma (float, default=None): The non-linear scale (h/Mpc). If this is not
            provided, then we used the fitting formula of Bartlett et al. 2024
        :neff (float, default=None): The effective slope of the linear power spectrum
            at the non-linear scale. If this is not provided, then we used the fitting
            formula of Bartlett et al. 2024
        :C (float, default=None): The curvature of the linear power spectrum
            at the non-linear scale. If this is not provided, then we used the fitting
            formula of Bartlett et al. 2024

    Returns:
        :A (np.ndarray): The correction factor for halofit

    Returns:


    Compute the non-linear power spectrum using halofit and a symbolic approximation
    to the linear power spectrum. The model sr-halofit of Bartlett et al. 2024 is
    obtained with the default parameters.

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :a (float): The scale factor to evaluate P(k) at
        :emulator (str, default='fiducial'): Which linear P(k)emulator to use from
            Bartlett et al. 2023. 'fiducial' uses the fiducial one, and 'max_precision'
            uses the most precise one.
        :extrapolate (bool, default=False): If using the Bartlett et al. 2023 linear
            fit, then if true we extrapolate outside range tested in paper. Otherwise,
            we use the E&H with baryons fit for this regime. Due to problems with the Colossus
            Eisenstein & Hu fit on large scales (typically a 2 percent offset), it is
            strongly recommended to use extrapolate=False.
        :which_params (str, default='Bartlett'): Which halofit parameters to use.
            Currently available: 'Takahashi' uses those from Takahashi et al. 2012, or
            'Bartlett' uses those from Bartlett et al. 2024.
        :add_correction (bool, default=True): Whether to add the correction 1+A from
            Bartlett et al. 2024 or not.

    Returns:
        :pk_lin (np.ndarray): The emulated non-linear P(k) [(Mpc/h)^3]


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit k / h Mpc-1
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): Amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): Energy per SFR of the galactic winds
        :A_AGN1 (Union[float, np.ndarray]): Energy per black-hole accretion rate of the kinetic black-hole feedback
        :A_SN2 (Union[float, np.ndarray]): Wind speed of the galactic winds
        :A_AGN2 (Union[float, np.ndarray]): Energy per unit black-hole accretion rate of the thermal model of the black-hole feedback
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :S (Union[float, np.ndarray)): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit h Mpc-1
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): Amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): Energy per unit SFR of the galactic winds
        :A_AGN1 (Union[float, np.ndarray]): Energy per unit blach-hole accretion rate, not used in the model, only here for consistency with other simulators.
        :A_SN2 (Union[float, np.ndarray]): Wind speed of the galactic winds
        :A_AGN2 (Union[float, np.ndarray]): Ejection speed/burstiness of the kinetic mode of the black-hole feedback
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :S (Union[float, np.ndarray)): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit h Mpc-1
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): Amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): Mass loading of the galactic winds
        :A_AGN1 (Union[float, np.ndarray]): Momentum flux of the QSO & jet-mode black-hole feedback
        :A_SN2 (Union[float, np.ndarray]): Wind speed of the galactic winds
        :A_AGN2 (Union[float, np.ndarray]): Jet speed of the jet-mode black-hole feedback
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :S (Union[float, np.ndarray)): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit h Mpc-1
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): Thermal energy injected in each SNII event
        :A_AGN1 (Union[float, np.ndarray]): Scaling of the black hole Bondi accretion rate
        :A_SN2 (Union[float, np.ndarray]): Metallicity dependence of the stellar feedback fraction per unit stellar mass
        :A_AGN2 (Union[float, np.ndarray]): Temperature jump of gas particles in AGN feedback events
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :S (Union[float, np.ndarray)): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum for a given hydro model.
    This is given by the ratio of the baryonic to non-baryonic power spectrum at ks provided.

    Args:
        :k (Union[float, np.ndarray]): Wavenumber in unit h Mpc-1
        :Omega_m (Union[float, np.ndarray]): Density of matter.
        :sigma_8 (Union[float, np.ndarray]): Amplitude of matter fluctuations
        :A_SN1 (Union[float, np.ndarray]): First supernova feedback parameter (see specific hydro model for details)
        :A_SN2 (Union[float, np.ndarray]): Second supernova feedback parameter (see specific hydro model for details)
        :A_AGN1 (Union[float, np.ndarray]): First AGN feedback parameter (see specific hydro model for details)
        :A_AGN2 (Union[float, np.ndarray]): Second AGN feedback parameter (see specific hydro model for details)
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)
        :hydro_model (str): Name of the hydro model to use ('Astrid', 'IllustrisTNG', 'SIMBA', 'Swift-EAGLE')

    Returns:
        :S (Union[float, np.ndarray]): Baryonic suppression of the matter power spectrum


    Compute nonlinear P(k) for the cosmology of interest

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :Om (Union[float, np.ndarray]): The z=0 total matter density parameter, Omega_m
        :Ob (Union[float, np.ndarray]): The z=0 baryonic density parameter, Omega_b
        :sigma8 (Union[float, np.ndarray]): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
        :logMc (Union[float, np.ndarray]): (Log10 of) mass fraction of hot gas in haloes
        :logeta (Union[float, np.ndarray]): (Log10 of) extent of ejected gas
        :logbeta (Union[float, np.ndarray]): (Log10 of) mass fraction of hot gas in haloes
        :logM1 (Union[float, np.ndarray]): (Log10 of) characteristic halo mass scale for central galaxy
        :logMinn (Union[float, np.ndarray]): (Log10 of) property describing density profile of hot gas in haloes
        :logthetainn (Union[float, np.ndarray]): (Log10 of) property describing density profile of hot gas in haloes
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :S (np.ndarray): Ratio of baryonic to non-baryonic power spectrum at ks provided.


    Calculate the typical error of the prediction of `Astrid`.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S of Astrid.


    Calculate the typical error of the prediction of `IllustrisTNG`.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S of IllustrisTNG.


    Calculate the typical error of the prediction of `SIMBA`.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S of SIMBA.


    Calculate the typical error of the prediction of `Swift_EAGLE`.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S of Swift-EAGLE.


    Calculate the typical error of the prediction for a given hydro model.

    Args:
        :k (Union[float, np.ndarray]): k values to evaluate P(k) at [h / Mpc]
        :a (Union[float, np.ndarray]): Scale factor, a = 1/(1+z)
        :hydro_model (str): Name of the hydro model to use ('Astrid', 'IllustrisTNG', 'SIMBA', 'Swift-EAGLE')

    Returns:
        :epsilon (Union[float, np.ndarray]): Covariance of the prediction and the actual S.


    Compute the emulated conversion As -> sigma8

    Args:
        :As (float): 10^9 times the amplitude of the primordial P(k)
        :Om (float): The z=0 total matter density parameter, Om
        :Ob (float): The z=0 baryonic density parameter, Ob
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS

    Returns:
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h


    Compute the emulated conversion sigma8 -> As

    Args:
        :sigma8 (float): The z=0 rms mass fluctuation in spheres of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Om
        :Ob (float): The z=0 baryonic density parameter, Ob
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS

    Returns:
        :As (float): 10^9 times the amplitude of the primordial P(k)


    Compute the emulated conversion As -> sigma8, using the most accurate expression

    Args:
        :As (float): 10^9 times the amplitude of the primordial P(k)
        :Om (float): The z=0 total matter density parameter, Om
        :Ob (float): The z=0 baryonic density parameter, Ob
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS

    Returns:
        :sigma8 (float): Root-mean-square density fluctuation when the linearly
            evolved field is smoothed with a top-hat filter of radius 8 Mpc/h


    Compute the emulated conversion sigma8 -> As, using the most accurate expression

    Args:
        :sigma8 (float): The z=0 rms mass fluctuation in spheres of radius 8 Mpc/h
        :Om (float): The z=0 total matter density parameter, Om
        :Ob (float): The z=0 baryonic density parameter, Ob
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS

    Returns:
        :As (float): 10^9 times the amplitude of the primordial P(k)


    Correction to the growth factor

    Args:
        :As (float): 10^9 times the amplitude of the primordial P(k)
        :Om (float): The z=0 total matter density parameter, Om
        :Ob (float): The z=0 baryonic density parameter, Ob
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS
        :a (float): The scale factor to evaluate P(k) at

    Returns:
        :result (float): correction to the growth factor


    Corrections to the present-day linear power spectrum

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :As (float): 10^9 times the amplitude of the primordial P(k)
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS
        :a (float): Scale factor to consider

    Returns:
        :result (np.ndarray): Corrections to the present-day linear power spectrum


    Compute the emulated logarithm of the ratio between the true linear power spectrum
    and the Eisenstein & Hu 1998 fit for LCDM as given in Bartlett et al. 2023.

    This calls the logF_fiducial function from symbolic_pofk.linear but extrapolates
    to all k considered

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :As (float): 10^9 times the amplitude of the primordial P(k)
        :Om (float): The z=0 total matter density parameter, Omega_m
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS
        :a (float): Scale factor to consider

    Returns:
        :logF (np.ndarray): The emulated logarithm of the ratio between the true linear power spectrum


    Compute the emulated linear matter power spectrum by combining the Eisenstein & Hu model, an approximation for the growth factor D,
    the fit from Bartlett et al. (2023), and corrections to both the present-day linear power spectrum and the growth factor.

    Args:
        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]
        :As (float): 10^9 times the amplitude of the primordial P(k)
        :Om (float): The z=0 total matter density parameter, Om
        :Ob (float): The z=0 baryonic density parameter, Omega_b
        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc
        :ns (float): Spectral tilt of primordial power spectrum
        :mnu (float): Sum of neutrino masses [eV / c^2]
        :w0 (float): Time independent part of the dark energy EoS
        :wa (float): Time dependent part of the dark energy EoS
        :a (float, default=1): The scale factor to evaluate P(k) at

    Returns:
        :pk_lin (np.ndarray): The emulated linear P(k) [(Mpc/h)^3]


    Modify a specific index in a tuple and return a new tuple.

    Args:
        :t (tuple): The input tuple.
        :i (int): The index to modify.
        :value (any): The new value to insert at index `i`.

    Returns:
        :ell (tuple): A new tuple with the specified index modified.


    Core implementation of Simpson's rule along a given axis.
    If x is None, assumes equally spaced data with spacing dx.

    Parameters
    ----------
    Args:
        :y (torch.Tensor): Function values to be integrated.
        :start (int): Starting index along the axis.
        :stop (int): Stopping index along the axis.
        :x (torch.Tensor, optional): Points at which y is sampled.
        :dx (float): Spacing between points if x is not provided.
        :axis (int): Axis along which to integrate.

    Returns:
        :result (torch.Tensor): Estimated integral using Simpson's rule.


    Integrate `y` along the given axis using Simpson's rule.

    Args:
        :y (torch.Tensor): Tensor to be integrated.
        :x (torch.Tensor, optional): Points at which `y` is sampled. If not given, uniform spacing `dx` is assumed.
        :dx (float, optional): Spacing between points if `x` is not provided. Default is 1.0.
        :axis (int, optional): Axis along which to integrate. Default is the last axis.

    Returns:
        :result (torch.Tensor): The estimated integral using Simpson's rule.

    Notes
    -----
    If the number of points along `axis` is even, Simpson's rule requires an adjustment
    since it only works with an odd number of points. The function handles this case by
    applying a correction to the last interval.


    Compute the hypergeometric series _2F_1(a, b; c; z) using a series expansion.

    Args:
        :a (float): First parameter of the hypergeometric function.
        :b (float): Second parameter of the hypergeometric function.
        :c (float): Third parameter of the hypergeometric function.
        :z :(torch.Tensor): Argument of the hypergeometric function.
        :max_iter (int): Maximum number of iterations for convergence.
        :tolerance (float): Tolerance for convergence.

    Returns:
        F (torch.Tensor): The computed hypergeometric series value.


    Compute the hypergeometric function _2F_1(a, b; c; z).

    Args:
        :a (float): First parameter of the hypergeometric function.
        :b (float): Second parameter of the hypergeometric function.
        :c (float): Third parameter of the hypergeometric function.
        :z :(torch.Tensor): Argument of the hypergeometric function.

    Returns:
        :result (torch.Tensor): The computed hypergeometric function value.


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 7),
            the 7 parameters are :
                Omega_m : Density of matter.
                sigma_8 : Amplitude of matter fluctuations
                A_SN1 : Energy per SFR of the galactic winds
                A_AGN1 : Energy per black-hole accretion rate of the kinetic black-hole feedback
                A_SN2 : Wind speed of the galactic winds
                A_AGN2 : Energy per unit black-hole accretion rate of the thermal model of the black-hole feedback
                a : Scale factor, a = 1/(1+z)

    Returns:
        :S (torch.Tensor): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 7),
            the 7 parameters are :
                Omega_m : Density of matter.
                sigma_8 : Amplitude of matter fluctuations
                A_SN1 : Energy per unit SFR of the galactic winds
                A_AGN1 : Energy per unit blach-hole accretion rate, not used in the model, only here for consistency with other simulators.
                A_SN2 : Wind speed of the galactic winds
                A_AGN2 : Ejection speed/burstiness of the kinetic mode of the black-hole feedback
                a : Scale factor, a = 1/(1+z)

    Returns:
        :S (torch.Tensor): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 7),
            the 7 parameters are :
                Omega_m : Density of matter.
                sigma_8 : Amplitude of matter fluctuations
                A_SN1 : Mass loading of the galactic winds
                A_AGN1 : Momentum flux of the QSO & jet-mode black-hole feedback
                A_SN2 : Wind speed of the galactic winds
                A_AGN2 : Jet speed of the jet-mode black-hole feedback
                a : Scale factor, a = 1/(1+z)

    Returns:
        :S (torch.Tensor): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum.

    The explanation of astrophysical parameters was taken from https://camels.readthedocs.io/en/latest/parameters.html

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 7),
            the 7 parameters are:
                Omega_m : Density of matter.
                sigma_8 : amplitude of matter fluctuations
                A_SN1 : Thermal energy injected in each SNII event
                A_AGN1 : Scaling of the black hole Bondi accretion rate
                A_SN2 : Metallicity dependence of the stellar feedback fraction per unit stellar mass
                A_AGN2 : Temperature jump of gas particles in AGN feedback events
                a : Scale factor, a = 1/(1+z)

    Returns:
        :S (torch.Tensor): Baryonic suppression of the matter power spectrum


    Compute the impact of baryonic physics on the matter power spectrum for a given hydro model.
    This is given by the ratio of the baryonic to non-baryonic power spectrum at ks provided.

    Args:
        :k_batch (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 7),
            the 7 parameters are:
                Omega_m : Density of matter.
                sigma_8 : Amplitude of matter fluctuations
                A_SN1 : First supernova feedback parameter (see specific hydro model for details)
                A_SN2 : Second supernova feedback parameter (see specific hydro model for details)
                A_AGN1 : First AGN feedback parameter (see specific hydro model for details)
                A_AGN2 : Second AGN feedback parameter (see specific hydro model for details)
                a : Scale factor, a = 1/(1+z)
        :hydro_model (str): Name of the hydro model to use ('Astrid', 'IllustrisTNG', 'SIMBA', 'Swift-EAGLE')

    Returns:
        :S (torch.Tensor): Baryonic suppression of the matter power spectrum


    Compute nonlinear P(k) for the cosmology of interest

     Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :theta_batch (torch.Tensor): tensor containing the parameters with shape (batch_size, 10),
            the 10 parameters are:
                sigma8 : Root-mean-square density fluctuation when the linearly
                    evolved field is smoothed with a top-hat filter of radius 8 Mpc/h
                Om : The z=0 total matter density parameter, Omega_m
                Ob : The z=0 baryonic density parameter, Omega_b
                logMc : (Log of) mass fraction of hot gas in haloes
                logeta : (Log of) extent of ejected gas
                logbeta : (Log of) mass fraction of hot gas in haloes
                logM1 : (Log of) characteristic halo mass scale for central galaxy
                logMinn : (Log of) property describing density profile of hot gas in haloes
                logthetainn : (Log of) property describing density profile of hot gas in haloes
                a : Scale factor, a = 1/(1+z)

    Returns:
        :S (torch.Tensor): Ratio of baryonic to non-baryonic power spectrum at ks provided.


    Calculate the typical error of the prediction of `Astrid`.

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :a (torch.Tensor): tensor containing the scale factor, a = 1/(1+z) with shape (batch_size, 1)

    Returns:
        :epsilon (torch.Tensor): Covariance of the prediction and the actual S of Astrid.


    Calculate the typical error of the prediction of `IllustrisTNG`.

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :a (torch.Tensor): tensor containing the scale factor, a = 1/(1+z) with shape (batch_size, 1)

    Returns:
        :epsilon (torch.Tensor): Covariance of the prediction and the actual S of IllustrisTNG.


    Calculate the typical error of the prediction of `SIMBA`.

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :a (torch.Tensor): tensor containing the scale factor, a = 1/(1+z) with shape (batch_size, 1)

    Returns:
        :epsilon (torch.Tensor): Covariance of the prediction and the actual S of SIMBA.


    Calculate the typical error of the prediction of `Swift_EAGLE`.

    Args:
        :k (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :a (torch.Tensor): tensor containing the scale factor, a = 1/(1+z) with shape (batch_size, 1)

    Returns:
        :epsilon (torch.Tensor): Covariance of the prediction and the actual S of Swift-EAGLE.


    Calculate the typical error of the prediction for a given hydro model.

    Args:
        :k_batch (torch.Tensor): tensor containing the k values to evaluate P(k) at [h / Mpc] with shape (n_k,1)
        :a_batch (torch.Tensor): tensor containing the scale factor, a = 1/(1+z) with shape (batch_size, 1)
        :hydro_model (str): Name of the hydro model to use ('Astrid', 'IllustrisTNG', 'SIMBA', 'Swift-EAGLE')

    Returns:
        :epsilon (torch.Tensor): Covariance of the prediction and the actual S.


    Return a dictionary containing the nominal rotation curve
    parameters.

    Parameters: Nothing

    Returns: params
      params :: dictionary
        params['theta0'], etc. : scalar
           The nominal rotation curve parameter


    This rotation curve does not support resampling. This function
    raises an error.

    Parameters:
      size : integer
        Unused

    Returns: Nothing


    Return circular orbit speed at a given Galactocentric radius.

    Parameters:
      R :: scalar or array of scalars
        Galactocentric radius (kpc)

      R0, theta0 :: scalars (optional)
        Solar Galactocentric radius (kpc) and circular orbit speed at
        R0 (km/s)

    Returns: theta
      theta :: scalar or array of scalars
        circular orbit speed at R (km/s)


    Return the LSR velocity at a given Galactic longitude and
    line-of-sight distance.

    Parameters:
      glong, glat :: scalar or array of scalars
        Galactic longitude and latitude (deg).

      dist :: scalar or array of scalars
        line-of-sight distance (kpc).

      R0, theta0 :: scalars (optional)
        Solar Galactocentric radius (kpc) and circular orbit speed at
        R0 (km/s)

    Returns: vlsr
      vlsr :: scalar or array of scalars
        LSR velocity (km/s).


Multiprocessing wrapper from KDE results.


    Return the parallax Galactocentric radius and distance given a
    parallax and parallax uncertainty. Generate distance posterior by
    resampling parallax and R0 within uncertainties. Peak of posterior
    and 68.3% minimum width Bayesian credible interval (BCI) are
    returned.

    Parameters:
      glong :: scalar or array of scalars
        Galactic longitude (deg).

      plx :: scalar or array of scalars
        Parallax in milli-arcseconds

      plx_err :: scalar or 1-D (optional)
        Parallax uncertainty in milli-arcseconds. If it is an array,
        it must have the same size as plx. Otherwise, this scalar
        uncertainty is applied to all plxs.

      dist_max :: scalar (optional)
        The maximum parallax distance to compute (kpc)

      R0, R0_err :: scalar (optional)
        Solar Galactocentric radius and uncertainty (kpc)

      pdf_bins :: integer (optional)
        number of bins used to calculate PDF

      num_samples :: integer (optional)
        Number of MC samples to use when generating PDF

      plot_pdf :: bool (optional)
        If True, plot each PDF. Filenames are
        plot_prefix+"{plx}_{err}.pdf".

      plot_prefix :: string (optional)
        The prefix for the plot filenames.

    Returns: output
      output["Rgal"] :: scalar or array of scalars
        Galactocentric radius (kpc).

      output["distance"] : scalar or array of scalars
        parallax distance (kpc)

      Each of these values is the mode of the posterior distribution.
      Also included in the dictionary for each of these parameters
      are param+"_err_neg" and param+"_err_pos", which define the
      68.3% Bayesian credible interval around the mode, and
      param+"_kde", which is the kernel density estimator fit to the
      posterior samples.

    Raises:
      ValueError : if glong and plx are not 1-D; or
                   if glong and plx are arrays and not the same size


    Return the Galactocentric radius of an object with a given
    Galacitic longitude, latitude, and distance.

    Parameters:
      glong, glat :: scalars or arrays of scalars
        Galactic longitude and latitude (deg).

      dist :: scalar or array of scalars
        line-of-sight distance (kpc).

      R0 :: scalar (optional)
        Galactocentric radius of the Sun.

    Returns: R
      Rgal :: scalar or array of scalars
        Galactocentric radius (kpc).


    Return the Galactocentric azimuth of an object with a given
    Galacitic longitude, latitude, and distance. Galactocentric
    azimuth is defined as zero in the direction of the Sun and
    increasing in the direction of the Solar orbit.

    Parameters:
      glong, glat :: scalars or arrays of scalars
        Galactic longitude and latitude (deg).

      dist :: scalar or array of scalars
        line-of-sight distance (kpc).

      R0 :: scalar (optional)
        Galactocentric radius of the Sun.

    Returns: az
      az :: scalar or array of scalars
        Galactocentric azimuth (degs).


    Return the line-of-sight distance of an object with a given
    Galactocentric azimuth, radius, and height above the plane.

    Parameters:
      az :: scalar or array of scalars
        Galactocentric azimuth (deg).

      Rgal :: scalar or array of scalars
        Galactocentric radius (kpc).

      Z :: scalar or array of scalars
        Height above the plane

      R0 :: scalar (optional)
        Galactocentric radius of the Sun.

    Returns: dist
      dist :: scalar or array of scalars
        Line-of-sight distance (kpc).


    Return the Galactic longitude of an object with a given
    Galacitocentric azimuth and radius. Galactic longitude is
    defined as zero in the direction of the Galactic Center and
    increasing in the direction of the Solar orbit.

    Parameters:
      az :: scalar or array of scalars
        Galactocentric azimuth (deg).

      Rgal :: scalar or array of scalars
        Galactocentric radius (kpc).

      Z :: scalar or array of scalars
        Height above the plane

      R0 :: scalar (optional)
        Galactocentric radius of the Sun.

    Returns: glong
      glong :: scalar or array of scalars
        Galactic longitude (degs).


    Return the "corrected" LSR velocity by updating the IAU-defined
    solar motion components.

    Parameters:
      glong :: scalar or array of scalars
        Galactic longitude (deg).

      glat :: scalar or array of scalars
        Galactic latitude (deg).

      vlsr :: scalar or array of scalars
             Measured LSR velocity (km/s).

      Ustd, Vstd, Wstd :: scalars (optional)
        IAU-defined solar motion parameters (km/s).

      Usun, Vsun, Wsun : scalars (optional)
        Updated solar motion parameters (km/s).

    Returns: corr_vlsr
      corr_vlsr :: scalar or array of scalars
        Corrected LSR velocity


    Return the Anderson+2012 kinematic distance uncertainties.

    Parameters:
      glong :: scalar or array of scalars
        Galactic longitude (deg).

      vlsr :: scalar or array of scalars
        Measured LSR velocity (km/s).

    Returns: near_err, far_err, tangent_err
      near_err :: scalar or array of scalars
        Anderson+2012 near distance uncertainty

      far_err :: scalar or array of scalars
        Anderson+2012 far distance uncertainty

      tangent_err :: scalar or array of scalars
        Anderson+2012 tangent distance uncertainty

    Raises:
      ValueError : if glong and vlsr are not 1-D; or
                   if glong and vlsr are arrays and
                   not the same size


    Fit a kernel density estimator (KDE) to the posterior given
    by a collection of samples. Return the mode (posterior peak)
    and the highest posterior density (HPD) determined by the minimum
    width Bayesian credible interval (BCI) containing a fraction of
    the posterior samples. The posterior should be well described by a
    single-modal distribution.

    Parameters:
      samples :: 1-D array of scalars
        The samples being fit with a KDE

      kdetype :: string
        Which KDE method to use
          'pyqt' uses pyqt_fit with boundary at 0
          'scipy' uses gaussian_kde with no boundary

      alpha :: scalar (optional)
        The fraction of samples included in the BCI.

      pdf_bins :: integer (optional)
        Number of bins used in calculating the PDF

    Returns: kde, mode, lower, upper
      kde :: scipy.gaussian_kde or pyqt_fit.1DKDE object
        The KDE calculated for this kinematic distance

      mode :: scalar
        The mode of the posterior

      lower :: scalar
        The lower bound of the BCI

      upper :: scalar
        The upper bound of the BCI


    Compute parallax distance and Galactocentric radius for a given
    Galactic longitude and parallax.

    Parameters:
      glong :: scalar or array of scalars
        Galactic longitude (deg).

      plx :: scalar or array of scalars
        Parallax (milli-arcseconds).

      plx_err :: scalar or 1-D (optional)
        Parallax uncertainty in milli-arcseconds. If it is an array,
        it must have the same size as plx. Otherwise, this scalar
        uncertainty is applied to all plxs.

      dist_max :: scalar (optional)
        The maximum parallax distance to compute (kpc)

      R0, R0_err :: scalar (optional)
        Solar Galactocentric radius and uncertainty (kpc)

    Returns: output
      output["Rgal"] :: scalar or array of scalars
        Galactocentric radius (kpc).

      output["distance"] :: scalar or array of scalars
        Parallax distance (kpc).


    Return a dictionary containing the nominal rotation curve
    parameters.

    Parameters: Nothing

    Returns: params
      params :: dictionary
        params['a1'], etc. : scalar
          The nominal rotation curve parameter


    Resample the Reid+2014 rotation curve parameters within their
    uncertainties assuming Gaussian probabilities.

    Parameters:
      size :: integer
        The number of random samples to generate. If None, generate
        only one sample and return a scalar.

    Returns: params
      params :: dictionary
        params['a1'], etc. : scalar or array of scalars
                             The re-sampled parameters


    Return circular orbit speed at a given Galactocentric radius.

    Parameters:
      R :: scalar or array of scalars
        Galactocentric radius (kpc)

      a1, a2, a3 :: scalars (optional)
        Reid+2014 rotation curve parameters

      R0 :: scalar (optional)
        Solar Galactocentric radius (kpc)

    Returns: theta
      theta :: scalar or array of scalars
        circular orbit speed at R (km/s)


    Return the LSR velocity at a given Galactic longitude and
    line-of-sight distance.

    Parameters:
      glong, glat :: scalar or array of scalars
        Galactic longitude and latitude (deg).

      dist :: scalar or array of scalars
        line-of-sight distance (kpc).

      a1, a2, a3 :: scalars (optional)
        Reid+2014 rotation curve parameters

      R0 :: scalar (optional)
        Solar Galactocentric radius (kpc)

    Returns: vlsr
      vlsr :: scalar or array of scalars
        LSR velocity (km/s).


    Return a dictionary containing the nominal rotation curve
    parameters.

    Parameters: Nothing

    Returns: params
      params :: dictionary
        params['a1'], etc. : scalar
           The nominal rotation curve parameter


    Return circular orbit speed at a given Galactocentric radius.

    Parameters:
      R :: scalar or array of scalars
        Galactocentric radius (kpc)

      a1, a2, a3 :: scalars (optional)
        Brand rotation curve parameters

      R0, theta0 :: scalars (optional)
        Solar Galactocentric radius (kpc) and circular orbit speed at
        R0 (km/s)

    Returns: theta
      theta :: scalar or array of scalars
        circular orbit speed at R (km/s)


    Return the LSR velocity at a given Galactic longitude and
    line-of-sight distance.

    Parameters:
      glong, glat :: scalar or array of scalars
        Galactic longitude and latitude (deg).

      dist :: scalar or array of scalars
        line-of-sight distance (kpc).

      a1, a2, a3 :: scalars (optional)
        Brand rotation curve parameters

      R0, theta0 :: scalars (optional)
        Solar Galactocentric radius (kpc) and circular orbit speed at
        R0 (km/s)

    Returns: vlsr
      vlsr :: scalar or array of scalars
        LSR velocity (km/s).


    Return the kinematic near, far, and tanget distance and distance
    uncertainties for a given Galactic longitude and LSR velocity
    assuming a given rotation curve. Generate distance posteriors by
    resampling within rotation curve parameter and velocity
    uncertainties. Peak of posterior and 68.3% minimum width Bayesian
    credible interval (BCI) are returned.

    Parameters:
      glong, glat :: scalar or array of scalars
        Galactic longitude and latitude (deg).

      velo :: scalar or array of scalars
        LSR velocity (km/s).

      velo_err :: scalar or array of scalars (optional)
        LSR velocity uncertainty (km/s). If it is an array, it must
        have the same size as velo. Otherwise, this scalar
        uncertainty is applied to all velos.

      rotcurve :: string (optional)
        rotation curve model

      rotcurve_dist_res :: scalar (optional)
        line-of-sight distance resolution when calculating kinematic
        distance with rotcurve_kd (kpc)

      rotcurve_dist_max :: scalar (optional)
        maximum line-of-sight distance when calculating kinematic
        distance with rotcurve_kd (kpc)

      pdf_bins :: integer (optional)
        number of bins used to calculate PDF

      num_samples :: integer (optional)
        Number of MC samples to use when generating PDF

      plot_pdf :: bool (optional)
        If True, plot each PDF. Filenames are:
        plot_prefix+"{glong}_{velo}.pdf".

      plot_prefix :: string (optional)
        The prefix for the plot filenames.

    Returns: output
      output["Rgal"] :: scalar or array of scalars
        Galactocentric radius (kpc).

      output["Rtan"] :: scalar or array of scalars
        Galactocentric radius of tangent point (kpc).

      output["near"] :: scalar or array of scalars
        kinematic near distance (kpc)

      output["far"] :: scalar or array of scalars
        kinematic far distance (kpc)

      output["distance"] :: scalar or array of scalars
        kinematic distance (near and far combined) (kpc)

      output["tangent"] :: scalar or array of scalars
        kinematic tangent distance (kpc)

      output["vlsr_tangent"] :: scalar or array of scalars
        LSR velocity of tangent point (km/s)

      Each of these values is the mode of the posterior distribution.
      Also included in the dictionary for each of these parameters
      are param+"_err_neg" and param+"_err_pos", which define the
      68.3% Bayesian credible interval around the mode, and
      param+"_kde", which is the kernel density estimator fit to the
      posterior samples.

    Raises:
      ValueError : if glong and velo are not the same shape
                   if velo_err is an array and not the same shape as
                   glong and velo


    Return the kinematic near, far, and tanget distance for a
    given Galactic longitude and LSR velocity assuming
    a given rotation curve.

    Parameters:
      glong, glat :: scalar or array of scalars
        Galactic longitude and latitude (deg). If it is an array, it must have the
        same shape as velo.

      velo :: scalar or array of scalars
        LSR velocity (km/s). If it is an array, it must have the same
        shape as glong.

      velo_err :: scalar or array of scalars (optional)
        LSR velocity uncertainty (km/s). If it is an array, it must
        have the same shape as velo. Otherwise, this scalar
        uncertainty is applied to all velos.

      velo_tol :: scalar (optional)
        LSR velocity tolerance to consider a match between velo and
        rotation curve velocity

      rotcurve :: string (optional)
        rotation curve model

      dist_res :: scalar (optional)
        line-of-sight distance resolution when calculating kinematic
        distance (kpc)

      dist_min :: scalar (optional)
        minimum line-of-sight distance when calculating kinematic
        distance (kpc)

      dist_max :: scalar (optional)
        maximum line-of-sight distance when calculating kinematic
        distance (kpc)

      resample :: bool (optional)
        if True, use resampled rotation curve parameters and LSR
        velocities.

      size :: integer (optional)
        if resample is True, generate this many samples

    Returns: output
      output["Rgal"] :: scalar or array of scalars
        Galactocentric radius (kpc).

      output["Rtan"] :: scalar or array of scalars
        Galactocentric radius of tangent point (kpc).

      output["near"] :: scalar or array of scalars
        kinematic near distance (kpc)

      output["far"] :: scalar or array of scalars
        kinematic far distance (kpc)

      output["tangent"] :: scalar or array of scalars
        kinematic tangent distance (kpc)

      output["vlsr_tangent"] :: scalar or array of scalars
        LSR velocity of tangent point (km/s)

      If glong and velo are scalars, each of these is a scalar.
      Otherwise they have the same shape as input glong and velo.

      If resample is True and size > 1, the samples are stored along
      the last axis of each element.

    Raises:
      ValueError : if glong and velo are not the same shape


    Resample the Reid+2019 rotation curve parameters within their
    uncertainties using the Wenger+2020 kernel density estimator
    to include parameter covariances.

    Parameters:
      size :: integer
        The number of random samples to generate. If None, generate
        only one sample and return a scalar.

    Returns: params
      params :: dictionary
        params['a1'], etc. : scalar or array of scalars
                             The re-sampled parameters


    Return circular orbit speed at a given Galactocentric radius.

    Parameters:
      R :: scalar or array of scalars
        Galactocentric radius (kpc)

      a2, a3 :: scalars (optional)
        Reid+2019 rotation curve parameters

      R0 :: scalar (optional)
        Solar Galactocentric radius (kpc)

    Returns: theta
      theta :: scalar or array of scalars
        circular orbit speed at R (km/s)


    Return the IAU-LSR velocity at a given Galactic longitude and
    line-of-sight distance.

    Parameters:
      glong, glat :: scalars or arrays of scalars
        Galactic longitude and latitude (deg).

      dist :: scalar or array of scalars
        line-of-sight distance (kpc).

      R0 :: scalar (optional)
        Solar Galactocentric radius (kpc)

      Usun, Vsun, Wsun, Upec, Vpec, a2, a3 :: scalars (optional)
        Reid+2019 rotation curve parameters

      Zsun :: scalar (optional)
        Height of sun above Galactic midplane (pc)

      roll :: scalar (optional)
        Roll of Galactic midplane relative to b=0 (deg)

      peculiar :: boolean (optional)
        If True, include HMSFR peculiar motion component

    Returns: vlsr
      vlsr :: scalar or array of scalars
        LSR velocity (km/s).


 from an array of mesh points and an index array of (quadratic) triangle connections,
    construct a bounding volume hierarchy (BVH) tree, which will be used to evaluate fields
    define on the mesh nodes.

    ARGS:
        points: an array of the (x,y) positions of the mesh nodes, dimension N x 2 for N nodes.
        connections: an array containing each triangle in the mesh; each triangle is represented
                     as 6 indices, corresponding to 6 points
    RETURNS:
        bvhtree: the BVH tree for the given mesh points and connections.


 create a BVH tree directly from a finite element mesh object.

    ARGS:
        mesh: a meshio object representing a finite element mesh
    RETURNS:
        bvhtree: the BVH tree for the given mesh points and connections.


create a BVH tree for ``mesh`` and pass it into to ``mesh.tree``


 find the index of the triangle in the mesh that contains the given point.

    ARGS:
        point: an array [x,y] corresponding to the query point..
        tree: the BVH tree for the mesh of interest.
    RETURNS:
        (int): the index of the triangle containing point in the mesh.


 evaluate a field sampled over a finite element mesh at a given point.

    ARGS:
        point: an [x,y] point, or an Nx2 array of points
        field: a real-valued field represented on a finite-element mesh
        tree: a BVH tree that stores the triangles of field's finite-element mesh.
    RETURNS:
        (float or vector): the field evaluated at point(s)


 resample a finite element field onto a new mesh

    ARGS:
        field: the finite element field to be resampled.
        mesh: the finite element mesh on which <field> is defined.
        newmesh: the new finite element mesh <field> should be sampled on.


 evaluate a field defined over a finite element mesh on a cartesian grid.

    ARGS:
        pointsx: a 1D array of x points
        pointsy: a 1D array of y points
        field: a real-valued field represented on a finite-element mesh
        tree: a BVH tree that stores the triangles of field's finite-element mesh

    RETURNS:
        (array): a 2D array corresponding to field, evaluated on the grid.


 compute the gradient of a real-valued finite element field with respect to x,y

    ARGS:
        field: a real-valued field represented on a finite-element mesh
        tris: an Nx6 array of triangles, each row contains the 6 indices of one (quadratic) triangle element
        points: an Mx2 array of [x,y] points representing the mesh nodes.
    RETURNS:
        (array): dimensions are dim(field) x 2; the last axis contains the x and y derivatives.


 plot a mesh and associated refractive index distribution
    Args:
    mesh: the mesh to be plotted. if None, we auto-compute a mesh using default values
    IOR_dict: dictionary that assigns each named region in the mesh to a refractive index value


 fragment the tool and the object, and return the fragments in the following order:
        intersection, object_fragment, tool_fragment.
        in some cases one of the later two may be empty


 update the primitive according to some args and return an Nx2 array of points.
            the default behavior is to manually pass in a points array. more specific primitives
            inheriting from Prim2D should implement their own update().


 make an Nx2 array of points for marking the primitive boundary,
            according to some args.


a quick check to see what this object looks like. generates a mesh with default parameters.


 this function computes the distance between the point (x,y) and the boundary of the primitive.
        negative distances -> inside the boundary, while positive -> outside. note that this doesn't need to be exact,
        the "distance" just needs to be positive outside the boundary, negative inside the boundary, and go to 0 as you approach the boundary.
        it is better if it works on vectorized inputs for x,y.

        ARGS:
            x (float or vector): x coordinate(s)  of point(s) to measure distance w/ boundary.
            y (float or vector): y coordinate(s)  of point(s) to measure distance w/ boundary.
        RETURNS:
            (float or vector): the distance(s) to the boundary from the point(s)


 this function computes the point on the boundary that is closest to a point (x,y).
        it is better if it works on vectorized inputs for x,y.

        ARGS:
            x (float or vector): x coordinate(s)  of point(s) to measure distance w/ boundary.
            y (float or vector): y coordinate(s)  of point(s) to measure distance w/ boundary.
        RETURNS:
            (tuple): a tuple containing:
                - (float or vector): the x coordinate(s) of the nearest boundary point
                - (float or vector): the y coordinate(s) of the nearest boundary point


 make points corresponding to the boundary of the primitive.

        ARGS:
            args1 (tuple): the arguments of make_points for the first primitive
            args2 (tuple): the arguments of make_points for the second primitive


 make points of prim2D at given z coord. should be implemented by
        inheriting classes.

        ARGS:
            z (float): the desired z coordinate

        RETURNS:
            (array): an Nx2 array of points corresponding to the boundary at z.


 for a point (x0,y0) inside the boundary at z0, compute a new point (x,y) at z, accounting
        for the z-variation of the Prim3D. should be implemented by inheriting classes, and
        should be vectorized if possible.

        ARGS:
            x0 (float or vector): reference x coordinate(s) inside the boundary at z0
            y0 (float or vector): reference y coordinate(s) inside the boundary at z0
            z0 (float): reference z coordiante
            z (float): new z coordinate
        RETURNS:
            (tuple): a tuple containing:
                - x1 (float or vector): the new x coordinate(s) at z
                - y1 (float or vector): the new y coordinate(s) at z


 update the mesh boundaries to the given z coordinate.

        ARGS:
            z (float):  desired z coordinate. note that all waveguides are
                        updated to z=0 upon initialization.