Added functionalities for plotting results from the empirical comparison

species/__init__.py

@@ -1,3 +1,5 @@
+from species.analysis.empirical import CompareSpectra
+
 from species.analysis.fit_model import FitModel
 
 from species.analysis.fit_spectrum import FitSpectrum
@@ -35,6 +37,9 @@
 
 from species.data.database import Database
 
+from species.plot.plot_empirical import plot_statistic, \
+                                        plot_empirical_spectra
+
 from species.plot.plot_color import plot_color_magnitude, \
                                     plot_color_color
 

species/analysis/empirical.py

@@ -0,0 +1,256 @@
+"""
+Module with functionalities for empirical spectral analysis.
+"""
+
+import configparser
+import os
+
+from typing import List, Optional, Tuple, Union
+
+import h5py
+import numpy as np
+
+from scipy.interpolate import interp1d
+from typeguard import typechecked
+
+from species.core import constants
+from species.data import database
+from species.read import read_object
+from species.util import dust_util, read_util
+
+
+class CompareSpectra:
+    """
+    Class for comparing a spectrum of an object with the spectra of a library.
+    """
+
+    @typechecked
+    def __init__(self,
+                 object_name: str,
+                 spec_name: str,
+                 spec_library: str) -> None:
+        """
+        Parameters
+        ----------
+        object_name : str
+            Object name as stored in the database with
+            :func:`~species.data.database.Database.add_object` or
+            :func:`~species.data.database.Database.add_companion`.
+        spec_name : str
+            Name of the spectrum that is stored at the object data of ``object_name``.
+        spec_library : str
+            Name of the spectral library ('irtf', 'spex', or 'kesseli+2017).
+
+        Returns
+        -------
+        NoneType
+            None
+        """
+
+        self.object_name = object_name
+        self.spec_name = spec_name
+        self.spec_library = spec_library
+
+        self.object = read_object.ReadObject(object_name)
+
+        config_file = os.path.join(os.getcwd(), 'species_config.ini')
+
+        config = configparser.ConfigParser()
+        config.read_file(open(config_file))
+
+        self.database = config['species']['database']
+
+    @typechecked
+    def spectral_type(self,
+                      tag: str,
+                      wavel_range: Optional[Tuple[Optional[float], Optional[float]]] = None,
+                      sptypes: Optional[List[str]] = None,
+                      av_ext: Optional[Union[List[float], np.array]] = None,
+                      rad_vel: Optional[Union[List[float], np.array]] = None) -> None:
+        """
+        Method for finding the best fitting empirical spectra from a selected library by
+        calculating for each spectrum the goodness-of-fit statistic from Cushing et al. (2008).
+
+        Parameters
+        ----------
+        tag : str
+            Database tag where for each spectrum from the spectral library the best-fit parameters
+            will be stored. So when testing a range of values for ``av_ext`` and ``rad_vel``, only
+            the parameters that minimize the goodness-of-fit statistic will be stored.
+        wavel_range : tuple(float, float), None
+            Wavelength range (um) that is used for the empirical comparison.
+        sptypes : list(str), None
+            List with spectral types to compare with. The list should only contains types, for
+            example ``sptypes=['M', 'L']``. All available spectral types in the ``spec_library``
+            are compared with if set to ``None``.
+        av_ext : list(float), np.array, None
+            List of A_V extinctions for which the goodness-of-fit statistic is tested. The
+            extinction is calculated with the empirical relation from Cardelli et al. (1989).
+        rad_vel : list(float), np.array, None
+            List of radial velocities (km s-1) for which the goodness-of-fit statistic is tested.
+
+        Returns
+        -------
+        NoneType
+            None
+        """
+
+        w_i = 1.
+
+        if av_ext is None:
+            av_ext = [0.]
+
+        if rad_vel is None:
+            rad_vel = [0.]
+
+        h5_file = h5py.File(self.database, 'r')
+
+        try:
+            h5_file[f'spectra/{self.spec_library}']
+
+        except KeyError:
+            h5_file.close()
+            species_db = database.Database()
+            species_db.add_spectrum(self.spec_library)
+            h5_file = h5py.File(self.database, 'r')
+
+        obj_spec = self.object.get_spectrum()[self.spec_name][0]
+        obj_res = self.object.get_spectrum()[self.spec_name][3]
+
+        name_list = []
+        spt_list = []
+        gk_list = []
+        ck_list = []
+        av_list = []
+        rv_list = []
+
+        print_message = ''
+
+        for i, item in enumerate(h5_file[f'spectra/{self.spec_library}']):
+            dset = h5_file[f'spectra/{self.spec_library}/{item}']
+            item_sptype = dset.attrs['sptype'].decode('utf-8')
+
+            if item_sptype == 'None':
+                continue
+
+            if sptypes is None or item_sptype[0] in sptypes:
+                spectrum = np.asarray(dset)
+
+                if wavel_range is not None:
+                    if wavel_range[0] is None:
+                        indices = np.where((spectrum[:, 0] < wavel_range[1]))[0]
+
+                    elif wavel_range[1] is None:
+                        indices = np.where((spectrum[:, 0] > wavel_range[0]))[0]
+
+                    else:
+                        indices = np.where((spectrum[:, 0] > wavel_range[0]) &
+                                           (spectrum[:, 0] < wavel_range[1]))[0]
+
+                    if len(indices) == 0:
+                        raise ValueError('The selected wavelength range does not cover any '
+                                         'wavelength points of the input spectrum. Please '
+                                         'use a broader range as argument of \'wavel_range\'.')
+
+                    spectrum = spectrum[indices, ]
+
+                empty_message = len(print_message)*' '
+                print(f'\r{empty_message}', end='')
+
+                print_message = f'Processing spectra... {item}'
+                print(f'\r{print_message}', end='')
+
+                for av_item in av_ext:
+                    for rv_item in rad_vel:
+                        ism_ext = dust_util.ism_extinction(av_item, 3.1, spectrum[:, 0])
+                        flux_scaling = 10.**(-0.4*ism_ext)
+
+                        wavel_shifted = spectrum[:, 0] + spectrum[:, 0] * 1e3*rv_item / constants.LIGHT
+
+                        flux_smooth = read_util.smooth_spectrum(wavel_shifted,
+                                                                spectrum[:, 1]*flux_scaling,
+                                                                spec_res=obj_res,
+                                                                force_smooth=True)
+
+                        interp_spec = interp1d(spectrum[:, 0],
+                                               flux_smooth,
+                                               kind='linear',
+                                               fill_value='extrapolate')
+
+                        indices = np.where((obj_spec[:, 0] > np.amin(spectrum[:, 0])) &
+                                           (obj_spec[:, 0] < np.amax(spectrum[:, 0])))[0]
+
+                        flux_resample = interp_spec(obj_spec[indices, 0])
+
+                        c_numer = w_i*obj_spec[indices, 1]*flux_resample/obj_spec[indices, 2]**2
+                        c_denom = w_i*flux_resample**2/obj_spec[indices, 2]**2
+
+                        c_k = np.sum(c_numer) / np.sum(c_denom)
+
+                        chi_sq = (obj_spec[indices, 1] - c_k*flux_resample) / obj_spec[indices, 2]
+                        g_k = np.sum(w_i * chi_sq**2)
+
+                        name_list.append(item)
+                        spt_list.append(item_sptype)
+                        gk_list.append(g_k)
+                        ck_list.append(c_k)
+                        av_list.append(av_item)
+                        rv_list.append(rv_item)
+
+        empty_message = len(print_message)*' '
+        print(f'\r{empty_message}', end='')
+
+        print('\rProcessing spectra... [DONE]')
+
+        h5_file.close()
+
+        name_list = np.asarray(name_list)
+        spt_list = np.asarray(spt_list)
+        gk_list = np.asarray(gk_list)
+        ck_list = np.asarray(ck_list)
+        av_list = np.asarray(av_list)
+        rv_list = np.asarray(rv_list)
+
+        sort_index = np.argsort(gk_list)
+
+        name_list = name_list[sort_index]
+        spt_list = spt_list[sort_index]
+        gk_list = gk_list[sort_index]
+        ck_list = ck_list[sort_index]
+        av_list = av_list[sort_index]
+        rv_list = rv_list[sort_index]
+
+        name_select = []
+        spt_select = []
+        gk_select = []
+        ck_select = []
+        av_select = []
+        rv_select = []
+
+        for i, item in enumerate(name_list):
+            if item not in name_select:
+                name_select.append(item)
+                spt_select.append(spt_list[i])
+                gk_select.append(gk_list[i])
+                ck_select.append(ck_list[i])
+                av_select.append(av_list[i])
+                rv_select.append(rv_list[i])
+
+        print('Best-fitting spectra:')
+
+        for i in range(10):
+            print(f'   - G = {gk_select[i]:.2e} -> {name_select[i]}, {spt_select[i]}, '
+                  f'A_V = {av_select[i]:.2f}, RV = {rv_select[i]:.0f} km/s')
+
+        species_db = database.Database()
+
+        species_db.add_empirical(tag=tag,
+                                 names=name_select,
+                                 sptypes=spt_select,
+                                 goodness_of_fit=gk_select,
+                                 flux_scaling=ck_select,
+                                 av_ext=av_select,
+                                 rad_vel=rv_select,
+                                 object_name=self.object_name,
+                                 spec_name=self.spec_name,
+                                 spec_library=self.spec_library)

species/data/database.py

@@ -20,7 +20,8 @@
 from species.core import box
 from species.data import ames_cond, ames_dusty, atmo, blackbody, btcond, btcond_feh, btnextgen, \
                          btsettl, btsettl_cifist, companions, drift_phoenix, dust, exo_rem, \
-                         filters, irtf, isochrones, leggett, petitcode, spex, vega, vlm_plx
+                         filters, irtf, isochrones, leggett, petitcode, spex, vega, vlm_plx, \
+                         kesseli2017
 from species.read import read_calibration, read_filter, read_model, read_planck
 from species.util import data_util, dust_util, read_util
 
@@ -1117,7 +1118,7 @@ def add_spectrum(self,
         Parameters
         ----------
         spec_library : str
-            Spectral library ('irtf' or 'spex').
+            Spectral library ('irtf', 'spex', 'kesseli+2017').
         sptypes : list(str)
             Spectral types ('F', 'G', 'K', 'M', 'L', 'T'). Currently only implemented for 'irtf'.
 
@@ -1144,6 +1145,9 @@ def add_spectrum(self,
         elif spec_library[0:5] == 'spex':
             spex.add_spex(self.input_path, h5_file)
 
+        elif spec_library[0:12] == 'kesseli+2017':
+            kesseli2017.add_kesseli2017(self.input_path, h5_file)
+
         h5_file.close()
 
     @typechecked
@@ -1762,3 +1766,79 @@ def get_samples(self,
                               ln_prob=ln_prob,
                               prob_sample=prob_sample,
                               median_sample=median_sample)
+
+    @typechecked
+    def add_empirical(self,
+                      tag: str,
+                      names: List[str],
+                      sptypes: List[str],
+                      goodness_of_fit: List[float],
+                      flux_scaling: List[float],
+                      av_ext: List[float],
+                      rad_vel: List[float],
+                      object_name: str,
+                      spec_name: str,
+                      spec_library: str) -> None:
+        """
+        Parameters
+        ----------
+        tag : str
+            Database tag where the results will be stored.
+        names : list(str)
+            Array with the names of the empirical spectra.
+        sptypes : list(str)
+            Array with the spectral types of ``names``.
+        goodness_of_fit : list(float)
+            Array with the goodness-of-fit values.
+        flux_scaling : list(float)
+            Array with the best-fit scaling values to match the library spectra with the data.
+        av_ext : list(float)
+            Array with the visual extinctions A_V.
+        rad_vel : list(float)
+            Array with the radial velocities (km s-1).
+        object_name : str
+            Object name as stored in the database with
+            :func:`~species.data.database.Database.add_object` or
+            :func:`~species.data.database.Database.add_companion`.
+        spec_name : str
+            Name of the spectrum that is stored at the object data of ``object_name``.
+        spec_library : str
+            Name of the spectral library that was used for the empirical comparison.
+
+        Returns
+        -------
+        NoneType
+            None
+        """
+
+        with h5py.File(self.database, 'a') as h5_file:
+
+            if 'results' not in h5_file:
+                h5_file.create_group('results')
+
+            if 'results/empirical' not in h5_file:
+                h5_file.create_group('results/empirical')
+
+            if f'results/empirical/{tag}' in h5_file:
+                del h5_file[f'results/empirical/{tag}']
+
+            dtype = h5py.special_dtype(vlen=str)
+
+            dset = h5_file.create_dataset(f'results/empirical/{tag}/names',
+                                          (np.size(names), ), dtype=dtype)
+
+            dset[...] = names
+
+            dset.attrs['object_name'] = str(object_name)
+            dset.attrs['spec_name'] = str(spec_name)
+            dset.attrs['spec_library'] = str(spec_library)
+
+            dset = h5_file.create_dataset(f'results/empirical/{tag}/sptypes',
+                                          (np.size(sptypes), ), dtype=dtype)
+
+            dset[...] = sptypes
+
+            h5_file.create_dataset(f'results/empirical/{tag}/goodness_of_fit', data=goodness_of_fit)
+            h5_file.create_dataset(f'results/empirical/{tag}/flux_scaling', data=flux_scaling)
+            h5_file.create_dataset(f'results/empirical/{tag}/av_ext', data=av_ext)
+            h5_file.create_dataset(f'results/empirical/{tag}/rad_vel', data=rad_vel)