class CosmicPopulation(Population):
    """Generate a cosmic FRB population."""

    def __init__(self,
                 n_srcs=1e4,
                 n_days=1,
                 name='cosmic',
                 repeaters=False,
                 generate=False):
        """Generate a popuation of FRBs.

        Args:
            n_srcs (int): Number of FRB sources to generate.
            n_days (float): Number of days over which FRBs are to be generated.
            name (str): Population name.
            repeaters (bool): Whether to generate a repeater population.
            generate (bool): Whether to create a population.

        Returns:
            Population: Population of FRBs

        """
        # Set up general population arguments
        Population.__init__(self)

        # Give population a name
        self.name = 'cosmic'
        if name:
            self.name = name

        # Set population arguments
        self.n_srcs = int(n_srcs)
        self.n_days = n_days
        self.repeaters = repeaters
        self.shape = (self.n_srcs,)

        # If wanting repeaters
        if self.repeaters:
            self.set_time()

        # Set up default models
        self.set_emission_range()
        self.set_dist()
        self.set_direction()
        self.set_lum()
        self.set_si()
        self.set_w()
        self.set_dm_mw()
        self.set_dm_igm()
        self.set_dm_host()
        self.set_dm()

        # Whether to start generating a Cosmic Population
        if generate:
            self.generate()

    def gen_index(self):
        """Generate indices for over each FRB source."""
        self.frbs.index = np.arange(self.n_srcs)

    def set_emission_range(self, low=100e6, high=10e9):
        """Set the emission range [Hz].

        The frequency range between which FRB sources should emit the given
        bolometric luminosity.

        Args:
            f_min (float): Lowest source frequency [Hz].
            f_max (float): Highest source frequency [Hz].
        """
        self.f_min = low
        self.f_max = high

    def gen_precalc(self):
        """Check whether pre-calculations have been run."""
        pc.DistanceTable(H_0=self.H_0, W_m=self.W_m, W_v=self.W_v)

    def set_dist(self, model='vol_co', **kwargs):
        """Set the number density model for calculating distances.

        Args:
            model (str): Number density model to use. Choice from
                ('vol_co', 'sfr', 'smd').
            z_max (float): Maximum redshift.
            H_0 (float): Hubble constant.
            W_m (float): Density parameter Ω_m.
            W_v (float): Cosmological constant Ω_Λ.
            alpha (float): Desired log N log S slope for a perfect,
                non-cosmological population.
        """
        # Option to use your own model
        if not isinstance(model, str):
            self.dist_func = lambda: model(**kwargs)
            return

        # I sometimes use 'constant' instead of 'vol_co'
        if model == 'constant':
            model = 'vol_co'

        # Check whether recognised number density model
        if model not in ['vol_co', 'sfr', 'smd']:
            raise ValueError('set_dist input not recognised')

        # Set number density model
        # Don't fear the lambda, merely delays executing the function
        self.dist_func = lambda: NumberDensity(model=model, **kwargs)

    def gen_dist(self):
        """Generate source distances."""
        n_model = self.dist_func()
        self.vol_co_max = n_model.vol_co_max
        self.frbs.z, self.frbs.dist_co = n_model.draw(self.n_srcs)

    def gen_gal_coords(self):
        """Generate galactic coordinates."""
        frbs = self.frbs
        # Get the proper distance
        dist_pr = frbs.dist_co/(1+frbs.z)
        # Convert into galactic coordinates
        frbs.gx, frbs.gy, frbs.gz = go.lb_to_xyz(frbs.gl, frbs.gb, dist_pr)

    def set_direction(self, model='uniform', **kwargs):
        """Set the model for generating the directions of the frb sources.

        Args:
            model (str): Choice from ('uniform').
        if model == 'uniform':
            min_ra (float): Minimum right ascenion [frac deg].
            max_ra (float): Maximum right ascenion [frac deg].
            min_dec (float): Minimum declination [frac deg].
            max_dec (float): Maximum declination [frac deg].
        """
        # Use your own function
        if not isinstance(model, str):
            self.direction_func = lambda: model(**kwargs)
            return

        # Or use a uniform distribution
        if model == 'uniform':
            self.direction_func = lambda: did.uniform(n_srcs=self.n_srcs,
                                                      **kwargs)
        else:
            raise ValueError('set_direction input not recognised')

    def gen_direction(self):
        """Generate the direction of frbs."""
        frbs = self.frbs
        # Calculate right ascenion and declination
        frbs.ra, frbs.dec = self.direction_func()
        # Convert to galactic lat/long coordinates
        frbs.gl, frbs.gb = go.radec_to_lb(frbs.ra, frbs.dec, frac=True)

    def set_dm_mw(self, model='ne2001', **kwargs):
        """Set the model for the Milky Way dispersion measure.

        Args:
            model (str): Option of 'ne2001'.
        """
        if not isinstance(model, str):
            self.dm_mw_func = lambda: model(**kwargs)
            return

        # Distribution from which to draw dm_mw
        if model == 'ne2001':
            self.dm_mw_func = lambda: pc.NE2001Table().lookup(self.frbs.gl,
                                                              self.frbs.gb)
        else:
            raise ValueError('set_dm_mw input not recognised')

    def gen_dm_mw(self):
        """Generate Milky Way dispersion measure."""
        self.frbs.dm_mw = self.dm_mw_func()

    def set_dm_igm(self, model='ioka', **kwargs):
        """Set intergalactic dispersion measure model.

        Args:
            model (str): Option of 'ioka'.
        if model == 'ioka':
            slope (float): Slope of the DM-z relationship.
            std (float): Spread around the DM-z relationship.
            spread_dist (str): 'normal' or 'lognormal'.
        """
        # Possibility to use your own function
        if not isinstance(model, str):
            self.dm_igm_func = lambda: model(**kwargs)
            return

        # Distribution from which to draw intergalactic dm
        if model == 'ioka':
            self.dm_igm_func = lambda: dmd.ioka(z=self.frbs.z, **kwargs)
        else:
            raise ValueError('set_dm_igm input not recognised')

    def gen_dm_igm(self):
        """Generate intergalactic dispersion measure."""
        self.frbs.dm_igm = self.dm_igm_func()

    def set_dm_host(self, model='gauss', **kwargs):
        """Set host galaxy dispersion measure.

        Args:
            model (str): Options from ('gauss', 'lognormal').

        if model in ('gauss', 'lognormal'):
            mean (float): Mean DM [pc/cm^3].
            std (float): Standard deviation DM [pc/cm^3].
        if model == 'constant':
            value (float): Value to adopt [pc/cm^3].
        """
        if not isinstance(model, str):
            self.dm_host_func = lambda: model(**kwargs)
            return

        # Distribution from which to draw host dispersion measure
        if model.startswith('gauss'):
            self.dm_host_func = lambda: dmd.gauss(z=self.frbs.z,
                                                  n_srcs=self.n_srcs,
                                                  **kwargs)
        elif model == 'lognormal':
            self.dm_host_func = lambda: dmd.lognormal(z=self.frbs.z,
                                                      n_srcs=self.n_srcs,
                                                      **kwargs)
        elif model == 'constant':
            self.dm_host_func = lambda: dmd.constant(n_srcs=self.n_srcs,
                                                     **kwargs)
        else:
            raise ValueError('set_dm_host input not recognised')

    def gen_dm_host(self):
        """Generate host dispersion measure."""
        self.frbs.dm_host = self.dm_host_func()

    def set_dm(self, mw=True, igm=True, host=True):
        """Set total dispersion measure.

        Args:
            mw (bool): Whether to include a Milky Way component
            igm (bool): Whether to include an IGM component
            host (bool): Whether to include a host galaxy component
        """
        # Which components to include
        self.dm_components = []
        if mw:
            self.dm_components.append(self.gen_dm_mw)
        if igm:
            self.dm_components.append(self.gen_dm_igm)
        if host:
            self.dm_components.append(self.gen_dm_host)

        # Save those components to execute at a later stage
        def run_dm():
            [c() for c in self.dm_components]
            return self.frbs.dm_mw + self.frbs.dm_igm + self.frbs.dm_host

        self.dm_func = run_dm

    def gen_dm(self):
        """Generate total dispersion measure."""
        self.frbs.dm = self.dm_func()

    def set_w(self, model='uniform', per_source='same', **kwargs):
        """Set intrinsic pulse widths model [ms].

        Args:
            model (str): Options from ('uniform', 'lognormal')
            per_source (str): Model for a single source burst
                distribution. Options from 'same' or 'different'
        If model == 'constant':
            value (float): Pulse width [ms].
        If model == 'uniform':
            low (float): Minimum pulse width [ms].
            high (float): Maximum pulse width [ms].
        If model == 'lognormal':
            mean (float): Mean pulse width [ms].
            std (float): Standard deviation pulse width [ms].

        """
        # Each burst from the same source: same or different widths?
        if per_source == 'same':
            self.w_shape = lambda: self.n_srcs
        elif per_source == 'different':
            self.w_shape = lambda: self.shape

        # Distribution from which to draw pulse widths

        # Find available distributions to draw from
        funcs = [d for d in dir(wd) if hasattr(getattr(wd, d), '__call__')]
        funcs.remove('calc_w_arr')

        # Set function
        if model in funcs:
            func = getattr(wd, model)

            # If you're getting fancy with combined distributions
            # See examples/adapting_population_parameters.py
            self._transpose_w = False
            for kw_value in kwargs.values():
                if isinstance(kw_value, (list, np.ndarray)):
                    self.w_shape = lambda: self.shape[::-1]
                    self._transpose_w = True

            self.w_func = lambda x: func(shape=x, z=self.frbs.z, **kwargs)
        else:
            raise ValueError('set_w input model not recognised')

    def gen_w(self):
        """Generate pulse widths [ms]."""
        shape = self.w_shape()
        self.frbs.w_int, self.frbs.w_arr = self.w_func(shape)

        # From combined distribution inputs
        if self._transpose_w:
            self.frbs.w_int = self.frbs.w_int.T
            self.frbs.w_arr = self.frbs.w_arr.T

    def set_si(self, model='gauss', per_source='same', **kwargs):
        """Set spectral index model.

        Args:
            model (str): Options from ('gauss')
            per_source (str): Model for a single source burst
                distribution. Options from ('same', 'different')
        If model == 'constant':
            value (float): Default spectal index.
        If model == 'gauss':
            mean (float): Mean spectral index
            std (float): Standard deviation spectral index

        """
        # Each burst from the same source: same or different si?
        if per_source == 'same':
            self.si_shape = lambda: self.n_srcs
        elif per_source == 'different':
            self.si_shape = lambda: self.shape

        # Find available distributions to draw from
        funcs = [d for d in dir(sid) if hasattr(getattr(sid, d), '__call__')]

        # Set function
        if model in funcs:
            func = getattr(sid, model)

            # If you're getting fancy with combined distributions
            self._transpose_si = False
            for kw_value in kwargs.values():
                if isinstance(kw_value, (list, np.ndarray)):
                    self.si_shape = lambda: self.shape[::-1]
                    self._transpose_si = True

            # Distribution from which to draw spectral indices
            self.si_func = lambda x: func(shape=x, **kwargs)
        else:
            raise ValueError('set_si input not recognised')

    def gen_si(self):
        """Generate spectral indices."""
        shape = self.si_shape()
        self.frbs.si = self.si_func(shape)

        if self._transpose_si:
            self.frbs.si = self.frbs.si.T

    def set_lum(self, model='powerlaw', per_source='same', **kwargs):
        """Set luminosity function [ergs/s].

        Args:
            model (str): Options from ('powerlaw')
            per_source (str): Model for a single source burst
                distribution. Options from ('same', 'different')

        If model == 'powerlaw':
            low (float): Minimum bolometric luminosity [ergs/s]
            high (float): Maximum bolometric luminosity [ergs/s]
            power (float): Power of luminosity function

        If model == 'constant':
            value (float): Value for standard candle [ergs/s]

        """
        # Each burst from the same source: same or different luminosities?
        if per_source == 'same':
            self.lum_shape = lambda: self.n_srcs
        elif per_source == 'different':
            self.lum_shape = lambda: self.shape

        # Find available distributions to draw from
        funcs = [d for d in dir(ld) if hasattr(getattr(ld, d), '__call__')]

        # Set function
        if model in funcs:
            func = getattr(ld, model)

            # Help out the user
            for s in ['slope', 'index']:
                if s in kwargs:
                    kwargs['power'] = kwargs.pop(s)

            # If you're getting fancy with combined distributions
            self._transpose_lum = False
            for kw_value in kwargs.values():
                if isinstance(kw_value, (list, np.ndarray)):
                    self.lum_shape = lambda: self.shape[::-1]
                    self._transpose_lum = True

            # Distribution from which to draw luminosities
            self.lum_func = lambda x: func(shape=x, **kwargs)
        else:
            raise ValueError('set_lum input not recognised')

    def gen_lum(self):
        """Generate luminosities [ergs/s]."""
        shape = self.lum_shape()
        self.frbs.lum_bol = self.lum_func(shape)

        # You need multiple luminosities if repeaters
        if self.repeaters and self.frbs.lum_bol.ndim == 1:
            repeat_lums = [self.frbs.lum_bol[..., np.newaxis]]*self.shape[1]
            self.frbs.lum_bol = np.concatenate(repeat_lums, axis=1)

        if self._transpose_lum:
            self.frbs.lum_bol = self.frbs.lum_bol.T

    def set_time(self, model='regular', **kwargs):
        """Set model from which to generate time stamps.

        Args:
            model (str): Options from ('single', 'regular', 'clustered',
                'poisson', 'cyclic')
        If model == 'regular':
            rate (float): Number of bursts per day
        If model == 'poisson':
            rate (float): Expected number of bursts per day
        If model == 'clustered':
            r (float): Rate parameter
            k (float): Shape parameter
        If model == 'cyclic':
            rate (float): Number of bursts per day
            period (float): Period of activity cycle (days)
            frac (float): Fraction of activity cycle a source is active
        """
        if not isinstance(model, str):
            # These lambda functions look complex, but aren't.
            # They merely stop the function from running immediately
            self.time_func = lambda: model(**kwargs)
            return

        # Find available distributions
        funcs = [d for d in dir(td) if hasattr(getattr(td, d), '__call__')]
        internal = ['gamma', 'iteratively_gen_times', '_weibull_dist',
                    '_poisson_dist']
        for f in internal:
            funcs.remove(f)

        # Set function
        if model in funcs:
            func = getattr(td, model)

            # Distribution from which to draw time stamps
            self.time_func = lambda: func(n_srcs=self.n_srcs,
                                          n_days=self.n_days,
                                          z=self.frbs.z,
                                          **kwargs)
        else:
            raise ValueError('set_time input not recognised')

    def gen_time(self):
        """Generate time stamps."""
        # Only relevant for repeaters
        if not self.repeaters:
            return

        pprint('Adding burst times')
        self.frbs.time = self.time_func()
        # Set size for all other parameters
        self.shape = self.frbs.time.shape
        pprint('Finished adding burst times')

    def generate(self):
        """Generate a full CosmicPopulation."""
        pprint(f'Generating {self.name} population')
        self.gen_index()
        self.gen_dist()
        self.gen_time()
        self.gen_direction()
        self.gen_gal_coords()
        self.gen_dm()
        self.gen_w()
        self.gen_lum()
        self.gen_si()
        pprint(f'Finished generating {self.name} population')

    @classmethod
    def simple(cls, n_srcs, n_days=1, repeaters=False, generate=False):
        """Set up a simple, local population."""
        pop = cls(n_srcs=n_srcs, n_days=n_days, name='simple',
                  repeaters=repeaters, generate=False)
        pop.set_dist(model='vol_co', z_max=0.01, alpha=-1.5,
                     H_0=67.74, W_m=0.3089, W_v=0.6911)
        pop.set_dm(mw=False, igm=False, host=False)
        pop.set_emission_range(low=10e7, high=10e9)
        pop.set_lum(model='constant', value=1e38)
        pop.set_w(model='constant', value=10)
        pop.set_si(model='constant', value=0)
        if pop.repeaters:
            pop.set_time(model='regular', rate=2)
        if generate:
            pop.generate()
        return pop

    @classmethod
    def complex(cls, n_srcs, n_days=1, repeaters=False, generate=False):
        """Set up a complex population."""
        pop = cls(n_srcs=n_srcs, n_days=n_days, name='complex',
                  repeaters=repeaters, generate=False)
        pop.set_dist(model='vol_co', z_max=1, alpha=-1.5,
                     H_0=67.74, W_m=0.3089, W_v=0.6911)
        pop.set_dm_host(model='gauss', mean=100, std=200)
        pop.set_dm_igm(model='ioka', slope=1000, std=None)
        pop.set_dm_mw(model='ne2001')
        pop.set_dm(mw=True, igm=True, host=True)
        pop.set_emission_range(low=10e7, high=10e9)
        pop.set_lum(model='powerlaw', low=1e40, high=1e45, power=0)
        pop.set_w(model='lognormal', mean=0.1, std=1)
        pop.set_si(model='gauss', mean=-1.4, std=1)
        if pop.repeaters:
            pop.set_time(model='poisson', rate=9)
        if generate:
            pop.generate()
        return pop

    @classmethod
    def optimal(cls, n_srcs, n_days=1, repeaters=False, generate=False):
        """Set up an optimal population."""
        pop = cls(n_srcs=n_srcs, n_days=n_days, name='optimal',
                  repeaters=repeaters, generate=False)
        pop.set_dist(model='vol_co', z_max=2.5, alpha=-2.2,
                     H_0=67.74, W_m=0.3089, W_v=0.6911)
        pop.set_dm_host(model='constant', value=50)
        pop.set_dm_igm(model='ioka', slope=1000, std=None)
        pop.set_dm_mw(model='ne2001')
        pop.set_dm(mw=True, igm=True, host=True)
        pop.set_emission_range(low=10e7, high=10e9)
        pop.set_lum(model='powerlaw', low=1e40, high=1e45, power=-0.8)
        pop.set_w(model='lognormal', mean=6.3e-3, std=.6)
        pop.set_si(model='constant', value=-0.4)
        if pop.repeaters:
            pop.set_time(model='poisson', rate=9)
        if generate:
            pop.generate()
        return pop class LargePopulation:
    """Method to deal with a large population.

    Basically splits it up, and merges the results.
    """

    def __init__(self, pop, *surveys, max_size=1e6, run=True):
        """Set arguments."""
        self.pop = pop
        self.base_name = pop.name
        self.surveys = surveys
        self.max_size = int(max_size)
        self.uids = None

        if run:
            self.run()
            self.merge()

    def run(self):
        """Run the generating and surveying of a large population."""
        pprint(f'Running a large {self.base_name} population')
        d = divmod(self.pop.n_srcs, self.max_size)
        sizes = [self.max_size for i in range(d[0])]
        if d[1] != 0:
            sizes.append(d[1])
        self.uids = [str(uuid.uuid4())[:8] for s in sizes]

        for i, n in enumerate(tqdm(sizes, desc='Subpopulations')):
            pop = self.pop
            pop.n_srcs = n
            pop.uid = self.uids[i]
            pop.generate()

            for surv in self.surveys:
                surv_pop = SurveyPopulation(pop, surv, scale_by_area=False)
                surv_pop.uid = pop.uid
                surv_pop.save()

    def merge(self):
        """Merge populations."""
        pprint('Merging populations')

        self.pops = []

        for s in self.surveys:
            pops = []
            files = [f'{self.base_name}_{s.name}_{uid}' for uid in self.uids]
            for f in files:
                pops.append(unpickle(f))

            # Main population
            mp = pops[0]

            # Merge each parameter
            for attr in mp.frbs.__dict__.keys():
                parm = getattr(mp.frbs, attr)
                if type(parm) is np.ndarray:
                    parms = []
                    for pop in pops:
                        parms.append(getattr(pop.frbs, attr))

                    try:
                        merged_parm = np.concatenate(parms, axis=0)
                    except ValueError:
                        # Check maximum size values should be padded to
                        max_size = max([p.shape[1] for p in parms])
                        new_parms = []

                        # Ensure matrices are the same shapes by padding them
                        for p in parms:
                            if p.shape[1] != max_size:
                                padded_p = np.zeros((p.shape[0], max_size))
                                padded_p[:] = np.nan
                                padded_p[:, :p.shape[1]] = p
                                new_parms.append(padded_p)
                            else:
                                new_parms.append(p)

                        merged_parm = np.concatenate(new_parms, axis=0)

                    setattr(mp.frbs, attr, merged_parm)

            # Add up detections
            for pop in pops[1:]:
                mp.source_rate.faint += pop.source_rate.faint
                mp.source_rate.late += pop.source_rate.late
                mp.source_rate.out += pop.source_rate.out
                mp.source_rate.det += pop.source_rate.det

                if mp.repeaters:
                    mp.burst_rate.faint += pop.burst_rate.faint
                    mp.burst_rate.late += pop.burst_rate.late
                    mp.burst_rate.out += pop.burst_rate.out
                    mp.burst_rate.det += pop.burst_rate.det
                    mp.burst_rate.pointing += pop.burst_rate.pointing

            # Recalculate detection rates
            mp.calc_rates(s)

            # Save the main population as one big population
            mp.uid = None
            mp.save()

            # Remove all of the smaller ones
            for f in files:
                p = paths.populations() + f'{f}.p'
                os.remove(p)

            self.pops.append(mp) def hist(parameter, bin_type='lin', n_bins=25, norm='max', edges=True,
         bins=None):
    """Bin up a parameter either in a lin or log space.

    Why is this not a standard option in numpy or matplotlib?

    Args:
        parameter (array): To be binned
        bin_type (str): Either 'lin', 'log' or 'ln'
        n_bins (int): Number of bins. Can be overriden internally
        norm (bool): Whether to normalise to 'max' or 'prob' or none

    Returns:
        tuple: bin centers, values per bin

    """
    if isinstance(parameter, list):
        parameter = np.array(parameter)

    if len(parameter) == 0:
        return np.nan, np.nan

    # Drop NaN-values
    parameter = parameter[~(np.isnan(parameter) | np.isinf(parameter))]

    # Determine number of bins
    if n_bins != 25:
        pass
    elif len(parameter) < 50:
        n_bins = 15
    elif len(parameter) > 500:
        n_bins = 50

    # Determine type of binning
    if bin_type == 'lin':
        _bins = n_bins
    elif bin_type == 'log':
        min_f = np.log10(np.min(parameter[parameter != 0]))
        max_f = np.log10(max(parameter))
        _bins = np.logspace(min_f, max_f, n_bins)
    elif bin_type == 'ln':
        min_f = np.log(np.min(parameter[parameter != 0]))
        max_f = np.log(max(parameter))
        _bins = np.logspace(min_f, max_f, n_bins, base=np.e)

    # Allow for custom bins
    if bins is not None:
        _bins = bins

    # Allow for probability weighting
    weights = None
    if norm == 'prob':
        weights = np.ones(len(parameter)) / len(parameter)

    # Bin
    n, bin_edges = np.histogram(parameter, bins=_bins, weights=weights)

    if norm == 'max':
        n = n/max(n)  # Normalise

    # Centre bins
    bins = (bin_edges[:-1] + bin_edges[1:]) / 2

    # Ensure there are edges on the outer bins of the histograms
    if edges:
        if bin_type == 'lin':
            bin_dif = np.diff(bins)[-1]
            bins = np.insert(bins, 0, bins[0] - bin_dif)
            bins = np.insert(bins, len(bins), bins[-1] + bin_dif)
            n = np.insert(n, 0, 0)
            n = np.insert(n, len(n), 0)
        else:
            bin_dif = np.diff(np.log10(bins))[-1]
            bins = np.insert(bins, 0, 10**(np.log10(bins[0])-bin_dif))
            bins = np.insert(bins, len(bins), 10**(np.log10(bins[-1])+bin_dif))
            n = np.insert(n, 0, 0)
            n = np.insert(n, len(n), 0)

    return bins, n class Survey:
    """
    Method containing survey parameters and functions.

    Args:
        name (str): Name of survey with which to observe population. Can
            either be a predefined survey present in frbpoppy or a path name to
            a new survey filename
        n_days (float): Time spent surveying [days]

    """

    def __init__(self,
                 name='perfect',
                 n_days=1):
        """Initializing."""
        # Set up parameters
        self.name = name
        self.n_days = n_days

        # Calculate some at a later stage
        self.beam_size = None
        self.beam_array = None
        self.pointings = None

        # Parse survey file
        self.read_survey_parameters()

        # Special treatment for perfect survey
        if self.name == 'perfect':
            self.beam_pattern = 'perfect'

        # Un-used
        self.strategy = 'regular'  # 'follow-up' not implemented
        if self.mount_type == 'transit':
            # Transit telescopes can't follow-up
            self.strategy = 'regular'

        # Set beam properties
        self.set_beam()

        # Set pointing properties
        self.set_pointings(mount_type=self.mount_type)

    def __str__(self):
        """Define how to print a survey object to a console."""
        s = 'Survey properties:'

        attributes = []
        for e in self.__dict__:
            attr = '\n\t{0:13.12}{1:.60}'.format(e, str(self.__dict__[e]))
            attributes.append(attr)

        s += ''.join(attributes)

        return s

    def read_survey_parameters(self):
        """Read in survey parameters."""
        # Read the survey file
        path = os.path.join(paths.surveys(), 'surveys.csv')
        df = pd.read_csv(path)
        df = df.set_index('survey')

        # Obtain the information relevant to a single survey
        survey = df[df.index == self.name].squeeze()

        # Parse parameters
        self.beta = survey['survey degradation factor']
        self.gain = survey['antenna gain (K/Jy)']
        self.t_obs = survey['integration time (s)']
        self.t_samp = survey['sampling time (ms)']
        self.T_rec = survey['receiver temperature (K)']
        self.central_freq = int(survey['centre frequency (MHz)'])
        self.bw = survey['bandwidth (MHz)']
        self.bw_chan = survey['channel bandwidth (MHz)']
        self.n_pol = survey['number of polarizations']
        self.beam_size_at_fwhm = survey['beam size (deg^2)']
        self.snr_limit = survey['signal-to-noise ratio [0-1]']
        self.max_w_eff = survey['maximum pulse width (ms)']
        self.latitude = survey['latitude (deg)']
        self.longitude = survey['longitude (deg)']
        self.mount_type = survey['mount type']
        self.ra_min = survey['minimum RA (deg)']
        self.ra_max = survey['maximum RA (deg)']
        self.dec_min = survey['minimum DEC (deg)']
        self.dec_max = survey['maximum DEC (deg)']
        self.gl_min = survey['minimum Galactic longitude (deg)']
        self.gl_max = survey['maximum Galactic longitude (deg)']
        self.gb_min = survey['minimum Galactic latitude (deg)']
        self.gb_max = survey['maximum Galactic latitude (deg)']

    def in_region(self, ra, dec, gl, gb):
        """
        Check if the given frbs are within the survey region.

        Args:
            ra, dec, gl, gb (float): Coordinates to check whether in region

        Returns:
            array: Boolean mask denoting whether frbs are within survey region

        """
        # Create mask with False if not in region
        mask = go.in_region(ra, dec, gl, gb,
                            ra_min=self.ra_min, ra_max=self.ra_max,
                            dec_min=self.dec_min, dec_max=self.dec_max,
                            gl_min=self.gl_min, gl_max=self.gl_max,
                            gb_min=self.gb_min, gb_max=self.gb_max)

        return mask

    def set_pointings(self, mount_type='tracking', n_pointings=None, ra=None,
                      dec=None):
        """Set pointing properties."""
        # If you want to use your own pointings
        if ra and dec:
            self.n_pointings = len(ra)
            self.point_func = lambda: (ra, dec)
            return

        # Calculate the number of pointings needed
        self.n_pointings = n_pointings
        if n_pointings is None:
            self.n_pointings = int(self.n_days*24*60*60 / self.t_obs)

        if mount_type == 'transit':
            self.point_func = lambda: pointings.transit(self.n_pointings,
                                                        self.latitude,
                                                        self.longitude,
                                                        self.t_obs)
        else:
            self.point_func = lambda: pointings.tracking(self.n_pointings,
                                                         ra_min=self.ra_min,
                                                         ra_max=self.ra_max,
                                                         dec_min=self.dec_min,
                                                         dec_max=self.dec_max,
                                                         gl_min=self.gl_min,
                                                         gl_max=self.gl_max,
                                                         gb_min=self.gb_min,
                                                         gb_max=self.gb_max)

    def gen_pointings(self):
        """Generate pointings."""
        self.pointings = self.point_func()

    def set_beam(self, model='perfect', size=None, random_loc=True,
                 n_sidelobes=0.5):
        """Set intensity profile.

        Set properties for int pro

        Args:
            model (str): Beam pattern. Choice from 'wsrt-apertif',
                'parkes-htru', 'chime-frb', 'gaussian', 'airy'.
            size (float): Beam size at FWHM [sq. deg].
                Defaults to that of the survey file
            random_loc (bool): Whether to calculate the precise or random
                location of each burst in the beam.
        if model in ('gaussian', 'airy'):
            n_sidelobes (int): Number of sidelobes to include. Defaults to
                cutting at the FWHM when set to 0.5.

        """
        # Set up beam properties
        self.beam_pattern = model
        self.n_sidelobes = n_sidelobes
        self.beam_array = None
        self.pixel_scale = None

        # Calculate beam properties
        if size is not None:
            self.beam_size_at_fwhm = size
        self.fwhm = 2*go.calc_sky_radius(self.beam_size_at_fwhm)

        # What should the maximum radius of the beam be?
        self.max_offset = bd.calc_max_offset(self.n_sidelobes, self.fwhm)
        self.beam_size = go.calc_sky_area(self.max_offset)

        beam_props = bd.get_beam_props(self.beam_pattern, self.fwhm)
        self.beam_size_array, self.pixel_scale, self.beam_array = beam_props

        if self.beam_size_array is not None:
            self.beam_size = self.beam_size_array
            # Not completely kosher -> a circle is not the same as a square,
            # but this is the smallest radius in which the full beam pattern
            # fits. The area of the circle extending beyond the beam pattern
            # will have an intensity of zero.
            self.max_offset = go.calc_sky_radius(self.beam_size)

        self.beam_func_oneoffs = lambda x: bd.int_pro_random(
                                 shape=x,
                                 fwhm=self.fwhm,
                                 pattern=self.beam_pattern,
                                 max_offset=self.max_offset,
                                 central_freq=self.central_freq,
                                 beam_array=self.beam_array,
                                 pixel_scale=self.pixel_scale)

        def int_pro(ra, dec, ra_p, dec_p, lst):
            return bd.int_pro_fixed(ra, dec, ra_p, dec_p, lst,
                                    pattern=self.beam_pattern,
                                    latitude=self.latitude,
                                    beam_array=self.beam_array,
                                    pixel_scale=self.pixel_scale,
                                    mount_type=self.mount_type)

        self.beam_func_rep = int_pro

    def calc_beam(self, repeaters=False, shape=None, ra=None, dec=None,
                  ra_p=None, dec_p=None, lst=None):
        """Calculate intensity profile."""
        if not repeaters and self.beam_pattern in ('airy', 'gaussian'):
            # What should the maximum radius of the beam be?
            self.max_offset = bd.calc_max_offset(self.n_sidelobes, self.fwhm)
            self.beam_size = go.calc_sky_area(self.max_offset)
        elif self.beam_pattern.startswith('perfect'):
            # What should the maximum radius of the beam be?
            self.max_offset = bd.calc_max_offset(self.n_sidelobes, self.fwhm)
            self.beam_size = go.calc_sky_area(self.max_offset)
        else:
            self.beam_size = self.beam_size_array

        if not repeaters:
            return self.beam_func_oneoffs(shape)
        else:
            return self.beam_func_rep(ra, dec, ra_p, dec_p, lst)

    def calc_dm_smear(self, dm):
        """
        Calculate delay in pulse across a channel due to dm smearing.

        Formula's based on 'Handbook of Pulsar Astronomy" by Duncan Lorimer
        & Michael Kramer, section A2.4. Note the power of the forefactor has
        changed due to the central frequency being given in MHz.

        Args:
            dm (array): Dispersion measure [pc/cm^3]

        Returns:
            t_dm (array): Time of delay [ms] at central band frequency

        """
        return 8.297616e6 * self.bw_chan * dm * (self.central_freq)**-3

    def calc_scat(self, dm):
        """Calculate scattering timescale for FRBs.

        Offset according to Lorimer et al. (doi:10.1093/mnrasl/slt098)

        Args:
            dm (array): Dispersion Measure

        Returns:
            array: Scattering timescales [ms]

        """
        freq = self.central_freq
        return go.scatter_bhat(dm, scindex=-3.86, offset=-9.5, freq=freq)

    def calc_Ts(self, gl, gb):
        """Set temperatures for frbs.

        Args:
            gl (array): Galactic longitude [deg]
            gb (array): Galactic latitude [deg]

        Returns:
            T_sky, T_sys [K]

        """
        # Special treatment for a perfect telescope
        if self.name.startswith('perfect'):
            T_sky = 0
            T_sys = self.T_rec
        else:
            T_sky = self.calc_T_sky(gl, gb)
            T_sys = self.T_rec + T_sky

        return T_sky, T_sys

    def calc_T_sky(self, gl, gb):
        """
        Calculate the sky temperature from the Haslam table.

        Afterwards scale to the survey frequency. The temperature sky map is
        given in the weird units of HealPix and despite looking up info on this
        coordinate system, I don't have the foggiest idea of how to transform
        these to galactic coordinates. I have therefore directly copied the
        following code from psrpoppy in the assumption Sam Bates managed to
        figure it out.

        Args:
            gl (array): Galactic longitude [deg]
            gb (array): Galactic latitude [deg]
        Returns:
            array: Sky temperature [K]

        """
        T_sky_list = go.load_T_sky()

        # ensure l is in range 0 -> 360
        B = gb
        L = np.copy(gl)
        L[L < 0.] += 360

        # convert from l and b to list indices
        j = B + 90.5
        j[j > 179] = 179

        nl = L - 0.5
        nl[L < 0.5] = 359
        i = nl / 4.

        index = 180*i.astype(int) + j.astype(int)
        T_sky_haslam = np.take(T_sky_list, index).astype(np.float32)

        # scale temperature
        # Assuming dominated by syncrotron radiation
        T_sky = T_sky_haslam * (self.central_freq/408.0)**(-2.6)

        return T_sky

    def calc_s_peak(self, si, lum_bol, z, dist_co, w_arr, w_eff,
                    f_low=100e6, f_high=10e9):
        """
        Calculate the mean spectral flux density.

        Following Lorimer et al, 2013, eq. 9., at the central frequency
        of the survey.

        Args:
            si (array): Spectral index
            lum_bol (array): Bolometric luminosity within emission band
            z (array): Redshift
            dist_co (array): Comoving distance [Gpc]
            w_arr (array): Pulse width at Earth [ms]
            w_eff (array): Pulse width at point of detection [ms]
            f_low (float): Source emission lower frequency limit [Hz].
            f_high (float): Source emission higher frequency limit [Hz].

        Returns:
            array: Mean spectral flux density [Jy]

        """
        # Limits observing bandwidth (as seen in rest frame source)
        f_1 = (self.central_freq - 0.5*self.bw)
        f_1 *= 1e6  # MHz -> Hz
        f_2 = (self.central_freq + 0.5*self.bw)
        f_2 *= 1e6  # MHz -> Hz

        # Spectral index
        if lum_bol.ndim == si.ndim:
            pass
        elif lum_bol.ndim > si.ndim:
            si = si[:, None]
        elif lum_bol.ndim < si.ndim:
            lum_bol = lum_bol[:, None]

        sp = si + 1
        sm = si - 1

        freq_frac = (f_2**sp - f_1**sp) / (f_2 - f_1)

        # Convert distance in Gpc to 10^25 metres
        dist = dist_co * 3.085678
        dist = dist.astype(np.float64)

        # Convert luminosity in 10^7 Watts such that s_peak will be in Janskys
        lum = lum_bol * 1e-31
        lum = lum.astype(np.float64)

        if lum.ndim > 1:
            z = z[:, None]
            dist = dist[:, None]

        nom = lum * (1+z)**sm * freq_frac
        den = 4*np.pi*dist**2 * (f_high**sp - f_low**sp)

        if nom.ndim == den.ndim:
            pass
        elif nom.ndim > 1:
            den = den[:, None]

        s_peak = nom/den

        # Add degradation factor due to pulse broadening (see Connor 2019)
        w_frac = (w_arr / w_eff)

        # Dimension conversions
        if w_frac.ndim == s_peak.ndim:
            pass
        elif s_peak.ndim == 1:
            s_peak = s_peak[:, None]
        elif w_frac.ndim == 1:
            w_frac = w_frac[:, None]

        s_peak = s_peak * w_frac

        return s_peak.astype(np.float32)

    def calc_w_eff(self, w_arr, t_dm, t_scat):
        """Calculate effective pulse width [ms].

        From Narayan (1987, DOI: 10.1086/165442), and also Cordes & McLaughlin
        (2003, DOI: 10.1086/378231). For details see p. 30 of Emily Petroff's
        thesis (2016), found here: http://hdl.handle.net/1959.3/417307

        Args:
            w_arr (array): Pulse width [ms]
            t_dm (array): Time delay due to dispersion [ms]
            t_scat (array): Time delay due to scattering [ms]

        Returns:
            array: Effective pulse width [ms]

        """
        if w_arr.ndim > 1:
            t_dm = t_dm[:, None]
            if isinstance(t_scat, np.ndarray):
                t_scat = t_scat[:, None]

        return np.sqrt(w_arr**2 + t_dm**2 + t_scat**2 + self.t_samp**2)

    def calc_snr(self, s_peak, w_arr, T_sys):
        """
        Caculate the SNR of several frbs.

        Radiometer equation for single pulse (Dewey et al., 1984), but
        adapted to allow for a degradation factor reducing the peak flux
        as a pulse is stretched

        Args:
            s_peak (array): Peak flux [Jy]
            w_arr (array): Pulse width at Earth [ms]
            T_sys (array): System temperature [K]

        Returns:
            array: Signal to noise ratio based on the radiometer
                equation for a single pulse.

        """
        # Dimension check
        if s_peak.ndim == w_arr.ndim:
            pass
        elif s_peak.ndim == 1:
            s_peak = s_peak[:, None]
        elif w_arr.ndim == 1:
            w_arr = w_arr[:, None]

        if (s_peak.ndim or w_arr.ndim) > 1:
            if isinstance(T_sys, np.ndarray):
                T_sys = T_sys[:, None]

        snr = s_peak*self.gain*np.sqrt(self.n_pol*self.bw*w_arr*1e3)
        snr /= (T_sys * self.beta)

        return snr

    def calc_fluence(self, s_peak, w_eff):
        """Calculate fluence [Jy ms].

        Args:
            s_peak (array): Peak flux
            w_eff (array): Effective pulse width

        Returns:
            array: Fluence in Jy*ms

        """
        if w_eff.ndim == s_peak.ndim:
            return s_peak * w_eff
        elif w_eff.ndim == 1:
            return s_peak * w_eff[:, None]
        elif s_peak.ndim == 1:
            return s_peak[:, None] * w_eff

    def calc_scint(self, t_scat, dist_co, gl, gb, snr):
        """
        Calculate scintillation effect on the signal to noise ratio.

        (Rather than adapting the flux, as the snr can change per survey
        attempt). Formulas based on 'Handbook of Pulsar Astronomy" by Duncan
        Lorimer & Michael Kramer, section 4.2. Test this before applying - no
        rigorous testing has been applied to this.

        Args:
            t_scat (array): Scattering timescale for frbs [ms]
                (can use self.calc_scat to calculate this)
            dist_co (array): Comoving distance array [Gpc]
            gl (array): Galactic longitude [deg]
            gb (array): Galactic latitude [deg]
            snr (array): Signal to Noise array to modify

        Returns:
            array: Signal to noise ratio modulation factors for scintillation

        """
        if not isinstance(t_scat, np.ndarray):
            m = 'Please ensure t_scat is calculated using self.calc_scat'
            raise ValueError(m)

        # Convert to seconds
        t_scat /= 1000.

        # Decorrelation bandwidth (eq. 4.39)
        decorr_bw = 1.16/(2*np.pi*t_scat)
        # Convert to MHz
        decorr_bw /= 1e6

        # Scintillation strength (eq. 4.33)
        u = np.sqrt(self.central_freq / decorr_bw)

        m = np.zeros_like(u)

        # Strong scintillation
        strong = (u < 1)
        m[strong] = np.sqrt(u[strong]**(5/3))   # (eq. 4.35)

        # Weak scintillation

        # Refractive scintillation (eq. 4.47)
        m_riss = u**-(1/3)

        # Taking the average kappa value
        kappa = 0.15

        t_diss, decorr_bw = go.ne2001_scint_time_bw(dist_co,
                                                    gl,
                                                    gb,
                                                    self.central_freq)

        # Following Cordes and Lazio (1991) (eq. 4.43)
        n_t = np.ones_like(t_diss)
        n_t[~np.isnan(t_diss)] = 1 + kappa * self.t_obs / t_diss

        n_f = np.ones_like(decorr_bw)
        n_f[~np.isnan(decorr_bw)] = 1 + kappa * self.bw / decorr_bw

        # Diffractive scintillation (eq. 4.41)
        m_diss = 1 / np.sqrt(n_t * n_f)

        # (eq. 4.48)
        weak = (u >= 1)
        m[weak] = np.sqrt(m_diss**2 + m_riss**2 + m_diss*m_riss)

        # Distribute the scintillation according to gaussian distribution
        snr = np.random.normal(snr, m*snr)

        return snr

    def calc_fluence_limit(self, w_eff=None):
        """Calculate the fluence limit.

        Read Keane, Petroff (2015) for more details on how this is calculated.

        Args:
            w_eff: Pulse width at which to calculate the fluence limit [ms].
                Default sets this to be at the maximum searched pulse width.

        Returns:
            float: Fluence limit for a maximum pulse width burst [Jy ms]

        """
        if not w_eff:
            w_eff = self.max_w_eff

        # Line of constant S/N
        self.s_peak_limit = self.snr_limit*self.T_rec*self.beta
        self.s_peak_limit /= self.gain*np.sqrt(self.n_pol*self.bw*1e3)

        # Line of constant fluence
        self.fluence_limit = self.s_peak_limit / np.sqrt(w_eff)
        self.fluence_limit *= w_eff

        return self.fluence_limit class SurveyPopulation(Population):
    """Class to create a survey population of FRBs."""

    def __init__(self, cosmic_pop, survey, scat=False, scin=False,
                 mute=False, scale_by_area=True):
        """
        Run a survey to detect FRB sources.

        Args:
            cosmic_pop (Population): Population class of FRB sources to observe
            survey (Survey): Survey class with which to observe
            scat (bool, optional): Whether to include scattering in signal to
                noise calculations.
            scin (bool, optional): Whether to apply scintillation to
                observations.
            mute (bool): Whether to suppress printing to terminal
            scale_by_area (bool): Whether to scale detection rates to the sky
                area visible to a survey. Only relevant for one-offs.
        """
        if not mute:
            pprint(f'Surveying {cosmic_pop.name} with {survey.name}')
        # Stops RuntimeWarnings about nan values
        np.warnings.filterwarnings('ignore')

        # Check whether CosmicPopulation has been generated
        try:
            if cosmic_pop.frbs.ra is None:
                m = 'You may have forgotten to generate your CosmicPopulation'
                raise ValueError(m)
        except AttributeError:
            m = 'You probably switched the population and survey in the'
            m += 'input of SurveyPopulation'
            raise ValueError(m)
        # Set up population
        Population.__init__(self)

        # Set attributes
        self.name = f'{cosmic_pop.name}_{survey.name}'
        self.vol_co_max = cosmic_pop.vol_co_max
        self.n_days = cosmic_pop.n_days
        self.repeaters = cosmic_pop.repeaters
        self.frbs = deepcopy(cosmic_pop.frbs)
        self.source_rate = Rates('source')
        if self.repeaters:
            self.burst_rate = Rates('burst')
        self.scat = scat
        self.scin = scin
        self.survey = survey
        self.scale_by_area = scale_by_area

        # Set survey attributes if not available
        if survey.n_days is None:
            survey.n_days = self.n_days

        # Calculations differ for repeaters
        if self.repeaters is True and scin is True:
            m = 'Scintillation is currently not implemented for '
            m += 'RepeaterPopulations'
            raise ValueError(m)

        # For convenience
        frbs = self.frbs
        sr = self.source_rate
        sr.tot = cosmic_pop.n_srcs
        if self.repeaters:
            br = self.burst_rate
            br.tot = self.frbs.time.size

            # Bursts which are too late have already been removed
            self.n_brst_pr_src = np.count_nonzero(~np.isnan(self.frbs.time), 1)
            br.late += br.tot - np.sum(self.n_brst_pr_src)
            sr.late += sr.tot - len(self.n_brst_pr_src)

        # Check whether source is in region
        region_mask = survey.in_region(frbs.ra, frbs.dec, frbs.gl, frbs.gb)
        frbs.apply(region_mask)

        # Keep track of detection numbers
        sr.out = np.sum(~region_mask)
        if self.repeaters:
            br.out = np.sum(self.n_brst_pr_src[~region_mask])
            self.n_brst_pr_src = self.n_brst_pr_src[region_mask]

        # Calculate dispersion measure across single channel
        frbs.t_dm = survey.calc_dm_smear(frbs.dm)

        # Set scattering timescale
        if scat:
            frbs.t_scat = survey.calc_scat(frbs.dm)

        # Calculate total temperature
        frbs.T_sky, frbs.T_sys = survey.calc_Ts(frbs.gl, frbs.gb)

        # Calculate effective pulse width
        frbs.w_eff = survey.calc_w_eff(frbs.w_arr, frbs.t_dm, frbs.t_scat)

        # Calculate peak flux density
        frbs.s_peak = survey.calc_s_peak(frbs.si,
                                         frbs.lum_bol,
                                         frbs.z,
                                         frbs.dist_co,
                                         frbs.w_arr,
                                         frbs.w_eff,
                                         f_low=cosmic_pop.f_min,
                                         f_high=cosmic_pop.f_max)

        # Calculations differ whether dealing with repeaters or not
        if self.repeaters:
            self.det_repeaters()
        else:
            self.det_oneoffs()

        # Prevent additional memory usage
        self.survey = None

    def det_oneoffs(self):
        """Detect one-off frbs."""
        frbs = self.frbs
        survey = self.survey

        # Account for beam offset
        int_pro, offset = survey.calc_beam(shape=frbs.s_peak.shape)
        frbs.s_peak *= int_pro
        frbs.offset = offset  # [deg]

        # Calculate fluence [Jy*ms]
        frbs.fluence = survey.calc_fluence(frbs.s_peak, frbs.w_eff)

        # Calculate Signal to Noise Ratio
        frbs.snr = survey.calc_snr(frbs.s_peak, frbs.w_arr, frbs.T_sys)

        # Add scintillation
        if self.scin:

            # Ensure scattering has been calculated
            if not isinstance(frbs.t_scat, np.ndarray):
                frbs.t_scat = survey.calc_scat(frbs.dm)

            # Calculate signal to noise ratio after scattering
            frbs.snr = survey.calc_scint(frbs.t_scat, frbs.dist_co, frbs.gl,
                                         frbs.gb, frbs.snr)

        # Check whether frbs would be above detection threshold
        snr_mask = (frbs.snr >= survey.snr_limit)
        frbs.apply(snr_mask)
        self.source_rate.faint = len(snr_mask) - np.count_nonzero(snr_mask)

        # Distant frbs are redshifted out of your observing time
        limit = 1/(1+frbs.z)
        rate_mask = np.random.random(len(frbs.z)) <= limit
        frbs.apply(rate_mask)
        self.source_rate.late = np.size(rate_mask)
        self.source_rate.late -= np.count_nonzero(rate_mask)

        self.source_rate.det = len(frbs.snr)

        # Calculate detection rates
        if self.scale_by_area:
            self.calc_rates(survey)

    def det_repeaters(self):
        """Detect repeating frbs."""
        frbs = self.frbs
        survey = self.survey
        br = self.burst_rate
        sr = self.source_rate

        # Set up a tuple of pointings if not given
        survey.gen_pointings()

        # t_obs in fractional days
        t_obs = survey.t_obs / 86400

        # Array with times of each pointing
        max_t = survey.n_days
        times = np.arange(0, max_t+t_obs, t_obs)  # [days]
        lsts = times*360*(24/23.9344696) % 360  # Local sidereal time [deg]
        lsts += np.random.uniform(0, 360)  # Add random offset

        # Only keep bursts within survey time
        time_mask = (frbs.time <= times[-1])
        # n_brst_pr_src = np.count_nonzero(~np.isnan(self.frbs.time), 1)
        frbs.apply(time_mask)

        # Keep track of losses
        br.late += np.sum(self.n_brst_pr_src - np.count_nonzero(time_mask, 1))
        sr.late += len(time_mask) - len(self.frbs.time)

        # Prepare for iterating over time
        max_n_pointings = len(times) - 1

        # Initialize some necessary arrays
        if frbs.w_eff.ndim == 2 or frbs.lum_bol.ndim == 2:
            sim_shape = frbs.time  # 2D
        else:
            sim_shape = frbs.lum_bol  # 1D

        frbs.fluence = np.full_like(sim_shape, np.nan)
        frbs.snr = np.full_like(sim_shape, np.nan)

        # Have to loop over the observing times
        ra_p = survey.pointings[0]
        dec_p = survey.pointings[1]
        lst = lsts[:-1]
        self.srcs_not_in_pointing = np.ones_like(frbs.index, dtype=bool)
        self.srcs_not_bright = np.ones_like(frbs.index, dtype=bool)

        # Parameters needed for for-loop
        keep = ([], [])
        for i in tqdm(np.arange(max_n_pointings), desc='Pointings'):
            xy = self._iter_pointings(ra_p[i % survey.n_pointings],
                                      dec_p[i % survey.n_pointings],
                                      lst[i],
                                      times[i],
                                      times[i+1])
            keep[0].extend(xy[0])
            keep[1].extend(xy[1])

        sr.pointing = np.sum(self.srcs_not_in_pointing)
        # Already outside of pointing, so unknown whether too faint
        self.srcs_not_bright[self.srcs_not_in_pointing] = False

        sr.faint = np.sum(self.srcs_not_bright)

        # Create SNR mask
        snr_mask = np.zeros_like(frbs.snr, dtype=bool)
        snr_mask[keep] = True
        frbs.apply(snr_mask)

        # Keep track of detections
        self.burst_rate.det = np.count_nonzero(snr_mask)
        self.source_rate.det = len(np.unique(keep[0]))

        # Reduce matrices' size
        frbs.clean_up()

        # Calculate detection rates
        self.calc_rates(survey)

    def _iter_pointings(self, ra_pt, dec_pt, lst, t_min, t_max):
        frbs = self.frbs
        survey = self.survey

        # Which frbs are within the pointing time?
        # Essential that each row is sorted from low to high!
        # Returns col, row index arrays
        t_ix = fast_where(frbs.time, t_min, t_max)

        # What's the intensity of them in the beam?
        int_pro, dx, dy = survey.calc_beam(repeaters=True,
                                           ra=frbs.ra[t_ix[0]],
                                           dec=frbs.dec[t_ix[0]],
                                           ra_p=ra_pt,
                                           dec_p=dec_pt,
                                           lst=lst)

        # If not an intensity of zero, they were inside the beam pattern
        p_ix = ~np.isnan(int_pro)

        # Time & position
        tp_ix = (t_ix[0][p_ix], t_ix[1][p_ix])

        # Time & not position
        tnp_ix = (t_ix[0][~p_ix], t_ix[1][~p_ix])

        # Number of bursts per source
        tp_unique, n_bursts = np.unique(tp_ix[0], return_counts=True)

        # Add to outside of pointing count
        self.burst_rate.pointing += len(tnp_ix[0])
        self.srcs_not_in_pointing[tp_unique] = False

        # Ensure relevant masks
        s_peak_ix = tp_unique
        not_ix = np.unique(tnp_ix[0])
        int_pro = int_pro[p_ix]

        if frbs.s_peak.ndim == 2:
            s_peak_ix = tp_ix
            not_ix = tnp_ix
        if frbs.s_peak.ndim == 1:
            int_pro = int_pro[tp_unique]

        # Apply intensities to those bursts' s_peak
        frbs.s_peak[s_peak_ix] *= int_pro
        frbs.s_peak[not_ix] = np.nan

        w_eff_ix = tp_unique
        if frbs.w_eff.ndim == 2:
            w_eff_ix = tp_ix

        # Ensure dimensionality is correct
        s_peak = frbs.s_peak[s_peak_ix]
        w_eff = frbs.w_eff[w_eff_ix]
        if frbs.w_eff.ndim < frbs.s_peak.ndim:
            w_eff = np.repeat(w_eff, n_bursts)
        elif frbs.w_eff.ndim > frbs.s_peak.ndim:
            s_peak = np.repeat(s_peak, n_bursts)

        # Calculate fluence [Jy*ms]
        frbs.fluence[s_peak_ix] = survey.calc_fluence(s_peak, w_eff)

        # Construct masks with correct dimension
        if frbs.w_arr.ndim == 2:
            w_arr = frbs.w_arr[tp_ix]
        elif frbs.w_arr.ndim == 1:
            w_arr = frbs.w_arr[tp_unique]

        # Ensure entries are repeated for the number of bursts
        s_peak = frbs.s_peak[s_peak_ix]
        if frbs.w_arr.ndim < frbs.s_peak.ndim:
            w_arr = np.repeat(w_arr, n_bursts)
        elif frbs.w_arr.ndim > frbs.s_peak.ndim:
            s_peak = np.repeat(s_peak, n_bursts)

        # And system temperatures
        if frbs.T_sys.ndim == 1:
            if frbs.s_peak.ndim == 2:
                T_sys = np.repeat(frbs.T_sys[tp_unique], n_bursts)
            elif frbs.s_peak.ndim == 1:
                T_sys = frbs.T_sys[tp_unique]
        elif frbs.T_sys.ndim == 0:
            T_sys = frbs.T_sys

        # Caculate Signal to Noise Ratio
        frbs.snr[s_peak_ix] = survey.calc_snr(s_peak, w_arr, T_sys)

        # Add scintillation
        if self.scin:
            # Not been fully tested with repeaters
            # Might break due to differing dimensionality of dist and snr
            t_scat = frbs.t_scat[tp_unique]
            dist_co = frbs.dist_co[tp_unique]
            gl = frbs.gl[tp_unique]
            gb = frbs.gb[tp_unique]
            snr = frbs.snr[s_peak_ix]
            new_snr = survey.calc_scint(t_scat, dist_co, gl, gb, snr)
            frbs.snr[s_peak_ix] = np.repeat(new_snr, n_bursts)

        # Only keep those in time, in position and above the snr limit
        snr_m = (frbs.snr[s_peak_ix] > survey.snr_limit)
        s_peak_ix = (tp_ix[0][snr_m], tp_ix[1][snr_m])

        self.burst_rate.faint += len(tp_ix[0]) - len(s_peak_ix[0])
        self.srcs_not_bright[np.unique(s_peak_ix[0])] = False

        return s_peak_ix[0], s_peak_ix[1]

    def calc_rates(self, survey):
        """Calculate the relative detection rates."""
        # Calculate scaling factors for rates
        area_sky = 4*np.pi*(180/np.pi)**2   # In sq. degrees
        f_area = survey.beam_size * self.source_rate.tot
        inside = self.source_rate.det + self.source_rate.late
        inside += self.source_rate.faint
        if inside > 0:
            f_area /= (inside*area_sky)
        else:
            f_area = 1

        # Saving scaling factors
        self.source_rate.days = self.n_days
        self.source_rate.name = self.name
        self.source_rate.vol = self.source_rate.tot
        self.source_rate.vol /= self.vol_co_max * (365.25/self.n_days)

        if self.repeaters:
            self.burst_rate.days = self.n_days
            self.burst_rate.name = self.name
            self.burst_rate.vol = self.burst_rate.tot
            self.burst_rate.vol /= self.vol_co_max * (365.25/self.n_days)

        # If oneoffs, you'll want to scale by the area
        if not self.repeaters:
            self.source_rate.f_area = f_area
            self.source_rate.scale_by_area() def pprint(*s, output=True):
    """Hack to make for more informative print statements."""
    f = inspect.stack()[1][1].split('/')[-1]
    m = '{:13.13} |'.format(f)

    if output:
        print(m, *s)
    else:
        lines = []
        for e in s:
            lines.append('\n'.join([f'{m} {f}' for f in e.split('\n')]))
        return '\n'.join(lines) def unpickle(filename=None, uid=None):
    """Quick function to unpickle a population.

    Args:
        filename (str, optional): Define the path to the pickled population,
            or give the population name
        uid (str, optional): Unique Identifier

    Returns:
        Population: Population class

    """
    # Find population file
    if os.path.isfile(filename):
        f = open(filename, 'rb')
    else:
        # Find standard population files
        try:
            name = filename.lower()
            if uid:
                name += f'_{uid}'
            p = paths.populations() + f'{name}.p'
            f = open(p, 'rb')
        except FileNotFoundError:
            s = 'Pickled population file "{0}" does not exist'.format(filename)
            raise FileNotFoundError(s)

    # Unpickle
    pop = pickle.load(f)
    f.close()
    return pop def plot_aa_style(cols=1):
    """Use a plotting style specifically for Astronomy & Astrophyiscs."""
    # Use a style file
    dir = os.path.abspath(os.path.dirname(__file__))
    style_file = os.path.join(dir, './aa.mplstyle')
    plt.style.use(style_file)

    # Check whether interactive backend is present
    if os.name == 'posix' and "DISPLAY" not in os.environ:
        # On a cluster
        plt.switch_backend('Agg')
        # Use just the basic latex package allowing \text to be used
        plt.rcParams['text.latex.preamble'] = r'\usepackage{amsmath}'

    if cols == 2:
        plt.rcParams["figure.figsize"] = (5.75373, 3.556) def rel_path(path):
    """Return the relative path name for this directory."""
    return os.path.join(os.path.abspath(os.path.dirname(__file__)), path) EXPECTED = {'parkes-htru': [9, 1549 / 0.551 / 24],  # N_frbs, N_days
            'wsrt-apertif': [9, 1100/24],  # 1100 hours
            'askap-fly': [20, 32840 / 8 / 24],
            'arecibo-palfa': [1, 24.1],
            'guppi': [0.4, 81],  # 0.4 is my own assumption
            'fast-crafts': [1, 1500/24],
            'chime-frb': [2, 1],  # My own assumption
            'askap-incoh': [9, 45]  # My own assumption
            }
import numpy as np
import matplotlib.pyplot as plt
import frbpoppy.direction_dists as did
import frbpoppy.galacticops as go
from copy import copy
from frbpoppy import CosmicPopulation, Survey, LargePopulation, SurveyPopulation, hist
from frbpoppy import unpickle, pprint
from tests.convenience import plot_aa_style, rel_path
from tests.rates.alpha_real import EXPECTED
           'puma-full',
           'chord',
           'ska1-low',
           'ska1-mid')
SIZE = 5e4


if MAKE:
    # Calculate the fraction of the sky that the survey covers
    surv_f_area = {}
    for name in SURVEYS:
        pop = CosmicPopulation.simple(5e5)
        pop.gen_direction()
        survey = Survey(name)
        mask = survey.in_region(pop.frbs.ra, pop.frbs.dec,
                                pop.frbs.gl, pop.frbs.gb)
        in_surv_region = np.sum(mask)
        tot_region = len(mask)

        area_sky = 4*np.pi*(180/np.pi)**2   # In sq. degrees
        f_area = (survey.beam_size/area_sky)*(tot_region/in_surv_region)

        surv_f_area[name] = f_area
        print(f'{name} covers {f_area*100}% of the sky')

    surv_pops = []

    for name in SURVEYS:
        # Set up survey
        survey = Survey(name)
        if name in ('parkes-htru', 'wsrt-apertif'):
            survey.set_beam(model=name)

        # Set up CosmicPopulation
        pop = CosmicPopulation.optimal(SIZE, generate=False)

        # Only generate FRBs in the survey region
        pop.set_direction(model='uniform',
                          min_ra=survey.ra_min,
                          max_ra=survey.ra_max,
                          min_dec=survey.dec_min,
                          max_dec=survey.dec_max)

        # Parkes also has galactic limits:
        if name == 'parkes-htru':
            pop.gen_index()
            pop.gen_dist()
            pop.gen_time()

            # Generate FRBs just within the galactic constraints
            pop.gen_direction()

            # Gather ra, dec coordinate limits
            lims = {'min_ra': survey.ra_min, 'max_ra': survey.ra_max,
                    'min_dec': survey.dec_min, 'max_dec': survey.dec_max}

            def sample(n_gen):
                ra, dec = did.uniform(n_srcs=n_gen, **lims)
                gl, gb = go.radec_to_lb(ra, dec, frac=True)
                coords = [ra, dec, gl, gb]
                return coords

            def accept(coords):
                return survey.in_region(*coords)

            coords = sample(int(SIZE))
            mask = accept(coords)
            reject, = np.where(~mask)
            while reject.size > 0:
                fill = sample(reject.size)
                mask = accept(fill)
                for i in range(len(coords)):
                    coords[i][reject[mask]] = fill[i][mask]
                reject = reject[~mask]

            # Assign the values
            frbs = pop.frbs
            frbs.ra, frbs.dec = coords[0], coords[1]
            frbs.gl, frbs.gb = coords[2], coords[3]

            # Continue with generation
            pop.gen_gal_coords()
            pop.gen_dm()
            pop.gen_w()
            pop.gen_lum()
            pop.gen_si()
        else:
            pop.generate()

        surv_pop = SurveyPopulation(pop, survey, scale_by_area=False)
        surv_pop.source_rate.f_area = surv_f_area[name]
        surv_pop.source_rate.scale_by_area()
        # surv_pop.save()
        surv_pops.append(surv_pop)

else:
    surv_pops = []
    for name in SURVEYS:
        surv_pops.append(unpickle(f'optimal_{name}'))

# Start plot
plot_aa_style(cols=2)
plt.rcParams["figure.figsize"] = (3.556*3, 3.556)
fig, axes = plt.subplots(1, 3)

for ax in axes.flatten():
    ax.set_aspect('auto')

# Get norm pop
y = 0
ys = []
names = []
rates = []
norm_sim_rate = surv_pops[0].source_rate.det
norm_real_rate = EXPECTED['parkes-htru'][0] / EXPECTED['parkes-htru'][1]
norm_rate = norm_sim_rate / norm_real_rate

for i, surv_pop in enumerate(surv_pops):

    name = surv_pop.name.split('_')[-1]

    pprint(name)