species.util.data_util — species 0.1.2 documentation

# Source code for species.util.data_util

    """
    Utility functions for data processing.
    """

    import warnings

    import numpy as np

    def update_sptype(sptypes):
        """
        Function to update a list with spectral types to two characters (e.g., M8, L3, or T1).

        Parameters
        ----------
        sptypes : numpy.ndarray
            Input spectral types.

        Returns
        -------
        numpy.ndarray
            Updated spectral types.
        """

        sptype_list = ['O', 'B', 'A', 'F', 'G', 'K', 'M', 'L', 'T', 'Y']

        for i, spt_item in enumerate(sptypes):
            if spt_item == 'None':
                pass

            elif spt_item == 'null':
                sptypes[i] = 'None'

            else:
                for list_item in sptype_list:
                    try:
                        sp_index = spt_item.index(list_item)
                        sptypes[i] = spt_item[sp_index:sp_index+2]

                    except ValueError:
                        pass

        return sptypes

    def update_filter(filter_in):
        """
        Function to update afilter ID from the Vizier Photometry viewer VOTable to the filter ID from
        the SVO Filter Profile Service.

        Parameters
        ----------
        filter_in : str
            Filter ID in the format of the Vizier Photometry viewer.

        Returns
        -------
        str
            Filter ID in the format of the SVO Filter Profile Service.
        """

        if filter_in[0:5] == b'2MASS':
            filter_out = str(b'2MASS/2MASS.'+filter_in[6:])

        elif filter_in[0:4] == b'WISE':
            filter_out = str(b'WISE/WISE.'+filter_in[5:])

        elif filter_in[0:10] == b'GAIA/GAIA2':
            filter_out = str(filter_in[0:9]+b'0'+filter_in[10:])

        else:
            filter_out = None

        return filter_out

    def sort_data(teff,
                  logg,
                  feh,
                  co,
                  fsed,
                  wavelength,
                  flux):
        """
        Parameters
        ----------
        teff : numpy.ndarray
        logg : numpy.ndarray
        feh : numpy.ndarray, None
        co : numpy.ndarray, None
        fsed : numpy.ndarray, None
        wavelength : numpy.ndarray
        flux : numpy.ndarray

        Returns
        -------
        tuple(numpy.ndarray, )
        """

        teff_unique = np.unique(teff)
        logg_unique = np.unique(logg)

        if feh is None and co is None and fsed is None:
            spectrum = np.zeros((teff_unique.shape[0],
                                 logg_unique.shape[0],
                                 wavelength.shape[0]))

        elif feh is not None and co is None and fsed is None:
            feh_unique = np.unique(feh)

            spectrum = np.zeros((teff_unique.shape[0],
                                 logg_unique.shape[0],
                                 feh_unique.shape[0],
                                 wavelength.shape[0]))

        elif feh is not None and co is not None and fsed is None:
            feh_unique = np.unique(feh)
            co_unique = np.unique(co)

            spectrum = np.zeros((teff_unique.shape[0],
                                 logg_unique.shape[0],
                                 feh_unique.shape[0],
                                 co_unique.shape[0],
                                 wavelength.shape[0]))

        elif feh is not None and co is None and fsed is not None:
            feh_unique = np.unique(feh)
            fsed_unique = np.unique(fsed)

            spectrum = np.zeros((teff_unique.shape[0],
                                 logg_unique.shape[0],
                                 feh_unique.shape[0],
                                 fsed_unique.shape[0],
                                 wavelength.shape[0]))

        else:
            feh_unique = np.unique(feh)
            co_unique = np.unique(co)
            fsed_unique = np.unique(fsed)

            spectrum = np.zeros((teff_unique.shape[0],
                                 logg_unique.shape[0],
                                 feh_unique.shape[0],
                                 co_unique.shape[0],
                                 fsed_unique.shape[0],
                                 wavelength.shape[0]))

        for i in range(teff.shape[0]):
            index_teff = np.argwhere(teff_unique == teff[i])[0]
            index_logg = np.argwhere(logg_unique == logg[i])[0]

            if feh is None and co is None and fsed is None:
                spectrum[index_teff, index_logg, :] = flux[i]

            elif feh is not None and co is None and fsed is None:
                index_feh = np.argwhere(feh_unique == feh[i])[0]
                spectrum[index_teff, index_logg, index_feh, :] = flux[i]

            elif feh is not None and co is not None and fsed is None:
                index_feh = np.argwhere(feh_unique == feh[i])[0]
                index_co = np.argwhere(co_unique == co[i])[0]
                spectrum[index_teff, index_logg, index_feh, index_co, :] = flux[i]

            elif feh is not None and co is None and fsed is not None:
                index_feh = np.argwhere(feh_unique == feh[i])[0]
                index_fsed = np.argwhere(fsed_unique == fsed[i])[0]
                spectrum[index_teff, index_logg, index_feh, index_fsed, :] = flux[i]

            else:
                index_feh = np.argwhere(feh_unique == feh[i])[0]
                index_co = np.argwhere(co_unique == co[i])[0]
                index_fsed = np.argwhere(fsed_unique == fsed[i])[0]
                spectrum[index_teff, index_logg, index_feh, index_co, index_fsed, :] = flux[i]

        if feh is None and co is None and fsed is None:
            sorted_data = [teff_unique, logg_unique, wavelength, spectrum]

        elif feh is not None and co is None and fsed is None:
            sorted_data = [teff_unique, logg_unique, feh_unique, wavelength, spectrum]

        elif feh is not None and co is not None and fsed is None:
            sorted_data = [teff_unique, logg_unique, feh_unique, co_unique, wavelength, spectrum]

        elif feh is not None and co is None and fsed is not None:
            sorted_data = [teff_unique, logg_unique, feh_unique, fsed_unique, wavelength, spectrum]

        else:
            sorted_data = [teff_unique, logg_unique, feh_unique, co_unique, fsed_unique, wavelength, spectrum]

        return sorted_data

    def write_data(model,
                   parameters,
                   database,
                   data_sorted):
        """
        Parameters
        ----------
        model : str
            Atmosphere model.
        parameters : tuple(str, )
            Model parameters.
        database: h5py._hl.files.File
            Database.
        data_sorted : tuple(numpy.ndarray, )

        Returns
        -------
        None
        """

        if 'models/'+model in database:
            del database['models/'+model]

        dset = database.create_group('models/'+model)

        dset.attrs['nparam'] = len(parameters)

        for i, item in enumerate(parameters):
            dset.attrs['parameter'+str(i)] = item

            database.create_dataset('models/'+model+'/'+item,
                                    data=data_sorted[i],
                                    dtype='f')

        database.create_dataset('models/'+model+'/wavelength',
                                data=data_sorted[len(parameters)],
                                dtype='f')

        database.create_dataset('models/'+model+'/flux',
                                data=data_sorted[len(parameters)+1],
                                dtype='f')

    def add_missing(model,
                    parameters,
                    database):
        """
        Parameters
        ----------
        model : str
            Atmosphere model.
        parameters : tuple(str, )
            Model parameters.
        database : h5py._hl.files.File
            Database.

        Returns
        -------
        None
        """

        grid_shape = []
        for item in parameters:
            grid_shape.append(database['models/'+model+'/'+item].shape[0])

        teff = np.asarray(database['models/'+model+'/teff'])
        flux = np.asarray(database['models/'+model+'/flux'])

        for i in range(grid_shape[0]):
            for j in range(grid_shape[1]):

                if len(parameters) == 2:
                    if np.count_nonzero(flux[i, j]) == 0:
                        try:
                            scaling = (teff[i+1]-teff[i])/(teff[i+1]-teff[i-1])
                            flux[i, j] = scaling*flux[i+1, j] + (1.-scaling)*flux[i-1, j]

                        except IndexError:
                            flux[i, j] = np.nan
                            warnings.warn(f'Interpolation is not possible at the edge of the '
                                          f'parameter grid. A NaN value is stored for Teff = '
                                          f'{teff[i]} K.')

                elif len(parameters) == 3:
                    for k in range(grid_shape[2]):
                        if np.count_nonzero(flux[i, j, k]) == 0:
                            try:
                                scaling = (teff[i+1]-teff[i])/(teff[i+1]-teff[i-1])
                                flux[i, j, k] = scaling*flux[i+1, j, k] + (1.-scaling)*flux[i-1, j, k]

                            except IndexError:
                                flux[i, j, k] = np.nan
                                warnings.warn(f'Interpolation is not possible at the edge of the '
                                              f'parameter grid. A NaN value is stored for Teff = '
                                              f'{teff[i]} K.')

                elif len(parameters) == 4:
                    for k in range(grid_shape[2]):
                        for m in range(grid_shape[3]):
                            if np.count_nonzero(flux[i, j, k, m]) == 0:
                                try:
                                    scaling = (teff[i+1]-teff[i])/(teff[i+1]-teff[i-1])
                                    flux[i, j, k, m] = scaling*flux[i+1, j, k, m] + (1.-scaling)*flux[i-1, j, k, m]

                                except IndexError:
                                    flux[i, j, k, m] = np.nan
                                    warnings.warn(f'Interpolation is not possible at the edge of the '
                                                  f'parameter grid. A NaN value is stored for Teff = '
                                                  f'{teff[i]} K.')

        del database['models/'+model+'/flux']

        database.create_dataset('models/'+model+'/flux',
                                data=flux,
                                dtype='f')