assembly_annotation.py

"""
AVIDbase: a biologically accurate structural dataset of nanobody-antigen complexes
    Copyright (C) 2026 Novartis Pharma AG

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import shutil
import argparse
from pathlib import Path
import numpy as np
import pandas as pd

from avidbase_functions.dataset_util.fetch import read_sabdab_summary, save_pdbs
from avidbase_functions.structure_processing.annotation import FindAssemblies, INPUT_TABLE_DICT, INPUT_TABLE_TYPES


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Generate an Excel table based on the given input PDB codes. SAbDab summary files can also be passed, in which case they are pre-filtered exclusively for X-ray crystal structures. Any structures not already present in the structure directory are downloaded from the RCSB PDB."
        )
    
    parser.add_argument(
        "-o", type=str, dest="out_dir", required=False,
        help="output directory for assembly input table", default=Path("output/intermediary_tables")
    )
    
    parser.add_argument(
        "-s", type=str, dest="struct_dir", required=False,
        help="structure output directory", default=Path("output/structures/CIF_files")
    )

    parser.add_argument(
        "--sabdab-files", type=str, nargs="*", dest="sabdab_files", required=False,
        help="(optional, accepts multiple filepaths) path(s) to SAbDab summary files for PDB codes to download. NOTE: by default, filters for method == X-RAY DIFFRACTION."
    )

    parser.add_argument(
        "--pdb-codes", type=str, dest="pdb_codes", required=False,
        help="(optional) comma-separated list of input PDB codes. Example: 1zvy,1jtp,1jtt"
    )

    parser.add_argument(
        "--overwrite", dest="overwrite", action="store_true",
        help="Overwrite all existing CIF files in the structure directory",
        default=False
    )

    parser.add_argument(
        "--verbose", dest="verbose", action="store_true",
        help="verbose",
        default=False
    )

    args = parser.parse_args()

    # Create structure directory if it doesn't exist
    struct_dir = Path(args.struct_dir)
    struct_dir.mkdir(exist_ok=True, parents=True)

    # Same for output directory
    out_dir = Path(args.out_dir)
    out_dir.mkdir(exist_ok=True, parents=True)

    # Empty structure directory if specified
    if args.overwrite is True:
        shutil.rmtree(struct_dir)

    pdb_list = []

    # Read unique PDB codes from SAbDab summary file
    if args.sabdab_files is not None:
        for f in args.sabdab_files:
            summary = read_sabdab_summary(Path(f))
            pdb_list.extend(sorted(summary.PDB.tolist()))

    # Read explicitly passed comma-separated PDB codes
    if args.pdb_codes is not None:
        pdb_list.extend(args.pdb_codes.lower().split(","))

    # Download CIF files from the PDB if any PDB codes were passed
    if len(pdb_list) > 0:
        save_pdbs(np.unique(pdb_list), out_dir=struct_dir, overwrite=args.overwrite) 

    # Automatically generate annotation input table
    out_df = pd.DataFrame(columns=list(INPUT_TABLE_DICT))
    
    for i, cif in enumerate(sorted([x for x in struct_dir.iterdir() if x.name.endswith(".cif")], key=lambda x: x.stem)):
        if args.verbose is True:
            print("Annotating:", cif.stem)

        fa = FindAssemblies(cif)
        fa.generate_input_table()

        # Output intermediary table
        out_df = pd.concat((out_df, fa.output_input_table()), axis=0)

    out_df.reset_index(drop=True).astype(INPUT_TABLE_TYPES).to_excel(out_dir / "annotation_input_table_UNMODIFIED.xlsx")