codeine 0.1.0__py3-none-any.whl

codeine/motifs/restriction.py

@@ -0,0 +1,105 @@
+"""
+Common restriction site motifs, obtained from New England BioLabs here:
+
+https://www.neb.com/en/tools-and-resources/selection-charts/alphabetized-list-of-recognition-specificities
+
+Retrieved: 2026-06-09
+"""
+
+from enum import Enum
+from typing import Tuple
+
+_COMPLEMENT = str.maketrans('ACGTacgt', 'TGCAtgca')
+
+
+def reverse_complement(seq: str) -> str:
+    """
+    Return the reverse complement of a DNA sequence.
+    """
+    return seq.translate(_COMPLEMENT)[::-1]
+
+
+class RestrictionSite(Enum):
+    """
+    Common restriction enzyme recognition sequences.
+
+    These can be passed directly to ``CodingSpace`` to forbid restriction
+    sites without specifying the DNA motifs manually. Both forward and reverse
+    sequences are passed.
+
+    Examples
+    --------
+    >>> CodingSpace(
+    ...     aa_seq,
+    ...     forbidden_motifs=[
+    ...         RestrictionSite.EcoRI,
+    ...         RestrictionSite.BsaI,
+    ...     ],
+    ... )
+
+    >>> RestrictionSite.EcoRI.forward
+    'GAATTC'
+    >>> RestrictionSite.EcoRI.reverse
+    'GAATTC'
+    """
+
+    # BioBricks
+    EcoRI = 'GAATTC'
+    XbaI = 'TCTAGA'
+    SpeI = 'ACTAGT'
+    PstI = 'CTGCAG'
+
+    # Cloning
+    BamHI = 'GGATCC'
+    HindIII = 'AAGCTT'
+    XhoI = 'CTCGAG'
+    SalI = 'GTCGAC'
+    KpnI = 'GGTACC'
+    SacI = 'GAGCTC'
+    NcoI = 'CCATGG'
+    NdeI = 'CATATG'
+    NotI = 'GCGGCCGC'
+    MluI = 'ACGCGT'
+    AgeI = 'ACCGGT'
+    AvrII = 'CCTAGG'
+    BglII = 'AGATCT'
+
+    # Golden Gate
+    BsaI = 'GGTCTC'
+    BsmBI = 'CGTCTC'
+    BbsI = 'GAAGAC'
+    SapI = 'GCTCTTC'
+
+    def __repr__(self):
+        return f'RestrictionSite.{self.name}'
+
+    def __str__(self):
+        if self.forward == self.reverse:
+            return f'{self.name} ({self.forward})'
+
+        return f'{self.name} ({self.forward} / {self.reverse})'
+
+    @property
+    def forward(self) -> str:
+        """
+        Forward recognition sequence.
+        """
+        return self.value
+
+    @property
+    def reverse(self) -> str:
+        """
+        Reverse-complemented recognition sequence.
+        """
+        return reverse_complement(self.value)
+
+    @property
+    def motifs(self) -> Tuple[str, ...]:
+        """
+        All motifs corresponding to this restriction site (forward and reverse).
+        Palindromic sites return a single motif.
+        """
+        if self.forward == self.reverse:
+            return self.forward,
+
+        return self.forward, self.reverse

codeine/motifs/validate.py

@@ -0,0 +1,117 @@
+
+from typing import List, Optional, Sequence, Union
+
+from codeine.motifs.restriction import RestrictionSite
+
+ForbiddenMotif = Union[str, RestrictionSite]
+ForbiddenMotifs = Union[ForbiddenMotif, Sequence[ForbiddenMotif]]
+
+
+def expand_and_validate_forbidden_motifs(
+        forbidden_motifs: ForbiddenMotifs,
+        rna: bool
+) -> List[str]:
+    """
+    Convert a set of forbidden motifs into a set of sequences to ban.
+
+    Parameters
+    ----------
+    forbidden_motifs
+        A sequence of either dna/rna strings or RestrictionSite objects.
+    rna
+        Whether to use RNA.
+
+    Returns
+    -------
+    A list of forbidden nucleotide sequences.
+    """
+    all_sequences = []
+
+    if isinstance(forbidden_motifs, (str, RestrictionSite)):
+        forbidden_motifs = [forbidden_motifs]
+
+    for motif in forbidden_motifs:
+        if isinstance(motif, RestrictionSite):
+            sequences = [*motif.motifs]
+
+        elif isinstance(motif, str):
+            if len(motif) == 0:
+                raise ValueError('Forbidden motifs cannot be empty.')
+
+            sequences = [motif]
+
+        else:
+            raise TypeError('Forbidden motifs must be strings or codeine.RestrictionSite.')
+
+        sequences = [seq.upper() for seq in sequences]
+        sequences = [seq.replace('T', 'U') if rna else seq.replace('U', 'T') for seq in sequences]
+
+        allowed = set('ACGU' if rna else 'ACGT')
+        for seq in sequences:
+            if not set(seq) <= allowed:
+                raise ValueError('Forbidden motifs must be nucleotide sequences.')
+
+        all_sequences += sequences
+
+    return sorted(set(all_sequences))
+
+
+def expand_and_validate_max_homopolymer(
+        max_homopolymer: int,
+        rna: bool = False
+) -> List[str]:
+    """
+    Convert a max homopolymer constraint into a set of banned sequences.
+
+    Parameters
+    ----------
+    max_homopolymer
+        The max length of homopolymer
+    rna
+        Whether to use RNA
+
+    Returns
+    -------
+    A list of forbidden nucleotide sequences.
+    """
+
+    if not isinstance(max_homopolymer, int):
+        raise TypeError('max_homopolymer must be an integer.')
+
+    if max_homopolymer < 1:
+        raise ValueError('max_homopolymer must be at least 1.')
+
+    nts = 'ACGU' if rna else 'ACGT'
+    return [nt * (max_homopolymer + 1) for nt in nts]
+
+
+def expand_and_validate_sequence_constraints(
+        forbidden_motifs: Optional[ForbiddenMotifs] = None,
+        max_homopolymer: Optional[int] = None,
+        rna: bool = False,
+):
+    """
+    Convert forbidden sequences and/or max homopolymer constraints into sets of banned sequences.
+
+    Parameters
+    ----------
+    forbidden_motifs
+        A sequence of either dna/rna strings or RestrictionSite objects.
+    max_homopolymer
+        The max allowed homopolymer.
+    rna
+        Whether to use RNA.
+
+    Returns
+    -------
+    A list of forbidden nucleotide sequences.
+    """
+    forbidden_sequences = []
+
+    if forbidden_motifs is not None:
+        forbidden_sequences += expand_and_validate_forbidden_motifs(forbidden_motifs, rna=rna)
+
+    if max_homopolymer is not None:
+        forbidden_sequences += expand_and_validate_max_homopolymer(max_homopolymer, rna=rna)
+
+    return sorted(set(forbidden_sequences))

codeine/space/coding.py

@@ -0,0 +1,490 @@
+import pickle
+import random
+
+from pathlib import Path
+from typing import Dict, Generator, List, Optional, Sequence, Tuple, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from codeine.space.mutation import MutationSpace
+
+from codeine.graph.base import CodonGraph, CodonRestriction
+from codeine.motifs.validate import expand_and_validate_sequence_constraints, ForbiddenMotifs
+from codeine.motifs.restriction import RestrictionSite
+from codeine.translation.tables import TranslationTable
+from codeine.translation.weights import CodonWeights
+from codeine.utils.display import format_forbidden_motifs, format_forbidden_motif,\
+    format_count, format_restrictions
+from codeine.utils.sampling import Seedable
+
+
+class CodingSpace:
+    """
+    Represents a space of valid coding sequences for a protein under constraints.
+    """
+    def __init__(
+        self,
+        aa_seq: str,
+        *,
+        translation_table: Optional[TranslationTable] = None,
+        rna: Optional[bool] = None,
+        codon_restrictions: Optional[Dict[int, CodonRestriction]] = None,
+        forbidden_motifs: Optional[ForbiddenMotifs] = None,
+        max_homopolymer: Optional[int] = None,
+        context_l: str = '',
+        context_r: str = '',
+        codon_weights: Optional[CodonWeights] = None,
+        seed: Optional[Seedable] = None,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        aa_seq
+            The amino acid sequence.
+        translation_table
+            The translation table to use. Leave blank to use standard table.
+        rna
+            Whether to use RNA. If false or blank, use DNA.
+        codon_restrictions
+            Any codon restrictions in the format e.g. ``{4: 'TCC'}`` or ``{5: ['AGT', 'AGC']}``. Positions are 1-based.
+        forbidden_motifs
+            Forbidden motifs, either as strings or as ``codeine.RestrictionSite``.
+        max_homopolymer
+            The maximum allowed length of nucleotide homopolymer
+        context_l
+            The context sequence to the left of the coding sequence.
+        context_r
+            The context sequence to the right of the coding sequence.
+        codon_weights
+            The codon weights to use. Leave blank to sample uniformly.
+        seed
+            Seed used to initialise the random number generator for sampling.
+        """
+
+        translation_table, codon_weights = self._resolve_tables(translation_table, codon_weights, rna)
+
+        graph = CodonGraph(
+            aa_seq,
+            codon_restrictions=codon_restrictions,
+            translation_table=translation_table,
+            weights=codon_weights,
+            context_l=context_l,
+            context_r=context_r,
+        )
+
+        view = graph.view(seed=seed)
+        self.view = view
+
+        self.forbidden_motifs = self._normalise_forbidden_motifs(forbidden_motifs)
+        self.max_homopolymer = max_homopolymer
+
+        self._update_forbidden_sequences()
+
+    @classmethod
+    def load(cls, path) -> 'CodingSpace':
+        """
+        Load a coding space from disc.
+        """
+        with Path(path).open('rb') as f:
+            return pickle.load(f)
+
+    def save(self, path) -> None:
+        """
+        Save this coding space to disc.
+        """
+        with Path(path).open('wb') as f:
+            pickle.dump(self, f)
+
+    def __getitem__(self, index: Union[int, slice]) -> Union[str, List[str]]:
+        """
+        Return one or more valid sequences.
+
+        Parameters
+        ----------
+        index
+            Zero-based sequence index or slice.
+
+        Returns
+        -------
+        str or List[str]
+            The indexed sequence, or a list of sequences for a slice.
+        """
+        return self.view[index]
+
+    def __iter__(self) -> Generator[str, None, None]:
+        """
+        Iterate over all valid sequences in this coding space.
+        Be aware that "all valid sequences" can be astronomically many!
+
+        Yields
+        ----------
+        All valid sequences in the coding space, in order.
+        """
+
+        yield from self.view
+
+    def __contains__(self, seq: str) -> bool:
+        """
+        Does the given seq exist in this space?
+
+        Returns
+        ----------
+        True if and only if this is a valid sequence in this space.
+        """
+
+        return seq in self.view
+
+    def __repr__(self) -> str:
+        molecule = 'RNA' if self.translation_table.rna else 'DNA'
+
+        lines = [
+            f'{type(self).__name__}',
+            '',
+            f'Translation table: {self.translation_table.table_id} ({self.translation_table.name})',
+            f'Molecule type: {molecule}',
+            '',
+            f'Amino acid sequence ({len(self.aa_seq)} aa):',
+            f'{self.aa_seq}',
+            '',
+        ]
+
+        if self.codon_restrictions:
+            lines += [
+                'Codon restrictions:',
+                *format_restrictions(
+                    self.codon_restrictions,
+                    label='restricted positions',
+                    max_lines=4,
+                ),
+                '',
+            ]
+
+        if self.forbidden_motifs:
+            motifs = self.forbidden_motifs
+
+            if isinstance(motifs, (str, RestrictionSite)):
+                motifs = [motifs]
+
+            lines += [
+                'Forbidden motifs:',
+                *format_forbidden_motifs(
+                    [
+                        format_forbidden_motif(motif, rna=self.translation_table.rna)
+                        for motif in motifs
+                    ],
+                    max_lines=4,
+                ),
+                '',
+            ]
+
+        if self.max_homopolymer is not None:
+            lines += [
+                'Maximum homopolymer length:',
+                f'  {self.max_homopolymer}',
+                '',
+            ]
+
+        if self.pinned_codons:
+            lines += [
+                'Temporary pins:',
+                *format_restrictions(
+                    self.pinned_codons,
+                    label='pinned positions',
+                    max_lines=4,
+                ),
+                '',
+            ]
+
+        lines.append(f'Num. valid coding sequences: {format_count(self.n_valid_sequences)}')
+
+        return '\n'.join(lines)
+
+    def sample(self, n: Optional[int] = None) -> str:
+        """
+        Sample one or more variants from this coding space.
+
+        Parameters
+        ----------
+        n
+            Number of sequences to sample. If omitted, return a single sequence.
+
+        Returns
+        -------
+        A sampled string sequence from this coding space.
+        """
+        return self.view.sample(n=n)
+
+    def enumerate(self) -> Generator[str, None, None]:
+        """
+        Generate all sequences in this space. If there are many (and often there are
+        astronomically many), one would not expect to reach the 'end'. However for smaller
+        sequence spaces, such as mutation spaces, it's quite possible to get there.
+
+        Yields
+        ------
+        str
+            A valid coding sequence.
+        """
+        yield from self.view.enumerate()
+
+    def contains(self, seq: str) -> bool:
+        """
+        Check whether a coding sequence is contained in this coding space.
+
+        Parameters
+        ----------
+        seq
+            The sequence to check.
+
+        Returns
+        -------
+        True if and only if the sequence is contained in this coding space.
+        """
+        return self.view.contains(seq)
+
+    def mutants(
+            self,
+            cds: str,
+            free_positions: Optional[Sequence[int]] = None,
+            min_nts: Optional[int] = None,
+            max_nts: Optional[int] = None,
+            min_codons: Optional[int] = None,
+            max_codons: Optional[int] = None,
+    ) -> 'MutationSpace':
+        """
+        Return a space of mutants relative to a given coding sequence, i.e. a space derived
+        from this one but which fixes the sequence on all but the specified positions.
+
+        Parameters
+        ----------
+        cds
+            The sequence to mutate.
+        free_positions
+            The positions that are allowed to vary.
+        min_nts
+            The min nucleotide (Hamming) distance relative to the reference sequence.
+        max_nts
+            The max nucleotide (Hamming) distance relative to the reference sequence.
+        min_codons
+            The min number of changed codons relative to the reference sequence.
+        max_codons
+            The max number of changed codons relative to the reference sequence.
+        """
+        cds = self.translation_table.normalise_sequence(cds)
+
+        if not self.contains(cds):
+            raise ValueError('CDS is not contained in this coding space.')
+
+        from codeine.space.mutation import MutationSpace
+
+        return MutationSpace(
+            space=self,
+            cds=cds,
+            free_positions=free_positions,
+            min_nts=min_nts,
+            max_nts=max_nts,
+            min_codons=min_codons,
+            max_codons=max_codons,
+        )
+
+    def pin_codons(self, pinned_codons: Dict[int, str]) -> None:
+        """
+        Pin temporary codons in this coding space.
+
+        Parameters
+        ----------
+        pinned_codons
+            A dict specifying which codons to pin, by position.
+        """
+        self.view.pin_codons(pinned_codons)
+
+    def unpin_codons(self, positions: Sequence[int]) -> None:
+        """
+        Remove temporary codon pins by position.
+
+        Parameters
+        ----------
+        positions
+            Positions to unpin.
+        """
+        self.view.unpin_codons(positions)
+
+    def set_pinned_codons(self, pinned_codons: Dict[int, str]) -> None:
+        """
+        Replace all temporary codon pins on this coding space.
+
+        Parameters
+        ----------
+        pinned_codons
+            A dict specifying which codons to pin, by position.
+        """
+        self.view.set_pinned_codons(pinned_codons)
+
+    def clear_pins(self) -> None:
+        """
+        Remove all temporary codon pins from this coding space.
+        """
+        self.view.clear_pins()
+
+    def set_forbidden_motifs(self, forbidden_motifs: ForbiddenMotifs) -> None:
+        """
+        Set the forbidden motifs for this coding space.
+
+        Parameters
+        ----------
+        forbidden_motifs
+            Motifs that should be forbidden in generated sequences.
+        """
+        self.forbidden_motifs = self._normalise_forbidden_motifs(forbidden_motifs)
+        self._update_forbidden_sequences()
+
+    def clear_forbidden_motifs(self) -> None:
+        """
+        Remove all forbidden motifs from this coding space.
+        """
+        self.set_forbidden_motifs(None)
+
+    def set_max_homopolymer(self, max_homopolymer: Optional[int]) -> None:
+        """
+        Set the maximum allowed homopolymer length.
+
+        Parameters
+        ----------
+        max_homopolymer
+            The longest allowed repeated run of one nucleotide, or None for no limit.
+        """
+        self.max_homopolymer = max_homopolymer
+        self._update_forbidden_sequences()
+
+    def clear_max_homopolymer(self) -> None:
+        """
+        Remove the maximum homopolymer constraint from this coding space.
+        """
+        self.set_max_homopolymer(None)
+
+    @property
+    def n_valid_sequences(self) -> int:
+        """
+        The number of valid sequences in this space.
+        """
+        return self.view.n_valid_sequences
+
+    @property
+    def aa_seq(self) -> str:
+        """
+        The amino acid sequence for this coding space.
+        """
+        return self.view.aa_seq
+
+    @property
+    def translation_table(self) -> TranslationTable:
+        """
+        The translation table being used in this space.
+        """
+        return self.view.translation_table
+
+    @property
+    def codon_weights(self) -> CodonWeights:
+        """
+        The codon weights being used in this space.
+        """
+        return self.view.codon_weights
+
+    @property
+    def codon_restrictions(self) -> Dict[int, CodonRestriction]:
+        """
+        The fixed codon restrictions from the underlying graph.
+        """
+        return self.view.codon_restrictions
+
+    @property
+    def context_l(self) -> str:
+        """
+        The left context sequence from the underlying graph.
+        """
+        return self.view.context_l
+
+    @property
+    def context_r(self) -> str:
+        """
+        The right context sequence from the underlying graph.
+        """
+        return self.view.context_r
+
+    @property
+    def pinned_codons(self) -> Dict[int, List[str]]:
+        """
+        Temporary codon pins currently applied to this coding space.
+        """
+        return self.view.pinned_codons
+
+    def _update_forbidden_sequences(self) -> None:
+        """
+        Rebuild concrete forbidden sequences and apply them to the view.
+        """
+        forbidden_sequences = expand_and_validate_sequence_constraints(
+            forbidden_motifs=self.forbidden_motifs,
+            max_homopolymer=self.max_homopolymer,
+            rna=self.translation_table.rna,
+        )
+        self.view.set_banned_sequences(forbidden_sequences)
+
+    @staticmethod
+    def _resolve_tables(
+            translation_table: Optional[TranslationTable],
+            codon_weights: Optional[CodonWeights],
+            rna: Optional[bool],
+    ) -> Tuple[TranslationTable, CodonWeights]:
+        """
+        Resolve user-submited (or not) translation table, codon weights, and RNA flag.
+        """
+
+        if rna is None:
+            if translation_table is not None and codon_weights is not None \
+                    and translation_table.rna != codon_weights.rna:
+                raise ValueError('Provided translation table and codon weights must have the same molecule type.')
+
+            if translation_table is not None:
+                rna = translation_table.rna
+            elif codon_weights is not None:
+                rna = codon_weights.rna
+            else:
+                rna = False
+
+        else:
+            if translation_table is not None and translation_table.rna != rna:
+                raise ValueError('Value for rna is inconsistent with the provided translation table.')
+
+            if codon_weights is not None and codon_weights.rna != rna:
+                raise ValueError('Value for rna is inconsistent with the provided codon weights.')
+
+        if translation_table is None:
+            translation_table = TranslationTable(table_id=1, rna=rna)
+
+        if codon_weights is None:
+            codon_weights = CodonWeights.uniform(table=translation_table)
+
+        return translation_table, codon_weights
+
+    def _normalise_forbidden_motifs(
+            self,
+            forbidden_motifs: Optional[ForbiddenMotifs],
+    ) -> Optional[ForbiddenMotifs]:
+        """
+        Normalise string forbidden motifs to the molecule type used by this coding
+        space. RestrictionSite objects are left unchanged.
+        """
+        if forbidden_motifs is None:
+            return None
+
+        if isinstance(forbidden_motifs, str):
+            return self.translation_table.normalise_sequence(forbidden_motifs)
+
+        if isinstance(forbidden_motifs, RestrictionSite):
+            return forbidden_motifs
+
+        return [
+            self.translation_table.normalise_sequence(motif)
+            if isinstance(motif, str)
+            else motif
+            for motif in forbidden_motifs
+        ]