Rewards#

Rewards assign weights to states in the state space. They are used behind the scenes to calculate different coalescent tree statistics but can also be passed directly to moment() to calculate more complex statistics.

class Reward[source]#

Bases: ABC

Base class for reward generation.

prod(*rewards: Reward)[source]#

Product of this reward with other rewards.

Parameters:

rewards (Reward) – Rewards to take the product with.

Return type:

ProductReward

Returns:

Product of the rewards.

sum(*rewards: Reward)[source]#

Sum of this reward with other rewards.

Parameters:

rewards (Reward) – Rewards to take the sum with.

Return type:

SumReward

Returns:

Sum of the rewards.

supports(state_space: Type[StateSpace])[source]#

Check if the reward supports the given state space.

Parameters:

state_space (Type[StateSpace]) – state space

Return type:

bool

Returns:

True if the reward supports the state space, False otherwise

static support(state_space: Type[StateSpace], rewards: Iterable[Reward])[source]#

Check if the rewards support the given state space.

Parameters:
  • state_space (Type[StateSpace]) – state space

  • rewards (Iterable[Reward]) – rewards

Return type:

bool

Returns:

True if the rewards support the state space, False otherwise

static requires_joint_state_space(rewards: Iterable[Reward])[source]#

Check whether any (possibly nested) reward can only be evaluated on the joint block-counting state space, i.e. supports it but neither the lineage- nor the block-counting state space (e.g. JointSFSReward).

Parameters:

rewards (Iterable[Reward]) – rewards

Return type:

bool

Returns:

True if some reward requires the joint state space

class LineageCountingReward[source]#

Bases: Reward, ABC

Base class for rewards that count lineages. Such rewards are compatible with LineageCountingStateSpace.

class BlockCountingReward[source]#

Bases: Reward, ABC

Base class for rewards that count blocks. Such rewards are compatible with BlockCountingStateSpace.

class JointBlockCountingReward[source]#

Bases: Reward, ABC

Base class for rewards that are compatible with JointBlockCountingStateSpace.

class TwoLocusBlockCountingReward[source]#

Bases: Reward, ABC

Base class for rewards that are compatible with TwoLocusBlockCountingStateSpace. This is deliberately not a JointBlockCountingReward: although the two-locus state space subclasses the joint one, its block axis encodes loci rather than populations, so generic joint-SFS rewards must not silently evaluate on it (and vice versa).

class JointSFSReward(config: Tuple[int, ...])[source]#

Bases: JointBlockCountingReward

Reward for a single bin of the joint (multi-population) site-frequency spectrum. The bin is identified by a descendant vector config = (k_0,...,k_{P-1}), and the reward of a state is the number of lineages (across all demes of residence and loci) whose descendant vector equals config.

__init__(config: Tuple[int, ...])[source]#

Initialize the reward.

Parameters:

config (Tuple[int, ...]) – The descendant vector identifying the joint SFS bin, i.e. the number of descendants subtended from each population.

class TwoLocusSFSReward(locus: int, count: int)[source]#

Bases: TwoLocusBlockCountingReward

Reward for one bin of the marginal site-frequency spectrum at a single locus in the two-locus block-counting state space. The reward of a state is the number of lineages that subtend exactly count samples at the given locus (i.e. whose two-locus descendant vector has component locus equal to count), regardless of how many they subtend at the other locus. The two-locus SFS is obtained as the cross-moment of two such rewards, one per locus.

__init__(locus: int, count: int)[source]#

Initialize the reward.

Parameters:
  • locus (int) – The locus index (0 or 1).

  • count (int) – The number of subtended samples at locus identifying the SFS bin.

class TreeHeightReward[source]#

Bases: LineageCountingReward, BlockCountingReward, JointBlockCountingReward

Reward for tree height. Note that when using multiple loci, this will provide the height of the locus with the highest tree.

class TotalTreeHeightReward[source]#

Bases: LineageCountingReward, BlockCountingReward

Reward based on tree height. When using multiple loci, this will provide the sum of the tree heights over all loci, regardless of whether they are linked or not.

class TotalBranchLengthReward[source]#

Bases: LineageCountingReward, BlockCountingReward, JointBlockCountingReward

Reward for total branch length. When using multiple loci, this will provide the sum of the total branch lengths over all loci, regardless of whether they are linked or not. Note that due to inherent limitation to rewards, we cannot determine the total branch length of the tree with the largest total branch length as done in TreeHeightReward.

class UnfoldedSFSReward(index: int)[source]#

Bases: SFSReward, BlockCountingReward

Reward for unfolded site frequency spectrum (SFS).

class FoldedSFSReward(index: int)[source]#

Bases: SFSReward, BlockCountingReward

Reward for folded site frequency spectrum (SFS).

class StateReward(state: int)[source]#

Bases: Reward

Reward for a specific state in the state space. This is useful for debugging or testing purposes.

__init__(state: int)[source]#

Initialize the reward.

Parameters:

state (int) – The state index to reward.

class LineageReward(n: int)[source]#

Bases: LineageCountingReward, JointBlockCountingReward

Reward for a specific number of lineages present in across all demes and loci. This reward can be used to, for example, track the individual coalescent times.

__init__(n: int)[source]#

Initialize the reward.

Parameters:

n (int) – The number of lineages to reward. Must be at least 2.

class DemeReward(pop: str)[source]#

Bases: LineageCountingReward, BlockCountingReward, JointBlockCountingReward

Reward fraction of lineages in a specific deme. Taking the product of this reward with another reward will result in a reward that only considers the specified deme. Use SumReward to marginalize over several demes.

__init__(pop: str)[source]#

Initialize the reward.

Parameters:

pop (str) – The population id to use.

class LocusReward(locus: int)[source]#

Bases: LineageCountingReward

Reward fraction of lineages in a specific locus. Taking the product of this reward with another reward will result in a reward that only considers the specified locus.

__init__(locus: int)[source]#

Initialize the reward.

Parameters:

locus (int) – The locus index to use.

class UnitReward[source]#

Bases: LineageCountingReward, BlockCountingReward, JointBlockCountingReward

Reward all states with 1 (including absorbing states).

class BlockCountingUnitReward[source]#

Bases: BlockCountingReward

Reward all states with 1 (including absorbing states), and only support block-counting state spaces. This reward can be used to force usage of the block-counting state space.

class TotalBranchLengthLocusReward(locus: int)[source]#

Bases: LocusReward, LineageCountingReward

Reward for total branch length per locus. This is needed as the taking the product of TotalBranchLengthReward and LocusReward will not work as expected (see CombinedReward).

class ProductReward(rewards: List[Reward])[source]#

Bases: CompositeReward

The product of multiple rewards.

class SumReward(rewards: List[Reward])[source]#

Bases: CompositeReward

The sum of multiple rewards.

class CombinedReward(rewards: List[Reward])[source]#

Bases: ProductReward

Class extending ProductReward to allow for more intuitive combination of rewards.

combinations: Dict[Tuple[Reward, Reward], Callable[[Reward, Reward], Reward]] = {(<class 'phasegen.rewards.TotalBranchLengthReward'>, <class 'phasegen.rewards.LocusReward'>): <function CombinedReward.<lambda>>}#

Dictionary of reward combinations

__init__(rewards: List[Reward])[source]#

Initialize the combined reward.

Parameters:

rewards (List[Reward]) – Rewards to combine

class CustomReward(func: ~typing.Callable[[~phasegen.state_space.StateSpace], ~numpy.ndarray], supports: ~typing.Callable[[~typing.Type[~phasegen.state_space.StateSpace]], bool] = <function CustomReward.<lambda>>)[source]#

Bases: Reward

Custom reward based on a user-defined function.

__init__(func: ~typing.Callable[[~phasegen.state_space.StateSpace], ~numpy.ndarray], supports: ~typing.Callable[[~typing.Type[~phasegen.state_space.StateSpace]], bool] = <function CustomReward.<lambda>>)[source]#

Initialize the custom reward.

Parameters:
  • func (Callable[[StateSpace], ndarray]) – The function to use to calculate the reward vector.

  • supports (Callable[[Type[StateSpace]], bool]) – The function to use to check if the reward supports the state space.

func: Callable[[StateSpace], ndarray]#

The function to calculate the reward vector

supports(state_space: Type[StateSpace])[source]#

Check if the reward supports the given state space.

Parameters:

state_space (Type[StateSpace]) – state space

Return type:

bool

Returns:

True if the reward supports the state space, False otherwise