codeine 0.1.0__py3-none-any.whl

codeine/graph/compile.py

@@ -0,0 +1,489 @@
+import math
+
+from typing import Dict, NamedTuple, Tuple, List, Optional, TYPE_CHECKING
+
+from codeine.constraints.banned import BannedTrackerStateId, BannedTrackerAdvanceResult
+from codeine.constraints.base import ConstraintState
+from codeine.graph.nodes import CodonNode, Node, ContextNode
+
+if TYPE_CHECKING:
+    from codeine.graph.view import CodonGraphView
+
+
+class TraversalState(NamedTuple):
+    """
+    The traversal state consists of the current node, plus a summary of the relevant
+    parts of how we got there. For example, it tracks whether we have seen parts of
+    banned sequences, and can track nucleotide or codon properties.
+
+    Different graph traversal histories that produce the same traversal state are
+    collapsed and are equivalent under this framework.
+    """
+    node: Node
+    banned_tracker_state_id: BannedTrackerStateId
+    constraint_state: ConstraintState
+
+
+class ChoiceResult(NamedTuple):
+    """
+    Cached result of taking one graph choice from one compiled state. The
+    "choice" is the graph edge label, i.e. a codon or a context sequence.
+
+    Each ChoiceResult is specific to its location in the graph. The descendant
+    counts and log mass are calculated iteratively by summing the values of
+    downstream nodes.
+    """
+    choice: str
+    descendant_count: int
+    descendant_log_mass: float
+    next_state_id: Optional[int]
+    is_coding: bool
+
+
+class CompiledView(NamedTuple):
+    """
+    Cached data for a compiled CodonGraphView, to speed up sampling and enumeration.
+    """
+    initial_state: TraversalState
+    initial_state_id: int
+    states: Tuple[TraversalState, ...]
+
+    # Compiled graph choices (lookup):
+    # state ID -> choice -> ChoiceResult
+    # Used for fast sequence validation and graph traversal in the graph view.
+    choices_by_state_id: Tuple[Dict[str, ChoiceResult], ...]
+
+    # Compiled graph choices (iteration):
+    # state ID -> ChoiceResults in graph order
+    # Used for fast sampling and sequence enumeration in the graph view.
+    choice_results_by_state_id: Tuple[Tuple[ChoiceResult, ...], ...]
+
+    n_valid_sequences: int
+
+
+class ViewCompiler:
+    """
+    Compile a CodonGraphView into cached choice, count, and sampling data.
+    """
+
+    def __init__(self, view: 'CodonGraphView') -> None:
+        self.view = view
+        self.graph = view.graph
+
+        self.banned_tracker = view.banned_tracker
+        self.has_banned_tracker = not self.banned_tracker.is_trivial
+
+        self.path_constraint = view.path_constraint
+        self.has_path_constraint = self.path_constraint is not None
+
+        self.state_ids: Dict[TraversalState, int] = {}
+        self.states: List[TraversalState] = []
+
+        # Dynamic-programming totals:
+        # state ID -> (descendant count, descendant log mass)
+        # Avoids repeatedly recomputing subtree sizes and probability masses.
+        self.totals_by_state_id: List[Optional[Tuple[int, float]]] = []
+
+        # Banned-sequence tracker transitions:
+        # (node, tracker state, choice) -> tracker result
+        # Avoids recomputing tracker advances during compilation.
+        self.banned_advance_cache: Dict[Tuple[Node, BannedTrackerStateId, str], BannedTrackerAdvanceResult] = {}
+
+        # Traversal transition table:
+        # state ID -> [(choice, child state ID), ...]
+        # Avoids rediscovering successor states during compilation.
+        self.child_results_by_state_id: List[Optional[List[Tuple[str, int]]]] = []
+
+        # Compiled graph choices (lookup):
+        # state ID -> choice -> ChoiceResult
+        # Used for fast sequence validation and graph traversal in the graph view.
+        self.choices_by_state_id: List[Optional[Dict[str, ChoiceResult]]] = []
+
+        # The cached log-ified codon weights, to avoid repeated log calculations
+        self.log_codon_weights = {
+            codon: math.log(weight)
+            for codon, weight in self.graph.cw.weights.items()
+            if weight > 0
+        }
+
+        # Cached version of the choices available at each node, taking into account fixed codons & pins
+        self.choices_by_node = {
+            node: tuple(self._get_choices_for_node(node))
+            for node in self.graph.nodes
+            if node is not self.graph.final_node
+        }
+
+    def compile(self) -> CompiledView:
+        """
+        Compile descendant counts, graph choices, and samplers.
+
+        Returns
+        -------
+        CompiledView
+            A compiled view.
+        """
+        initial_state = self._initial_state()
+        initial_state_id = self._get_or_register_state_id(initial_state)
+
+        self._compile_from(initial_state_id)
+
+        choices_by_state_id = tuple(
+            choices or {}
+            for choices in self.choices_by_state_id
+        )
+
+        choice_results_by_state_id = tuple(
+            tuple(choices.values()) if choices else ()
+            for choices in self.choices_by_state_id
+        )
+
+        initial_total = self.totals_by_state_id[initial_state_id]
+        assert initial_total is not None
+
+        return CompiledView(
+            initial_state=initial_state,
+            initial_state_id=initial_state_id,
+            states=tuple(self.states),
+            n_valid_sequences=initial_total[0],
+            choices_by_state_id=choices_by_state_id,
+            choice_results_by_state_id=choice_results_by_state_id,
+        )
+
+    def _get_or_register_state_id(self, state: TraversalState) -> int:
+        """
+        Return the stable integer ID for a traversal state, creating one if needed.
+        """
+        state_id = self.state_ids.get(state)
+
+        if state_id is None:
+            state_id = len(self.states)
+            self.state_ids[state] = state_id
+            self.states.append(state)
+            self.totals_by_state_id.append(None)
+            self.choices_by_state_id.append(None)
+            self.child_results_by_state_id.append(None)
+
+        return state_id
+
+    def _initial_state(self) -> TraversalState:
+        """
+        Return the starting traversal state.
+
+        The initial state starts at the graph's initial node, with fresh banned-
+        sequence and path-constraint states if those systems are active.
+
+        Returns
+        -------
+        TraversalState
+            Starting state for graph compilation.
+        """
+        if self.has_path_constraint:
+            constraint_state = self.path_constraint.initial_state
+        else:
+            constraint_state = ()
+
+        banned_tracker_state_id = 0
+
+        return TraversalState(self.graph.initial_node, banned_tracker_state_id, constraint_state)
+
+    def _compile_from(self, initial_state_id: int) -> None:
+        """
+        Compile every reachable traversal state starting from an initial state ID.
+
+        Uses an explicit depth-first stack so that each non-final state is compiled
+        only after its child states have been compiled.
+
+        Parameters
+        ----------
+        initial_state_id
+            ID of the state from which graph compilation should begin.
+        """
+        stack = [(initial_state_id, False)]
+
+        while stack:
+            state_id, expanded = stack.pop()
+            state = self.states[state_id]
+            node = state.node
+
+            if self.totals_by_state_id[state_id] is not None:
+                continue
+
+            if node is self.graph.final_node:
+                self._compile_final_state(state_id)
+                continue
+
+            if not expanded:
+                stack.append((state_id, True))
+                stack.extend(self._uncompiled_children(state_id))
+                continue
+
+            self._compile_state(state_id)
+
+    def _compile_final_state(self, state_id: int) -> None:
+        """
+        Compile a terminal traversal state.
+
+        By the time a terminal state is reached, all graph choices have already
+        been processed, including the right context. Choices rejected by the
+        banned-sequence tracker or path constraint would not have reached this
+        state.
+
+        The only remaining decision is whether the final path-constraint state is
+        acceptable. If so, this terminal state contributes one complete sequence
+        with log mass 0.0. Otherwise, it contributes no sequences.
+
+        Parameters
+        ----------
+        state_id
+            ID of the terminal traversal state being compiled.
+        """
+        state = self.states[state_id]
+
+        if not self.has_path_constraint or self.path_constraint.is_satisfied(state.constraint_state):
+            total = (1, 0.0)
+        else:
+            total = (0, -math.inf)
+
+        self.totals_by_state_id[state_id] = total
+        self.choices_by_state_id[state_id] = {}
+
+    def _compile_state(self, state_id: int) -> None:
+        """
+        Compile one non-final traversal state.
+
+        For each outgoing graph choice, combine the previously compiled child state
+        with the contribution from the current node to produce a ChoiceResult.
+        The total descendant count and log mass are then cached for the current state.
+
+        Parameters
+        ----------
+        state_id
+            ID of the traversal state being compiled.
+        """
+        state = self.states[state_id]
+        node = state.node
+        choice_results = {}
+        descendant_count = 0
+        descendant_log_masses = []
+        is_coding = isinstance(node, CodonNode)
+        child_results = self.child_results_by_state_id[state_id] or ()
+
+        for choice, child_id in child_results:
+            child_state = self.states[child_id]
+            child = child_state.node
+            child_total = self.totals_by_state_id[child_id]
+
+            if child_total is None:
+                continue
+
+            child_count, subtree_log_mass = child_total
+
+            if child_count == 0:
+                continue
+
+            choice_log_mass = self._accumulate_log_mass(node, choice, subtree_log_mass)
+
+            if choice_log_mass == -math.inf:
+                continue
+
+            result = ChoiceResult(
+                choice=choice,
+                descendant_count=child_count,
+                descendant_log_mass=choice_log_mass,
+                next_state_id=None if child is self.graph.final_node else child_id,
+                is_coding=is_coding,
+            )
+
+            choice_results[choice] = result
+            descendant_count += child_count
+            descendant_log_masses.append(choice_log_mass)
+
+        descendant_log_mass = self._sum_log_masses(descendant_log_masses)
+
+        total = (descendant_count, descendant_log_mass)
+
+        self.choices_by_state_id[state_id] = choice_results
+        self.totals_by_state_id[state_id] = total
+
+    def _uncompiled_children(self, state_id: int) -> List[Tuple[int, bool]]:
+        """
+        Return child state IDs reached by taking each outgoing graph choice.
+
+        Choices rejected by the banned-sequence tracker or path constraint are skipped.
+        Only child states that have not yet been compiled are returned.
+
+        Parameters
+        ----------
+        state_id
+            ID of the traversal state whose children should be discovered.
+
+        Returns
+        -------
+        list of tuple
+            Stack entries for child state IDs still needing compilation.
+        """
+        children = []
+        child_results = []
+
+        state = self.states[state_id]
+        node = state.node
+
+        for choice in self.choices_by_node[node]:
+            child = node.transitions.get(choice)
+
+            if child is None:
+                continue
+
+            advance = self._advance_banned_tracker(state.banned_tracker_state_id, node, choice)
+
+            if advance.banned:
+                continue
+
+            if self.has_path_constraint:
+                next_constraint_state = self.path_constraint.advance(state.constraint_state, node.pos, choice)
+
+                if next_constraint_state is None:
+                    continue
+            else:
+                next_constraint_state = ()
+
+            child_state = TraversalState(child, advance.state_id, next_constraint_state)
+            child_id = self._get_or_register_state_id(child_state)
+            child_results.append((choice, child_id))
+
+            if self.totals_by_state_id[child_id] is None:
+                children.append((child_id, False))
+
+        self.child_results_by_state_id[state_id] = child_results
+
+        return children
+
+    def _get_choices_for_node(self, node: Node) -> List[str]:
+        """
+        Return the graph choices available from a node in this view.
+
+        Codon nodes respect any pinned codons defined by the view. Context nodes
+        always have a single fixed sequence.
+
+        Parameters
+        ----------
+        node
+            Graph node whose available choices are required.
+
+        Returns
+        -------
+        list of str
+            The choices that may be taken from this node in the current view.
+        """
+        if isinstance(node, CodonNode):
+            if node.pos in self.view.pinned_codons:
+                return self.view.pinned_codons[node.pos]
+            else:
+                return node.codons
+
+        elif isinstance(node, ContextNode):
+            return [node.sequence]
+
+    def _advance_banned_tracker(
+            self,
+            banned_tracker_state_id: BannedTrackerStateId,
+            node: Node,
+            choice: str,
+    ) -> BannedTrackerAdvanceResult:
+        """
+        Advance the banned-sequence tracker after taking one graph choice.
+
+        Results are cached because the same tracker transition may be encountered
+        from many traversal states during compilation.
+
+        Parameters
+        ----------
+        banned_tracker_state_id
+            Current banned-sequence tracker state.
+        node
+            Current graph node.
+        choice
+            Graph choice taken from the current node.
+
+        Returns
+        -------
+        BannedTrackerAdvanceResult
+            Whether the choice enters a banned state and the resulting tracker
+            state.
+        """
+        key = (node, banned_tracker_state_id, choice)
+
+        if key in self.banned_advance_cache:
+            return self.banned_advance_cache[key]
+
+        if not self.has_banned_tracker:
+            result = BannedTrackerAdvanceResult(banned=False, state_id=banned_tracker_state_id)
+        else:
+            step = (node.pos, choice)
+            result = self.banned_tracker.advance(step, banned_tracker_state_id)
+
+        self.banned_advance_cache[key] = result
+        return result
+
+    def _accumulate_log_mass(
+            self,
+            node: Node,
+            choice: str,
+            subtree_log_mass: float,
+    ) -> float:
+        """
+        Accumulate the log probability mass contributed by one graph choice.
+
+        The subtree log mass has already been computed for the child state. Codon
+        nodes contribute the log of their codon weight, whereas context nodes
+        contribute no additional mass.
+
+        Parameters
+        ----------
+        node
+            The graph node from which the choice is taken.
+        choice
+            The outgoing graph choice.
+        subtree_log_mass
+            The total log mass reachable from the child state.
+
+        Returns
+        -------
+        float
+            The total log mass reachable after taking this choice.
+        """
+        if isinstance(node, CodonNode):
+            codon_log_weight = self.log_codon_weights.get(choice)
+
+            if codon_log_weight is None:
+                return -math.inf
+
+            return codon_log_weight + subtree_log_mass
+
+        return subtree_log_mass
+
+    def _sum_log_masses(self, log_masses: List[float]) -> float:
+        """
+        Combine several subtree log masses into a single log mass.
+
+        The calculation is performed using the log-sum-exp trick to avoid numerical
+        underflow when the subtree probabilities are extremely small.
+
+        Parameters
+        ----------
+        log_masses
+            Log-space masses to sum.
+
+        Returns
+        -------
+        float
+            The log of the summed masses, or -inf if no finite masses exist.
+        """
+        if not log_masses:
+            return -math.inf
+
+        max_log_mass = max(log_masses)
+
+        total_relative_mass = sum(math.exp(log_mass - max_log_mass) for log_mass in log_masses)
+
+        return max_log_mass + math.log(total_relative_mass)

codeine/graph/nodes.py

@@ -0,0 +1,111 @@
+from typing import Sequence
+
+
+class Node:
+    """
+    Basic CodonGraph node.
+    """
+
+    def __init__(self) -> None:
+        """
+        Constructor for the Node class.
+        """
+        self.parents = set()
+        self.transitions = {}
+        self.pos = None
+
+
+class ContextNode(Node):
+    """
+    Basic class representing a sequence context node on the codon graph.
+    This refers to the sequence either to the left of or to the right of
+    the coding sequence, and can be empty.
+    """
+
+    def __init__(self, pos: int, sequence: str) -> None:
+        """
+        Constructor for the ContextNode class.
+
+        Parameters
+        ----------
+        pos
+            The graph position. Left context is 0; right context is len(aa_seq) + 1.
+        sequence
+            The context sequence contained on this node.
+        """
+        super().__init__()
+
+        # Basic info.
+        self.pos = pos
+        self.sequence = sequence
+
+        # Set an ID for this node.
+        self.id = f'context-{pos}'
+
+    def __repr__(self) -> str:
+        return (
+            f'ContextNode('
+            f'id={self.id}'
+            f', pos={self.pos}'
+            f')'
+        )
+
+
+class CodonNode(Node):
+    """
+    Basic class representing a codon node on the codon graph.
+    """
+
+    def __init__(self, pos: int, aa: str, codons: Sequence[str]) -> None:
+        """
+        Constructor for the CodonNode class.
+
+        Parameters
+        ----------
+        pos
+            The aa position. Positioning is 1-based.
+        aa
+            The aa identity.
+        codons
+            The possible codons for this node.
+        """
+        super().__init__()
+
+        # Basic info. Positioning is 1-based.
+        self.pos = pos
+        self.aa = aa
+
+        # Set an ID for this node.
+        self.id = f'{aa}{pos}'
+
+        # Initialise the basic attributes.
+        self.codons = tuple(codons)
+
+    def __repr__(self) -> str:
+        codons = ','.join(self.codons)
+
+        return (
+            f'CodonNode('
+            f'id={self.id}'
+            f', pos={self.pos}'
+            #f', codons=[{codons}]'
+            f')'
+        )
+
+
+class EndNode(Node):
+    """
+    Final node for the codon graph.
+
+    Marks successful completion of a graph walk.
+    """
+
+    def __init__(self) -> None:
+        """
+        Constructor for the EndNode class.
+        """
+        super().__init__()
+        self.id = 'end'
+
+    def __repr__(self) -> str:
+        return (f'EndNode(id={self.id})')