<commit_before>#!/usr/bin/env python
# -*- coding:utf-8 -*-

# @Filename: utils.py
# @License: BSD 3-clause (http://www.opensource.org/licenses/BSD-3-Clause)


import numpy as np
import kcorrect.template


def sed_ab(wave=None):
    """Calculate AB source spectrum

    Parameters
    ----------

    wave : ndarray of np.float32
        wavelength in Angstrom

    Returns
    -------

    ab : ndarray of np.float32
        AB flux in erg cm^{-2} s^{-1} Angstrom^{-1}
"""
    # This is 3631 Jy * c / wavelength^2
    # Keeping cspeed in order unity units to avoid overflows
    cspeed = 2.99792   # speed of light in 10^{18} A / s
    ab = 3.631e-2 / wave**2 * cspeed
    return(ab)


def read_basel(filename=None):
    """Read Basel-style spectrum file

    Parameters
    ----------

    filename : str
        file to read from

    Returns
    -------

    info : dict
        meta-data for model

    wave : ndarray of np.float32
        [1221] wavelengths in nm

    flux : ndarray of np.float32
        [N, 1221] fluxes in erg/cm^2/s/Hz

    Notes
    -----

    1221 hardcoded as number of spectral elements.

    Model parameters stored as arrays in info attribute are:
       'modelno'
       'teff'
       'logg'
       'mh'
       'vturb'
       'xh'
"""
    fp = open(filename, "r")
    fpstr = fp.read()
    fp.close

    words = fpstr.split()

    nwave = 1221

    nunits = (len(words) - nwave) // (nwave + 6)

    wavestr = words[0:nwave]
    wave = np.array([np.float32(x) for x in wavestr], dtype=np.float32)

    flux = np.zeros((nunits, nwave), dtype=np.float32)
    modelno = np.zeros(nunits, dtype=np.int32)
    teff = np.zeros(nunits, dtype=np.float32)
    logg = np.zeros(nunits, dtype=np.float32)
    mh = np.zeros(nunits, dtype=np.float32)
    vturb = np.zeros(nunits, dtype=np.float32)
    xh = np.zeros(nunits, dtype=np.float32)

    for unit in range(nunits):
        start = nwave + unit * (nwave + 6)

        modelno[unit] = np.int32(words[start])
        teff[unit] = np.float32(words[start + 1])
        logg[unit] = np.float32(words[start + 2])
        mh[unit] = np.float32(words[start + 3])
        vturb[unit] = np.float32(words[start + 4])
        xh[unit] = np.float32(words[start + 5])

        fluxstr = words[start + 6:start + 6 + nwave]
        flux[unit, :] = np.array([np.float32(x) for x in fluxstr],
                                 dtype=np.float32)

    info = dict()
    info['modelno'] = modelno
    info['teff'] = teff
    info['logg'] = logg
    info['mh'] = mh
    info['vturb'] = vturb
    info['xh'] = xh

    return(info, wave, flux)<commit_msg>add ab fix and asinh stuff<commit_after>#!/usr/bin/env python
# -*- coding:utf-8 -*-

# @Filename: utils.py
# @License: BSD 3-clause (http://www.opensource.org/licenses/BSD-3-Clause)


import numpy as np
import kcorrect.template


def sed_ab(wave=None):
    """Calculate AB source spectrum

    Parameters
    ----------

    wave : ndarray of np.float32
        wavelength in Angstrom

    Returns
    -------

    ab : ndarray of np.float32
        AB flux in erg cm^{-2} s^{-1} Angstrom^{-1}
"""
    # This is 3631 Jy * c / wavelength^2
    # Keeping cspeed in order unity units to avoid overflows
    cspeed = 2.99792   # speed of light in 10^{18} A / s
    ab = 3.631e-2 / wave**2 * cspeed
    return(ab)


def read_basel(filename=None):
    """Read Basel-style spectrum file

    Parameters
    ----------

    filename : str
        file to read from

    Returns
    -------

    info : dict
        meta-data for model

    wave : ndarray of np.float32
        [1221] wavelengths in nm

    flux : ndarray of np.float32
        [N, 1221] fluxes in erg/cm^2/s/Hz

    Notes
    -----

    1221 hardcoded as number of spectral elements.

    Model parameters stored as arrays in info attribute are:
       'modelno'
       'teff'
       'logg'
       'mh'
       'vturb'
       'xh'
"""
    fp = open(filename, "r")
    fpstr = fp.read()
    fp.close

    words = fpstr.split()

    nwave = 1221

    nunits = (len(words) - nwave) // (nwave + 6)

    wavestr = words[0:nwave]
    wave = np.array([np.float32(x) for x in wavestr], dtype=np.float32)

    flux = np.zeros((nunits, nwave), dtype=np.float32)
    modelno = np.zeros(nunits, dtype=np.int32)
    teff = np.zeros(nunits, dtype=np.float32)
    logg = np.zeros(nunits, dtype=np.float32)
    mh = np.zeros(nunits, dtype=np.float32)
    vturb = np.zeros(nunits, dtype=np.float32)
    xh = np.zeros(nunits, dtype=np.float32)

    for unit in range(nunits):
        start = nwave + unit * (nwave + 6)

        modelno[unit] = np.int32(words[start])
        teff[unit] = np.float32(words[start + 1])
        logg[unit] = np.float32(words[start + 2])
        mh[unit] = np.float32(words[start + 3])
        vturb[unit] = np.float32(words[start + 4])
        xh[unit] = np.float32(words[start + 5])

        fluxstr = words[start + 6:start + 6 + nwave]
        flux[unit, :] = np.array([np.float32(x) for x in fluxstr],
                                 dtype=np.float32)

    info = dict()
    info['modelno'] = modelno
    info['teff'] = teff
    info['logg'] = logg
    info['mh'] = mh
    info['vturb'] = vturb
    info['xh'] = xh

    return(info, wave, flux)


def sdss_ab_correct(maggies=None, ivar=None,
                    abfix=[-0.036, 0.012, 0.010, 0.028, 0.040]):
    """Return AB maggies based on SDSS standard maggies

    Parameters
    ----------

    maggies : ndarray of np.float32
        array of fluxes in standard SDSS system

    ivar : ndarray of np.float32
        inverse variances in standard SDSS system (optional)

    abfix : ndarray or list of np.float32
        [5]-array of magnitude offset to apply to standard values

    Returns
    -------

    ab_maggies : ndarray of np.float32
        array of fluxes converted to AB

    ab_ivar : ndarray of np.float32
        inverse variances converted to AB (if ivar input)

    Notes
    -----

    Uses the AB conversions produced by <NAME>, in his
    message sdss-calib/1152

      u(AB,2.5m) = u(database, 2.5m) - 0.036
      g(AB,2.5m) = g(database, 2.5m) + 0.012
      r(AB,2.5m) = r(database, 2.5m) + 0.010
      i(AB,2.5m) = i(database, 2.5m) + 0.028
      z(AB,2.5m) = z(database, 2.5m) + 0.040

    Unless an alternative is specified with the abfix parameter.
"""
    maggies = np.float32(maggies)
    if(ivar is not None):
        ivar = np.float32(ivar)

        if(maggies.shape != ivar.shape):
            raise("maggies and ivar must be the same shape")

    abfix = np.array(abfix, dtype=np.float32)
    if(abfix.shape != (5,)):
        raise("abfix must have 5 values")
    abfac = 10.**(- 0.4 * abfix)

    if(maggies.ndim == 1):
        if(maggies.size != 5):
            raise("sdss_ab_correct expects 5 SDSS maggies values (ugriz)")
        ab_maggies = maggies * abfac
        if(ivar is not None):
            ab_ivar = ivar / abfac**2
    else:
        ab_maggies = np.zeros(maggies.shape, dtype=np.float32)
        for i, cabfac in enumerate(abfac):
            ab_maggies[..., i] = maggies[..., i] * cabfac
        if(ivar is not None):
            ab_ivar = np.zeros(maggies.shape, dtype=np.float32)
            for i, cabfac in enumerate(abfac):
                ab_ivar[..., i] = ivar[..., i] / cabfac**2

    if(ivar is not None):
        return(ab_maggies, ab_ivar)
    else:
        return(ab_maggies)


def sdss_asinh_to_maggies(mag=None, mag_err=None, extinction=None):
    """Calculate maggies and ivar based on catalog parameters

    Parameters
    ----------

    mag : ndarray of np.float32
        [N, 5] or [5] array of asinh magnitudes from SDSS

    mag_err : ndarray of np.float32
        [N, 5] or [5] array of asinh magnitude errors from SDSS

    extinction : ndarray or list of np.float32
        [N, 5] or [5] array of Galactic extinctions from SDSS

    Returns
    -------

    maggies : ndarray of np.float32
        [N, 5] or [5] array of maggies

    ivar : ndarray of np.float32
        [N, 5] or [5] array of inverse variances
"""
    mag = np.float32(mag)
    mag_err = np.float32(mag_err)
    extinction = np.float32(extinction)

    if(mag.shape != mag_err.shape):
        raise("mag and mag_err must be the same shape")

    if(mag.shape != extinction.shape):
        raise("mag and extinction must be the same shape")

    if(mag.shape[-1] != 5):
        raise("mag must have 5 values per object")
{'fraction_non_alphanumeric': 0.13646106215765658, 'fraction_numerical': 0.037747082638723506, 'mean_word_length': 2.1246279761904763, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '205579', 'n_tokens_mistral': 3238, 'n_tokens_neox': 2925, 'n_words': 791}