Type checks added in FitPlanck

species/analysis/fit_planck.py

@@ -6,6 +6,8 @@
 import math
 import warnings
 
+from typing import Optional, Union, List, Tuple, Dict
+
 from multiprocessing import Pool, cpu_count
 
 import emcee
@@ -18,13 +20,16 @@
 except:
     warnings.warn('PyMultiNest could not be imported.')
 
+from typeguard import typechecked
+
 from species.analysis import photometry
 from species.data import database
 from species.read import read_object, read_planck
 
 
-def lnprior(param,
-            bounds):
+@typechecked
+def lnprior(param: np.ndarray,
+            bounds: dict) -> np.float64:
     """
     Internal function for the prior probability.
 
@@ -56,13 +61,14 @@ def lnprior(param,
     return ln_prior
 
 
-def lnlike(param,
-           bounds,
-           objphot,
-           synphot,
-           distance,
-           spectrum,
-           n_planck):
+@typechecked
+def lnlike(param: np.ndarray,
+           bounds: dict,
+           objphot: list,
+           synphot: list,
+           distance: float,
+           spectrum: Optional[dict],
+           n_planck: int) -> np.float64:
     """
     Internal function for the likelihood function.
 
@@ -81,9 +87,10 @@ def lnlike(param,
         calculation of synthetic photometry from the model spectra.
     distance : float
         Distance (pc).
-    spectrum : numpy.ndarray, None
-        Spectrum array with the wavelength (um), flux (W m-2 um-1), and error
-        (W m-2 um-1). Not used if set to None.
+    spectrum : dict, None
+        Dictionary with the spectra, covariance matrix, inverse of the covariance matrix, and the
+        spectral resolution. The spectrum contains columns with wavelength (um), flux (W m-2 um-1),
+        and error (W m-2 um-1). Not used if set to ``None``.
     n_planck : int
         Number of Planck components.
 
@@ -146,13 +153,14 @@ def lnlike(param,
     return -0.5*chisq
 
 
-def lnprob(param,
-           bounds,
-           objphot,
-           synphot,
-           distance,
-           spectrum,
-           n_planck):
+@typechecked
+def lnprob(param: np.ndarray,
+           bounds: dict,
+           objphot: list,
+           synphot: list,
+           distance: float,
+           spectrum: Optional[dict],
+           n_planck: int) -> np.float64:
     """
     Internal function for the posterior probability.
 
@@ -171,9 +179,10 @@ def lnprob(param,
         calculation of synthetic photometry from the model spectra.
     distance : float
         Distance (pc).
-    spectrum : numpy.ndarray, None
-        Spectrum array with the wavelength (um), flux (W m-2 um-1), and error
-        (W m-2 um-1). Not used if set to None.
+    spectrum : dict, None
+        Dictionary with the spectra, covariance matrix, inverse of the covariance matrix, and the
+        spectral resolution. The spectrum contains columns with wavelength (um), flux (W m-2 um-1),
+        and error (W m-2 um-1). Not used if set to ``None``.
     n_planck : int
         Number of Planck components.
 
@@ -210,24 +219,27 @@ class FitPlanck:
     temperatures and radii to decrease and increase, respectively.
     """
 
+    @typechecked
     def __init__(self,
-                 object_name,
-                 filters,
-                 bounds,
-                 inc_phot=True,
-                 inc_spec=True):
+                 object_name: str,
+                 filters: Optional[List[str]],
+                 bounds: Union[Dict[str, Tuple[float, float]],
+                               Dict[str, List[Tuple[float, float]]]],
+                 inc_phot: bool = True,
+                 inc_spec: bool = True) -> None:
         """
         Parameters
         ----------
         object_name : str
             Object name in the database.
-        filters : tuple(str, )
+        filters : tuple(str, ), None
             Filter names for which the photometry is selected. All available photometric data of
-            the object are used if set to None.
+            the object are used if set to ``None``.
         bounds : dict
             Parameter boundaries for 'teff' and 'radius'. The values should be provided either as
-            float or as list of floats such that multiple Planck functions can be combined,
-            e.g. ``{'teff': [(1000., 2000.), (500., 1500.)], 'radius': [(0.5, 1.5), (1.5, 2.0)]}``.
+            tuple (with two float) or as list of tuples (with two floats) such that multiple Planck
+            functions can be combined, e.g. ``{'teff': [(1000., 2000.), (500., 1500.)],
+            'radius': [(0.5, 1.5), (1.5, 2.0)]}``.
         inc_phot : bool
             Include photometric data with the fit.
         inc_spec : bool
@@ -292,11 +304,12 @@ def __init__(self,
         else:
             self.spectrum = None
 
+    @typechecked
     def run_mcmc(self,
-                 tag,
-                 guess,
-                 nwalkers=200,
-                 nsteps=1000):
+                 tag: str,
+                 guess: Optional[Union[Dict[str, float], Dict[str, List[float]]]],
+                 nwalkers: int = 200,
+                 nsteps: int = 1000) -> None:
         """
         Function to run the MCMC sampler.
 
@@ -377,10 +390,11 @@ def run_mcmc(self,
                                distance=self.distance[0],
                                spec_labels=None)
 
+    @typechecked
     def run_multinest(self,
-                      tag,
-                      n_live_points=4000,
-                      output='multinest/'):
+                      tag: str,
+                      n_live_points: int = 4000,
+                      output: str = 'multinest/') -> None:
         """
         Function to run the ``PyMultiNest`` wrapper of the ``MultiNest`` sampler. While
         ``PyMultiNest`` can be installed with ``pip`` from the PyPI repository, ``MultiNest``
@@ -424,7 +438,10 @@ def run_multinest(self,
         for i, item in enumerate(self.modelpar):
             cube_index[item] = i
 
-        def lnprior_multinest(cube, n_dim, n_param):
+        @typechecked
+        def lnprior_multinest(cube,
+                              n_dim: int,
+                              n_param: int) -> None:
             """
             Function to transform the unit cube into the parameter cube. It is not clear how to
             pass additional arguments to the function, therefore it is placed here.
@@ -462,7 +479,10 @@ def lnprior_multinest(cube, n_dim, n_param):
                         (self.bounds[f'radius_{i}'][1]-self.bounds[f'radius_{i}'][0]) * \
                         cube[cube_index[f'radius_{i}']]
 
-        def lnlike_multinest(cube, n_dim, n_param):
+        @typechecked
+        def lnlike_multinest(cube,
+                             n_dim: int,
+                             n_param: int) -> np.float64:
             """
             Function for the logarithm of the likelihood, computed from the parameter cube.
 

species/data/companions.py

@@ -104,9 +104,9 @@ def get_data():
                                      'Paranal/SPHERE.IRDIS_D_H23_3': (17.95, 0.17),  # Keppler et al. 2018
                                      'Paranal/SPHERE.IRDIS_D_K12_1': [(16.78, 0.31), (16.72, 0.05)],  # Stolker et al. in prep.
                                      'Paranal/SPHERE.IRDIS_D_K12_2': [(16.23, 0.32), (16.38, 0.06)], # Stolker et al. in prep.
-                                     'Paranal/NACO.Lp': (14.08, 0.33),  # Stolker et al. in prep.
-                                     'Paranal/NACO.NB405': (13.91, 0.34),  # Stolker et al. in prep.
-                                     'Paranal/NACO.Mp': (13.64, 0.22),  # Stolker et al. in prep.
+                                     # 'Paranal/NACO.Lp': (14.08, 0.33),  # Stolker et al. in prep.
+                                     # 'Paranal/NACO.NB405': (13.91, 0.34),  # Stolker et al. in prep.
+                                     # 'Paranal/NACO.Mp': (13.64, 0.22),  # Stolker et al. in prep.
                                      'Keck/NIRC2.Lp': (14.64, 0.18)}},  # Wang et al. 2020
 
             'PDS 70 c': {'distance': (113.43, 0.52),

species/plot/plot_mcmc.py

@@ -168,17 +168,35 @@ def plot_posterior(tag,
 
     print(f'Plotting the posterior: {output}...', end='', flush=True)
 
-    if inc_luminosity and 'teff' in box.parameters and 'radius' in box.parameters:
+    if inc_luminosity:
         ndim += 1
 
-        teff_index = np.argwhere(np.array(box.parameters) == 'teff')[0]
-        radius_index = np.argwhere(np.array(box.parameters) == 'radius')[0]
+        if 'teff' in box.parameters and 'radius' in box.parameters:
+            teff_index = np.argwhere(np.array(box.parameters) == 'teff')[0]
+            radius_index = np.argwhere(np.array(box.parameters) == 'radius')[0]
 
-        luminosity = 4. * np.pi * (samples[..., radius_index]*constants.R_JUP)**2 * \
-            constants.SIGMA_SB * samples[..., teff_index]**4. / constants.L_SUN
+            luminosity = 4. * np.pi * (samples[..., radius_index]*constants.R_JUP)**2 * \
+                constants.SIGMA_SB * samples[..., teff_index]**4. / constants.L_SUN
 
-        samples = np.append(samples, np.log10(luminosity), axis=-1)
-        box.parameters.append('luminosity')
+            samples = np.append(samples, np.log10(luminosity), axis=-1)
+            box.parameters.append('luminosity')
+
+        elif 'teff_0' in box.parameters and 'radius_0' in box.parameters:
+            luminosity = 0.
+
+            for i in range(100):
+                teff_index = np.argwhere(np.array(box.parameters) == f'teff_{i}')
+                radius_index = np.argwhere(np.array(box.parameters) == f'radius_{i}')
+
+                if len(teff_index) > 0 and len(radius_index) > 0:
+                    luminosity += 4. * np.pi * (samples[..., radius_index[0]]*constants.R_JUP)**2 \
+                        * constants.SIGMA_SB * samples[..., teff_index[0]]**4. / constants.L_SUN
+
+                else:
+                    break
+
+            samples = np.append(samples, np.log10(luminosity), axis=-1)
+            box.parameters.append('luminosity')
 
     labels = plot_util.update_labels(box.parameters)
 

species/read/read_model.py

@@ -490,19 +490,19 @@ def get_model(self,
                              f'interpolated at {model_param} because zeros are stored in the '
                              f'database.')
 
-        spec_box = box.create_box(boxtype='model',
-                                  model=self.model,
-                                  wavelength=wavelength,
-                                  flux=flux[is_finite],
-                                  parameters=model_param,
-                                  quantity=quantity)
-
-        if 'radius' in spec_box.parameters:
-            spec_box.parameters['luminosity'] = 4. * np.pi * (spec_box.parameters['radius'] * \
-                constants.R_JUP)**2 * constants.SIGMA_SB * spec_box.parameters['teff']**4. / \
+        model_box = box.create_box(boxtype='model',
+                                   model=self.model,
+                                   wavelength=wavelength,
+                                   flux=flux[is_finite],
+                                   parameters=model_param,
+                                   quantity=quantity)
+
+        if 'radius' in model_box.parameters:
+            model_box.parameters['luminosity'] = 4. * np.pi * (model_box.parameters['radius'] * \
+                constants.R_JUP)**2 * constants.SIGMA_SB * model_box.parameters['teff']**4. / \
                 constants.L_SUN  # (Lsun)
 
-        return spec_box
+        return model_box
 
     def get_data(self,
                  model_param):
@@ -588,19 +588,19 @@ def get_data(self,
 
         h5_file.close()
 
-        spec_box = box.create_box(boxtype='model',
-                                  model=self.model,
-                                  wavelength=wl_points,
-                                  flux=flux,
-                                  parameters=model_param,
-                                  quantity='flux')
+        model_box = box.create_box(boxtype='model',
+                                   model=self.model,
+                                   wavelength=wl_points,
+                                   flux=flux,
+                                   parameters=model_param,
+                                   quantity='flux')
 
-        if 'radius' in spec_box.parameters:
-            spec_box.parameters['luminosity'] = 4. * np.pi * (spec_box.parameters['radius'] * \
-                constants.R_JUP)**2 * constants.SIGMA_SB * spec_box.parameters['teff']**4. / \
+        if 'radius' in model_box.parameters:
+            model_box.parameters['luminosity'] = 4. * np.pi * (model_box.parameters['radius'] * \
+                constants.R_JUP)**2 * constants.SIGMA_SB * model_box.parameters['teff']**4. / \
                 constants.L_SUN  # (Lsun)
 
-        return spec_box
+        return model_box
 
     def get_flux(self,
                  model_param,

species/read/read_planck.py

@@ -165,12 +165,19 @@ def get_spectrum(self,
                                     model_param[f'teff_{i}'],
                                     scaling)  # (W m-2 um-1)
 
-        return box.create_box(boxtype='model',
-                              model='planck',
-                              wavelength=wavel_points,
-                              flux=flux,
-                              parameters=model_param,
-                              quantity='flux')
+        model_box = box.create_box(boxtype='model',
+                                   model='planck',
+                                   wavelength=wavel_points,
+                                   flux=flux,
+                                   parameters=model_param,
+                                   quantity='flux')
+
+        if 'radius' in model_box.parameters:
+            model_box.parameters['luminosity'] = 4. * np.pi * (model_box.parameters['radius'] * \
+                constants.R_JUP)**2 * constants.SIGMA_SB * model_box.parameters['teff']**4. / \
+                constants.L_SUN  # (Lsun)
+
+        return model_box
 
     def get_flux(self,
                  model_param,

species/util/plot_util.py

@@ -135,7 +135,7 @@ def update_labels(param):
 
     if 'radius' in param:
         index = param.index('radius')
-        param[index] = r'$R$ ($\mathregular{R_{Jup}}$)'
+        param[index] = r'$R$ ($\mathregular{R_{J}}$)'
 
     if 'luminosity' in param:
         index = param.index('luminosity')
@@ -169,6 +169,22 @@ def update_labels(param):
         elif item[0:6] == 'error_':
             param[i] = rf'$b_\mathregular{{{item[6:]}}}$'
 
+    for i in range(100):
+        if f'teff_{i}' in param:
+            index = param.index(f'teff_{i}')
+            param[index] = rf'$T_\mathregular{{{i+1}}}$ (K)'
+
+        else:
+            break
+
+    for i in range(100):
+        if f'radius_{i}' in param:
+            index = param.index(f'radius_{i}')
+            param[index] = rf'$R_\mathregular{{{i+1}}}$ ' + r'($\mathregular{R_{J}}$)'
+
+        else:
+            break
+
     return param
 
 
@@ -255,7 +271,7 @@ def quantity_unit(param,
         quantity.append('radius')
 
         if object_type == 'planet':
-            unit.append(r'$R_\mathregular{Jup}}$')
+            unit.append(r'$R_\mathregular{J}}$')
         elif object_type == 'star':
             unit.append(r'$R_\mathregular{\odot}}$')
 
@@ -270,7 +286,7 @@ def quantity_unit(param,
         quantity.append('mass')
 
         if object_type == 'planet':
-            unit.append(r'$M_\mathregular{Jup}$')
+            unit.append(r'$M_\mathregular{J}$')
         elif object_type == 'star':
             unit.append(r'$M_\mathregular{\odot}$')