Skip to content

Annotation mixin

annotation_mixin

Mixin for antibody sequence annotation using ANARCI/riot_na.

Classes

AnnotationMixin

Mixin for ANARCI sequence annotation capabilities.

Source code in src/antibody_training_esm/datasets/mixins/annotation_mixin.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class AnnotationMixin:
    """Mixin for ANARCI sequence annotation capabilities."""

    logger: logging.Logger

    def annotate_sequence(
        self, sequence_id: str, sequence: str, chain: str
    ) -> dict[str, str] | None:
        """
        Annotate a single sequence using ANARCI (IMGT numbering).

        This method wraps riot_na.create_riot_aa() to extract CDR/FWR regions.

        Args:
            sequence_id: Unique identifier for the sequence
            sequence: Protein sequence to annotate
            chain: Chain type ("H" for heavy, "L" for light)

        Returns:
            Dictionary with keys: FWR1, CDR1, FWR2, CDR2, FWR3, CDR3, FWR4
            Returns None if annotation fails
        """
        try:
            # Import riot_na here to avoid dependency issues
            from riot_na import create_riot_aa

            # Run ANARCI annotation
            result = create_riot_aa(sequence_id, sequence, chain=chain)

            if result is None:
                self.logger.warning(
                    f"ANARCI annotation failed for {sequence_id} ({chain} chain)"
                )
                return None

            # Extract regions
            annotations = {
                "FWR1": result.get("FWR1", ""),
                "CDR1": result.get("CDR1", ""),
                "FWR2": result.get("FWR2", ""),
                "CDR2": result.get("CDR2", ""),
                "FWR3": result.get("FWR3", ""),
                "CDR3": result.get("CDR3", ""),
                "FWR4": result.get("FWR4", ""),
            }

            # Validate annotations (should not be empty)
            if not any(annotations.values()):
                self.logger.warning(
                    f"All annotations empty for {sequence_id} ({chain} chain)"
                )
                return None

            return annotations

        except Exception as e:
            self.logger.error(f"Error annotating {sequence_id} ({chain} chain): {e}")
            return None

    def annotate_all(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Annotate all sequences in a DataFrame.

        Adds annotation columns for heavy and light chains.

        Args:
            df: DataFrame with VH_sequence and optionally VL_sequence columns

        Returns:
            DataFrame with annotation columns added
        """
        self.logger.info(f"Annotating {len(df)} sequences...")

        # Annotate heavy chains
        if "VH_sequence" in df.columns:
            self.logger.info("Annotating VH sequences...")
            vh_annotations: pd.Series = df.apply(
                lambda row: self.annotate_sequence(
                    row.get("id", f"seq_{row.name}"), row["VH_sequence"], "H"
                )
                if pd.notna(row["VH_sequence"])
                else None,
                axis=1,
            )

            # Extract annotation fields
            for field in ["FWR1", "CDR1", "FWR2", "CDR2", "FWR3", "CDR3", "FWR4"]:
                df[f"VH_{field}"] = vh_annotations.apply(
                    lambda x, f=field: x.get(f, "") if x else ""
                )

        # Annotate light chains (if present)
        if "VL_sequence" in df.columns:
            self.logger.info("Annotating VL sequences...")
            vl_annotations: pd.Series = df.apply(
                lambda row: self.annotate_sequence(
                    row.get("id", f"seq_{row.name}"), row["VL_sequence"], "L"
                )
                if pd.notna(row["VL_sequence"])
                else None,
                axis=1,
            )

            # Extract annotation fields
            for field in ["FWR1", "CDR1", "FWR2", "CDR2", "FWR3", "CDR3", "FWR4"]:
                df[f"VL_{field}"] = vl_annotations.apply(
                    lambda x, f=field: x.get(f, "") if x else ""
                )

        self.logger.info("Annotation complete")
        return df
Functions
annotate_sequence(sequence_id, sequence, chain)

Annotate a single sequence using ANARCI (IMGT numbering).

This method wraps riot_na.create_riot_aa() to extract CDR/FWR regions.

Parameters:

Name Type Description Default
sequence_id str

Unique identifier for the sequence

 required
sequence str

Protein sequence to annotate

 required
chain str

Chain type ("H" for heavy, "L" for light)

 required

Returns:

Type Description
dict[str, str] | None

Dictionary with keys: FWR1, CDR1, FWR2, CDR2, FWR3, CDR3, FWR4
dict[str, str] | None

Returns None if annotation fails
Source code in src/antibody_training_esm/datasets/mixins/annotation_mixin.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def annotate_sequence(
    self, sequence_id: str, sequence: str, chain: str
) -> dict[str, str] | None:
    """
    Annotate a single sequence using ANARCI (IMGT numbering).

    This method wraps riot_na.create_riot_aa() to extract CDR/FWR regions.

    Args:
        sequence_id: Unique identifier for the sequence
        sequence: Protein sequence to annotate
        chain: Chain type ("H" for heavy, "L" for light)

    Returns:
        Dictionary with keys: FWR1, CDR1, FWR2, CDR2, FWR3, CDR3, FWR4
        Returns None if annotation fails
    """
    try:
        # Import riot_na here to avoid dependency issues
        from riot_na import create_riot_aa

        # Run ANARCI annotation
        result = create_riot_aa(sequence_id, sequence, chain=chain)

        if result is None:
            self.logger.warning(
                f"ANARCI annotation failed for {sequence_id} ({chain} chain)"
            )
            return None

        # Extract regions
        annotations = {
            "FWR1": result.get("FWR1", ""),
            "CDR1": result.get("CDR1", ""),
            "FWR2": result.get("FWR2", ""),
            "CDR2": result.get("CDR2", ""),
            "FWR3": result.get("FWR3", ""),
            "CDR3": result.get("CDR3", ""),
            "FWR4": result.get("FWR4", ""),
        }

        # Validate annotations (should not be empty)
        if not any(annotations.values()):
            self.logger.warning(
                f"All annotations empty for {sequence_id} ({chain} chain)"
            )
            return None

        return annotations

    except Exception as e:
        self.logger.error(f"Error annotating {sequence_id} ({chain} chain): {e}")
        return None
annotate_all(df)

Annotate all sequences in a DataFrame.

Adds annotation columns for heavy and light chains.

Parameters:

Name Type Description Default
df DataFrame

DataFrame with VH_sequence and optionally VL_sequence columns

 required

Returns:

Type Description
DataFrame

DataFrame with annotation columns added
Source code in src/antibody_training_esm/datasets/mixins/annotation_mixin.py 
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def annotate_all(self, df: pd.DataFrame) -> pd.DataFrame:
    """
    Annotate all sequences in a DataFrame.

    Adds annotation columns for heavy and light chains.

    Args:
        df: DataFrame with VH_sequence and optionally VL_sequence columns

    Returns:
        DataFrame with annotation columns added
    """
    self.logger.info(f"Annotating {len(df)} sequences...")

    # Annotate heavy chains
    if "VH_sequence" in df.columns:
        self.logger.info("Annotating VH sequences...")
        vh_annotations: pd.Series = df.apply(
            lambda row: self.annotate_sequence(
                row.get("id", f"seq_{row.name}"), row["VH_sequence"], "H"
            )
            if pd.notna(row["VH_sequence"])
            else None,
            axis=1,
        )

        # Extract annotation fields
        for field in ["FWR1", "CDR1", "FWR2", "CDR2", "FWR3", "CDR3", "FWR4"]:
            df[f"VH_{field}"] = vh_annotations.apply(
                lambda x, f=field: x.get(f, "") if x else ""
            )

    # Annotate light chains (if present)
    if "VL_sequence" in df.columns:
        self.logger.info("Annotating VL sequences...")
        vl_annotations: pd.Series = df.apply(
            lambda row: self.annotate_sequence(
                row.get("id", f"seq_{row.name}"), row["VL_sequence"], "L"
            )
            if pd.notna(row["VL_sequence"])
            else None,
            axis=1,
        )

        # Extract annotation fields
        for field in ["FWR1", "CDR1", "FWR2", "CDR2", "FWR3", "CDR3", "FWR4"]:
            df[f"VL_{field}"] = vl_annotations.apply(
                lambda x, f=field: x.get(f, "") if x else ""
            )

    self.logger.info("Annotation complete")
    return df