import abc
from typing import Dict
import numpy as np
from class_registry import AutoRegister
from popsynth.utils.logging import setup_logger
from popsynth.utils.meta import Parameter, ParameterMeta
from popsynth.utils.registry import selection_registry
log = setup_logger(__name__)
[docs]class SelectionParameter(Parameter):
pass
[docs]class SelectionProbability(object,
metaclass=AutoRegister(selection_registry,
base_type=ParameterMeta)):
[docs] def __init__(
self,
name: str = "name",
use_obs_value: bool = False,
use_distance: bool = False,
use_luminosity: bool = False,
use_flux: bool = False,
) -> None:
"""
Base class for selections on a population.
:param name: Name of the selection
:type name: str
:param use_obs_value: If `True`, make selection on
observed_value. `False` by default.
:type use_obs_value: bool
:param use_distance: If `True` make selection on distance.
`False` by default.
:type use_distance: bool
:param use_luminosity: If `True` make selection on luminosity.
`False` by default.
:type use_luminosity: bool
:param use_flux: If `True` make selection on flux. `False` by default.
:type use_flux: bool
"""
self._name = name # type: str
self._parameter_storage = {}
self._observed_flux = None # type: np.ndarray
self._observed_value = None # type: np.ndarray
self._distance = None # type: np.ndarray
self._luminosity = None # type: np.ndarray
self._selection = None # type: np.ndarray
self._is_sampled: bool = False
self._use_obs_value: bool = use_obs_value
self._use_distance: bool = use_distance
self._use_luminosity: bool = use_luminosity
self._use_flux: bool = use_flux
def __add__(self, other):
log.debug(f"adding selection from {other.name} to {self.name}")
new_selection = np.logical_and(self._selection, other.selection)
out = DummySelection()
out._selection = new_selection
return out
# def __radd__(self, other):
# return self.__add__(other)
[docs] def set_luminosity(self, luminosity: np.ndarray) -> None:
"""
Set the luminosity of the selection.
"""
self._luminosity = luminosity # type: np.ndarray
[docs] def set_distance(self, distance: np.ndarray) -> None:
"""
Set the distance of the selection.
"""
self._distance = distance # type: np.ndarray
[docs] def set_observed_flux(self, observed_flux: np.ndarray) -> None:
"""
Set the observed flux of the selection.
"""
self._observed_flux = observed_flux # type: np.ndarray
[docs] def set_observed_value(self, observed_value: np.ndarray) -> None:
"""
Set the observed value of the selection.
"""
self._observed_value = observed_value # type: np.ndarray
[docs] @abc.abstractclassmethod
def draw(self, size: int) -> None:
pass
[docs] def reset(self):
"""
Reset the selector.
"""
if self._is_sampled:
log.info(f"Selection: {self.name} is being reset")
self._is_sampled = False
self._selection = None
[docs] def select(self, size: int):
if not self._is_sampled:
self.draw(size)
else:
log.warning(f"selecting with {self.name} more than once!")
log.debug(f"{self.name} sampled {sum(self._selection)} objects")
self._is_sampled = True
@property
def parameters(self) -> Dict[str, float]:
out = {}
for k, v in self._parameter_storage.items():
if v is not None:
out[k] = v
return out
@property
def selection(self) -> np.ndarray:
if self._selection is not None:
return self._selection
else:
log.error(f"selector: {self.name} as not be sampled yet!")
raise RuntimeError()
@property
def n_selected(self) -> int:
return sum(self._selection)
@property
def n_non_selected(self) -> int:
return sum(~self._selection)
@property
def n_objects(self) -> int:
return self._selection.shape[0]
@property
def selection_index(self) -> np.ndarray:
return np.where(self._selection)[0]
@property
def non_selection_index(self) -> np.ndarray:
return np.where(~self._selection)[0]
@property
def name(self) -> str:
return self._name
[docs]class DummySelection(SelectionProbability):
_selection_name = "DummySelection"
[docs] def __init__(self):
"""
A Dummy selection for testing.
"""
super(DummySelection, self).__init__(name="dummy")
[docs] def draw(self, size=1):
pass