aspartik.b3.operators

class DeltaExchange:

@dataclass(slots=True)
class DeltaExchange(Operator):
    """Scales a multidimensional parameter without changing its sum

    This operator is analogous to BEAST's `DeltaExchangeOperator`.  It picks
    two random dimensions from a parameter, samples a random delta from
    `distribution`, and increments one of them by delta and decrements the
    other one.
    """

    param: SimpleSequence[float]
    """
    The multidimensional parameter to edit.  Two random dimensions ones will be
    changed for each proposal.
    """
    factor: float
    """
    The move size is a random value between 0 and 1 multiplied by `factor`.
    """
    rng: RNG
    weight: float = 1

    def __post_init__(self):
        if len(self.param) <= 1:
            raise ValueError(f"`param` must have at least two dimensions")

    def propose(self) -> Proposal:
        rng = self.rng

        delta = rng.random_float() * self.factor

        dim_1 = rng.random_int(0, len(self.param))
        dim_2 = dim_1
        while dim_2 == dim_1:
            dim_1 = rng.random_int(0, len(self.param))

        self.param[dim_1] -= delta
        self.param[dim_2] += delta

        # The move is symmetrical, so the Hastings ratio is 0
        return Proposal.Hastings(0)
#

Scales a multidimensional parameter without changing its sum

This operator is analogous to BEAST's DeltaExchangeOperator. It picks two random dimensions from a parameter, samples a random delta from distribution, and increments one of them by delta and decrements the other one.

param: aspartik.utils.typing.SimpleSequence

#

The multidimensional parameter to edit. Two random dimensions ones will be changed for each proposal.

factor: float

#

The move size is a random value between 0 and 1 multiplied by factor.

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        rng = self.rng

        delta = rng.random_float() * self.factor

        dim_1 = rng.random_int(0, len(self.param))
        dim_2 = dim_1
        while dim_2 == dim_1:
            dim_1 = rng.random_int(0, len(self.param))

        self.param[dim_1] -= delta
        self.param[dim_2] += delta

        # The move is symmetrical, so the Hastings ratio is 0
        return Proposal.Hastings(0)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class NodeSlide:

@dataclass(slots=True)
class NodeSlide(Operator):
    """Slides the age of a random internal node between its parent and children

    This operator is similar to BEAST2's `EpochFlexOperator`: it will only
    affect the age of the selected node without altering the tree topology (a
    node cannot slide above its parent).
    """

    tree: Tree
    """The tree to edit."""
    distribution: Distribution
    """
    The distribution which will sample the new node height on the interval
    between its parent and the closest child.
    """
    rng: RNG
    weight: float = 1

    def __post_init__(self):
        assert_two_internals(self)

    def propose(self) -> Proposal:
        tree = self.tree

        node, parent = tree.random_nonroot_internal(self.rng)
        left, right = tree.children_of(node)

        oldest = tree.height_of(parent)
        youngest = max(tree.height_of(left), tree.height_of(right))

        new_height = sample_range(youngest, oldest, self.distribution, self.rng)

        tree.set_height(node, new_height)

        return Proposal.Hastings(0.0)
#

Slides the age of a random internal node between its parent and children

This operator is similar to BEAST2's EpochFlexOperator: it will only affect the age of the selected node without altering the tree topology (a node cannot slide above its parent).

tree: aspartik.b3.Tree

#

The tree to edit.

distribution: aspartik.stats.distributions.Distribution

#

The distribution which will sample the new node height on the interval between its parent and the closest child.

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        tree = self.tree

        node, parent = tree.random_nonroot_internal(self.rng)
        left, right = tree.children_of(node)

        oldest = tree.height_of(parent)
        youngest = max(tree.height_of(left), tree.height_of(right))

        new_height = sample_range(youngest, oldest, self.distribution, self.rng)

        tree.set_height(node, new_height)

        return Proposal.Hastings(0.0)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class ParamScale:

@dataclass(slots=True)
class ParamScale(Operator):
    """Scales one parameter

    This operator is analogous to BEAST2's `ScaleOperator`, except it only
    works for parameters.

    Note that this operator doesn't have the upper/lower bounds BEAST2's analog
    has.  Instead the `Bound` prior should be used to put limits on the
    parameter values.
    """

    param: Real
    """The parameter to scale."""
    factor: float
    """
    The scale ratio will be sampled from `(factor, 1 / factor)`.  So, the
    smaller the factor, the larger the moves proposed by this operator are.
    Also, this means that `factor` must be within `(0, 1)`.
    """
    distribution: Distribution
    """The distribution from which to sample the scaling factor."""
    rng: RNG
    weight: float = 1

    def __post_init__(self):
        assert_factor(self)

    def propose(self) -> Proposal:
        low, high = self.factor, 1 / self.factor
        scale = sample_range(low, high, self.distribution, self.rng)

        self.param.scale(scale)

        ratio = -log(scale)
        return Proposal.Hastings(ratio)
#

Scales one parameter

This operator is analogous to BEAST2's ScaleOperator, except it only works for parameters.

Note that this operator doesn't have the upper/lower bounds BEAST2's analog has. Instead the Bound prior should be used to put limits on the parameter values.

param: aspartik.b3.parameters.Real

#

The parameter to scale.

factor: float

#

The scale ratio will be sampled from (factor, 1 / factor). So, the smaller the factor, the larger the moves proposed by this operator are. Also, this means that factor must be within (0, 1).

distribution: aspartik.stats.distributions.Distribution

#

The distribution from which to sample the scaling factor.

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        low, high = self.factor, 1 / self.factor
        scale = sample_range(low, high, self.distribution, self.rng)

        self.param.scale(scale)

        ratio = -log(scale)
        return Proposal.Hastings(ratio)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class RandomWalk:

@dataclass(slots=True)
class RandomWalk(Operator):
    """
    Adds or subtracts values to a parameter irrespective of its value

    The delta is sampled uniformly from `[0, window)` and is negated half of
    the time.

    This parameter doesn't have bounds.  Use the `Bound` prior to restrict
    parameter values.
    """

    param: Real
    window: float
    rng: RNG
    lower: float = 0
    upper: float = math.inf
    boundary: Literal["reflect"] = "reflect"
    weight: float = 1

    _dist: Uniform = field(default_factory=lambda: Uniform(0, 1), init=False)

    def propose(self) -> Proposal:
        lower, upper = self.lower, self.upper
        rng = self.rng
        param = self.param

        diff = sample_range(0, self.window, self._dist, rng)
        if rng.random_bool():
            diff *= -1

        new_value = float(param) + diff

        if new_value < self.lower:
            match self.boundary:
                case "reflect":
                    if self.upper == math.inf:
                        new_value = lower + (lower - new_value)
                    else:
                        # TODO: ping-pong reflection
                        pass

        if new_value > self.upper:
            # TODO
            if self.lower == math.inf:
                new_value = upper + (upper - new_value)
            else:
                pass

        self.param.set(new_value)

        return Proposal.Hastings(0.0)
#

Adds or subtracts values to a parameter irrespective of its value

The delta is sampled uniformly from [0, window) and is negated half of the time.

This parameter doesn't have bounds. Use the Bound prior to restrict parameter values.

param: aspartik.b3.parameters.Real

#

window: float

#

rng: aspartik.rng.RNG

#

lower: float

#

upper: float

#

boundary: typing.Literal

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        lower, upper = self.lower, self.upper
        rng = self.rng
        param = self.param

        diff = sample_range(0, self.window, self._dist, rng)
        if rng.random_bool():
            diff *= -1

        new_value = float(param) + diff

        if new_value < self.lower:
            match self.boundary:
                case "reflect":
                    if self.upper == math.inf:
                        new_value = lower + (lower - new_value)
                    else:
                        # TODO: ping-pong reflection
                        pass

        if new_value > self.upper:
            # TODO
            if self.lower == math.inf:
                new_value = upper + (upper - new_value)
            else:
                pass

        self.param.set(new_value)

        return Proposal.Hastings(0.0)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class RootSlide:

@dataclass(slots=True)
class RootSlide(Operator):
    """
    Scales the height of the root node

    The height will be scaled between `factor` and `1 / factor` sampled with
    `distribution`.  If the move will put the root below one of its children
    the operation gets rejected.
    """

    tree: Tree
    """The tree to edit."""
    factor: float
    distribution: Distribution
    """The distribution to draw the height move distance from"""
    rng: RNG
    weight: float = 1

    def __post_init__(self):
        assert_factor(self)

    def propose(self) -> Proposal:
        tree = self.tree
        rng = self.rng
        root = tree.root

        low, high = self.factor, 1 / self.factor
        scale = sample_range(low, high, self.distribution, rng)

        old_height = tree.height_of(root)

        left, right = tree.children_of(root)
        lower_height = max(tree.height_of(left), tree.height_of(right))

        new_height = (old_height - lower_height) * scale + lower_height

        tree.set_height(root, new_height)

        return Proposal.Hastings(-log(scale))
#

Scales the height of the root node

The height will be scaled between factor and 1 / factor sampled with distribution. If the move will put the root below one of its children the operation gets rejected.

tree: aspartik.b3.Tree

#

The tree to edit.

factor: float

#

distribution: aspartik.stats.distributions.Distribution

#

The distribution to draw the height move distance from

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        tree = self.tree
        rng = self.rng
        root = tree.root

        low, high = self.factor, 1 / self.factor
        scale = sample_range(low, high, self.distribution, rng)

        old_height = tree.height_of(root)

        left, right = tree.children_of(root)
        lower_height = max(tree.height_of(left), tree.height_of(right))

        new_height = (old_height - lower_height) * scale + lower_height

        tree.set_height(root, new_height)

        return Proposal.Hastings(-log(scale))
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class SubtreePruneRegraft:

@dataclass(slots=True)
class SubtreePruneRegraft(Operator):
    """
    Fixed height subtree and regraft move.

    This operator was described in [Hoehna et al 2008][l], section 3.2.7.  The
    move selects a random node `i` and its parent `i_parent`.  It then selects
    a random edge whose height overlaps with the height of `i_parent`.
    `i_parent` is spliced into the middle of this edge.

    [l]: https://alexeidrummond.org/assets/publications/2008-hoehna-evalution.pdf
    """

    tree: Tree
    """The tree to edit."""
    rng: RNG
    weight: float = 1

    def __post_init__(self):
        assert_two_internals(self)

    def propose(self) -> Proposal:
        rng = self.rng
        tree = self.tree
        root = tree.root

        node = tree.random_node(rng)
        parent = tree.parent_of(node)
        while node == root or parent == root:
            node = tree.random_node(rng)
            parent = tree.parent_of(node)

        assert parent is not None
        grandparent = tree.parent_of(parent)
        assert grandparent is not None

        sibling = tree.other_child(parent, node)

        parent_height = tree.height_of(parent)

        edge = tree.random_intersecting_edge(parent_height, rng)

        if edge is None:
            return Proposal.Reject()

        other, other_parent = tree.edge_nodes(edge)

        # Tree update
        grandparent_to_parent = tree.edge_index(parent)
        parent_to_sibling = tree.edge_index(sibling)
        other_parent_to_other = tree.edge_index(other)

        tree.update_edge(grandparent_to_parent, sibling)
        tree.update_edge(other_parent_to_other, parent)
        tree.update_edge(parent_to_sibling, other)

        tree.validate()

        return Proposal.Hastings(0.0)
#

Fixed height subtree and regraft move.

This operator was described in Hoehna et al 2008, section 3.2.7. The move selects a random node i and its parent i_parent. It then selects a random edge whose height overlaps with the height of i_parent. i_parent is spliced into the middle of this edge.

tree: aspartik.b3.Tree

#

The tree to edit.

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        rng = self.rng
        tree = self.tree
        root = tree.root

        node = tree.random_node(rng)
        parent = tree.parent_of(node)
        while node == root or parent == root:
            node = tree.random_node(rng)
            parent = tree.parent_of(node)

        assert parent is not None
        grandparent = tree.parent_of(parent)
        assert grandparent is not None

        sibling = tree.other_child(parent, node)

        parent_height = tree.height_of(parent)

        edge = tree.random_intersecting_edge(parent_height, rng)

        if edge is None:
            return Proposal.Reject()

        other, other_parent = tree.edge_nodes(edge)

        # Tree update
        grandparent_to_parent = tree.edge_index(parent)
        parent_to_sibling = tree.edge_index(sibling)
        other_parent_to_other = tree.edge_index(other)

        tree.update_edge(grandparent_to_parent, sibling)
        tree.update_edge(other_parent_to_other, parent)
        tree.update_edge(parent_to_sibling, other)

        tree.validate()

        return Proposal.Hastings(0.0)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class SubtreeSlide:

@dataclass(slots=True)
class SubtreeSlide(Operator):
    """"""

    tree: Tree
    """The tree to edit."""
    distribution: Sample[float]
    """
    The distribution which will sample the new node height on the interval
    between its parent and the closest child.
    """
    rng: RNG
    weight: float = 1

    def propose(self) -> Proposal:
        """
        If there are no non-root internal nodes, the operator will bail with
        `Proposal.Reject`.
        """

        rng = self.rng
        tree = self.tree
        root = tree.root

        ratio = 0.0

        # automatically fail on trees without non-root internal nodes
        if tree.num_internals == 1:
            return Proposal.Reject()

        node, sibling, parent, grandparent = family(tree, rng)
        delta = self.distribution.sample(rng)
        old_height = tree.height_of(parent)
        new_height = old_height + delta

        if delta > 0:
            if grandparent is not None and tree.height_of(grandparent) < new_height:
                # two nodes whose edge intersects `new_height`
                new_child = parent
                new_parent = grandparent

                while tree.height_of(new_parent) < new_height:
                    new_child = new_parent
                    new_parent = tree.parent_of(new_child)
                    if new_parent is None:
                        break

                parent_to_sibling = tree.edge_index(sibling)
                grandparent_to_parent = tree.edge_index(parent)

                if new_parent is None:
                    # new_child was the root
                    tree.update_edge(parent_to_sibling, new_child)
                    tree.update_edge(grandparent_to_parent, sibling)
                    tree.set_root(parent)
                else:
                    new_parent_to_new_child = tree.edge_index(new_child)

                    tree.update_edge(parent_to_sibling, new_child)
                    tree.update_edge(grandparent_to_parent, sibling)
                    tree.update_edge(new_parent_to_new_child, parent)

                num_reverse_sources = len(intersections(tree, new_child, old_height))
                ratio = -log(num_reverse_sources)
        else:
            if tree.height_of(node) > new_height:
                return Proposal.Reject()

            if tree.height_of(sibling) > new_height:
                # topological changes

                destinations = intersections(tree, sibling, new_height)
                random_idx = rng.random_int(0, len(destinations))
                new_child = destinations[random_idx]
                new_parent = tree.parent_of(new_child)

                parent_to_sibling = tree.edge_index(sibling)
                new_parent_to_new_child = tree.edge_index(new_child)

                if parent == tree.root:
                    tree.update_edge(parent_to_sibling, new_child)
                    tree.update_edge(new_parent_to_new_child, parent)

                    tree.set_root(sibling)
                else:
                    grandparent_to_parent = tree.edge_index(parent)

                    tree.update_edge(parent_to_sibling, new_child)
                    tree.update_edge(grandparent_to_parent, sibling)
                    tree.update_edge(new_parent_to_new_child, parent)

                ratio = -log(len(destinations))

        tree.set_height(parent, new_height)

        tree.validate()

        return Proposal.Hastings(ratio)
#

tree: aspartik.b3.Tree

#

The tree to edit.

distribution: aspartik.stats.distributions.Sample

#

The distribution which will sample the new node height on the interval between its parent and the closest child.

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        """
        If there are no non-root internal nodes, the operator will bail with
        `Proposal.Reject`.
        """

        rng = self.rng
        tree = self.tree
        root = tree.root

        ratio = 0.0

        # automatically fail on trees without non-root internal nodes
        if tree.num_internals == 1:
            return Proposal.Reject()

        node, sibling, parent, grandparent = family(tree, rng)
        delta = self.distribution.sample(rng)
        old_height = tree.height_of(parent)
        new_height = old_height + delta

        if delta > 0:
            if grandparent is not None and tree.height_of(grandparent) < new_height:
                # two nodes whose edge intersects `new_height`
                new_child = parent
                new_parent = grandparent

                while tree.height_of(new_parent) < new_height:
                    new_child = new_parent
                    new_parent = tree.parent_of(new_child)
                    if new_parent is None:
                        break

                parent_to_sibling = tree.edge_index(sibling)
                grandparent_to_parent = tree.edge_index(parent)

                if new_parent is None:
                    # new_child was the root
                    tree.update_edge(parent_to_sibling, new_child)
                    tree.update_edge(grandparent_to_parent, sibling)
                    tree.set_root(parent)
                else:
                    new_parent_to_new_child = tree.edge_index(new_child)

                    tree.update_edge(parent_to_sibling, new_child)
                    tree.update_edge(grandparent_to_parent, sibling)
                    tree.update_edge(new_parent_to_new_child, parent)

                num_reverse_sources = len(intersections(tree, new_child, old_height))
                ratio = -log(num_reverse_sources)
        else:
            if tree.height_of(node) > new_height:
                return Proposal.Reject()

            if tree.height_of(sibling) > new_height:
                # topological changes

                destinations = intersections(tree, sibling, new_height)
                random_idx = rng.random_int(0, len(destinations))
                new_child = destinations[random_idx]
                new_parent = tree.parent_of(new_child)

                parent_to_sibling = tree.edge_index(sibling)
                new_parent_to_new_child = tree.edge_index(new_child)

                if parent == tree.root:
                    tree.update_edge(parent_to_sibling, new_child)
                    tree.update_edge(new_parent_to_new_child, parent)

                    tree.set_root(sibling)
                else:
                    grandparent_to_parent = tree.edge_index(parent)

                    tree.update_edge(parent_to_sibling, new_child)
                    tree.update_edge(grandparent_to_parent, sibling)
                    tree.update_edge(new_parent_to_new_child, parent)

                ratio = -log(len(destinations))

        tree.set_height(parent, new_height)

        tree.validate()

        return Proposal.Hastings(ratio)
#

If there are no non-root internal nodes, the operator will bail with Proposal.Reject.

class BeastNarrowExchange:

@dataclass(slots=True)
class BeastNarrowExchange(Operator):
    """
    Narrow exchange operator compatible with BEASTs `narrowExchange`
    """

    tree: Tree
    rng: RNG
    weight: float = 1

    def propose(self):
        rng = self.rng
        tree = self.tree
        root = tree.root

        node = tree.random_node(rng)
        while node == root or tree.parent_of(node) == root:
            node = tree.random_node(rng)

        parent = tree.parent_of(node)
        assert parent is not None
        grandparent = tree.parent_of(parent)
        assert grandparent is not None
        uncle = tree.other_child(grandparent, parent)

        if tree.height_of(uncle) < tree.height_of(parent):
            tree.swap_parents(node, uncle)
            return Proposal.Hastings(0.0)
        else:
            return Proposal.Reject()
#

Narrow exchange operator compatible with BEASTs narrowExchange

tree: aspartik.b3.Tree

#

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self)

    def propose(self):
        rng = self.rng
        tree = self.tree
        root = tree.root

        node = tree.random_node(rng)
        while node == root or tree.parent_of(node) == root:
            node = tree.random_node(rng)

        parent = tree.parent_of(node)
        assert parent is not None
        grandparent = tree.parent_of(parent)
        assert grandparent is not None
        uncle = tree.other_child(grandparent, parent)

        if tree.height_of(uncle) < tree.height_of(parent):
            tree.swap_parents(node, uncle)
            return Proposal.Hastings(0.0)
        else:
            return Proposal.Reject()
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class BeastWideExchange:

@dataclass(slots=True)
class BeastWideExchange(Operator):
    """
    Wide exchange operator compatible with BEASTs `wideExchange`
    """

    tree: Tree
    rng: RNG
    weight: float = 1

    def propose(self):
        rng = self.rng
        tree = self.tree
        root = tree.root

        node_a = tree.random_node(rng)
        while node_a == root:
            node_a = tree.random_node(rng)

        node_b = tree.random_node(rng)
        while node_b == root or node_b == node_a:
            node_b = tree.random_node(rng)

        node_a_parent = tree.parent_of(node_a)
        node_b_parent = tree.parent_of(node_b)
        assert node_a_parent is not None
        assert node_b_parent is not None

        if (
            node_a_parent != node_b_parent
            and node_a != node_b_parent
            and node_b != node_a_parent
            and tree.height_of(node_a) < tree.height_of(node_b_parent)
            and tree.height_of(node_b) < tree.height_of(node_a_parent)
        ):
            tree.swap_parents(node_a, node_b)
            return Proposal.Hastings(0.0)
        else:
            return Proposal.Reject()
#

Wide exchange operator compatible with BEASTs wideExchange

tree: aspartik.b3.Tree

#

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self)

    def propose(self):
        rng = self.rng
        tree = self.tree
        root = tree.root

        node_a = tree.random_node(rng)
        while node_a == root:
            node_a = tree.random_node(rng)

        node_b = tree.random_node(rng)
        while node_b == root or node_b == node_a:
            node_b = tree.random_node(rng)

        node_a_parent = tree.parent_of(node_a)
        node_b_parent = tree.parent_of(node_b)
        assert node_a_parent is not None
        assert node_b_parent is not None

        if (
            node_a_parent != node_b_parent
            and node_a != node_b_parent
            and node_b != node_a_parent
            and tree.height_of(node_a) < tree.height_of(node_b_parent)
            and tree.height_of(node_b) < tree.height_of(node_a_parent)
        ):
            tree.swap_parents(node_a, node_b)
            return Proposal.Hastings(0.0)
        else:
            return Proposal.Reject()
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class NarrowExchange:

@dataclass(slots=True)
class NarrowExchange(Operator):
    """Exchanges the parents of two neighbouring nodes

    This operator is analogous to BEAST2's `Exchange` operator with `isNarrow`
    set to true.  It finds a grandparent (internal node both of whose children
    are also internal) with two kids: `parent` and `uncle` (uncle is younger
    than the parent).  And one of the children of `parent` is swapped with
    `uncle`.
    """

    tree: Tree
    rng: RNG
    weight: float = 1

    def __post_init__(self):
        assert_two_internals(self)

    def propose(self) -> Proposal:
        tree = self.tree

        num_grandparents_before = tree.num_grandparents()
        if num_grandparents_before == 0:
            # no grandparents to pick `grandparent` from
            return Proposal.Reject()

        while True:
            grandparent = tree.random_internal(self.rng)
            if tree.is_grandparent(grandparent):
                break

        left, right = tree.children_of(grandparent)
        if tree.height_of(left) > tree.height_of(right):
            parent, uncle = left, right
        elif tree.height_of(right) > tree.height_of(left):
            parent, uncle = right, left
        else:
            return Proposal.Reject()

        # guaranteed because `grandparent` is a grandparent
        assert isinstance(parent, Internal)
        assert isinstance(uncle, Internal)

        before = int(tree.is_grandparent(parent)) + int(tree.is_grandparent(uncle))

        if self.rng.random_bool(0.5):
            child = tree.children_of(parent)[0]
        else:
            child = tree.children_of(parent)[1]

        tree.swap_parents(uncle, child)

        after = int(tree.is_grandparent(parent)) + int(tree.is_grandparent(uncle))
        num_grandparents_after = num_grandparents_before - before + after
        ratio = log(num_grandparents_before / num_grandparents_after)

        return Proposal.Hastings(ratio)
#

Exchanges the parents of two neighbouring nodes

This operator is analogous to BEAST2's Exchange operator with isNarrow set to true. It finds a grandparent (internal node both of whose children are also internal) with two kids: parent and uncle (uncle is younger than the parent). And one of the children of parent is swapped with uncle.

tree: aspartik.b3.Tree

#

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        tree = self.tree

        num_grandparents_before = tree.num_grandparents()
        if num_grandparents_before == 0:
            # no grandparents to pick `grandparent` from
            return Proposal.Reject()

        while True:
            grandparent = tree.random_internal(self.rng)
            if tree.is_grandparent(grandparent):
                break

        left, right = tree.children_of(grandparent)
        if tree.height_of(left) > tree.height_of(right):
            parent, uncle = left, right
        elif tree.height_of(right) > tree.height_of(left):
            parent, uncle = right, left
        else:
            return Proposal.Reject()

        # guaranteed because `grandparent` is a grandparent
        assert isinstance(parent, Internal)
        assert isinstance(uncle, Internal)

        before = int(tree.is_grandparent(parent)) + int(tree.is_grandparent(uncle))

        if self.rng.random_bool(0.5):
            child = tree.children_of(parent)[0]
        else:
            child = tree.children_of(parent)[1]

        tree.swap_parents(uncle, child)

        after = int(tree.is_grandparent(parent)) + int(tree.is_grandparent(uncle))
        num_grandparents_after = num_grandparents_before - before + after
        ratio = log(num_grandparents_before / num_grandparents_after)

        return Proposal.Hastings(ratio)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class WideExchange:

@dataclass(slots=True)
class WideExchange(Operator):
    """Exchanges the parent of two random nodes

    This operator is analogous to BEAST2's `Exchange` operator with `isNarrow`
    set to false.  It picks two random nodes in the tree (they could be either
    leaves or internals) and swaps their parents.

    If a randomly selected move is impossible (a parent would be younger than
    its child) the operator aborts with `Proposal.Reject`.
    """

    tree: Tree
    rng: RNG
    weight: float = 1

    def propose(self) -> Proposal:
        tree = self.tree

        root = tree.root

        i = tree.random_node(self.rng)
        while i == root:
            i = tree.random_node(self.rng)

        j = None
        while j is None or j == i or j == root:
            j = tree.random_node(self.rng)
        assert j is not None

        i_parent = tree.parent_of(i)
        if i_parent is None:
            return Proposal.Reject()
        j_parent = tree.parent_of(j)
        if j_parent is None:
            return Proposal.Reject()

        # Abort if j and i are parent-child or if one of the parents would be
        # younger than its new child or if the two selected nodes.
        if (
            j != i_parent
            and i != j_parent
            and tree.height_of(j) < tree.height_of(i_parent)
            and tree.height_of(i) < tree.height_of(j_parent)
        ):
            tree.swap_parents(i, j)

            return Proposal.Hastings(0.0)
        else:
            return Proposal.Reject()
#

Exchanges the parent of two random nodes

This operator is analogous to BEAST2's Exchange operator with isNarrow set to false. It picks two random nodes in the tree (they could be either leaves or internals) and swaps their parents.

If a randomly selected move is impossible (a parent would be younger than its child) the operator aborts with Proposal.Reject.

tree: aspartik.b3.Tree

#

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        tree = self.tree

        root = tree.root

        i = tree.random_node(self.rng)
        while i == root:
            i = tree.random_node(self.rng)

        j = None
        while j is None or j == i or j == root:
            j = tree.random_node(self.rng)
        assert j is not None

        i_parent = tree.parent_of(i)
        if i_parent is None:
            return Proposal.Reject()
        j_parent = tree.parent_of(j)
        if j_parent is None:
            return Proposal.Reject()

        # Abort if j and i are parent-child or if one of the parents would be
        # younger than its new child or if the two selected nodes.
        if (
            j != i_parent
            and i != j_parent
            and tree.height_of(j) < tree.height_of(i_parent)
            and tree.height_of(i) < tree.height_of(j_parent)
        ):
            tree.swap_parents(i, j)

            return Proposal.Hastings(0.0)
        else:
            return Proposal.Reject()
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class TreeScale:

@dataclass(slots=True)
class TreeScale(Operator):
    """Scales the age of the entire tree

    This parameter is analogous to BEAST2's `ScaleOperator` when it's used on a
    tree.  It will scale all internal nodes by a random scale which is randomly
    picked depending on `factor` and `distribution`.
    """

    tree: Tree
    """The tree to scale."""
    factor: float
    """
    The scaling ratio will be sampled from `(factor, 1 / factor)`.  So, the
    factor must be between 0 and 1 and the smaller it is the larger the steps
    will be.
    """
    distribution: Distribution
    """Distribution from which the scale is sampled."""
    rng: RNG
    weight: float = 1

    def __post_init__(self):
        assert_factor(self)

    def propose(self) -> Proposal:
        tree = self.tree
        rng = self.rng
        root = tree.root

        low, high = self.factor, 1 / self.factor
        scale = sample_range(low, high, self.distribution, rng)

        try:
            tree.scale(scale)
        except:
            return Proposal.Reject()

        ratio = log(scale) * (tree.num_internals - 2)
        return Proposal.Hastings(ratio)
#

Scales the age of the entire tree

This parameter is analogous to BEAST2's ScaleOperator when it's used on a tree. It will scale all internal nodes by a random scale which is randomly picked depending on factor and distribution.

tree: aspartik.b3.Tree

#

The tree to scale.

factor: float

#

The scaling ratio will be sampled from (factor, 1 / factor). So, the factor must be between 0 and 1 and the smaller it is the larger the steps will be.

distribution: aspartik.stats.distributions.Distribution

#

Distribution from which the scale is sampled.

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        tree = self.tree
        rng = self.rng
        root = tree.root

        low, high = self.factor, 1 / self.factor
        scale = sample_range(low, high, self.distribution, rng)

        try:
            tree.scale(scale)
        except:
            return Proposal.Reject()

        ratio = log(scale) * (tree.num_internals - 2)
        return Proposal.Hastings(ratio)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class UpDown:

@dataclass(slots=True)
class UpDown(Operator):
    """
    Scales `up` by factor and `down` by inverse factor on each step

    See `factor` documentation for how the scaling factor is sampled.  Note
    that it can be both more and less than 1, so `up` can sometimes be scaled
    down (and so `down` would be scaled up on those steps).

    Though, when using the uniform distribution, scaling factor will be biased
    towards values `> 1`, so `up` will be scaled up more often than not (not
    accounting for rejections).
    """

    up: Scalable
    """The parameter to scale up."""
    down: Scalable
    """The parameter to scale down."""
    factor: float
    """
    The scale ratio will be sampled from `(factor, 1 / factor)`.  So, the
    smaller the factor, the larger the moves proposed by this operator are.
    This also means that `factor` must be within `(0, 1)`.
    """
    distribution: Distribution
    """The distribution from which to sample the scaling factor."""
    rng: RNG
    weight: float = 1

    def __post_init__(self):
        assert_factor(self)

    def propose(self) -> Proposal:
        low, high = self.factor, 1 / self.factor
        scale = sample_range(low, high, self.distribution, self.rng)

        try:
            num_scaling_up = self.up.scale(scale)
            num_scaling_down = self.down.scale(1 / scale)
        except:
            return Proposal.Reject()

        ratio = log(scale) * (num_scaling_up - num_scaling_down - 2)
        return Proposal.Hastings(ratio)
#

Scales up by factor and down by inverse factor on each step

See factor documentation for how the scaling factor is sampled. Note that it can be both more and less than 1, so up can sometimes be scaled down (and so down would be scaled up on those steps).

Though, when using the uniform distribution, scaling factor will be biased towards values > 1, so up will be scaled up more often than not (not accounting for rejections).

up: aspartik.b3.parameters.Scalable

#

The parameter to scale up.

down: aspartik.b3.parameters.Scalable

#

The parameter to scale down.

factor: float

#

The scale ratio will be sampled from (factor, 1 / factor). So, the smaller the factor, the larger the moves proposed by this operator are. This also means that factor must be within (0, 1).

distribution: aspartik.stats.distributions.Distribution

#

The distribution from which to sample the scaling factor.

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        low, high = self.factor, 1 / self.factor
        scale = sample_range(low, high, self.distribution, self.rng)

        try:
            num_scaling_up = self.up.scale(scale)
            num_scaling_down = self.down.scale(1 / scale)
        except:
            return Proposal.Reject()

        ratio = log(scale) * (num_scaling_up - num_scaling_down - 2)
        return Proposal.Hastings(ratio)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class WilsonBalding:

@dataclass(slots=True)
class WilsonBalding(Operator):
    """A version of a subtree regraft move

    Introduced in [this paper][paper], it picks a random subtree and inserts it
    in-between two other nodes.

    [paper]: https://doi.org/10.1093/genetics/161.3.1307
    """

    tree: Tree
    rng: RNG
    weight: float = 1

    def propose(self) -> Proposal:
        tree = self.tree
        rng = self.rng

        # pick a random non-root node
        while True:
            i_parent = tree.random_internal(rng)
            i_grandparent = tree.parent_of(i_parent)
            if i_grandparent is not None:
                break
        i_parent_height = tree.height_of(i_parent)

        i, i_brother = tree.children_of(i_parent)
        if rng.random_bool():
            i, i_brother = i_brother, i

        # Pick a node j_parent, such that it's above i_parent and one of its
        # children is below i_parent
        while True:
            j_parent = tree.random_internal(rng)
            j, j_brother = tree.children_of(j_parent)
            if rng.random_bool():
                j = j_brother

            if tree.height_of(j_parent) > i_parent_height > tree.height_of(j):
                break

        before = tree.height_of(i_grandparent) - max(
            tree.height_of(i), tree.height_of(i_brother)
        )
        after = tree.height_of(j_parent) - max(tree.height_of(i), tree.height_of(j))
        ratio = log(after / before)

        i_grandparent_to_i_parent = tree.edge_index(i_parent)
        j_parent_to_j = tree.edge_index(j)
        i_parent_to_i_brother = tree.edge_index(i_brother)

        # Cut out i_parent and replace it with a direct edge from grandparent
        # to i_brother
        tree.update_edge(i_grandparent_to_i_parent, i_brother)
        # Hook up i_parent to j_parent.  It's fine because we checked that
        # i_parent is lower than j_parent when selecting j
        tree.update_edge(j_parent_to_j, i_parent)
        # Replace i_brother edge from i_parent with j.  Once again, we've
        # enforced i_parent being above j earlier
        tree.update_edge(i_parent_to_i_brother, j)

        return Proposal.Hastings(ratio)
#

A version of a subtree regraft move

Introduced in this paper, it picks a random subtree and inserts it in-between two other nodes.

tree: aspartik.b3.Tree

#

rng: aspartik.rng.RNG

#

weight: float

#

Influences the probability of the operator being picked

On each step MCMC picks a random operator from the list passed to it. It uses this value to weight them. So, the larger it is, the more often the operator will be picked, and visa versa. This value is read once on startup. Therefore, if it's changed mid-execution the old cached value will still be used.

def propose(self) -> aspartik.b3.Proposal

    def propose(self) -> Proposal:
        tree = self.tree
        rng = self.rng

        # pick a random non-root node
        while True:
            i_parent = tree.random_internal(rng)
            i_grandparent = tree.parent_of(i_parent)
            if i_grandparent is not None:
                break
        i_parent_height = tree.height_of(i_parent)

        i, i_brother = tree.children_of(i_parent)
        if rng.random_bool():
            i, i_brother = i_brother, i

        # Pick a node j_parent, such that it's above i_parent and one of its
        # children is below i_parent
        while True:
            j_parent = tree.random_internal(rng)
            j, j_brother = tree.children_of(j_parent)
            if rng.random_bool():
                j = j_brother

            if tree.height_of(j_parent) > i_parent_height > tree.height_of(j):
                break

        before = tree.height_of(i_grandparent) - max(
            tree.height_of(i), tree.height_of(i_brother)
        )
        after = tree.height_of(j_parent) - max(tree.height_of(i), tree.height_of(j))
        ratio = log(after / before)

        i_grandparent_to_i_parent = tree.edge_index(i_parent)
        j_parent_to_j = tree.edge_index(j)
        i_parent_to_i_brother = tree.edge_index(i_brother)

        # Cut out i_parent and replace it with a direct edge from grandparent
        # to i_brother
        tree.update_edge(i_grandparent_to_i_parent, i_brother)
        # Hook up i_parent to j_parent.  It's fine because we checked that
        # i_parent is lower than j_parent when selecting j
        tree.update_edge(j_parent_to_j, i_parent)
        # Replace i_brother edge from i_parent with j.  Once again, we've
        # enforced i_parent being above j earlier
        tree.update_edge(i_parent_to_i_brother, j)

        return Proposal.Hastings(ratio)
#

Proposes a new MCMC step

It is presumed that the operator will store all the references to parameters and trees it wants to edit and will change them accordingly. If a move cannot be proposed for any reason Proposal.Reject should be returned. MCMC will deal with rolling back the state.

class EpochScale:

#

Scales a random epoch in a tree

This parameter is analogous to BEAST2's ScaleOperator when it's used on a tree. It will scale the full tree (so, for now, only its internal nodes, since leaves all have the height of 0).

factor

#

The scaling ratio will be sampled from (factor, 1 / factor). So, the factor must be between 0 and 1 and the smaller it is the larger the steps will be.

tree

#

weight

#

distribution

#

rng

#

def propose(self, /)

#

class SubtreeLeap:

#

Moves a node a distance, changing the topology randomly

First, a distance delta is sampled from the distribution. The operator selects a random node and all edges delta away from that node (down if the delta is negative or up and down if it's positive). One of those edges is randomly selected and the node is spliced into it. If the delta is above the root, the node will become the new root.

tree

#

The tree to edit.

distribution

#

The distribution to draw the height move distance from

rng

#

weight

#

def propose(self, /)

#