RNA Secondary Structure

multimolecule.utils.rna.secondary_structure

RnaSecondaryStructureTopology

Bases: UndirectedGraph

Source code in multimolecule/utils/rna/secondary_structure/topology.py

class RnaSecondaryStructureTopology(UndirectedGraph):
    def __init__(
        self,
        sequence: str,
        secondary_structure: str | Tensor,
        device: torch.device | None = None,
        **kwargs,
    ):

        length = len(sequence)
        if isinstance(secondary_structure, torch.Tensor):
            pairs = secondary_structure
            if device is None:
                device = pairs.device
            elif pairs.device != device:
                pairs = pairs.to(device)
        elif isinstance(secondary_structure, str):
            if len(sequence) != len(secondary_structure):
                raise ValueError("sequence and secondary_structure must have the same length")
            pairs = torch.as_tensor(dot_bracket_to_pairs(secondary_structure), device=device)
            if device is None:
                device = pairs.device
        else:
            raise TypeError("secondary_structure must be a str or torch.Tensor")
        pairs = normalize_pairs(pairs)
        if device is None:
            device = pairs.device
        nested_pairs, pseudoknot_pairs = split_pseudoknot_pairs(pairs)

        if length > 1:
            backbone_idx = torch.arange(length - 1, dtype=torch.long, device=device)
            backbone_edges = torch.stack((backbone_idx, backbone_idx + 1), dim=1)
        else:
            backbone_edges = torch.empty((0, 2), dtype=torch.long, device=device)

        edge_index_parts: List[Tensor] = []
        edge_type_parts: List[Tensor] = []
        if len(backbone_edges):
            edge_index_parts.append(backbone_edges)
            edge_type_parts.append(
                torch.full((len(backbone_edges), 1), EdgeType.BACKBONE.value, dtype=torch.long, device=device)
            )
        if len(nested_pairs):
            edge_index_parts.append(nested_pairs)
            edge_type_parts.append(
                torch.full(
                    (len(nested_pairs), 1), EdgeType.NESTED_PAIRS.value, dtype=torch.long, device=nested_pairs.device
                )
            )
        if len(pseudoknot_pairs):
            edge_index_parts.append(pseudoknot_pairs)
            edge_type_parts.append(
                torch.full(
                    (len(pseudoknot_pairs), 1),
                    EdgeType.PSEUDOKNOT_PAIR.value,
                    dtype=torch.long,
                    device=pseudoknot_pairs.device,
                )
            )

        if edge_index_parts:
            edge_index = torch.cat(edge_index_parts, dim=0)
            edge_types = torch.cat(edge_type_parts, dim=0)
        else:
            edge_index = torch.empty((0, 2), dtype=torch.long, device=device)
            edge_types = torch.empty((0, 1), dtype=torch.long, device=device)

        self.sequence = sequence
        self.secondary_structure = secondary_structure
        self._all_pairs = pairs
        self._nested_pairs = nested_pairs
        self._pseudoknot_pairs = pseudoknot_pairs
        super().__init__(edge_index=edge_index, edge_features={"type": edge_types}, device=device, **kwargs)

    @property
    def all_pairs(self) -> Tensor:
        return self._all_pairs

    @property
    def nested_pairs(self) -> Tensor:
        return self._nested_pairs

    @property
    def pseudoknot_pairs(self) -> Tensor:
        return self._pseudoknot_pairs

    def pairs(self, view: StructureView | str | None = None) -> Tensor:
        view = StructureView.parse(view)
        if view == StructureView.ALL:
            return self._all_pairs
        if view == StructureView.NESTED:
            return self._nested_pairs
        if view == StructureView.PSEUDOKNOT:
            return self._pseudoknot_pairs
        raise ValueError("view must be 'all', 'nested', or 'pseudoknot'")

    @cached_property
    def crossing_events(self) -> Tensor:
        if self._all_pairs.shape[0] < 2:
            return self._all_pairs.new_empty((0, 2))
        return crossing_events(self._all_pairs)

    @cached_property
    def crossing_arcs(self) -> Tensor:
        if self._all_pairs.shape[0] < 2:
            return self._all_pairs.new_empty((0, 2))
        return crossing_arcs(self._all_pairs)

    @cached_property
    def crossing_pairs(self) -> Tensor:
        if self._all_pairs.shape[0] < 2:
            return self._all_pairs.new_empty((0, 2))
        return crossing_pairs(self._all_pairs)

    @cached_property
    def crossing_nucleotides(self) -> Tensor:
        if self._all_pairs.shape[0] < 2:
            return self._all_pairs.new_empty((0, 2))
        return crossing_nucleotides(self._all_pairs)

    @cached_property
    def tiers(self) -> List[Tier]:
        if len(self) == 0:
            return []
        tiers = StemSegments._tier_pairs(self._all_pairs)
        out: List[Tier] = []
        for idx, tier_pairs in enumerate(tiers):
            if not isinstance(tier_pairs, Tensor):
                tier_pairs = torch.tensor(tier_pairs, dtype=torch.long, device=self._all_pairs.device)
            out.append(Tier(idx, tier_pairs, len(self)))
        return out

    @cached_property
    def nested_stem_segments(self) -> StemSegments:
        return StemSegments(self._nested_pairs, 0)

    @cached_property
    def all_stem_segments(self) -> List[StemSegments]:
        return [tier.stem_segments for tier in self.tiers]

    def tier_stem_segments(self, tier: int) -> StemSegments:
        return self._tier_item(tier).stem_segments

    @cached_property
    def nested_helix_segments(self) -> HelixSegments:
        return HelixSegments(self._nested_pairs, 0)

    @cached_property
    def all_helix_segments(self) -> List[HelixSegments]:
        return [tier.helix_segments for tier in self.tiers]

    def tier_helix_segments(self, tier: int) -> HelixSegments:
        return self._tier_item(tier).helix_segments

    @cached_property
    def nested_stem_graph(self) -> DirectedGraph:
        return self.nested_stem_segments.graph

    @cached_property
    def nested_stem_edges(self) -> List[StemEdge]:
        return self.nested_stem_segments.edges

    def tier_stem_edges(self, tier: int) -> List[StemEdge]:
        return self._tier_item(tier).stem_segments.edges

    @cached_property
    def all_stem_edges(self) -> List[StemEdge]:
        return [edge for stems in self.all_stem_segments for edge in stems.edges]

    @cached_property
    def nested_helix_graph(self) -> DirectedGraph:
        return self.nested_helix_segments.graph

    @cached_property
    def nested_helix_edges(self) -> List[StemEdge]:
        return self.nested_helix_segments.edges

    def tier_helix_edges(self, tier: int) -> List[StemEdge]:
        return self._tier_item(tier).helix_segments.edges

    @cached_property
    def all_helix_edges(self) -> List[StemEdge]:
        return [edge for segments in self.all_helix_segments for edge in segments.edges]

    def tier_stem_graph(self, tier: int) -> DirectedGraph:
        return self._tier_item(tier).stem_segments.graph

    @cached_property
    def all_stem_graphs(self) -> List[DirectedGraph]:
        return [stems.graph for stems in self.all_stem_segments]

    def tier_helix_graph(self, tier: int) -> DirectedGraph:
        return self._tier_item(tier).helix_segments.graph

    @cached_property
    def all_helix_graphs(self) -> List[DirectedGraph]:
        return [segments.graph for segments in self.all_helix_segments]

    @cached_property
    def nested_loop_segments(self) -> LoopSegments:
        return LoopSegments(self._nested_pairs, len(self), self._nested_pairs.device, 0)

    @cached_property
    def all_loop_segments(self) -> List[LoopSegments]:
        return [tier.loop_segments for tier in self.tiers]

    def loop_segments_by_tier(self, view: StructureView | str | None = None) -> List[LoopSegments]:
        return [tier.loop_segments for tier in self._tiers_for_view(view)]

    def loop_segment_contexts_by_tier(self, view: StructureView | str | None = None) -> List[List[LoopSegmentContext]]:
        return [
            loop_segments.contexts(self._pseudoknot_pairs) for loop_segments in self.loop_segments_by_tier(view=view)
        ]

    def loop_segments(
        self,
        tier: int | None = None,
        *,
        view: StructureView | str | None = None,
    ) -> LoopSegments:
        if view is not None and tier is not None:
            raise ValueError("view and tier are mutually exclusive")
        if view is None:
            if tier is None:
                return self.nested_loop_segments
            return self._tier_item(tier).loop_segments
        view = StructureView.parse(view)
        if view == StructureView.NESTED:
            return self.nested_loop_segments
        raise ValueError(
            "view must be 'nested' when requesting a single LoopSegments; "
            "use loop_segments_by_tier for 'all' or 'pseudoknot'"
        )

    def loop_segment_contexts(
        self,
        tier: int | None = None,
        *,
        view: StructureView | str | None = None,
    ) -> List[LoopSegmentContext]:
        if view is not None and tier is not None:
            raise ValueError("view and tier are mutually exclusive")
        if view is None:
            if tier is None:
                return self.nested_loop_segments.contexts(self._pseudoknot_pairs)
            return self._tier_item(tier).loop_segments.contexts(self._pseudoknot_pairs)
        view = StructureView.parse(view)
        if view == StructureView.NESTED:
            return self.nested_loop_segments.contexts(self._pseudoknot_pairs)
        raise ValueError(
            "view must be 'nested' when requesting a single LoopSegments; "
            "use loop_segment_contexts_by_tier for 'all' or 'pseudoknot'"
        )

    def _tiers_for_view(self, view: StructureView | str | None) -> List[Tier]:
        view = StructureView.parse(view)
        if view == StructureView.ALL:
            return self.tiers
        if view == StructureView.NESTED:
            return self.tiers[:1]
        if view == StructureView.PSEUDOKNOT:
            return self.tiers[1:]
        raise ValueError("view must be 'all', 'nested', or 'pseudoknot'")

    def stems(self, view: StructureView | str | None = None) -> List[StemSegment]:
        tiers = self._tiers_for_view(view)
        return [segment for tier in tiers for segment in tier.stem_segments.segments]

    def helices(self, view: StructureView | str | None = None) -> List[HelixSegment]:
        tiers = self._tiers_for_view(view)
        return [segment for tier in tiers for segment in tier.helix_segments.segments]

    @cached_property
    def _partner_map_all(self) -> Tensor:
        return self._partner_map_from_pairs(self._all_pairs, len(self))

    @staticmethod
    def _partner_map_from_pairs(pairs: Tensor, length: int) -> Tensor:
        pair_map = pairs.new_full((length,), -1, dtype=torch.long)
        if pairs.numel() == 0:
            return pair_map
        left = torch.minimum(pairs[:, 0], pairs[:, 1]).to(torch.long)
        right = torch.maximum(pairs[:, 0], pairs[:, 1]).to(torch.long)
        pair_map[left] = right
        pair_map[right] = left
        return pair_map

    def partner_at(
        self,
        pos: int,
        tier: int | None = None,
        *,
        view: StructureView | str | None = None,
    ) -> int | None:
        if pos >= len(self):
            raise IndexError("position is out of range")
        if view is not None and tier is not None:
            raise ValueError("view and tier are mutually exclusive")
        if view is not None:
            view = StructureView.parse(view)
            if view == StructureView.ALL:
                partner = int(self._partner_map_all[pos].item())
                return None if partner < 0 else partner
            pair_map = self._partner_map_from_pairs(self.pairs(view=view), len(self))
            partner = int(pair_map[pos].item())
            return None if partner < 0 else partner
        if tier is None:
            partner = int(self._partner_map_all[pos].item())
            return None if partner < 0 else partner
        pair_map = self._partner_map_from_pairs(self._tier_item(tier).pairs, len(self))
        partner = int(pair_map[pos].item())
        return None if partner < 0 else partner

    def paired_positions(self, view: StructureView | str | None = None) -> List[int]:
        view = StructureView.parse(view)
        pairs = self.pairs(view=view)
        if isinstance(pairs, Tensor):
            return StemSegments._paired_positions_from_pairs(pairs)
        pairs_list = ensure_pairs_list(pairs)
        if not pairs_list:
            return []
        positions = {i for i, j in pairs_list}
        positions.update(j for i, j in pairs_list)
        return sorted(positions)

    def unpaired_positions(self, view: StructureView | str | None = None) -> List[int]:
        view = StructureView.parse(view)
        if view == StructureView.ALL:
            pair_map = self._partner_map_all
        else:
            pair_map = self._partner_map_from_pairs(self.pairs(view=view), len(self))
        if pair_map.numel() == 0:
            return []
        return [idx for idx, partner in enumerate(pair_map.tolist()) if partner < 0]

    def noncanonical_pairs(
        self,
        view: StructureView | str | None = None,
        *,
        unsafe: bool = True,
    ) -> Tensor:
        pairs = self.pairs(view=view)
        return noncanonical_pairs(pairs, self.sequence, unsafe=unsafe)

    @cached_property
    def _stem_pair_map(self) -> Dict[Pair, StemSegment]:
        return StemSegments.pair_map_by_tiers(ensure_pairs_list(self._all_pairs))

    @cached_property
    def _helix_pair_map(self) -> Dict[Pair, HelixSegment]:
        mapping: Dict[Pair, HelixSegment] = {}
        for segments in self.all_helix_segments:
            mapping.update(segments.pair_map)
        return mapping

    def stem_at(
        self,
        pos: int,
        tier: int | None = None,
        *,
        view: StructureView | str | None = None,
    ) -> StemSegment | None:
        if view is not None and tier is not None:
            raise ValueError("view and tier are mutually exclusive")
        if view is not None:
            view = StructureView.parse(view)
            partner = self.partner_at(pos, view=view)
            if partner is None:
                return None
            pair = (min(pos, partner), max(pos, partner))
            if view == StructureView.NESTED:
                return self.tier_stem_segments(0).pair_map.get(pair)
            if view == StructureView.PSEUDOKNOT:
                segment = self._stem_pair_map.get(pair)
                return segment if segment is not None and segment.tier > 0 else None
            return self._stem_pair_map.get(pair)
        partner = self.partner_at(pos, tier=tier)
        if partner is None:
            return None
        pair = (min(pos, partner), max(pos, partner))
        if tier is None:
            return self._stem_pair_map.get(pair)
        return self.tier_stem_segments(tier).pair_map.get(pair)

    def helix_at(
        self,
        pos: int,
        tier: int | None = None,
        *,
        view: StructureView | str | None = None,
    ) -> HelixSegment | None:
        if view is not None and tier is not None:
            raise ValueError("view and tier are mutually exclusive")
        if view is not None:
            view = StructureView.parse(view)
            partner = self.partner_at(pos, view=view)
            if partner is None:
                return None
            pair = (min(pos, partner), max(pos, partner))
            if view == StructureView.NESTED:
                return self.tier_helix_segments(0).pair_map.get(pair)
            if view == StructureView.PSEUDOKNOT:
                segment = self._helix_pair_map.get(pair)
                return segment if segment is not None and segment.tier > 0 else None
            return self._helix_pair_map.get(pair)
        partner = self.partner_at(pos, tier=tier)
        if partner is None:
            return None
        pair = (min(pos, partner), max(pos, partner))
        if tier is None:
            return self._helix_pair_map.get(pair)
        return self.tier_helix_segments(tier).pair_map.get(pair)

    @staticmethod
    def _select_preferred_region(
        regions: Sequence[StemSegment | HelixSegment | LoopSegment],
        tier_preference: StructureView | str | None,
    ) -> StemSegment | HelixSegment | LoopSegment | None:
        if not regions:
            return None
        if isinstance(tier_preference, StructureView):
            tier_preference = tier_preference.value
        if tier_preference is None:
            return regions[0]
        tier_preference = str(tier_preference).lower()
        if tier_preference == "first":
            return regions[0]
        if tier_preference == "last":
            return regions[-1]
        if tier_preference == "nested":
            return next((region for region in regions if region.tier == 0), None)
        if tier_preference == "pseudoknot":
            return next((region for region in regions if region.tier > 0), None)
        raise ValueError("tier_preference must be 'nested', 'pseudoknot', 'first', or 'last'")

    def region_at(
        self,
        pos: int,
        *,
        view: StructureView | str | None = None,
        paired: str = "stem",
        tier_preference: StructureView | str | None = None,
        unpaired: str = "segment",
    ) -> StemSegment | HelixSegment | LoopSegment | Loop | None:
        if unpaired not in ("segment", "loop"):
            raise ValueError("unpaired must be 'segment' or 'loop'")
        if unpaired == "loop":
            if paired not in ("stem", "helix"):
                raise ValueError("paired must be 'stem' or 'helix'")
            selected_view = view
            if tier_preference is not None:
                if isinstance(tier_preference, StructureView):
                    prefer_str = tier_preference.value
                else:
                    prefer_str = str(tier_preference).lower()
                if prefer_str not in ("nested", "pseudoknot", "first", "last"):
                    raise ValueError("tier_preference must be 'nested', 'pseudoknot', 'first', or 'last'")
                if selected_view is None:
                    if prefer_str == "nested":
                        selected_view = StructureView.NESTED
                    elif prefer_str == "pseudoknot":
                        selected_view = StructureView.PSEUDOKNOT
            partner = self.partner_at(pos, view=selected_view)
            if partner is not None:
                if paired == "helix":
                    return self.helix_at(pos, view=selected_view)
                return self.stem_at(pos, view=selected_view)
            return self.loop_at(pos, view=selected_view)
        regions = self.regions_at(pos, view=view, paired=paired)
        if not regions:
            return None
        if tier_preference is None:
            for region in regions:
                if not isinstance(region, LoopSegment):
                    return region
            return regions[0]
        return self._select_preferred_region(regions, tier_preference)

    def regions_at(
        self,
        pos: int,
        *,
        view: StructureView | str | None = None,
        paired: str = "stem",
    ) -> List[StemSegment | HelixSegment | LoopSegment]:
        if pos >= len(self):
            raise IndexError("position is out of range")
        view = StructureView.parse(view)
        if paired not in ("stem", "helix"):
            raise ValueError("paired must be 'stem' or 'helix'")

        segments: List[StemSegment | HelixSegment | LoopSegment] = []
        for tier in self._tiers_for_view(view):
            pair_map = self._partner_map_from_pairs(tier.pairs, len(self))
            partner = int(pair_map[pos].item())
            segment: StemSegment | HelixSegment | LoopSegment | None = None
            if partner >= 0:
                pair = (min(pos, partner), max(pos, partner))
                if paired == "helix":
                    segment = tier.helix_segments.pair_map.get(pair)
                else:
                    segment = tier.stem_segments.pair_map.get(pair)
                if segment is not None:
                    segments.append(segment)
                continue
            segment = tier.loop_segments.segment_at(pos)
            if segment is not None:
                segments.append(segment)
        return segments

    def loop_at(
        self,
        pos: int,
        *,
        view: StructureView | str | None = None,
        mode: LoopView | str | None = None,
    ) -> Loop | None:
        if pos >= len(self):
            raise IndexError("position is out of range")
        return self.loops(view=view, mode=mode).loop_at(pos)

    def loop_segment_at(
        self,
        pos: int,
        *,
        view: StructureView | str | None = None,
        tier_preference: StructureView | str | None = None,
    ) -> LoopSegment | None:
        if pos >= len(self):
            raise IndexError("position is out of range")
        view = StructureView.parse(view)
        segments: List[LoopSegment] = []
        for tier in self._tiers_for_view(view):
            segment = tier.loop_segments.segment_at(pos)
            if segment is not None:
                segments.append(segment)
        selected = self._select_preferred_region(segments, tier_preference)
        return selected if isinstance(selected, LoopSegment) else None

    def loop_contexts(self) -> List[LoopContext]:
        """
        Return nested-loop contexts annotated with pseudoknot pairs.

        Loops are defined on nested pairs only; pseudoknot pairs are reported
        as inside or crossing each loop's spans.
        """
        loops = self.nested_loops
        pairs = self._pseudoknot_pairs
        if not loops:
            return []
        if pairs.numel() == 0:
            return [LoopContext(loop, pairs, pairs) for loop in loops]

        left = pairs[:, 0].view(1, -1)
        right = pairs[:, 1].view(1, -1)
        out: List[LoopContext] = []
        for loop in loops:
            if not loop.spans:
                empty = pairs.new_empty((0, 2))
                out.append(LoopContext(loop, empty, empty))
                continue
            starts = torch.tensor(
                [span.start for span in loop.spans],
                device=pairs.device,
                dtype=pairs.dtype,
            ).view(-1, 1)
            stops = torch.tensor(
                [span.stop for span in loop.spans],
                device=pairs.device,
                dtype=pairs.dtype,
            ).view(-1, 1)
            in_left = (left >= starts) & (left <= stops)
            in_right = (right >= starts) & (right <= stops)
            left_mask = in_left.any(dim=0)
            right_mask = in_right.any(dim=0)
            inside_mask = left_mask & right_mask
            crossing_mask = left_mask ^ right_mask
            out.append(LoopContext(loop, pairs[inside_mask], pairs[crossing_mask]))
        return out

    def _resolve_loop_inputs(
        self,
        *,
        view: StructureView | str | None,
        mode: LoopView | str | None,
        pairs: Tensor | np.ndarray | Pairs | None,
    ) -> Tuple[Tensor | np.ndarray | Pairs, LoopView]:
        if mode is None:
            mode = LoopView.Topological
        mode = LoopView.parse(mode)
        if pairs is None:
            if mode == LoopView.Nested:
                if view is None:
                    view = StructureView.NESTED
                else:
                    view = StructureView.parse(view)
                    if view != StructureView.NESTED:
                        raise ValueError("mode 'nested' requires view='nested'")
            pairs = self.pairs(view=view)
        elif view is not None:
            raise ValueError("view and pairs are mutually exclusive")
        return pairs, mode

    def _tier_item(self, tier: int) -> Tier:
        if tier < 0 or tier >= len(self.tiers):
            raise IndexError("tier is out of range")
        return self.tiers[tier]

    @cached_property
    def nested_loops(self) -> Loops:
        return Loops.from_pairs(self._nested_pairs, len(self))

    @cached_property
    def all_loops(self) -> Loops:
        return self.loops(view=StructureView.ALL)

    @cached_property
    def taxonomy_loops(self) -> Loops:
        return Loops.taxonomy_from_pairs(self._all_pairs, len(self))

    def loops(
        self,
        view: StructureView | str | None = None,
        *,
        pairs: Tensor | np.ndarray | Pairs | None = None,
        mode: LoopView | str | None = None,
    ) -> Loops:
        pairs_input = pairs
        pairs, mode = self._resolve_loop_inputs(view=view, mode=mode, pairs=pairs)
        if mode == LoopView.Taxonomy:
            return Loops.taxonomy_from_pairs(pairs, len(self))
        if pairs_input is None:
            view = StructureView.parse(view)
            if view == StructureView.ALL and self.crossing_pairs.numel() > 0:
                return self.nested_loops
        return Loops.from_pairs(pairs, len(self))

    def loop_helix_graph(
        self,
        pairs: Tensor | np.ndarray | Pairs | None = None,
        *,
        view: StructureView | str | None = None,
        mode: LoopView | str | None = None,
        tier: int | None = None,
    ) -> LoopHelixGraph:
        pairs, mode = self._resolve_loop_inputs(view=view, mode=mode, pairs=pairs)
        pairs = normalize_pairs(pairs)
        pairs_list = ensure_pairs_list(pairs)
        device = pairs.device if isinstance(pairs, Tensor) else None
        if device is None:
            device = self._all_pairs.device
        if mode == LoopView.Taxonomy:
            loops = Loops.taxonomy_from_pairs(pairs, len(self))
        else:
            loops = Loops.from_pairs(pairs, len(self))
        tier_indices = None if tier is None else (tier,)
        helices = HelixSegments.segments_by_tier(pairs_list, tiers=tier_indices, device=device)

        loop_count = len(loops)
        loop_nodes = tuple(range(loop_count))
        helix_nodes = tuple(range(loop_count, loop_count + len(helices)))
        graph = UndirectedGraph(device=device)
        graph.add_nodes(loop_nodes)
        graph.add_nodes(helix_nodes)

        helix_key_to_idx = {(*helix.key, helix.tier): idx for idx, helix in enumerate(helices)}
        for loop_idx, loop in enumerate(loops):
            for helix in loop.anchor_helices:
                key = (*helix.key, helix.tier)
                helix_idx = helix_key_to_idx.get(key)
                if helix_idx is None:
                    continue
                graph.add_edge(loop_idx, loop_count + helix_idx, attr={"tier": helix.tier})

        return LoopHelixGraph(graph, loops, tuple(helices), loop_nodes, helix_nodes)

    def loop_span_sequences(self, loop: Loop) -> Tuple[str, ...]:
        return tuple(self.sequence[span.start : span.stop + 1] for span in loop.spans)

    def loop_sequence(self, loop: Loop, *, joiner: str = "") -> str:
        return joiner.join(self.loop_span_sequences(loop))

    def loop_anchor_pair_sequences(self, loop: Loop) -> Tuple[str, ...]:
        return tuple(self.sequence[i] + self.sequence[j] for i, j in loop.anchor_pairs)

    def stem_strands(self, stem: StemSegment, *, orientation: str = "5p-3p") -> Tuple[str, str]:
        return self._duplex_strands(stem.start_5p, stem.stop_5p, stem.start_3p, stem.stop_3p, orientation=orientation)

    def helix_strands(self, helix: HelixSegment, *, orientation: str = "5p-3p") -> Tuple[str, str]:
        return self._duplex_strands(
            helix.start_5p,
            helix.stop_5p,
            helix.start_3p,
            helix.stop_3p,
            orientation=orientation,
        )

    def _duplex_strands(
        self,
        start_5p: int,
        stop_5p: int,
        start_3p: int,
        stop_3p: int,
        *,
        orientation: str,
    ) -> Tuple[str, str]:
        if orientation not in ("5p-3p", "index"):
            raise ValueError("orientation must be '5p-3p' or 'index'")
        strand_5p = self.sequence[start_5p : stop_5p + 1]
        strand_3p = self.sequence[stop_3p : start_3p + 1]
        if orientation == "5p-3p":
            strand_3p = strand_3p[::-1]
        return strand_5p, strand_3p

    def annotate_positions(
        self,
        *,
        view: StructureView | str | None = None,
        paired: str = "stem",
        mode: LoopView | str | None = None,
        tier_preference: StructureView | str | None = None,
    ) -> List[Dict[str, object]]:
        """
        Return per-position annotations for paired/loop context.

        Each entry includes partner, paired flag, tier, segment/loop objects,
        and indices where available. Negative indexing follows normal Python rules.
        """
        if paired not in ("stem", "helix"):
            raise ValueError("paired must be 'stem' or 'helix'")
        view = StructureView.parse(view)
        loops = self.loops(view=view, mode=mode)
        loop_index_map = loops._position_to_loop_index
        segments = self.helices(view=view) if paired == "helix" else self.stems(view=view)
        segment_index_map = {segment: idx for idx, segment in enumerate(segments)}

        annotations: List[Dict[str, object]] = []
        for pos in range(len(self)):
            partner = self.partner_at(pos, view=view)
            is_paired = partner is not None
            segment: StemSegment | HelixSegment | None = None
            segment_index: int | None = None
            segment_type: str | None = None
            loop_segment: LoopSegment | None = None
            loop: Loop | None = None
            loop_index: int | None = None
            loop_kind: LoopType | None = None
            tier: int | None = None

            if is_paired:
                if paired == "helix":
                    segment = self.helix_at(pos, view=view)
                    segment_type = "helix"
                else:
                    segment = self.stem_at(pos, view=view)
                    segment_type = "stem"
                if segment is not None:
                    segment_index = segment_index_map.get(segment)
                    tier = segment.tier
            else:
                loop_segment = self.loop_segment_at(pos, view=view, tier_preference=tier_preference)
                if loop_segment is not None:
                    tier = loop_segment.tier
                if pos < len(loop_index_map):
                    idx = loop_index_map[pos]
                    if idx != -1:
                        loop_index = idx
                        loop = loops[idx]
                        loop_kind = loop.kind

            annotations.append(
                {
                    "pos": pos,
                    "partner": partner,
                    "paired": is_paired,
                    "tier": tier,
                    "segment_type": segment_type,
                    "segment_index": segment_index,
                    "segment": segment,
                    "loop_segment": loop_segment,
                    "loop_index": loop_index,
                    "loop": loop,
                    "loop_kind": loop_kind,
                }
            )
        return annotations

    def __repr__(self) -> str:
        parts = [
            f"length={len(self)}",
            f"pairs={int(self._all_pairs.shape[0])}",
            f"nested={int(self._nested_pairs.shape[0])}",
            f"pseudoknot={int(self._pseudoknot_pairs.shape[0])}",
        ]
        tiers = self.__dict__.get("tiers")
        if tiers is not None:
            parts.append(f"tiers={len(tiers)}")
        parts.append(f"device='{self._all_pairs.device}'")
        return f"{self.__class__.__name__}({', '.join(parts)})"

    def __len__(self) -> int:
        return len(self.sequence)

loop_contexts

Python

loop_contexts() -> List[LoopContext]

Return nested-loop contexts annotated with pseudoknot pairs.

Loops are defined on nested pairs only; pseudoknot pairs are reported as inside or crossing each loop’s spans.

Source code in multimolecule/utils/rna/secondary_structure/topology.py

def loop_contexts(self) -> List[LoopContext]:
    """
    Return nested-loop contexts annotated with pseudoknot pairs.

    Loops are defined on nested pairs only; pseudoknot pairs are reported
    as inside or crossing each loop's spans.
    """
    loops = self.nested_loops
    pairs = self._pseudoknot_pairs
    if not loops:
        return []
    if pairs.numel() == 0:
        return [LoopContext(loop, pairs, pairs) for loop in loops]

    left = pairs[:, 0].view(1, -1)
    right = pairs[:, 1].view(1, -1)
    out: List[LoopContext] = []
    for loop in loops:
        if not loop.spans:
            empty = pairs.new_empty((0, 2))
            out.append(LoopContext(loop, empty, empty))
            continue
        starts = torch.tensor(
            [span.start for span in loop.spans],
            device=pairs.device,
            dtype=pairs.dtype,
        ).view(-1, 1)
        stops = torch.tensor(
            [span.stop for span in loop.spans],
            device=pairs.device,
            dtype=pairs.dtype,
        ).view(-1, 1)
        in_left = (left >= starts) & (left <= stops)
        in_right = (right >= starts) & (right <= stops)
        left_mask = in_left.any(dim=0)
        right_mask = in_right.any(dim=0)
        inside_mask = left_mask & right_mask
        crossing_mask = left_mask ^ right_mask
        out.append(LoopContext(loop, pairs[inside_mask], pairs[crossing_mask]))
    return out

annotate_positions

Python

annotate_positions(
    *,
    view: StructureView | str | None = None,
    paired: str = "stem",
    mode: LoopView | str | None = None,
    tier_preference: StructureView | str | None = None
) -> List[Dict[str, object]]

Return per-position annotations for paired/loop context.

Each entry includes partner, paired flag, tier, segment/loop objects, and indices where available. Negative indexing follows normal Python rules.

Source code in multimolecule/utils/rna/secondary_structure/topology.py

def annotate_positions(
    self,
    *,
    view: StructureView | str | None = None,
    paired: str = "stem",
    mode: LoopView | str | None = None,
    tier_preference: StructureView | str | None = None,
) -> List[Dict[str, object]]:
    """
    Return per-position annotations for paired/loop context.

    Each entry includes partner, paired flag, tier, segment/loop objects,
    and indices where available. Negative indexing follows normal Python rules.
    """
    if paired not in ("stem", "helix"):
        raise ValueError("paired must be 'stem' or 'helix'")
    view = StructureView.parse(view)
    loops = self.loops(view=view, mode=mode)
    loop_index_map = loops._position_to_loop_index
    segments = self.helices(view=view) if paired == "helix" else self.stems(view=view)
    segment_index_map = {segment: idx for idx, segment in enumerate(segments)}

    annotations: List[Dict[str, object]] = []
    for pos in range(len(self)):
        partner = self.partner_at(pos, view=view)
        is_paired = partner is not None
        segment: StemSegment | HelixSegment | None = None
        segment_index: int | None = None
        segment_type: str | None = None
        loop_segment: LoopSegment | None = None
        loop: Loop | None = None
        loop_index: int | None = None
        loop_kind: LoopType | None = None
        tier: int | None = None

        if is_paired:
            if paired == "helix":
                segment = self.helix_at(pos, view=view)
                segment_type = "helix"
            else:
                segment = self.stem_at(pos, view=view)
                segment_type = "stem"
            if segment is not None:
                segment_index = segment_index_map.get(segment)
                tier = segment.tier
        else:
            loop_segment = self.loop_segment_at(pos, view=view, tier_preference=tier_preference)
            if loop_segment is not None:
                tier = loop_segment.tier
            if pos < len(loop_index_map):
                idx = loop_index_map[pos]
                if idx != -1:
                    loop_index = idx
                    loop = loops[idx]
                    loop_kind = loop.kind

        annotations.append(
            {
                "pos": pos,
                "partner": partner,
                "paired": is_paired,
                "tier": tier,
                "segment_type": segment_type,
                "segment_index": segment_index,
                "segment": segment,
                "loop_segment": loop_segment,
                "loop_index": loop_index,
                "loop": loop,
                "loop_kind": loop_kind,
            }
        )
    return annotations

annotate

Python

annotate(
    structure: RnaSecondaryStructureTopology,
) -> Tuple[str, str]

Return both structural and functional annotations for this structure.

Source code in multimolecule/utils/rna/secondary_structure/bprna.py

def annotate(structure: RnaSecondaryStructureTopology) -> Tuple[str, str]:
    """
    Return both structural and functional annotations for this structure.
    """
    segment_data = BpRnaSecondaryStructureTopology(structure.sequence, topology=structure)
    return segment_data.structural_annotation, segment_data.functional_annotation

annotate_function

Python

annotate_function(
    structure: RnaSecondaryStructureTopology,
) -> str

Return a bpRNA-like functional annotation string (knot/function array) for this structure.

Labels: K for bases involved in pseudoknot pairs, N otherwise.

Source code in multimolecule/utils/rna/secondary_structure/bprna.py

def annotate_function(structure: RnaSecondaryStructureTopology) -> str:
    """
    Return a bpRNA-like functional annotation string (knot/function array) for this structure.

    Labels: K for bases involved in pseudoknot pairs, N otherwise.
    """
    segment_data = BpRnaSecondaryStructureTopology(structure.sequence, topology=structure)
    return segment_data.functional_annotation

annotate_structure

Python

annotate_structure(
    structure: RnaSecondaryStructureTopology,
) -> str

Return a bpRNA-like structural annotation string (structure array) for this structure.

Labels: S (stems), H (hairpins), B (bulges), I (internals), M (multibranch), X (external), E (end).

Source code in multimolecule/utils/rna/secondary_structure/bprna.py

def annotate_structure(structure: RnaSecondaryStructureTopology) -> str:
    """
    Return a bpRNA-like structural annotation string (structure array) for this structure.

    Labels: S (stems), H (hairpins), B (bulges), I (internals), M (multibranch),
    X (external), E (end).
    """
    segment_data = BpRnaSecondaryStructureTopology(structure.sequence, topology=structure)
    return segment_data.structural_annotation

noncanonical_pairs

Python

noncanonical_pairs(
    pairs: Tensor | ndarray | Pairs,
    sequence: str,
    unsafe: bool = False,
) -> Tensor | ndarray | PairsList

Return subset of base pairs that are non-canonical (backend-aware).

Non-ACGU bases are treated as non-canonical.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	Base pairs as a (n, 2) tensor or array.	required
sequence	str	RNA sequence string.	required
unsafe	bool	Ignore invalid self-pairs when True; raise when False.	False

Returns:

Type	Description
Tensor | ndarray | PairsList	Non-canonical base pairs using the same backend as input (list inputs return list of tuples).

Examples:

Torch input

>>> import torch
>>> noncanonical_pairs(torch.tensor([[0, 3]]), "ACGU").tolist()
[]
>>> noncanonical_pairs(torch.tensor([[0, 3]]), "ACGA").tolist()
[[0, 3]]
>>> noncanonical_pairs(torch.tensor([[0, 1]]), "GU").tolist()
[]

NumPy input

>>> import numpy as np
>>> noncanonical_pairs(np.array([[0, 3]]), "ACGU").tolist()
[]
>>> noncanonical_pairs(np.array([[0, 3]]), "ACGA").tolist()
[[0, 3]]
>>> noncanonical_pairs(np.array([[0, 1]]), "GU").tolist()
[]
>>> noncanonical_pairs(np.array([[0, 1]]), "AX").tolist()
[[0, 1]]

Source code in multimolecule/utils/rna/secondary_structure/noncanonical.py

def noncanonical_pairs(
    pairs: Tensor | np.ndarray | Pairs, sequence: str, unsafe: bool = False
) -> Tensor | np.ndarray | PairsList:
    """
    Return subset of base pairs that are non-canonical (backend-aware).

    Non-ACGU bases are treated as non-canonical.

    Args:
        pairs: Base pairs as a (n, 2) tensor or array.
        sequence: RNA sequence string.
        unsafe: Ignore invalid self-pairs when True; raise when False.

    Returns:
        Non-canonical base pairs using the same backend as input (list inputs return list of tuples).

    Examples:
        Torch input
        >>> import torch
        >>> noncanonical_pairs(torch.tensor([[0, 3]]), "ACGU").tolist()
        []
        >>> noncanonical_pairs(torch.tensor([[0, 3]]), "ACGA").tolist()
        [[0, 3]]
        >>> noncanonical_pairs(torch.tensor([[0, 1]]), "GU").tolist()
        []

        NumPy input
        >>> import numpy as np
        >>> noncanonical_pairs(np.array([[0, 3]]), "ACGU").tolist()
        []
        >>> noncanonical_pairs(np.array([[0, 3]]), "ACGA").tolist()
        [[0, 3]]
        >>> noncanonical_pairs(np.array([[0, 1]]), "GU").tolist()
        []
        >>> noncanonical_pairs(np.array([[0, 1]]), "AX").tolist()
        [[0, 1]]
    """
    if isinstance(pairs, Tensor):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_noncanonical_pairs(pairs, sequence, unsafe)
    if isinstance(pairs, np.ndarray):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_noncanonical_pairs(pairs, sequence, unsafe)
    if isinstance(pairs, Sequence):
        if not pairs:
            return []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        return list(map(tuple, _numpy_noncanonical_pairs(pairs, sequence, unsafe).tolist()))
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

noncanonical_pairs_set

Python

noncanonical_pairs_set(
    pairs: Tensor | ndarray | Pairs,
    sequence: str,
    unsafe: bool = False,
) -> set[Pair]

Return non-canonical base pairs as a set of (i, j) tuples.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	Base pairs as a tensor, array, or sequence of (i, j) tuples.	required
sequence	str	RNA sequence string.	required
unsafe	bool	Ignore invalid self-pairs when True; raise when False.	False

Returns:

Type	Description
set[Pair]	A set of (i, j) tuples for non-canonical pairs.

Examples:

>>> sorted(noncanonical_pairs_set([(0, 3), (1, 2)], "ACGA"))
[(0, 3)]

Source code in multimolecule/utils/rna/secondary_structure/noncanonical.py

def noncanonical_pairs_set(
    pairs: Tensor | np.ndarray | Pairs,
    sequence: str,
    unsafe: bool = False,
) -> set[Pair]:
    """
    Return non-canonical base pairs as a set of ``(i, j)`` tuples.

    Args:
        pairs: Base pairs as a tensor, array, or sequence of (i, j) tuples.
        sequence: RNA sequence string.
        unsafe: Ignore invalid self-pairs when True; raise when False.

    Returns:
        A set of (i, j) tuples for non-canonical pairs.

    Examples:
        >>> sorted(noncanonical_pairs_set([(0, 3), (1, 2)], "ACGA"))
        [(0, 3)]
    """
    if isinstance(pairs, Tensor):
        if pairs.numel() == 0:
            return set()
        pairs = pairs.detach().cpu().numpy()
    elif isinstance(pairs, np.ndarray):
        if pairs.size == 0:
            return set()
    elif isinstance(pairs, Sequence):
        if not pairs:
            return set()
        pairs = np.asarray(pairs, dtype=int)
    else:
        raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

    if pairs.size == 0:
        return set()
    if pairs.ndim != 2 or pairs.shape[1] != 2:
        raise ValueError("pairs must be a sequence/array of (i, j) index tuples")

    pairs = pairs.astype(int, copy=False)
    pairs = _numpy_normalize_pairs_low_high(pairs)
    mask = noncanonical_pair_mask(pairs[:, 0], pairs[:, 1], sequence, unsafe)
    if not np.any(mask):
        return set()
    return {(int(pairs[idx, 0]), int(pairs[idx, 1])) for idx in np.flatnonzero(mask)}

contact_map_to_dot_bracket

Python

contact_map_to_dot_bracket(
    contact_map: Tensor | ndarray,
    unsafe: bool = False,
    *,
    threshold: float = 0.5
) -> str

Convert a contact map (NumPy or Torch) to a dot-bracket notation string.

Examples:

Torch input

>>> import torch
>>> contact_map_tensor = torch.tensor([[0, 0, 0, 1],
...                                    [0, 0, 1, 0],
...                                    [0, 1, 0, 0],
...                                    [1, 0, 0, 0]])
>>> contact_map_to_dot_bracket(contact_map_tensor)
'(())'

NumPy input

>>> import numpy as np
>>> contact_map = np.array([[0, 0, 0, 1],
...                          [0, 0, 1, 0],
...                          [0, 1, 0, 0],
...                          [1, 0, 0, 0]])
>>> contact_map_to_dot_bracket(contact_map)
'(())'

List input

>>> contact_map_to_dot_bracket([[0, 1], [1, 0]])
'()'

Source code in multimolecule/utils/rna/secondary_structure/notations.py

def contact_map_to_dot_bracket(
    contact_map: Tensor | np.ndarray, unsafe: bool = False, *, threshold: float = 0.5
) -> str:
    """
    Convert a contact map (NumPy or Torch) to a dot-bracket notation string.

    Examples:
        Torch input
        >>> import torch
        >>> contact_map_tensor = torch.tensor([[0, 0, 0, 1],
        ...                                    [0, 0, 1, 0],
        ...                                    [0, 1, 0, 0],
        ...                                    [1, 0, 0, 0]])
        >>> contact_map_to_dot_bracket(contact_map_tensor)
        '(())'

        NumPy input
        >>> import numpy as np
        >>> contact_map = np.array([[0, 0, 0, 1],
        ...                          [0, 0, 1, 0],
        ...                          [0, 1, 0, 0],
        ...                          [1, 0, 0, 0]])
        >>> contact_map_to_dot_bracket(contact_map)
        '(())'

        List input
        >>> contact_map_to_dot_bracket([[0, 1], [1, 0]])
        '()'
    """
    return pairs_to_dot_bracket(
        contact_map_to_pairs(contact_map, unsafe=unsafe, threshold=threshold), length=len(contact_map), unsafe=unsafe
    )

contact_map_to_pairs

Python

contact_map_to_pairs(
    contact_map: Tensor,
    unsafe: bool = False,
    *,
    threshold: float = 0.5
) -> Tensor

Python

contact_map_to_pairs(
    contact_map: ndarray,
    unsafe: bool = False,
    *,
    threshold: float = 0.5
) -> ndarray

Python

contact_map_to_pairs(
    contact_map: Sequence,
    unsafe: bool = False,
    *,
    threshold: float = 0.5
) -> PairsList

Python

contact_map_to_pairs(
    contact_map: Tensor | ndarray | Sequence,
    unsafe: bool = False,
    *,
    threshold: float = 0.5
) -> Tensor | ndarray | PairsList

Convert a contact map to a list of base pairs.

If contact_map is a torch tensor, returns a (K, 2) torch.LongTensor. Otherwise, returns a numpy (K, 2) int array (list inputs return a list of tuples).

For integer/bool contact maps, any non-zero entry is treated as a contact and the map is expected to represent a binary (symmetric) adjacency matrix.

For floating-point contact maps, values are interpreted as pairing probabilities in [0, 1] (or logits/scores in unsafe mode), and pairs are decoded using a greedy NMS-style one-to-one matching that prioritizes higher scores above threshold.

Examples:

Torch input

>>> import torch
>>> contact_map_tensor = torch.tensor([[0, 0, 0, 1],
...                                   [0, 0, 1, 0],
...                                   [0, 1, 0, 0],
...                                   [1, 0, 0, 0]])
>>> contact_map_to_pairs(contact_map_tensor).tolist()
[[0, 3], [1, 2]]

NumPy input

>>> import numpy as np
>>> contact_map_array = np.array([[0, 0, 0, 1],
...                               [0, 0, 1, 0],
...                               [0, 1, 0, 0],
...                               [1, 0, 0, 0]])
>>> contact_map_to_pairs(contact_map_array).tolist()
[[0, 3], [1, 2]]
>>> contact_map_to_pairs(np.array([[0.0, 0.8], [0.8, 0.0]]), threshold=0.5).tolist()
[[0, 1]]

List input

>>> contact_map_to_pairs([[0, 1], [1, 0]])
[(0, 1)]

Source code in multimolecule/utils/rna/secondary_structure/notations.py

def contact_map_to_pairs(
    contact_map: Tensor | np.ndarray | Sequence, unsafe: bool = False, *, threshold: float = 0.5
) -> Tensor | np.ndarray | PairsList:
    """
    Convert a contact map to a list of base pairs.

    If ``contact_map`` is a torch tensor, returns a ``(K, 2)`` torch.LongTensor.
    Otherwise, returns a numpy ``(K, 2)`` int array (list inputs return a list of tuples).

    For integer/bool contact maps, any non-zero entry is treated as a contact and the map is
    expected to represent a binary (symmetric) adjacency matrix.

    For floating-point contact maps, values are interpreted as pairing probabilities in ``[0, 1]``
    (or logits/scores in ``unsafe`` mode), and pairs are decoded using a greedy NMS-style
    one-to-one matching that prioritizes higher scores above ``threshold``.

    Examples:
        Torch input
        >>> import torch
        >>> contact_map_tensor = torch.tensor([[0, 0, 0, 1],
        ...                                   [0, 0, 1, 0],
        ...                                   [0, 1, 0, 0],
        ...                                   [1, 0, 0, 0]])
        >>> contact_map_to_pairs(contact_map_tensor).tolist()
        [[0, 3], [1, 2]]

        NumPy input
        >>> import numpy as np
        >>> contact_map_array = np.array([[0, 0, 0, 1],
        ...                               [0, 0, 1, 0],
        ...                               [0, 1, 0, 0],
        ...                               [1, 0, 0, 0]])
        >>> contact_map_to_pairs(contact_map_array).tolist()
        [[0, 3], [1, 2]]
        >>> contact_map_to_pairs(np.array([[0.0, 0.8], [0.8, 0.0]]), threshold=0.5).tolist()
        [[0, 1]]

        List input
        >>> contact_map_to_pairs([[0, 1], [1, 0]])
        [(0, 1)]
    """
    if isinstance(contact_map, Tensor):
        if contact_map.ndim != 2 or contact_map.shape[0] != contact_map.shape[1]:
            raise ValueError("Contact map must be a square 2D matrix.")
        if contact_map.is_floating_point():
            return _torch_contact_map_to_pairs_float(contact_map, unsafe=unsafe, threshold=threshold)
        return _torch_contact_map_to_pairs_binary(contact_map, unsafe=unsafe)
    if isinstance(contact_map, np.ndarray):
        if contact_map.ndim != 2 or contact_map.shape[0] != contact_map.shape[1]:
            raise ValueError("Contact map must be a square 2D matrix.")
        if np.issubdtype(contact_map.dtype, np.floating):
            return _numpy_contact_map_to_pairs_float(contact_map, unsafe=unsafe, threshold=threshold)
        return _numpy_contact_map_to_pairs_binary(contact_map, unsafe=unsafe)
    if isinstance(contact_map, Sequence):
        contact_map = np.asarray(contact_map)
        if contact_map.ndim != 2 or contact_map.shape[0] != contact_map.shape[1]:
            raise ValueError("Contact map must be a square 2D matrix.")
        if np.issubdtype(contact_map.dtype, np.floating):
            pairs = _numpy_contact_map_to_pairs_float(contact_map, unsafe=unsafe, threshold=threshold)
        else:
            pairs = _numpy_contact_map_to_pairs_binary(contact_map, unsafe=unsafe)
        return [tuple(pair) for pair in pairs.tolist()]
    raise TypeError("contact_map must be a torch.Tensor, numpy.ndarray, or sequence")

dot_bracket_to_contact_map

Python

dot_bracket_to_contact_map(dot_bracket: str) -> ndarray

Convert a dot-bracket notation string to a numpy contact map.

Examples:

>>> dot_bracket_to_contact_map('(())').astype(int).tolist()
[[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]]

Source code in multimolecule/utils/rna/secondary_structure/notations.py

def dot_bracket_to_contact_map(dot_bracket: str) -> np.ndarray:
    """
    Convert a dot-bracket notation string to a numpy contact map.

    Examples:
        >>> dot_bracket_to_contact_map('(())').astype(int).tolist()
        [[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]]
    """
    return pairs_to_contact_map(dot_bracket_to_pairs(dot_bracket), length=len(dot_bracket))

dot_bracket_to_pairs

Python

dot_bracket_to_pairs(dot_bracket: str) -> ndarray

Convert a dot-bracket notation string to a list of base-pair indices.

Parameters:

Name	Type	Description	Default
dot_bracket	str	Dot-bracket notation. Supports pseudoknots via multiple bracket types, including (), [], {}, <>, and A-Z/a-z. Unpaired tokens (., +, _, ,) are treated as unpaired positions.	required

Returns:

Type	Description
ndarray	A numpy array of shape (n, 2) with pairs (i, j) where 0 <= i < j < len(dot_bracket).

Raises:

Type	Description
ValueError	On unmatched or invalid symbols.

Examples:

>>> dot_bracket_to_pairs("((.))").tolist()
[[0, 4], [1, 3]]
>>> dot_bracket_to_pairs("([)]").tolist()
[[0, 2], [1, 3]]
>>> dot_bracket_to_pairs("...").tolist()
[]

Source code in multimolecule/utils/rna/secondary_structure/notations.py

def dot_bracket_to_pairs(dot_bracket: str) -> np.ndarray:
    """
    Convert a dot-bracket notation string to a list of base-pair indices.

    Args:
        dot_bracket: Dot-bracket notation. Supports pseudoknots via multiple
            bracket types, including (), [], {}, <>, and A-Z/a-z. Unpaired
            tokens (`.`, `+`, `_`, `,`) are treated as unpaired positions.

    Returns:
        A numpy array of shape (n, 2) with pairs ``(i, j)`` where ``0 <= i < j < len(dot_bracket)``.

    Raises:
        ValueError: On unmatched or invalid symbols.

    Examples:
        >>> dot_bracket_to_pairs("((.))").tolist()
        [[0, 4], [1, 3]]
        >>> dot_bracket_to_pairs("([)]").tolist()
        [[0, 2], [1, 3]]
        >>> dot_bracket_to_pairs("...").tolist()
        []
    """
    stacks: defaultdict[str, List[int]] = defaultdict(list)
    pairs: PairsList = []
    for i, symbol in enumerate(dot_bracket):
        if symbol in _DOT_BRACKET_PAIR_TABLE:
            stacks[symbol].append(i)
        elif symbol in _REVERSE_DOT_BRACKET_PAIR_TABLE:
            opener = _REVERSE_DOT_BRACKET_PAIR_TABLE[symbol]
            try:
                j = stacks[opener].pop()
            except IndexError:
                raise ValueError(f"Unmatched symbol {symbol} at position {i} in sequence {dot_bracket}") from None
            pairs.append((j, i))
        elif symbol not in _UNPAIRED_TOKENS:
            raise ValueError(f"Invalid symbol {symbol} at position {i} in sequence {dot_bracket}")
    for symbol, stack in stacks.items():
        if stack:
            raise ValueError(f"Unmatched symbol {symbol} at position {stack[0]} in sequence {dot_bracket}")
    if not pairs:
        return np.empty((0, 2), dtype=int)
    pairs.sort()
    return np.asarray(pairs, dtype=int)

pairs_to_contact_map

Python

pairs_to_contact_map(
    pairs: Tensor,
    length: int | None = None,
    unsafe: bool = False,
) -> Tensor

Python

pairs_to_contact_map(
    pairs: ndarray,
    length: int | None = None,
    unsafe: bool = False,
) -> ndarray

Python

pairs_to_contact_map(
    pairs: PairsList,
    length: int | None = None,
    unsafe: bool = False,
) -> List[List[bool]]

Python

pairs_to_contact_map(
    pairs: Tensor | ndarray | Pairs,
    length: int | None = None,
    unsafe: bool = False,
) -> Tensor | ndarray | List[List[bool]]

Convert base pairs to a symmetric contact map.

If pairs is a torch tensor, returns a boolean torch.Tensor on the same device. Otherwise, returns a numpy boolean array. If length is None, it is inferred as max(pairs) + 1.

Examples:

Torch input

>>> import torch
>>> contact_map_tensor = pairs_to_contact_map(torch.tensor([[0, 3], [1, 2]]), length=4)
>>> contact_map_tensor.to(torch.int).tolist()
[[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]]

NumPy input

>>> import numpy as np
>>> contact_map = pairs_to_contact_map(np.array([(0, 3), (1, 2)]), length=4)
>>> contact_map.astype(int).tolist()
[[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]]
>>> pairs_to_contact_map(np.array([(0, 2)])).astype(int).tolist()
[[0, 0, 1], [0, 0, 0], [1, 0, 0]]

List input

>>> pairs_to_contact_map([(0, 2)])
[[False, False, True], [False, False, False], [True, False, False]]

Source code in multimolecule/utils/rna/secondary_structure/notations.py

def pairs_to_contact_map(
    pairs: Tensor | np.ndarray | Pairs,
    length: int | None = None,
    unsafe: bool = False,
) -> Tensor | np.ndarray | List[List[bool]]:
    """
    Convert base pairs to a symmetric contact map.

    If ``pairs`` is a torch tensor, returns a boolean torch.Tensor on the same device.
    Otherwise, returns a numpy boolean array.
    If ``length`` is None, it is inferred as ``max(pairs) + 1``.

    Examples:
        Torch input
        >>> import torch
        >>> contact_map_tensor = pairs_to_contact_map(torch.tensor([[0, 3], [1, 2]]), length=4)
        >>> contact_map_tensor.to(torch.int).tolist()
        [[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]]

        NumPy input
        >>> import numpy as np
        >>> contact_map = pairs_to_contact_map(np.array([(0, 3), (1, 2)]), length=4)
        >>> contact_map.astype(int).tolist()
        [[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]]
        >>> pairs_to_contact_map(np.array([(0, 2)])).astype(int).tolist()
        [[0, 0, 1], [0, 0, 0], [1, 0, 0]]

        List input
        >>> pairs_to_contact_map([(0, 2)])
        [[False, False, True], [False, False, False], [True, False, False]]
    """
    if isinstance(pairs, Tensor):
        if pairs.numel() == 0:
            return _torch_pairs_to_contact_map(pairs.view(0, 2), length, unsafe)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_pairs_to_contact_map(pairs, length, unsafe)
    if isinstance(pairs, np.ndarray):
        if pairs.size == 0:
            return _numpy_pairs_to_contact_map(pairs.reshape(0, 2), length, unsafe)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_pairs_to_contact_map(pairs, length, unsafe)
    if isinstance(pairs, Sequence):
        pairs = np.asarray(pairs, dtype=int)
        if pairs.size == 0:
            return _numpy_pairs_to_contact_map(pairs.reshape(0, 2), length, unsafe)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        return _numpy_pairs_to_contact_map(pairs, length, unsafe).tolist()
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

pairs_to_dot_bracket

Python

pairs_to_dot_bracket(
    pairs: Tensor | ndarray | Pairs,
    length: int,
    unsafe: bool = False,
) -> str

Convert base pairs to a dot-bracket string (backend-aware input, string output).

Torch inputs are accepted and internally converted to NumPy for string building. In safe mode, tiers are assigned using an exact minimal-tier coloring. In unsafe mode, a greedy tiering is used for speed and may use more bracket types.

Examples:

Torch input

>>> import torch
>>> pairs_to_dot_bracket(torch.tensor([[0, 2], [1, 3]]), length=4)
'([)]'

NumPy input

>>> import numpy as np
>>> pairs_to_dot_bracket(np.array([(0, 3), (1, 2)]), length=4)
'(())'

List input

>>> pairs_to_dot_bracket([(0, 3), (1, 2)], length=4)
'(())'

Source code in multimolecule/utils/rna/secondary_structure/notations.py

def pairs_to_dot_bracket(
    pairs: Tensor | np.ndarray | Pairs,
    length: int,
    unsafe: bool = False,
) -> str:
    """
    Convert base pairs to a dot-bracket string (backend-aware input, string output).

    Torch inputs are accepted and internally converted to NumPy for string building.
    In safe mode, tiers are assigned using an exact minimal-tier coloring.
    In unsafe mode, a greedy tiering is used for speed and may use more bracket types.

    Examples:
        Torch input
        >>> import torch
        >>> pairs_to_dot_bracket(torch.tensor([[0, 2], [1, 3]]), length=4)
        '([)]'

        NumPy input
        >>> import numpy as np
        >>> pairs_to_dot_bracket(np.array([(0, 3), (1, 2)]), length=4)
        '(())'

        List input
        >>> pairs_to_dot_bracket([(0, 3), (1, 2)], length=4)
        '(())'
    """
    # Always operate in NumPy for string construction
    if isinstance(pairs, Tensor):
        pairs = pairs.detach().cpu().numpy()
    elif isinstance(pairs, np.ndarray):
        pass
    elif isinstance(pairs, Sequence):
        pairs = np.asarray(list(pairs), dtype=int)
    else:
        raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")
    if pairs.size == 0:
        return _numpy_pairs_to_dot_bracket(pairs, length, unsafe)
    if pairs.ndim != 2 or pairs.shape[1] != 2:
        raise ValueError("pairs must be an array-like with shape (n, 2)")
    return _numpy_pairs_to_dot_bracket(pairs, length, unsafe)

normalize_pairs

Python

normalize_pairs(pairs: Tensor) -> Tensor

Python

normalize_pairs(pairs: ndarray) -> ndarray

Python

normalize_pairs(pairs: Pairs) -> PairsList

Python

normalize_pairs(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | PairsList

Normalize base-pair indices to unique, sorted (i < j) pairs.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | PairsList	Normalized pairs using the same backend as input.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> normalize_pairs(torch.tensor([[3, 1], [1, 3]])).tolist()
[[1, 3]]

NumPy input

>>> import numpy as np
>>> normalize_pairs(np.array([[3, 1], [1, 3]])).tolist()
[[1, 3]]

List input

>>> normalize_pairs([(3, 1), (1, 3), (2, 0)])
[(0, 2), (1, 3)]

Source code in multimolecule/utils/rna/secondary_structure/pairs.py

def normalize_pairs(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | PairsList:
    """
    Normalize base-pair indices to unique, sorted (i < j) pairs.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Normalized pairs using the same backend as input.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> normalize_pairs(torch.tensor([[3, 1], [1, 3]])).tolist()
        [[1, 3]]

        NumPy input
        >>> import numpy as np
        >>> normalize_pairs(np.array([[3, 1], [1, 3]])).tolist()
        [[1, 3]]

        List input
        >>> normalize_pairs([(3, 1), (1, 3), (2, 0)])
        [(0, 2), (1, 3)]
    """
    if isinstance(pairs, Tensor):
        if pairs.numel() == 0:
            return pairs.view(0, 2).to(dtype=torch.long)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_normalize_pairs(pairs)
    if isinstance(pairs, np.ndarray):
        if pairs.size == 0:
            return np.empty((0, 2), dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_normalize_pairs(pairs)
    if isinstance(pairs, Sequence):
        if not pairs:
            return []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.size == 0:
            return []
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        return list(map(tuple, _numpy_normalize_pairs(pairs).tolist()))
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

pairs_to_duplex_segment_list

Python

pairs_to_duplex_segment_list(
    pairs: Tensor | ndarray | Pairs,
) -> List[Segment]

Convert base pairs to bulge-tolerant segments while preserving actual pairs.

Segments grow by stepping to the next paired 5’ position and the previous paired 3’ position; if those positions pair to each other, they belong to the same segment.

Source code in multimolecule/utils/rna/secondary_structure/pairs.py

def pairs_to_duplex_segment_list(pairs: Tensor | np.ndarray | Pairs) -> List[Segment]:
    """
    Convert base pairs to bulge-tolerant segments while preserving actual pairs.

    Segments grow by stepping to the next paired 5' position and the previous paired 3' position;
    if those positions pair to each other, they belong to the same segment.
    """
    if isinstance(pairs, Tensor):
        pairs = pairs.detach().cpu().numpy()
    if isinstance(pairs, Sequence):
        pairs = np.asarray(pairs, dtype=int)
    if isinstance(pairs, np.ndarray):
        if pairs.size == 0:
            return []
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_pairs_to_duplex_segment_list(pairs)
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

pairs_to_helix_segment_arrays

Python

pairs_to_helix_segment_arrays(
    pairs: Tensor,
) -> Tuple[Tensor, Tensor, Tensor]

Python

pairs_to_helix_segment_arrays(
    pairs: ndarray,
) -> Tuple[ndarray, ndarray, ndarray]

Python

pairs_to_helix_segment_arrays(
    pairs: Pairs,
) -> Tuple[List[int], List[int], List[int]]

Python

pairs_to_helix_segment_arrays(
    pairs: Tensor | ndarray | Pairs,
) -> (
    Tuple[Tensor, Tensor, Tensor]
    | Tuple[ndarray, ndarray, ndarray]
    | Tuple[List[int], List[int], List[int]]
)

Convert base pairs to strict stacked segments (no bulges).

Source code in multimolecule/utils/rna/secondary_structure/pairs.py

def pairs_to_helix_segment_arrays(
    pairs: Tensor | np.ndarray | Pairs,
) -> Tuple[Tensor, Tensor, Tensor] | Tuple[np.ndarray, np.ndarray, np.ndarray] | Tuple[List[int], List[int], List[int]]:
    """
    Convert base pairs to strict stacked segments (no bulges).
    """
    if isinstance(pairs, Tensor):
        if pairs.numel() == 0:
            empty = pairs.new_empty((0,), dtype=torch.long)
            return empty, empty, empty
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_pairs_to_helix_segment_arrays(pairs)
    if isinstance(pairs, np.ndarray):
        if pairs.size == 0:
            empty_arr = np.empty((0,), dtype=int)
            return empty_arr, empty_arr, empty_arr
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_pairs_to_helix_segment_arrays(pairs)
    if isinstance(pairs, Sequence):
        if not pairs:
            empty_arr = np.empty((0,), dtype=int)
            return empty_arr, empty_arr, empty_arr
        pairs_array = np.asarray(pairs, dtype=int)
        if pairs_array.size == 0:
            empty_arr = np.empty((0,), dtype=int)
            return empty_arr, empty_arr, empty_arr
        if pairs_array.ndim != 2 or pairs_array.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        start_i, start_j, lengths = _numpy_pairs_to_helix_segment_arrays(pairs_array)
        return start_i.tolist(), start_j.tolist(), lengths.tolist()
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

pairs_to_stem_segment_arrays

Python

pairs_to_stem_segment_arrays(
    pairs: Tensor,
) -> Tuple[Tensor, Tensor, Tensor]

Python

pairs_to_stem_segment_arrays(
    pairs: ndarray,
) -> Tuple[ndarray, ndarray, ndarray]

Python

pairs_to_stem_segment_arrays(
    pairs: Pairs,
) -> Tuple[List[int], List[int], List[int]]

Python

pairs_to_stem_segment_arrays(
    pairs: Tensor | ndarray | Pairs,
) -> (
    Tuple[Tensor, Tensor, Tensor]
    | Tuple[ndarray, ndarray, ndarray]
    | Tuple[List[int], List[int], List[int]]
)

Convert base pairs to stem segments that allow bulges/internal loops.

Source code in multimolecule/utils/rna/secondary_structure/pairs.py

def pairs_to_stem_segment_arrays(
    pairs: Tensor | np.ndarray | Pairs,
) -> Tuple[Tensor, Tensor, Tensor] | Tuple[np.ndarray, np.ndarray, np.ndarray] | Tuple[List[int], List[int], List[int]]:
    """
    Convert base pairs to stem segments that allow bulges/internal loops.
    """
    if isinstance(pairs, Tensor):
        if pairs.numel() == 0:
            empty = pairs.new_empty((0,), dtype=torch.long)
            return empty, empty, empty
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_pairs_to_stem_segment_arrays(pairs)
    if isinstance(pairs, np.ndarray):
        if pairs.size == 0:
            empty_arr = np.empty((0,), dtype=int)
            return empty_arr, empty_arr, empty_arr
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_pairs_to_stem_segment_arrays(pairs)
    if isinstance(pairs, Sequence):
        if not pairs:
            empty_arr = np.empty((0,), dtype=int)
            return empty_arr, empty_arr, empty_arr
        pairs_array = np.asarray(pairs, dtype=int)
        if pairs_array.size == 0:
            empty_arr = np.empty((0,), dtype=int)
            return empty_arr, empty_arr, empty_arr
        if pairs_array.ndim != 2 or pairs_array.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        start_i, start_j, lengths = _numpy_pairs_to_stem_segment_arrays(pairs_array)
        return start_i.tolist(), start_j.tolist(), lengths.tolist()
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

segment_arrays_to_pairs

Python

segment_arrays_to_pairs(
    segments: Tuple[Tensor, Tensor, Tensor] | List[Segment],
    mask: Tensor | None = None,
    empty: ndarray | None = None,
) -> Tensor | ndarray

Convert segments back to base pairs.

Source code in multimolecule/utils/rna/secondary_structure/pairs.py

def segment_arrays_to_pairs(
    segments: Tuple[Tensor, Tensor, Tensor] | List[Segment],
    mask: Tensor | None = None,
    empty: np.ndarray | None = None,
) -> Tensor | np.ndarray:
    """
    Convert segments back to base pairs.
    """
    if isinstance(segments, tuple):
        if len(segments) != 3:
            raise ValueError("segments must be (start_i, start_j, lengths)")
        start_i, start_j, lengths = segments
        if isinstance(start_i, Tensor):
            if mask is None:
                mask = torch.ones_like(start_i, dtype=torch.bool)
            return _torch_segment_arrays_to_pairs(start_i, start_j, lengths, mask)
        if isinstance(start_i, np.ndarray):
            if empty is None:
                empty = np.empty((0, 2), dtype=int)
            if mask is None:
                mask_array = np.ones_like(start_i, dtype=bool)
            elif (isinstance(mask, np.ndarray) and mask.dtype == bool) or (
                isinstance(mask, (list, tuple)) and mask and isinstance(mask[0], bool)
            ):
                mask_array = np.asarray(mask, dtype=bool)
            elif isinstance(mask, (list, tuple, np.ndarray)):
                mask_array = np.zeros_like(start_i, dtype=bool)
                mask_array[np.asarray(mask, dtype=int)] = True
            else:
                raise TypeError("mask must be a boolean array or sequence of indices")
            return _numpy_segment_arrays_to_pairs(start_i, start_j, lengths, mask_array, empty)
        if isinstance(start_i, (list, tuple)):
            start_i = np.asarray(start_i, dtype=int)
            start_j = np.asarray(start_j, dtype=int)
            lengths = np.asarray(lengths, dtype=int)
            if empty is None:
                empty = np.empty((0, 2), dtype=int)
            if mask is None:
                mask_array = np.ones_like(start_i, dtype=bool)
            elif (isinstance(mask, np.ndarray) and mask.dtype == bool) or (
                isinstance(mask, (list, tuple)) and mask and isinstance(mask[0], bool)
            ):
                mask_array = np.asarray(mask, dtype=bool)
            elif isinstance(mask, (list, tuple, np.ndarray)):
                mask_array = np.zeros_like(start_i, dtype=bool)
                mask_array[np.asarray(mask, dtype=int)] = True
            else:
                raise TypeError("mask must be a boolean array or sequence of indices")
            return _numpy_segment_arrays_to_pairs(start_i, start_j, lengths, mask_array, empty)
        raise TypeError("segments tuple must contain torch.Tensor, numpy.ndarray, or list values")
    assert not isinstance(segments, tuple)
    if empty is None:
        empty = np.empty((0, 2), dtype=int)
    if mask is not None:
        if (isinstance(mask, np.ndarray) and mask.dtype == bool) or (
            isinstance(mask, (list, tuple)) and mask and isinstance(mask[0], bool)
        ):
            segments = [seg for seg, keep in zip(segments, mask) if keep]
        elif isinstance(mask, (list, tuple, np.ndarray)):
            segments = [segments[int(idx)] for idx in mask]
        else:
            raise TypeError("mask must be a boolean array or sequence of indices")
    return segment_list_to_pairs(segments, empty)

sort_pairs

Python

sort_pairs(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | PairsList

Sort base-pair indices by (i, j) without normalization or de-duplication.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | PairsList	Sorted pairs using the same backend as input.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> sort_pairs(torch.tensor([[2, 5], [0, 1]])).tolist()
[[0, 1], [2, 5]]

NumPy input

>>> import numpy as np
>>> sort_pairs(np.array([[2, 5], [0, 1]])).tolist()
[[0, 1], [2, 5]]

List input

>>> sort_pairs([(2, 5), (0, 1)])
[(0, 1), (2, 5)]

Source code in multimolecule/utils/rna/secondary_structure/pairs.py

def sort_pairs(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | PairsList:
    """
    Sort base-pair indices by (i, j) without normalization or de-duplication.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Sorted pairs using the same backend as input.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> sort_pairs(torch.tensor([[2, 5], [0, 1]])).tolist()
        [[0, 1], [2, 5]]

        NumPy input
        >>> import numpy as np
        >>> sort_pairs(np.array([[2, 5], [0, 1]])).tolist()
        [[0, 1], [2, 5]]

        List input
        >>> sort_pairs([(2, 5), (0, 1)])
        [(0, 1), (2, 5)]
    """
    if isinstance(pairs, Tensor):
        if pairs.numel() == 0:
            return pairs.view(0, 2)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_sort_pairs(pairs)
    if isinstance(pairs, np.ndarray):
        if pairs.size == 0:
            return np.empty((0, 2), dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_sort_pairs(pairs)
    if isinstance(pairs, Sequence):
        if not pairs:
            return []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        return list(map(tuple, _numpy_sort_pairs(pairs).tolist()))
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

crossing_arcs

Python

crossing_arcs(pairs: Tensor) -> Tensor

Python

crossing_arcs(pairs: ndarray) -> ndarray

Python

crossing_arcs(pairs: Pairs) -> List[Tuple[Pair, Pair]]

Python

crossing_arcs(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | List[Tuple[Pair, Pair]]

Return pair-level crossing arcs as ((i, j), (k, l)) where i < k < j < l.

Pair inputs are expected to be normalized.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | List[Tuple[Pair, Pair]]	Crossing arcs using the same backend as input. Tensor/NumPy inputs
Tensor | ndarray | List[Tuple[Pair, Pair]]	return shape (n, 2, 2).

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> crossing_arcs(torch.tensor([[0, 2], [1, 3]])).tolist()
[[[0, 2], [1, 3]]]

NumPy input

>>> import numpy as np
>>> crossing_arcs(np.array([[0, 2], [1, 3]])).tolist()
[[[0, 2], [1, 3]]]

List input

>>> crossing_arcs([(0, 2), (1, 3)])
[((0, 2), (1, 3))]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def crossing_arcs(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | List[Tuple[Pair, Pair]]:
    """
    Return pair-level crossing arcs as ((i, j), (k, l)) where i < k < j < l.

    Pair inputs are expected to be normalized.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Crossing arcs using the same backend as input. Tensor/NumPy inputs
        return shape (n, 2, 2).

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> crossing_arcs(torch.tensor([[0, 2], [1, 3]])).tolist()
        [[[0, 2], [1, 3]]]

        NumPy input
        >>> import numpy as np
        >>> crossing_arcs(np.array([[0, 2], [1, 3]])).tolist()
        [[[0, 2], [1, 3]]]

        List input
        >>> crossing_arcs([(0, 2), (1, 3)])
        [((0, 2), (1, 3))]
    """
    if isinstance(pairs, Tensor):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_crossing_arcs(pairs)
    if isinstance(pairs, np.ndarray):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_crossing_arcs(pairs)
    if isinstance(pairs, Sequence):
        if not pairs:
            return []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        events = _numpy_crossing_arcs(pairs)
        return [(tuple(pair_a), tuple(pair_b)) for pair_a, pair_b in events.tolist()]
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

crossing_events

Python

crossing_events(pairs: Tensor) -> Tensor

Python

crossing_events(pairs: ndarray) -> ndarray

Python

crossing_events(
    pairs: Pairs,
) -> List[
    Tuple[
        Tuple[int, int, int, int], Tuple[int, int, int, int]
    ]
]

Python

crossing_events(
    pairs: Tensor | ndarray | Pairs,
) -> (
    Tensor
    | ndarray
    | List[
        Tuple[
            Tuple[int, int, int, int],
            Tuple[int, int, int, int],
        ]
    ]
)

Return crossing events between stem segments.

Each stem is encoded as (start_5p, stop_5p, start_3p, stop_3p).

For pair-level crossing arcs, use crossing_arcs.

Pair inputs are expected to be normalized.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | List[Tuple[Tuple[int, int, int, int], Tuple[int, int, int, int]]]	Crossing events using the same backend as input. Tensor/NumPy inputs
Tensor | ndarray | List[Tuple[Tuple[int, int, int, int], Tuple[int, int, int, int]]]	return shape (n, 2, 4).

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> crossing_events(torch.tensor([[0, 2], [1, 3]])).tolist()
[[[0, 0, 2, 2], [1, 1, 3, 3]]]

NumPy input

>>> import numpy as np
>>> crossing_events(np.array([[0, 2], [1, 3]])).tolist()
[[[0, 0, 2, 2], [1, 1, 3, 3]]]

List input

>>> crossing_events([(0, 2), (1, 3)])
[((0, 0, 2, 2), (1, 1, 3, 3))]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def crossing_events(
    pairs: Tensor | np.ndarray | Pairs,
) -> Tensor | np.ndarray | List[Tuple[Tuple[int, int, int, int], Tuple[int, int, int, int]]]:
    """
    Return crossing events between stem segments.

    Each stem is encoded as (start_5p, stop_5p, start_3p, stop_3p).

    For pair-level crossing arcs, use ``crossing_arcs``.

    Pair inputs are expected to be normalized.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Crossing events using the same backend as input. Tensor/NumPy inputs
        return shape (n, 2, 4).

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> crossing_events(torch.tensor([[0, 2], [1, 3]])).tolist()
        [[[0, 0, 2, 2], [1, 1, 3, 3]]]

        NumPy input
        >>> import numpy as np
        >>> crossing_events(np.array([[0, 2], [1, 3]])).tolist()
        [[[0, 0, 2, 2], [1, 1, 3, 3]]]

        List input
        >>> crossing_events([(0, 2), (1, 3)])
        [((0, 0, 2, 2), (1, 1, 3, 3))]
    """
    if isinstance(pairs, Tensor):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_crossing_events(pairs)
    if isinstance(pairs, np.ndarray):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_crossing_events(pairs)
    if isinstance(pairs, Sequence):
        if not pairs:
            return []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        events = _numpy_crossing_events(pairs)
        return [(tuple(stem_a), tuple(stem_b)) for stem_a, stem_b in events.tolist()]
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

crossing_mask

Python

crossing_mask(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | List[bool]

Return a boolean mask for pairs that cross any other pair.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | List[bool]	Boolean mask for the input pairs using the same backend as input.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> crossing_mask(torch.tensor([[0, 2], [1, 3]])).tolist()
[True, True]

NumPy input

>>> import numpy as np
>>> crossing_mask(np.array([[0, 2], [1, 3]])).tolist()
[True, True]

List input

>>> crossing_mask([(0, 2), (1, 3)])
[True, True]
>>> crossing_mask([(0, 3), (1, 2)])
[False, False]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def crossing_mask(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | List[bool]:
    """
    Return a boolean mask for pairs that cross any other pair.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Boolean mask for the input pairs using the same backend as input.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> crossing_mask(torch.tensor([[0, 2], [1, 3]])).tolist()
        [True, True]

        NumPy input
        >>> import numpy as np
        >>> crossing_mask(np.array([[0, 2], [1, 3]])).tolist()
        [True, True]

        List input
        >>> crossing_mask([(0, 2), (1, 3)])
        [True, True]
        >>> crossing_mask([(0, 3), (1, 2)])
        [False, False]
    """
    if isinstance(pairs, Tensor):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_crossing_mask(pairs)
    if isinstance(pairs, np.ndarray):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_crossing_mask(pairs)
    if isinstance(pairs, Sequence):
        if not pairs:
            return []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        return _numpy_crossing_mask(pairs).tolist()
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

crossing_nucleotides

Python

crossing_nucleotides(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | List[int]

Return nucleotide indices involved in any crossing pair.

Pair inputs are expected to be normalized.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | List[int]	Unique nucleotide indices using the same backend as input.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> crossing_nucleotides(torch.tensor([[0, 2], [1, 3]])).tolist()
[0, 1, 2, 3]

NumPy input

>>> import numpy as np
>>> crossing_nucleotides(np.array([[0, 2], [1, 3]])).tolist()
[0, 1, 2, 3]

List input

>>> crossing_nucleotides([(0, 2), (1, 3)])
[0, 1, 2, 3]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def crossing_nucleotides(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | List[int]:
    """
    Return nucleotide indices involved in any crossing pair.

    Pair inputs are expected to be normalized.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Unique nucleotide indices using the same backend as input.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> crossing_nucleotides(torch.tensor([[0, 2], [1, 3]])).tolist()
        [0, 1, 2, 3]

        NumPy input
        >>> import numpy as np
        >>> crossing_nucleotides(np.array([[0, 2], [1, 3]])).tolist()
        [0, 1, 2, 3]

        List input
        >>> crossing_nucleotides([(0, 2), (1, 3)])
        [0, 1, 2, 3]
    """
    if isinstance(pairs, Tensor):
        crossing = crossing_pairs(pairs)
        if crossing.numel() == 0:
            return torch.empty((0,), dtype=torch.long, device=pairs.device)
        return torch.unique(crossing.view(-1)).to(torch.long)
    if isinstance(pairs, np.ndarray):
        crossing = crossing_pairs(pairs)
        if crossing.size == 0:
            return np.empty((0,), dtype=int)
        return np.unique(crossing.reshape(-1))
    if isinstance(pairs, Sequence):
        crossing = crossing_pairs(pairs)
        if not crossing:
            return []
        return np.unique(np.asarray(crossing, dtype=int).reshape(-1)).tolist()
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

crossing_pairs

Python

crossing_pairs(pairs: ndarray) -> ndarray

Python

crossing_pairs(pairs: Pairs) -> PairsList

Python

crossing_pairs(pairs: Tensor) -> Tensor

Python

crossing_pairs(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | PairsList

Return pairs from segments that cross any other segment (no-heuristic PK).

Pairs are expected to be normalized (unique, sorted with i < j). Use normalize_pairs if you need to normalize raw inputs.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | PairsList	Crossing pairs using the same backend as input.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> crossing_pairs(torch.tensor([[0, 2], [1, 3]])).tolist()
[[0, 2], [1, 3]]

NumPy input

>>> import numpy as np
>>> crossing_pairs(np.array([[0, 2], [1, 3]])).tolist()
[[0, 2], [1, 3]]

List input

>>> crossing_pairs([(0, 2), (1, 3)])
[(0, 2), (1, 3)]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def crossing_pairs(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | PairsList:
    """
    Return pairs from segments that cross any other segment (no-heuristic PK).

    Pairs are expected to be normalized (unique, sorted with i < j).
    Use ``normalize_pairs`` if you need to normalize raw inputs.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Crossing pairs using the same backend as input.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> crossing_pairs(torch.tensor([[0, 2], [1, 3]])).tolist()
        [[0, 2], [1, 3]]

        NumPy input
        >>> import numpy as np
        >>> crossing_pairs(np.array([[0, 2], [1, 3]])).tolist()
        [[0, 2], [1, 3]]

        List input
        >>> crossing_pairs([(0, 2), (1, 3)])
        [(0, 2), (1, 3)]
    """
    if isinstance(pairs, Tensor):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_crossing_pairs(pairs)
    if isinstance(pairs, np.ndarray):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_crossing_pairs(pairs)
    if isinstance(pairs, Sequence):
        if not pairs:
            return []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        return list(map(tuple, _numpy_crossing_pairs(pairs).tolist()))
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

has_pseudoknot

Python

has_pseudoknot(pairs: Tensor | ndarray | Pairs) -> bool

Return True if any pseudoknot pairs are present under segment-MWIS split.

Pair inputs are expected to be normalized.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
bool	True if pseudoknot pairs exist, otherwise False.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> has_pseudoknot(torch.tensor([[0, 2], [1, 3]]))
True

NumPy input

>>> import numpy as np
>>> has_pseudoknot(np.array([[0, 2], [1, 3]]))
True

List input

>>> has_pseudoknot([(0, 2), (1, 3)])
True
>>> has_pseudoknot([(0, 3), (1, 2)])
False

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def has_pseudoknot(pairs: Tensor | np.ndarray | Pairs) -> bool:
    """
    Return True if any pseudoknot pairs are present under segment-MWIS split.

    Pair inputs are expected to be normalized.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        True if pseudoknot pairs exist, otherwise False.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> has_pseudoknot(torch.tensor([[0, 2], [1, 3]]))
        True

        NumPy input
        >>> import numpy as np
        >>> has_pseudoknot(np.array([[0, 2], [1, 3]]))
        True

        List input
        >>> has_pseudoknot([(0, 2), (1, 3)])
        True
        >>> has_pseudoknot([(0, 3), (1, 2)])
        False
    """
    _, pseudoknot_pairs = split_pseudoknot_pairs(pairs)
    if isinstance(pseudoknot_pairs, Tensor):
        return bool(pseudoknot_pairs.numel())
    if isinstance(pseudoknot_pairs, np.ndarray):
        return bool(pseudoknot_pairs.size)
    if isinstance(pseudoknot_pairs, Sequence):
        return bool(pseudoknot_pairs)
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

nested_pairs

Python

nested_pairs(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | PairsList

Return primary pairs from the segment-MWIS split.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | PairsList	Primary pairs using the same backend as input.

This is equivalent to split_pseudoknot_pairs(pairs)[0] and expects normalized unique pairs.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> nested_pairs(torch.tensor([[0, 2], [1, 3]])).tolist()
[[1, 3]]

NumPy input

>>> import numpy as np
>>> nested_pairs(np.array([[0, 2], [1, 3]])).tolist()
[[1, 3]]

List input

>>> nested_pairs([(0, 2), (1, 3)])
[(1, 3)]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def nested_pairs(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | PairsList:
    """
    Return primary pairs from the segment-MWIS split.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Primary pairs using the same backend as input.

    This is equivalent to ``split_pseudoknot_pairs(pairs)[0]`` and expects
    normalized unique pairs.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> nested_pairs(torch.tensor([[0, 2], [1, 3]])).tolist()
        [[1, 3]]

        NumPy input
        >>> import numpy as np
        >>> nested_pairs(np.array([[0, 2], [1, 3]])).tolist()
        [[1, 3]]

        List input
        >>> nested_pairs([(0, 2), (1, 3)])
        [(1, 3)]
    """
    primary, _ = split_pseudoknot_pairs(pairs)
    return primary

pseudoknot_nucleotides

Python

pseudoknot_nucleotides(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | List[int]

Return nucleotide indices involved in any pseudoknot pair.

Pair inputs are expected to be normalized.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | List[int]	Unique nucleotide indices using the same backend as input (sequence inputs return Python lists).

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> pseudoknot_nucleotides(torch.tensor([[0, 2], [1, 3]])).tolist()
[0, 2]

NumPy input

>>> import numpy as np
>>> pseudoknot_nucleotides(np.array([[0, 2], [1, 3]])).tolist()
[0, 2]

List input

>>> pseudoknot_nucleotides([(0, 2), (1, 3)])
[0, 2]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def pseudoknot_nucleotides(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | List[int]:
    """
    Return nucleotide indices involved in any pseudoknot pair.

    Pair inputs are expected to be normalized.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Unique nucleotide indices using the same backend as input (sequence inputs return Python lists).

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> pseudoknot_nucleotides(torch.tensor([[0, 2], [1, 3]])).tolist()
        [0, 2]

        NumPy input
        >>> import numpy as np
        >>> pseudoknot_nucleotides(np.array([[0, 2], [1, 3]])).tolist()
        [0, 2]

        List input
        >>> pseudoknot_nucleotides([(0, 2), (1, 3)])
        [0, 2]
    """
    if isinstance(pairs, Tensor):
        pseudoknot_pairs_data = pseudoknot_pairs(pairs)
        if pseudoknot_pairs_data.numel() == 0:
            return torch.empty((0,), dtype=torch.long, device=pairs.device)
        return torch.unique(pseudoknot_pairs_data.view(-1)).to(torch.long)
    if isinstance(pairs, np.ndarray):
        pseudoknot_pairs_data = pseudoknot_pairs(pairs)
        if pseudoknot_pairs_data.size == 0:
            return np.empty((0,), dtype=int)
        return np.unique(pseudoknot_pairs_data.reshape(-1))
    if isinstance(pairs, Sequence):
        pseudoknot_pairs_data = pseudoknot_pairs(pairs)
        if not pseudoknot_pairs_data:
            return []
        return np.unique(np.asarray(pseudoknot_pairs_data, dtype=int).reshape(-1)).tolist()
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

pseudoknot_pairs

Python

pseudoknot_pairs(pairs: ndarray) -> ndarray

Python

pseudoknot_pairs(pairs: Pairs) -> PairsList

Python

pseudoknot_pairs(pairs: Tensor) -> Tensor

Python

pseudoknot_pairs(
    pairs: Tensor | ndarray | Pairs,
) -> Tensor | ndarray | PairsList

Return pseudoknot pairs from segments not selected by MWIS.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tensor | ndarray | PairsList	Pseudoknot pairs using the same backend as input.

This is equivalent to split_pseudoknot_pairs(pairs)[1] and expects normalized unique pairs.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Tie-breaks for equal total base pairs: (1) minimize unpaired-within-span, (2) minimize total span, (3) minimize number of segments, (4) deterministic order fallback.

Examples:

Torch input

>>> import torch
>>> pseudoknot_pairs(torch.tensor([[0, 2], [1, 3]])).tolist()
[[0, 2]]

NumPy input

>>> import numpy as np
>>> pseudoknot_pairs(np.array([[0, 2], [1, 3]])).tolist()
[[0, 2]]

List input

>>> pseudoknot_pairs([(0, 2), (1, 3)])
[(0, 2)]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def pseudoknot_pairs(pairs: Tensor | np.ndarray | Pairs) -> Tensor | np.ndarray | PairsList:
    """
    Return pseudoknot pairs from segments not selected by MWIS.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        Pseudoknot pairs using the same backend as input.

    This is equivalent to ``split_pseudoknot_pairs(pairs)[1]`` and expects
    normalized unique pairs.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Tie-breaks for equal total base pairs: (1) minimize unpaired-within-span,
    (2) minimize total span, (3) minimize number of segments, (4) deterministic
    order fallback.

    Examples:
        Torch input
        >>> import torch
        >>> pseudoknot_pairs(torch.tensor([[0, 2], [1, 3]])).tolist()
        [[0, 2]]

        NumPy input
        >>> import numpy as np
        >>> pseudoknot_pairs(np.array([[0, 2], [1, 3]])).tolist()
        [[0, 2]]

        List input
        >>> pseudoknot_pairs([(0, 2), (1, 3)])
        [(0, 2)]
    """
    _, pseudoknot = split_pseudoknot_pairs(pairs)
    return pseudoknot

pseudoknot_tiers

Python

pseudoknot_tiers(
    pairs: Tensor | ndarray | Pairs, unsafe: bool = False
) -> List[Tensor] | List[ndarray] | Tiers

Return dot-bracket tiers as non-crossing groups of pairs.

Pairs are expected to be normalized (unique, sorted with i < j). Use normalize_pairs if you need to normalize raw inputs.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required
unsafe	bool	Use greedy tiering for speed instead of minimal coloring.	False

Returns:

Type	Description
List[Tensor] | List[ndarray] | Tiers	A list of tiers. Each tier is a list/array/tensor of pairs.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> tiers = pseudoknot_tiers(torch.tensor([[0, 2], [1, 3]]))
>>> [tier.tolist() for tier in tiers]
[[[0, 2]], [[1, 3]]]

NumPy input

>>> import numpy as np
>>> tiers = pseudoknot_tiers(np.array([[0, 2], [1, 3]]))
>>> [tier.tolist() for tier in tiers]
[[[0, 2]], [[1, 3]]]

List input

>>> pseudoknot_tiers([(0, 3), (1, 2)])
[[(0, 3), (1, 2)]]

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def pseudoknot_tiers(
    pairs: Tensor | np.ndarray | Pairs, unsafe: bool = False
) -> List[Tensor] | List[np.ndarray] | Tiers:
    """
    Return dot-bracket tiers as non-crossing groups of pairs.

    Pairs are expected to be normalized (unique, sorted with i < j).
    Use ``normalize_pairs`` if you need to normalize raw inputs.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.
        unsafe: Use greedy tiering for speed instead of minimal coloring.

    Returns:
        A list of tiers. Each tier is a list/array/tensor of pairs.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> tiers = pseudoknot_tiers(torch.tensor([[0, 2], [1, 3]]))
        >>> [tier.tolist() for tier in tiers]
        [[[0, 2]], [[1, 3]]]

        NumPy input
        >>> import numpy as np
        >>> tiers = pseudoknot_tiers(np.array([[0, 2], [1, 3]]))
        >>> [tier.tolist() for tier in tiers]
        [[[0, 2]], [[1, 3]]]

        List input
        >>> pseudoknot_tiers([(0, 3), (1, 2)])
        [[(0, 3), (1, 2)]]
    """
    if isinstance(pairs, Tensor):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        return _torch_tiers_from_pairs(pairs, unsafe=unsafe)
    if isinstance(pairs, np.ndarray):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        return _numpy_tiers_from_pairs(pairs, unsafe=unsafe)
    if isinstance(pairs, Sequence):
        if not pairs:
            return []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        tiers = _numpy_tiers_from_pairs(pairs, unsafe=unsafe)
        return [list(map(tuple, tier.tolist())) for tier in tiers]
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

split_crossing_pairs

Python

split_crossing_pairs(
    pairs: Tensor,
) -> Tuple[Tensor, Tensor]

Python

split_crossing_pairs(
    pairs: ndarray,
) -> Tuple[ndarray, ndarray]

Python

split_crossing_pairs(
    pairs: PairsList,
) -> Tuple[PairsList, PairsList]

Python

split_crossing_pairs(
    pairs: Tensor | ndarray | Pairs,
) -> Tuple[
    Tensor | ndarray | PairsList,
    Tensor | ndarray | PairsList,
]

Split pairs into non-crossing pairs and crossing pairs (no-heuristic).

Pairs are expected to be normalized (unique, sorted with i < j). Use normalize_pairs if you need to normalize raw inputs.

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tuple[Tensor | ndarray | PairsList, Tensor | ndarray | PairsList]	(non_crossing_pairs, crossing_pairs) using the same backend as input.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> primary, crossing = split_crossing_pairs(torch.tensor([[0, 2], [1, 3]]))
>>> primary.tolist(), crossing.tolist()
([], [[0, 2], [1, 3]])

NumPy input

>>> import numpy as np
>>> primary, crossing = split_crossing_pairs(np.array([[0, 3], [1, 2]]))
>>> primary.tolist(), crossing.tolist()
([[0, 3], [1, 2]], [])

List input

>>> split_crossing_pairs([(0, 2), (1, 3)])
([], [(0, 2), (1, 3)])
>>> split_crossing_pairs([(0, 3), (1, 2)])
([(0, 3), (1, 2)], [])

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def split_crossing_pairs(
    pairs: Tensor | np.ndarray | Pairs,
) -> Tuple[Tensor | np.ndarray | PairsList, Tensor | np.ndarray | PairsList]:
    """
    Split pairs into non-crossing pairs and crossing pairs (no-heuristic).

    Pairs are expected to be normalized (unique, sorted with i < j).
    Use ``normalize_pairs`` if you need to normalize raw inputs.

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        (non_crossing_pairs, crossing_pairs) using the same backend as input.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> primary, crossing = split_crossing_pairs(torch.tensor([[0, 2], [1, 3]]))
        >>> primary.tolist(), crossing.tolist()
        ([], [[0, 2], [1, 3]])

        NumPy input
        >>> import numpy as np
        >>> primary, crossing = split_crossing_pairs(np.array([[0, 3], [1, 2]]))
        >>> primary.tolist(), crossing.tolist()
        ([[0, 3], [1, 2]], [])

        List input
        >>> split_crossing_pairs([(0, 2), (1, 3)])
        ([], [(0, 2), (1, 3)])
        >>> split_crossing_pairs([(0, 3), (1, 2)])
        ([(0, 3), (1, 2)], [])
    """
    if isinstance(pairs, Tensor):
        if pairs.numel() == 0:
            empty = pairs.view(0, 2)
            return empty, empty
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        if pairs.shape[0] < 2:
            return _torch_sort_pairs(pairs), pairs.new_empty((0, 2))
        mask = _torch_crossing_mask(pairs)
        if not bool(mask.any().item()):
            return _torch_sort_pairs(pairs), pairs.new_empty((0, 2))
        primary = _torch_sort_pairs(pairs[~mask])
        crossing = _torch_sort_pairs(pairs[mask])
        return primary, crossing
    if isinstance(pairs, np.ndarray):
        if pairs.size == 0:
            empty = pairs.reshape(0, 2)
            return empty, empty
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        if pairs.shape[0] < 2:
            return _numpy_sort_pairs(pairs), pairs[:0]
        mask = _numpy_crossing_mask(pairs)
        if not mask.any():
            return _numpy_sort_pairs(pairs), pairs[:0]
        primary = _numpy_sort_pairs(pairs[~mask])
        crossing = _numpy_sort_pairs(pairs[mask])
        return primary, crossing
    if isinstance(pairs, Sequence):
        if not pairs:
            return [], []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        mask = _numpy_crossing_mask(pairs)
        if not mask.any():
            return list(map(tuple, _numpy_sort_pairs(pairs).tolist())), []
        primary = _numpy_sort_pairs(pairs[~mask]).tolist()
        crossing = _numpy_sort_pairs(pairs[mask]).tolist()
        return list(map(tuple, primary)), list(map(tuple, crossing))
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")

split_pseudoknot_pairs

Python

split_pseudoknot_pairs(
    pairs: ndarray,
) -> Tuple[ndarray, ndarray]

Python

split_pseudoknot_pairs(
    pairs: Pairs,
) -> Tuple[PairsList, PairsList]

Python

split_pseudoknot_pairs(
    pairs: Tensor,
) -> Tuple[Tensor, Tensor]

Python

split_pseudoknot_pairs(
    pairs: Tensor | ndarray | Pairs,
) -> Tuple[
    Tensor | ndarray | PairsList,
    Tensor | ndarray | PairsList,
]

Split base pairs into primary and pseudoknot pairs using segment-level MWIS.

Pairs are expected to be normalized (unique, sorted with i < j). Use normalize_pairs if you need to normalize raw inputs.

Tie-breaks order for equal total base pairs is lexicographic on:

1. minimize unpaired-within-span
2. minimize total span
3. minimize number of segments
4. deterministic segment order (prefer 3’ segments)

Parameters:

Name	Type	Description	Default
pairs	Tensor | ndarray | Pairs	torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.	required

Returns:

Type	Description
Tuple[Tensor | ndarray | PairsList, Tensor | ndarray | PairsList]	(nested_pairs, pseudoknot_pairs) using the same backend as input.

Raises:

Type	Description
ValueError	If pairs has invalid shape for the selected backend.
TypeError	If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

Examples:

Torch input

>>> import torch
>>> primary, pseudoknot_pairs = split_pseudoknot_pairs(torch.tensor([[0, 2], [1, 3]]))
>>> primary.tolist(), pseudoknot_pairs.tolist()
([[1, 3]], [[0, 2]])

NumPy input

>>> import numpy as np
>>> primary, pseudoknot_pairs = split_pseudoknot_pairs(np.array([[0, 2], [1, 3]]))
>>> primary.tolist(), pseudoknot_pairs.tolist()
([[1, 3]], [[0, 2]])

List input

>>> split_pseudoknot_pairs([(0, 2), (1, 3)])
([(1, 3)], [(0, 2)])
>>> split_pseudoknot_pairs([(0, 3), (1, 2)])
([(0, 3), (1, 2)], [])

Source code in multimolecule/utils/rna/secondary_structure/pseudoknot.py

def split_pseudoknot_pairs(
    pairs: Tensor | np.ndarray | Pairs,
) -> Tuple[Tensor | np.ndarray | PairsList, Tensor | np.ndarray | PairsList]:
    """
    Split base pairs into primary and pseudoknot pairs using segment-level MWIS.

    Pairs are expected to be normalized (unique, sorted with i < j).
    Use ``normalize_pairs`` if you need to normalize raw inputs.

    Tie-breaks order for equal total base pairs is lexicographic on:

    1. minimize unpaired-within-span
    2. minimize total span
    3. minimize number of segments
    4. deterministic segment order (prefer 3' segments)

    Args:
        pairs: torch.Tensor, numpy.ndarray, or array-like with shape (n, 2) and 0-based indices.

    Returns:
        (nested_pairs, pseudoknot_pairs) using the same backend as input.

    Raises:
        ValueError: If pairs has invalid shape for the selected backend.
        TypeError: If pairs is not a torch.Tensor, numpy.ndarray, or array-like with shape (n, 2).

    Examples:
        Torch input
        >>> import torch
        >>> primary, pseudoknot_pairs = split_pseudoknot_pairs(torch.tensor([[0, 2], [1, 3]]))
        >>> primary.tolist(), pseudoknot_pairs.tolist()
        ([[1, 3]], [[0, 2]])

        NumPy input
        >>> import numpy as np
        >>> primary, pseudoknot_pairs = split_pseudoknot_pairs(np.array([[0, 2], [1, 3]]))
        >>> primary.tolist(), pseudoknot_pairs.tolist()
        ([[1, 3]], [[0, 2]])

        List input
        >>> split_pseudoknot_pairs([(0, 2), (1, 3)])
        ([(1, 3)], [(0, 2)])
        >>> split_pseudoknot_pairs([(0, 3), (1, 2)])
        ([(0, 3), (1, 2)], [])
    """
    if isinstance(pairs, Tensor):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a torch.Tensor with shape (n, 2)")
        primary, pseudoknot = _torch_split_pseudoknot_pairs(pairs)
        return _torch_sort_pairs(primary), _torch_sort_pairs(pseudoknot)
    if isinstance(pairs, np.ndarray):
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be a numpy.ndarray with shape (n, 2)")
        primary, pseudoknot = _numpy_split_pseudoknot_pairs(pairs)
        return _numpy_sort_pairs(primary), _numpy_sort_pairs(pseudoknot)
    if isinstance(pairs, Sequence):
        if not pairs:
            return [], []
        pairs = np.asarray(pairs, dtype=int)
        if pairs.ndim != 2 or pairs.shape[1] != 2:
            raise ValueError("pairs must be an array-like with shape (n, 2)")
        primary, pseudoknot = _numpy_split_pseudoknot_pairs(pairs)
        primary = _numpy_sort_pairs(primary)
        pseudoknot = _numpy_sort_pairs(pseudoknot)
        return list(map(tuple, primary.tolist())), list(map(tuple, pseudoknot.tolist()))
    raise TypeError("pairs must be a torch.Tensor, numpy.ndarray, or sequence of (i, j) pairs")