Source code for gwtoolbox.tools_earth

from astropy import units as u
import numpy as np
from .sources_kHz import *
from .detectors_earth import *
from .functions_earth import *
from .cosmology import Cosmology
from .constants import *
from .parameters import *

def set_cosmology(cosmoID=None, H0=None, Om0=None):
    cosmos_class = Cosmology()
    cosmos = cosmos_class.set_cosmo(cosmoID,H0,Om0)
    return cosmos

class Tools:
    """
    This is a class with tools to manipulate GW detectors and sources on Earth.
    
    Args:
      detID (int): detector id
      popID (int): population id
      cosmos (class): cosmological model
      det_setup (Optional[list of floats]): parameters to modify the detector noise
 
    Attributes:
      detector (class): detector setup class
      ant_pat (array of dtype float): detector shape
      noise (array of dtype float): detector noise
      
    """
  
    def __init__(self, detID, popID, cosmos, det_setup=None):
        """
        Parameters:
          detID (int): detector id
          popID (int): population id
          cosmos (class): cosmological model
          det_setup (Optional[list of dtype floats]): parameters to modify the detector noise
        """
        self.cosmos = cosmos
        self.detID = detID
        self.popID = popID
        if self.detID in [-1,0,1]:
            self.detector = LigoLike(self.detID)
        elif self.detID in [-2,2]:
            self.detector = ETLike(self.detID)
        self.ant_pat = self.detector.ante_pattern
        self.noise = self.detector.noise_curve(det_setup)

        if self.popID == 0:
            self.population_class = BHB(self.cosmos)
            z_ranges = [1.e-5,3.]
            m1_ranges = [5.,50.]
            m2_ranges = m1_ranges
            chi_ranges = [-0.5,0.5]
            self.param_ranges = [z_ranges,m1_ranges,m2_ranges,chi_ranges]
            self.param_init = [0.1,20.,20.,0.]
        elif self.popID == 1:
            self.population_class = DNS(self.cosmos)
            z_ranges = [1.e-5,1.]
            m1_ranges = [0.5,2.]
            m2_ranges = m1_ranges
            chi_ranges = [-0.5,0.5]
            self.param_ranges = [z_ranges,m1_ranges,m2_ranges,chi_ranges]
            self.param_init = [0.1,1.,1.,0.]
        elif self.popID == 2:
            self.population_class = BHNS(self.cosmos)
            z_ranges = [1.e-5,1.]
            m1_ranges = [5.,12.]
            m2_ranges = [0.5,2.]
            chi_ranges = [-0.5,0.5]
            self.param_ranges = [z_ranges,m1_ranges,m2_ranges,chi_ranges]
            self.param_init = [0.1,10.,1.,0.]

    def total_number(self, time_obs=TIME_OBS_EARTH, rho_cri=RHO_CRIT_EARTH):
        """
        Return the expected number of detections.
        
        Parameters:
          time_obs (float): Observation time, in unit of minute
          rho_cri (float): The SNR threshold of detection
          
        Returns:
          (float): expected number of detections
        """
        self.ndet = self.population_class.tot_num(time_obs,rho_cri,self.ant_pat,self.noise,self.param_ranges)
        return self.ndet

    def list_params(self, time_obs=TIME_OBS_EARTH, rho_cri=RHO_CRIT_EARTH, size=None):
        """
        Return an array of final parameters.
        
        Parameters:
          time_obs (float): Observation time, in unit of minute
          rho_cri (float): The SNR threshold of detection
          size (Optional[int]): sample size
          
        Returns:
          (list of arrays of dtype float): samples
        """
        if size == None:
            self.number = SAMPLE_SIZE
        else:
            self.number = size

        self.samples = self.population_class.MCMCsample(self.number,self.param_init,self.param_ranges,time_obs,rho_cri,self.ant_pat,self.noise)

        return self.samples
    
    def list_param_errors(self):
        """
        Return array of final parameter errors.
        It can be used only after list_params is used.
        
        Returns:
          (list of arrays of dtype float): errors m1,m2,chi
        """
        errors = self.population_class.errors_FIM(self.number, self.samples, self.noise)

        return errors