import numpy as np
import scipy.special as sf
import scipy.stats as stats
from popsynth.utils.configuration import popsynth_config
from popsynth.utils.logging import setup_logger
from popsynth.utils.progress_bar import progress_bar
from .selection_probability import SelectionParameter, SelectionProbability
log = setup_logger(__name__)
[docs]class UnitySelection(SelectionProbability):
_selection_name = "UnitySelection"
[docs] def __init__(self, name="unity"):
"""
A selection that returns all selected.
:param name: Name of the selection
"""
super(UnitySelection, self).__init__(name=name)
[docs] def draw(self, size: int) -> None:
self._selection = np.ones(size, dtype=int).astype(bool)
[docs]class BernoulliSelection(SelectionProbability):
_selection_name = "BernoulliSelection"
probability = SelectionParameter(vmin=0, vmax=1, default=0.5)
[docs] def __init__(self) -> None:
"""
A Bernoulli selection with ``probability`` as a parameter.
:param probability: Probability for each Bernoulli trial
:type probability: :class:`SelectionParameter`
"""
super(BernoulliSelection, self).__init__(name="Bernoulli")
[docs] def draw(self, size: int) -> None:
if popsynth_config.show_progress:
self._selection = np.zeros(size, dtype=int).astype(
bool) # type: np.ndarray
for i in progress_bar(range(size), desc=f"Selecting {self.name}"):
# see if we detect the distance
if stats.bernoulli.rvs(self.probability) == 1:
self._selection[i] = 1
else:
self._selection = stats.bernoulli.rvs(
self.probability, size=size).astype(bool) # type: np.ndarray
[docs]class BoxSelection(SelectionProbability):
_selection_name = "BoxSelection"
vmin = SelectionParameter()
vmax = SelectionParameter()
[docs] def __init__(
self,
name: str = "box selection",
use_obs_value: bool = False,
use_distance: bool = False,
use_luminosity: bool = False,
use_flux: bool = False,
):
"""
A box selection on observed_value, distance,
luminosity or flux.
: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
:param vmin: Minimum value of selection
:type vmin: :class:`SelectionParameter`
:param vmax: Maximum value of selection
:type vmax: :class:`SelectionParameter`
"""
super(BoxSelection, self).__init__(
name=name,
use_distance=use_distance,
use_luminosity=use_luminosity,
use_obs_value=use_obs_value,
use_flux=use_flux,
)
[docs] def draw(self, size: int) -> np.ndarray:
if self._use_distance:
values = self._distance
if self._use_obs_value:
values = self._observed_value
if self._use_luminosity:
values = self._luminosity
if self._use_flux:
values = self._observed_flux
self._selection = (values >= self.vmin) & (values <= self.vmax)
[docs]class LowerBound(SelectionProbability):
_selection_name = "LowerBound"
boundary = SelectionParameter()
[docs] def __init__(
self,
name: str = "Hard selection",
use_obs_value: bool = False,
use_distance: bool = False,
use_luminosity: bool = False,
use_flux: bool = False,
):
"""
A hard, lower bound selection on obs_value, distance,
luminosity or flux.
Selects values >= ``boundary``.
: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
:param boundary: Value of the selection boundary
:type boundary: :class:`SelectionParameter`
"""
super(LowerBound, self).__init__(
name=name,
use_distance=use_distance,
use_luminosity=use_luminosity,
use_obs_value=use_obs_value,
use_flux=use_flux,
)
[docs] def draw(self, size: int) -> None:
if self._use_distance:
values = self._distance
if self._use_obs_value:
values = self._observed_value
if self._use_luminosity:
values = self._luminosity
if self._use_flux:
values = self._observed_flux
self._selection = values >= self.boundary
[docs]class UpperBound(SelectionProbability):
_selection_name = "UpperBound"
boundary = SelectionParameter()
[docs] def __init__(
self,
name: str = "Hard selection",
use_obs_value: bool = False,
use_distance: bool = False,
use_luminosity: bool = False,
use_flux: bool = False,
):
"""
A hard, upper bound selection on obs_value, distance,
luminosity or flux.
Selects values <= ``boundary``.
: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
:param boundary: Value of the selection boundary
:type boundary: :class:`SelectionParameter`
"""
super(UpperBound, self).__init__(
name=name,
use_distance=use_distance,
use_luminosity=use_luminosity,
use_obs_value=use_obs_value,
use_flux=use_flux,
)
[docs] def draw(self, size: int) -> None:
if self._use_distance:
values = self._distance
if self._use_obs_value:
values = self._observed_value
if self._use_luminosity:
values = self._luminosity
if self._use_flux:
values = self._observed_flux
self._selection = values <= self.boundary
[docs]class SoftSelection(SelectionProbability):
_selection_name = "SoftSelection"
boundary = SelectionParameter()
strength = SelectionParameter(vmin=0)
[docs] def __init__(
self,
name: str = "Soft Selection",
use_obs_value: bool = False,
use_distance: bool = False,
use_luminosity: bool = False,
use_flux: bool = False,
) -> None:
"""
A soft selection using an inverse logit function either on the
log or linear value of the observed_value, distance,
luminosity or flux.
: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
:param boundary: Center of the inverse logit
:type boundary: :class:`SelectionParameter`
:param strength: Width of the logit
:type strength: :class:`SelectionParameter`
"""
super(SoftSelection, self).__init__(
name=name,
use_distance=use_distance,
use_luminosity=use_luminosity,
use_obs_value=use_obs_value,
use_flux=use_flux,
)
[docs] def draw(self, size: int, use_log=True) -> None:
if self._use_distance:
values = self._distance
if self._use_obs_value:
values = self._observed_value
if self._use_luminosity:
values = self._luminosity
if self._use_flux:
values = self._observed_flux
if not use_log:
probs = sf.expit(self.strength *
(values - self.boundary)) # type: np.ndarray
else:
probs = sf.expit(self.strength *
(np.log10(values) -
np.log10(self.boundary))) # type: np.ndarray
self._selection = stats.bernoulli.rvs(probs,
size=len(values)).astype(bool)
__all__ = [
"UnitySelection",
"BernoulliSelection",
"BoxSelection",
"UpperBound",
"UpperBound",
"SoftSelection",
]