mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

mlmm/fix_altloc.py

@@ -0,0 +1,464 @@
+#!/usr/bin/env python3
+"""
+fix_altloc.py - Drop alternate locations from PDB files
+
+What it does
+------------
+1) Blank the PDB altLoc column (column 17, 1-based) with a single space.
+   - This is a 1-character replacement (no shifting / no reformatting).
+2) If the same atom appears multiple times due to alternate locations
+   (altLoc like A/B/... or custom labels like H/L),
+   keep the "best" one by the default rule:
+      - Highest occupancy first
+      - If tied (or occupancy missing), keep the earliest one in the file
+
+Handled records
+---------------
+- ATOM / HETATM
+- ANISOU is also handled: ANISOU lines are kept only if the corresponding
+  ATOM/HETATM line (same serial) is kept.
+
+Notes
+-----
+- Atom serial numbers are NOT renumbered (gaps may remain).
+- CONECT and other connectivity/annotation records are NOT updated.
+
+Usage
+-----
+  mlmm fix-altloc -i input.pdb -o output.pdb
+  mlmm fix-altloc -i ./dir -o ./dir_clean --recursive
+  mlmm fix-altloc -i ./dir --inplace --recursive
+"""
+
+import shutil
+from pathlib import Path
+from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple
+
+import click
+
+COORD_RECORDS = ("ATOM  ", "HETATM")
+ANISOU_RECORD = "ANISOU"
+
+# PDB fixed columns (0-based Python indices)
+ALTLOC_IDX = 16                 # column 17 (1-based)
+SERIAL_SLICE = slice(6, 11)     # columns 7-11 (1-based), width 5
+OCC_SLICE = slice(54, 60)       # columns 55-60 (1-based), width 6
+
+
+def split_newline(line: str) -> Tuple[str, str]:
+    """Split a line into (core, newline) while preserving the newline exactly."""
+    if line.endswith("\r\n"):
+        return line[:-2], "\r\n"
+    if line.endswith("\n"):
+        return line[:-1], "\n"
+    if line.endswith("\r"):
+        return line[:-1], "\r"
+    return line, ""
+
+
+def ensure_len(core: str, n: int) -> str:
+    """Right-pad with spaces to guarantee at least n characters (no shifting)."""
+    return core if len(core) >= n else core.ljust(n)
+
+
+def blank_altloc(line: str) -> str:
+    """
+    Blank the altLoc field (column 17, 1-based) with a single space.
+
+    IMPORTANT: This does NOT remove characters; it replaces exactly one character,
+    so the fixed-width PDB formatting is preserved.
+    """
+    core, nl = split_newline(line)
+    core = ensure_len(core, ALTLOC_IDX + 1)  # make sure core[ALTLOC_IDX] exists
+    core = core[:ALTLOC_IDX] + " " + core[ALTLOC_IDX + 1:]
+    return core + nl
+
+
+def atom_serial_5(line: str) -> str:
+    """Return the 5-character atom serial field exactly as it appears (cols 7-11)."""
+    core, _ = split_newline(line)
+    core = ensure_len(core, SERIAL_SLICE.stop)
+    return core[SERIAL_SLICE]
+
+
+def parse_occupancy(line: str) -> Optional[float]:
+    """
+    Parse occupancy from columns 55-60 (1-based).
+    Returns None if missing/unparseable.
+    """
+    core, _ = split_newline(line)
+    core = ensure_len(core, OCC_SLICE.stop)
+    s = core[OCC_SLICE].strip()
+    if not s:
+        return None
+    try:
+        return float(s)
+    except ValueError:
+        return None
+
+
+def atom_identity_key(line: str) -> Tuple[str, str, str, str, str, str, str]:
+    """
+    Build a key to identify the "same atom" while IGNORING altLoc.
+
+    Fields used (fixed columns, classic PDB):
+      - record name (ATOM/HETATM)     cols 1-6
+      - atom name                      cols 13-16
+      - residue name                   cols 18-20
+      - chain ID                       col 22
+      - residue sequence number        cols 23-26
+      - insertion code                 col 27
+      - segID (non-standard, common)   cols 73-76
+
+    segID is included to reduce accidental merging in MD-style PDBs where chain ID may be blank.
+    """
+    core, _ = split_newline(line)
+    core = ensure_len(core, 76)
+
+    record = core[0:6]
+    atom_name = core[12:16]
+    res_name = core[17:20]
+    chain_id = core[21:22]
+    res_seq = core[22:26]
+    i_code = core[26:27]
+    seg_id = core[72:76]  # 73-76 (1-based), optional/non-standard
+
+    return (record, atom_name, res_name, chain_id, res_seq, i_code, seg_id)
+
+
+def process_block(lines: List[str]) -> List[str]:
+    """
+    Two-pass processing for a block (either the whole file if no MODEL,
+    or the content between MODEL and ENDMDL):
+
+    Pass 1: determine the best coordinate line per atom key
+            by (occupancy desc, first-appearance asc).
+    Pass 2: output only the chosen coordinate lines (altLoc blanked),
+            and keep only ANISOU lines whose serial is chosen (altLoc blanked).
+            All other records are passed through unchanged.
+
+    Handling of different atom counts between altLoc states:
+    --------------------------------------------------------
+    When different altLoc states have different atoms (e.g., altLoc A has
+    atoms N,CA,CB,CG while altLoc B has N,CA,CB,CD), this function:
+    - For DUPLICATE atoms (same identity key, e.g., N,CA,CB): selects the best
+      one based on occupancy
+    - For UNIQUE atoms (only in one altLoc, e.g., CG in A, CD in B): keeps ALL
+      of them in the output
+
+    This ensures the output structure contains all unique atoms from all altLoc
+    states, with duplicates resolved to the best conformer.
+    """
+    # key -> (occ_val_for_compare, line_index, serial5)
+    best: Dict[Tuple[str, str, str, str, str, str, str], Tuple[float, int, str]] = {}
+
+    for idx, line in enumerate(lines):
+        if line.startswith(COORD_RECORDS):
+            key = atom_identity_key(line)
+            occ = parse_occupancy(line)
+            occ_val = occ if occ is not None else float("-inf")
+            serial = atom_serial_5(line)
+
+            if key not in best:
+                best[key] = (occ_val, idx, serial)
+            else:
+                best_occ, best_idx, _best_serial = best[key]
+                # Prefer higher occupancy; if tied, prefer earlier line (smaller idx)
+                if (occ_val > best_occ) or (occ_val == best_occ and idx < best_idx):
+                    best[key] = (occ_val, idx, serial)
+
+    chosen_serials: Set[str] = set(v[2] for v in best.values())
+
+    out: List[str] = []
+    for idx, line in enumerate(lines):
+        if line.startswith(COORD_RECORDS):
+            key = atom_identity_key(line)
+            # Keep only the selected "best" line for this key
+            if key in best and best[key][1] == idx:
+                out.append(blank_altloc(line))
+            continue
+
+        if line.startswith(ANISOU_RECORD):
+            serial = atom_serial_5(line)
+            if serial in chosen_serials:
+                out.append(blank_altloc(line))
+            continue
+
+        out.append(line)
+
+    return out
+
+
+def process_stream(lines: Iterable[str]) -> Iterator[str]:
+    """
+    Handle MODEL/ENDMDL blocks:
+      - If MODEL records exist, apply the selection independently within each MODEL block.
+      - Text outside MODEL blocks is processed as a single block.
+    """
+    buffer: List[str] = []
+    in_model = False
+
+    for line in lines:
+        if line.startswith("MODEL "):
+            # Flush anything accumulated before this MODEL
+            if buffer:
+                for x in process_block(buffer):
+                    yield x
+                buffer = []
+            in_model = True
+            yield line
+            continue
+
+        if in_model and line.startswith("ENDMDL"):
+            # Process the model contents, then emit ENDMDL
+            for x in process_block(buffer):
+                yield x
+            buffer = []
+            in_model = False
+            yield line
+            continue
+
+        buffer.append(line)
+
+    # Flush remaining lines at EOF
+    if buffer:
+        for x in process_block(buffer):
+            yield x
+
+
+def clean_pdb_file(in_path: Path, out_path: Path) -> None:
+    """Process a PDB file and write the cleaned output."""
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    with in_path.open("r", newline="") as fin, out_path.open("w", newline="") as fout:
+        for out_line in process_stream(fin):
+            fout.write(out_line)
+
+
+def collect_pdb_files(input_path: Path, recursive: bool) -> List[Path]:
+    """Collect *.pdb files from a file or directory (optionally recursive)."""
+    if input_path.is_file():
+        return [input_path]
+    pattern = "**/*.pdb" if recursive else "*.pdb"
+    return sorted([p for p in input_path.glob(pattern) if p.is_file()])
+
+
+# =============================================================================
+# Public API for programmatic use
+# =============================================================================
+
+def has_altloc(pdb_path: Path) -> bool:
+    """
+    Check if a PDB file contains any non-blank altLoc characters (column 17, 1-based).
+
+    Returns True if at least one ATOM/HETATM record has a non-space character
+    in the altLoc column. Returns False if no altLoc is found.
+    """
+    try:
+        with open(pdb_path, "r", encoding="utf-8", errors="ignore") as fh:
+            for line in fh:
+                if line.startswith(COORD_RECORDS):
+                    # altLoc is at column 17 (1-based), which is index 16 (0-based)
+                    if len(line) > ALTLOC_IDX:
+                        altloc_char = line[ALTLOC_IDX]
+                        if altloc_char != " " and altloc_char != "":
+                            return True
+        return False
+    except Exception:
+        return False
+
+
+def fix_altloc_file(
+    in_path: str | Path,
+    out_path: str | Path,
+    *,
+    overwrite: bool = False,
+    skip_if_no_altloc: bool = True,
+) -> bool:
+    """
+    Fix alternate locations in a PDB file.
+
+    Parameters
+    ----------
+    in_path : str | Path
+        Input PDB file path.
+    out_path : str | Path
+        Output PDB file path.
+    overwrite : bool
+        If True, overwrite existing output file. Default False.
+    skip_if_no_altloc : bool
+        If True, skip processing if no altLoc is detected. Default True.
+
+    Returns
+    -------
+    bool
+        True if the file was processed (altloc found and fixed),
+        False if skipped (no altloc detected).
+
+    Raises
+    ------
+    FileExistsError
+        If output file exists and overwrite=False.
+    FileNotFoundError
+        If input file does not exist.
+    """
+    in_path = Path(in_path)
+    out_path = Path(out_path)
+
+    if not in_path.exists():
+        raise FileNotFoundError(f"Input file not found: {in_path}")
+
+    if out_path.exists() and not overwrite:
+        raise FileExistsError(f"Output file exists: {out_path}")
+
+    if skip_if_no_altloc and not has_altloc(in_path):
+        return False
+
+    clean_pdb_file(in_path, out_path)
+    return True
+
+
+# =============================================================================
+# CLI
+# =============================================================================
+
+def _run_fix_altloc(
+    input_path: Path,
+    out: Optional[Path],
+    recursive: bool,
+    inplace: bool,
+    overwrite: bool,
+    force: bool,
+) -> None:
+    """Core business logic for fix-altloc (called from Click CLI)."""
+    pdb_files = collect_pdb_files(input_path, recursive)
+    if not pdb_files:
+        raise click.ClickException(f"No .pdb files found in: {input_path}")
+
+    skip_if_no_altloc = not force
+    processed_count = 0
+    skipped_count = 0
+
+    # In-place mode
+    if inplace:
+        for in_path in pdb_files:
+            if skip_if_no_altloc and not has_altloc(in_path):
+                skipped_count += 1
+                continue
+
+            bak_path = in_path.with_suffix(in_path.suffix + ".bak")
+            if not bak_path.exists():
+                shutil.copy2(in_path, bak_path)
+
+            tmp_path = in_path.with_suffix(in_path.suffix + ".tmp")
+            clean_pdb_file(in_path, tmp_path)
+            tmp_path.replace(in_path)
+            processed_count += 1
+
+        if processed_count > 0:
+            click.echo(f"[fix-altloc] Processed {processed_count} file(s) in-place.")
+        if skipped_count > 0:
+            click.echo(f"[fix-altloc] Skipped {skipped_count} file(s) (no altLoc detected).")
+        return
+
+    # File input
+    if input_path.is_file():
+        in_path = input_path
+
+        if skip_if_no_altloc and not has_altloc(in_path):
+            click.echo(f"[fix-altloc] Skipped {in_path} (no altLoc detected).")
+            return
+
+        if out is None:
+            out_path = in_path.with_name(in_path.stem + "_clean.pdb")
+        else:
+            if out.suffix.lower() == ".pdb":
+                out_path = out
+            else:
+                out.mkdir(parents=True, exist_ok=True)
+                out_path = out / in_path.name
+
+        if out_path.exists() and not overwrite:
+            raise click.ClickException(f"Output exists: {out_path} (use --overwrite to overwrite)")
+
+        clean_pdb_file(in_path, out_path)
+        click.echo(f"[fix-altloc] Fixed altLoc → {out_path}")
+        return
+
+    # Directory input
+    in_dir = input_path
+    out_dir = out if out is not None else in_dir.with_name(in_dir.name + "_clean")
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    for in_path in pdb_files:
+        if skip_if_no_altloc and not has_altloc(in_path):
+            skipped_count += 1
+            continue
+
+        rel = in_path.relative_to(in_dir)
+        out_path = out_dir / rel
+
+        if out_path.exists() and not overwrite:
+            raise click.ClickException(f"Output exists: {out_path} (use --overwrite to overwrite)")
+
+        clean_pdb_file(in_path, out_path)
+        processed_count += 1
+
+    if processed_count > 0:
+        click.echo(f"[fix-altloc] Processed {processed_count} file(s) → {out_dir}")
+    if skipped_count > 0:
+        click.echo(f"[fix-altloc] Skipped {skipped_count} file(s) (no altLoc detected).")
+
+
+@click.command(
+    name="fix-altloc",
+    help=(
+        "Blank PDB altLoc column (col 17) without shifting, and keep one altLoc "
+        "per atom by default rule: highest occupancy, then earliest appearance."
+    ),
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+@click.option(
+    "-i", "--input", "input_path",
+    type=click.Path(exists=True, path_type=Path),
+    required=True,
+    help="Input PDB file or directory.",
+)
+@click.option(
+    "-o", "--out",
+    type=click.Path(path_type=Path),
+    default=None,
+    help="Output file (if input is a file) or output directory (if input is a directory).",
+)
+@click.option(
+    "--recursive/--no-recursive",
+    default=False, show_default=True,
+    help="When input is a directory, process *.pdb recursively (including subdirectories).",
+)
+@click.option(
+    "--inplace/--no-inplace",
+    default=False, show_default=True,
+    help="Overwrite input file(s) in place (creates .bak next to each file).",
+)
+@click.option(
+    "--overwrite/--no-overwrite",
+    default=False, show_default=True,
+    help="Allow overwriting existing output files.",
+)
+@click.option(
+    "--force/--no-force",
+    default=False, show_default=True,
+    help="Process files even if no altLoc is detected (default: skip files without altLoc).",
+)
+def cli(
+    input_path: Path,
+    out: Optional[Path],
+    recursive: bool,
+    inplace: bool,
+    overwrite: bool,
+    force: bool,
+) -> None:
+    _run_fix_altloc(input_path, out, recursive, inplace, overwrite, force)
+
+
+if __name__ == "__main__":
+    cli()