# Source code for xga.samples.general

#  This code is a part of X-ray: Generate and Analyse (XGA), a module designed for the XMM Cluster Survey (XCS).

from warnings import warn

import numpy as np
from astropy.cosmology import Planck15
from astropy.units import Quantity, Unit
from tqdm import tqdm

from .base import BaseSample
from ..exceptions import NoValidObservationsError, PeakConvergenceFailedError
from ..sources.general import PointSource, ExtendedSource

[docs]class ExtendedSample(BaseSample):
"""
The sample class for exploring general extended sources without the extra information required to
analyse more specific X-ray extended sources (like galaxy clusters).

:param np.ndarray ra: The right-ascensions of the extended sources, in degrees.
:param np.ndarray dec: The declinations of the extended sources, in degrees.
:param np.ndarray redshift: The redshifts of the extended sources, optional. Default is None.
:param np.ndarray name: The names of the extended sources, optional. If no names are supplied
then they will be constructed from the supplied coordinates.
:param Quantity custom_region_radius: Custom analysis region radius(ii) for these sources, optional. Either
pass a scalar astropy quantity, or a non-scalar astropy quantity with length equal to the number of sources.
:param bool use_peak: Whether peak positions should be found and used.
:param Quantity peak_lo_en: The lower energy bound for the RateMap to calculate peak
position from. Default is 0.5keV.
:param Quantity peak_hi_en: The upper energy bound for the RateMap to calculate peak
position from. Default is 2.0keV.
:param float back_inn_rad_factor: This factor is multiplied by an analysis region radius, and gives the inner
radius for the background region. Default is 1.05.
:param float back_out_rad_factor: This factor is multiplied by an analysis region radius, and gives the outer
radius for the background region. Default is 1.5.
:param cosmology: An astropy cosmology object for use throughout analysis of the source.
:param bool no_prog_bar: Should a source declaration progress bar be shown during setup.
:param bool psf_corr: Should images be PSF corrected with default settings during sample setup.
:param str peak_find_method: Which peak finding method should be used (if use_peak is True). Default
is hierarchical, simple may also be passed.
"""
def __init__(self, ra: np.ndarray, dec: np.ndarray, redshift: np.ndarray = None, name: np.ndarray = None,
custom_region_radius: Quantity = None, use_peak: bool = True,
peak_lo_en: Quantity = Quantity(0.5, "keV"), peak_hi_en: Quantity = Quantity(2.0, "keV"),
load_fits: bool = False, no_prog_bar: bool = False, psf_corr: bool = False,
peak_find_method: str = "hierarchical"):
"""
The init method of the ExtendedSample class.
"""
raise TypeError("Please pass None or a quantity object for custom_region_radius, rather than an "
"array or list.")
# People might pass a single value for custom_region_radius, in which case we turn it into
#  a non-scalar quantity
" then there must be one entry per object passed to this sample object.")

# I don't like having this here, but it does avoid a circular import problem
from xga.sas import evselect_image, eexpmap, emosaic

# Using the super defines BaseSources and stores them in the self._sources dictionary
no_prog_bar=no_prog_bar)

evselect_image(self, peak_lo_en, peak_hi_en)
eexpmap(self, peak_lo_en, peak_hi_en)
emosaic(self, "image", peak_lo_en, peak_hi_en)
emosaic(self, "expmap", peak_lo_en, peak_hi_en)

del self._sources
self._sources = {}

# A list to store the inds of any declarations that failed due to NoValidObservations, so some
#  attributes can be cleaned up later
final_names = []
# This records which sources had a failed peak finding attempt, for a warning at the end of the declaration
failed_peak_find = []
with tqdm(desc="Setting up Extended Sources", total=len(self._accepted_inds), disable=no_prog_bar) as dec_lb:
for ind in range(len(self._accepted_inds)):
r, d = ra[self._accepted_inds[ind]], dec[self._accepted_inds[ind]]
if redshift is None:
z = None
else:
z = redshift[self._accepted_inds[ind]]
n = self._names[ind]

try:
# Declare a generic extended source, telling it is that it is part of a sample with in_sample=True
self._sources[n] = ExtendedSource(r, d, z, n, cr, use_peak, peak_lo_en, peak_hi_en,
if isinstance(cr, Quantity):
# I know this will write to this over and over, but it seems a bit silly to check
#  whether this has been set yet when all radii should be forced to be the same unit
self._cr_unit = cr.unit
else:
self._cr_unit = Unit('')
final_names.append(n)
except PeakConvergenceFailedError:
warn("The peak finding algorithm has not converged for {}, using user "
"supplied coordinates".format(n))
# Have to re-declare the source if peak finding failed
self._sources[n] = ExtendedSource(r, d, z, n, cr, False, peak_lo_en, peak_hi_en,
final_names.append(n)
except NoValidObservationsError:
self._failed_sources[n] = "CleanedNoMatch"

dec_lb.update(1)
self._names = final_names

# I've cleaned the observations, and its possible some of the data has been thrown away,
#  so I should regenerate the mosaic images/expmaps
emosaic(self, "image", peak_lo_en, peak_hi_en)
emosaic(self, "expmap", peak_lo_en, peak_hi_en)

# I don't offer the user choices as to the configuration for PSF correction at the moment
if psf_corr:
# Trying to see if this stops a circular import issue I've been having
from ..imagetools.psf import rl_psf
rl_psf(self, lo_en=peak_lo_en, hi_en=peak_hi_en)

# It is possible (especially if someone is using the Sample classes as a way to check whether things have
#  XMM data) that no sources will have been declared by this point, in which case it should fail now
if len(self._sources) == 0:
raise NoValidObservationsError(
"No Extended Sources have been declared, none of the sample passed the cleaning steps.")

# Put all the warnings for there being no XMM data in one - I think it's neater. Wait until after the check
#  to make sure that are some sources because in that case this warning is redundant.
no_data = [name for name in self._failed_sources if self._failed_sources[name] == 'NoMatch' or
self._failed_sources[name] == 'Failed ObsClean']
# If there are names in that list, then we do the warning
if len(no_data) != 0:
warn("The following do not appear to have any XMM data, and will not be included in the "
"sample (can also check .failed_names); {n}".format(n=', '.join(no_data)), stacklevel=2)

# We also do a combined warning for those clusters that had a failed peak finding attempt, if there are any
if len(failed_peak_find) != 0:
warn("Peak finding did not converge for the following; {n}, using user "
"supplied coordinates".format(n=', '.join(failed_peak_find)), stacklevel=2)

# This shows a warning that tells the user how to see any suppressed warnings that occurred during source
#  declarations, but only if there actually were any.
self._check_source_warnings()

@property
"""
Property getter for the radii of the custom analysis regions that can be used for analysis of the
extended sources in this sample. Users are not required to pass a custom analysis region so this
may be NaN.

:return: A non-scalar Quantity of the custom source radii passed in by the user.
:rtype: Quantity
"""

@property
"""
Property getter for the unit which the custom analysis radii values are stored in.

:return: The unit that the custom radii are stored in.
:rtype: Unit
"""
return self._cr_unit

def _del_data(self, key: int):
"""
Specific to the ExtendedSample class, this deletes the extra data stored during the initialisation
of this type of sample.

:param int key: The index or name of the source to delete.
"""

[docs]class PointSample(BaseSample):
"""
The sample class for general point sources, without the extra information required to analyse more specific
X-ray point sources.

:param np.ndarray ra: The right-ascensions of the point sources, in degrees.
:param np.ndarray dec: The declinations of the point sources, in degrees.
:param np.ndarray redshift: The redshifts of the point sources, optional. Default is None.
:param np.ndarray name: The names of the point sources, optional. If no names are supplied
then they will be constructed from the supplied coordinates.
:param Quantity point_radius: The point source analysis region radius(ii) for this sample. Either
pass a scalar astropy quantity, or a non-scalar astropy quantity with length equal to the number of sources.
:param bool use_peak: Whether peak positions should be found and used. For PointSample the 'simple' peak
finding method is the only one available.
:param Quantity peak_lo_en: The lower energy bound for the RateMap to calculate peak
position from. Default is 0.5keV.
:param Quantity peak_hi_en: The upper energy bound for the RateMap to calculate peak
position from. Default is 2.0keV.
:param float back_inn_rad_factor: This factor is multiplied by an analysis region radius, and gives the inner
radius for the background region. Default is 1.05.
:param float back_out_rad_factor: This factor is multiplied by an analysis region radius, and gives the outer
radius for the background region. Default is 1.5.
:param cosmology: An astropy cosmology object for use throughout analysis of the source.
:param bool no_prog_bar: Should a source declaration progress bar be shown during setup.
:param bool psf_corr: Should images be PSF corrected with default settings during sample setup.
"""
def __init__(self, ra: np.ndarray, dec: np.ndarray, redshift: np.ndarray = None, name: np.ndarray = None,
point_radius: Quantity = Quantity(30, 'arcsec'), use_peak: bool = False,
peak_lo_en: Quantity = Quantity(0.5, "keV"), peak_hi_en: Quantity = Quantity(2.0, "keV"),
cosmology=Planck15, load_fits: bool = False, no_prog_bar: bool = False, psf_corr: bool = False):
"""
The init method of the PointSample class.
"""

# Strongly enforce that its a quantity, this also means that it should be guaranteed that all radii have
#  a single unit
raise TypeError("Please pass a quantity object for point_radius, rather than an array or list.")
# People might pass a single value for point_radius, in which case we turn it into a non-scalar quantity
raise ValueError("If you pass a set of radii (rather than a single radius) to point_radius then there"
" must be one entry per object passed to this sample object.")

# I don't like having this here, but it does avoid a circular import problem
from xga.sas import evselect_image, eexpmap, emosaic

# Using the super defines BaseSources and stores them in the self._sources dictionary
no_prog_bar=no_prog_bar)

evselect_image(self, peak_lo_en, peak_hi_en)
eexpmap(self, peak_lo_en, peak_hi_en)
emosaic(self, "image", peak_lo_en, peak_hi_en)
emosaic(self, "expmap", peak_lo_en, peak_hi_en)

del self._sources
self._sources = {}

# A list to store the inds of any declarations that failed due to NoValidObservations, so some
#  attributes can be cleaned up later
final_names = []
# This records which sources had a failed peak finding attempt, for a warning at the end of the declaration
failed_peak_find = []
with tqdm(desc="Setting up Point Sources", total=len(self._accepted_inds), disable=no_prog_bar) as dec_lb:
for ind in range(len(self._accepted_inds)):
r, d = ra[self._accepted_inds[ind]], dec[self._accepted_inds[ind]]
if redshift is None:
z = None
else:
z = redshift[self._accepted_inds[ind]]
n = self._names[ind]

# Observation cleaning goes on automatically in PointSource, so if a NoValidObservations error is
#  thrown I have to catch it and not add that source to this sample.
try:
self._sources[n] = PointSource(r, d, z, n, pr, use_peak, peak_lo_en, peak_hi_en,
# I know this will write to this over and over, but it seems a bit silly to check whether this has
#  been set yet when all radii should be forced to be the same unit
self._pr_unit = pr.unit
final_names.append(n)
except PeakConvergenceFailedError:
warn("The peak finding algorithm has not converged for {}, using user "
"supplied coordinates".format(n))
self._sources[n] = PointSource(r, d, z, n, pr, False, peak_lo_en, peak_hi_en,
final_names.append(n)
except NoValidObservationsError:
self._failed_sources[n] = "CleanedNoMatch"

dec_lb.update(1)
self._names = final_names

# I've cleaned the observations, and its possible some of the data has been thrown away,
#  so I should regenerate the mosaic images/expmaps
emosaic(self, "image", peak_lo_en, peak_hi_en)
emosaic(self, "expmap", peak_lo_en, peak_hi_en)

# I don't offer the user choices as to the configuration for PSF correction at the moment
if psf_corr:
# Trying to see if this stops a circular import issue I've been having
from ..imagetools.psf import rl_psf
rl_psf(self, lo_en=peak_lo_en, hi_en=peak_hi_en)

# It is possible (especially if someone is using the Sample classes as a way to check whether things have
#  XMM data) that no sources will have been declared by this point, in which case it should fail now
if len(self._sources) == 0:
raise NoValidObservationsError(
"No Point Sources have been declared, none of the sample passed the cleaning steps.")

# Put all the warnings for there being no XMM data in one - I think it's neater. Wait until after the check
#  to make sure that are some sources because in that case this warning is redundant.
no_data = [name for name in self._failed_sources if self._failed_sources[name] == 'NoMatch' or
self._failed_sources[name] == 'Failed ObsClean']
# If there are names in that list, then we do the warning
if len(no_data) != 0:
warn("The following do not appear to have any XMM data, and will not be included in the "
"sample (can also check .failed_names); {n}".format(n=', '.join(no_data)), stacklevel=2)

# We also do a combined warning for those clusters that had a failed peak finding attempt, if there are any
if len(failed_peak_find) != 0:
warn("Peak finding did not converge for the following; {n}, using user "
"supplied coordinates".format(n=', '.join(failed_peak_find)), stacklevel=2)

# This shows a warning that tells the user how to see any suppressed warnings that occurred during source
#  declarations, but only if there actually were any.
self._check_source_warnings()

@property
"""
Property getter for the radii of the regions used for analysis of the point sources in this sample.

:return: A non-scalar Quantity of the point source radii used for analysis of the point sources in
this sample.
:rtype: Quantity
"""

@property
"""
Property getter for the unit which the point radii values are stored in.

:return: The unit that the point radii are stored in.
:rtype: Unit
"""
return self._pr_unit

def _del_data(self, key: int):
"""
Specific to the PointSample class, this deletes the extra data stored during the initialisation
of this type of sample.

:param int key: The index or name of the source to delete.
"""