uniaf3 0.2.0__py3-none-any.whl

uniaf3/__init__.py

@@ -0,0 +1,5 @@
+"""Unified data processing for AlphaFold3-like models."""
+
+from importlib.metadata import version
+
+__version__ = version("uniaf3")

uniaf3/adapters/__init__.py

@@ -0,0 +1,119 @@
+"""Adapters to convert between UniAF3Config and model-specific configs.
+
+Each model has a ``to_*`` and ``from_*`` function pair in its own module under
+``uniaf3.adapters``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from uniaf3.adapters.alphafold3 import from_alphafold3, to_alphafold3
+from uniaf3.adapters.alphafold3_server import (
+    from_alphafold3_server,
+    to_alphafold3_server,
+)
+from uniaf3.adapters.boltz import from_boltz, to_boltz
+from uniaf3.adapters.chai import from_chai, to_chai
+from uniaf3.adapters.protenix import from_protenix, to_protenix
+from uniaf3.schema import (
+    AF3Config,
+    AF3ServerConfig,
+    AnyConfig,
+    AnyConfigList,
+    BoltzConfig,
+    ChaiConfig,
+    ProtenixConfig,
+    UniAF3Config,
+)
+
+__all__ = [
+    "from_alphafold3",
+    "from_alphafold3_server",
+    "from_boltz",
+    "from_chai",
+    "from_protenix",
+    "from_uniaf3",
+    "to_alphafold3",
+    "to_alphafold3_server",
+    "to_boltz",
+    "to_chai",
+    "to_protenix",
+    "to_uniaf3",
+]
+
+
+def to_uniaf3(
+    conf: AnyConfig, *, msa_dir: str | Path = "."
+) -> UniAF3Config | list[UniAF3Config]:
+    """Convert any supported model config to UniAF3Config.
+
+    Args:
+        conf: A config object from any supported model format.
+        msa_dir: Directory to save MSA files (used by Boltz and Chai).
+
+    Returns:
+        The equivalent UniAF3Config.
+
+    Raises:
+        TypeError: If the config type is not recognized.
+
+    """
+    if isinstance(conf, UniAF3Config):
+        return conf
+    if isinstance(conf, AF3ServerConfig):
+        return from_alphafold3_server(conf)
+    if isinstance(conf, AF3Config):
+        return from_alphafold3(conf)
+    if isinstance(conf, BoltzConfig):
+        return from_boltz(conf, msa_dir=msa_dir)
+    if isinstance(conf, ChaiConfig):
+        return from_chai(conf, msa_dir=msa_dir)
+    if isinstance(conf, ProtenixConfig):
+        return from_protenix(conf)
+    raise TypeError(f"Unsupported config type: {type(conf)}")
+
+
+def from_uniaf3(
+    conf: UniAF3Config | list[UniAF3Config],
+    target: type[AnyConfig],
+    *,
+    name: str = "uniaf3_job",
+    msa_dir: str | Path = ".",
+    strict: bool = False,
+) -> AnyConfig | AnyConfigList:
+    """Convert a UniAF3Config to a specific model config.
+
+    Args:
+        conf: The UniAF3Config to convert.
+        target: The target config class.
+        name: Job name for models that require one (AF3, Protenix).
+        msa_dir: Directory to save MSA CSV files (used by Boltz and Chai).
+        strict: If True, raise errors for unsupported features.
+
+    Returns:
+        The target model config.
+
+    Raises:
+        TypeError: If the target type is not recognized.
+
+    """
+    if target is UniAF3Config:
+        return conf
+    if target is AF3Config:
+        if isinstance(conf, list):
+            return [to_alphafold3(c, name=name, strict=strict) for c in conf]
+        return to_alphafold3(conf, name=name, strict=strict)
+    if target is AF3ServerConfig:
+        return to_alphafold3_server(conf, name=name, strict=strict)
+    if target is BoltzConfig:
+        if isinstance(conf, list):
+            return [to_boltz(c, msa_dir=msa_dir, strict=strict) for c in conf]
+        return to_boltz(conf, msa_dir=msa_dir, strict=strict)
+    if target is ChaiConfig:
+        if isinstance(conf, list):
+            return [to_chai(c, msa_dir=msa_dir, strict=strict) for c in conf]
+        return to_chai(conf, msa_dir=msa_dir, strict=strict)
+    if target is ProtenixConfig:
+        return to_protenix(conf, name=name, strict=strict)
+    raise TypeError(f"Unsupported target type: {target}")

uniaf3/adapters/_helpers.py

@@ -0,0 +1,23 @@
+"""Shared helpers for adapter modules."""
+
+from __future__ import annotations
+
+import warnings
+
+
+def ensure_list(val: str | list[str]) -> list[str]:
+    """Normalize id field to a list."""
+    return val if isinstance(val, list) else [val]
+
+
+def warn_lossy_conversion(msg: str):
+    """Emit a warning for lossy conversion behavior."""
+    warnings.warn(f"Lossy conversion: {msg}", UserWarning, stacklevel=3)
+
+
+def err_unsupported_feature(strict: bool, msg: str):
+    """Help handle unsupported features based on the strict flag."""
+    if strict:
+        raise ValueError(msg)
+    else:
+        warnings.warn(f"Skipping unsupported feature: {msg}", UserWarning, stacklevel=3)

uniaf3/adapters/alphafold3.py

@@ -0,0 +1,356 @@
+"""Adapter for converting between UniAF3Config and AlphaFold3 config."""
+
+from __future__ import annotations
+
+from uniaf3.adapters._helpers import (
+    err_unsupported_feature,
+    warn_lossy_conversion,
+)
+from uniaf3.schema.alphafold3 import (
+    AF3DNA,
+    AF3RNA,
+    AF3BondedAtom,
+    AF3Config,
+    AF3Ligand,
+    AF3NucleotideModification,
+    AF3Protein,
+    AF3ProteinModification,
+    AF3SequenceEntry,
+    AF3Template,
+)
+from uniaf3.schema.base import (
+    Atom,
+    AuxiliaryParams,
+    CovalentBond,
+    Glycan,
+    Ligand,
+    Polymer,
+    PolymerType,
+    ProteinSeq,
+    SequenceModification,
+    StructuralTemplate,
+    UniAF3Config,
+)
+
+
+def to_alphafold3(
+    config: UniAF3Config,
+    name: str = "uniaf3_job",
+    strict: bool = False,
+) -> AF3Config:
+    """Convert a UniAF3Config to an AlphaFold3 config.
+
+    Args:
+        config: UniAF3Config pydantic object.
+        name: Job name for the AF3 config.
+        strict: If True, raise errors when encountering unsupported features.
+            If False, skip unsupported features with warnings.
+
+    """
+    sequences: list[AF3SequenceEntry] = []
+    for seq in config.sequences:
+        if isinstance(seq, Glycan):
+            # TODO: AF3 does not have a native glycan type. Glycans must be
+            # represented as multi-CCD ligands. This requires knowing the
+            # component CCD codes, which the chai_str notation may not directly
+            # map to.
+            err_unsupported_feature(
+                strict,
+                f"Glycans are not directly supported in AF3: {seq}",
+            )
+            continue
+
+        if isinstance(seq, ProteinSeq):
+            mods = (
+                [
+                    AF3ProteinModification(ptmType=m.ccd, ptmPosition=m.position)
+                    for m in seq.modifications
+                ]
+                if seq.modifications
+                else None
+            )
+            protein = AF3Protein(
+                id=seq.id,
+                sequence=seq.sequence,
+                modifications=mods,
+                description=seq.description,
+                unpairedMsaPath=seq.unpaired_msa,
+                pairedMsaPath=seq.paired_msa,
+            )
+            # Templates
+            if seq.templates:
+                af3_templates = []
+                for tmpl in seq.templates:
+                    if tmpl.query_chains or tmpl.template_chains:
+                        warn_lossy_conversion(
+                            "UniAF3Config.sequences[*].templates.{query_chains,template_chains} are not represented by AF3Config.sequences[*].protein.templates."
+                        )
+                    if (
+                        tmpl.boltz_enable_force
+                        or tmpl.boltz_template_threshold is not None
+                    ):
+                        warn_lossy_conversion(
+                            "UniAF3Config.sequences[*].templates.{boltz_enable_force,boltz_template_threshold} are not represented by AF3Config.sequences[*].protein.templates."
+                        )
+                    # TODO: extract a single chain from the template structure
+                    af3_templates.append(
+                        AF3Template(
+                            mmcifPath=tmpl.path,
+                            queryIndices=tmpl.query_idx or [],
+                            templateIndices=tmpl.template_idx or [],
+                        )
+                    )
+                protein.templates = af3_templates
+            sequences.append(AF3SequenceEntry(protein=protein))
+
+        elif isinstance(seq, Polymer):
+            if seq.polymer_type == PolymerType.Protein:
+                mods = (
+                    [
+                        AF3ProteinModification(ptmType=m.ccd, ptmPosition=m.position)
+                        for m in seq.modifications
+                    ]
+                    if seq.modifications
+                    else None
+                )
+                protein = AF3Protein(
+                    id=seq.id,
+                    sequence=seq.sequence,
+                    modifications=mods,
+                    description=seq.description,
+                )
+                sequences.append(AF3SequenceEntry(protein=protein))
+            elif seq.polymer_type == PolymerType.DNA:
+                mods = (
+                    [
+                        AF3NucleotideModification(
+                            modificationType=m.ccd, basePosition=m.position
+                        )
+                        for m in seq.modifications
+                    ]
+                    if seq.modifications
+                    else None
+                )
+                dna = AF3DNA(
+                    id=seq.id,
+                    sequence=seq.sequence,
+                    modifications=mods,
+                    description=seq.description,
+                )
+                sequences.append(AF3SequenceEntry(dna=dna))
+            elif seq.polymer_type == PolymerType.RNA:
+                mods = (
+                    [
+                        AF3NucleotideModification(
+                            modificationType=m.ccd, basePosition=m.position
+                        )
+                        for m in seq.modifications
+                    ]
+                    if seq.modifications
+                    else None
+                )
+                rna = AF3RNA(
+                    id=seq.id,
+                    sequence=seq.sequence,
+                    modifications=mods,
+                    description=seq.description,
+                )
+                sequences.append(AF3SequenceEntry(rna=rna))
+
+        elif isinstance(seq, Ligand):
+            if seq.ccd:
+                lig = AF3Ligand(
+                    id=seq.id,
+                    ccdCodes=seq.ccd,
+                    description=seq.description,
+                )
+            elif seq.smiles:
+                lig = AF3Ligand(
+                    id=seq.id,
+                    smiles=seq.smiles,
+                    description=seq.description,
+                )
+            else:
+                continue
+            sequences.append(AF3SequenceEntry(ligand=lig))
+
+    # Bonded atom pairs (only covalent bonds)
+    bonded_atom_pairs: list[tuple[AF3BondedAtom, AF3BondedAtom]] = []
+    if config.covalent_bonds:
+        for r in config.covalent_bonds:
+            # AF3 requires atom names to be given
+            if r.atom1.atom_name is None or r.atom2.atom_name is None:
+                err_unsupported_feature(
+                    strict,
+                    f"AF3 bondedAtomPairs require atom names, but got: {r}",
+                )
+                continue
+            a1: AF3BondedAtom = (
+                r.atom1.chain_id,
+                r.atom1.residue_idx,
+                r.atom1.atom_name,
+            )
+            a2: AF3BondedAtom = (
+                r.atom2.chain_id,
+                r.atom2.residue_idx,
+                r.atom2.atom_name,
+            )
+            bonded_atom_pairs.append((a1, a2))
+
+    if config.contact_restraints or config.pocket_restraints:
+        err_unsupported_feature(
+            strict,
+            "AF3 does not support contact or pocket restraints.",
+        )
+
+    return AF3Config(
+        name=name,
+        modelSeeds=config.aux.seeds,
+        sequences=sequences,
+        bondedAtomPairs=bonded_atom_pairs or None,
+    )
+
+
+def from_alphafold3(config: AF3Config) -> UniAF3Config:
+    """Convert an AlphaFold3 config to a UniAF3Config.
+
+    Args:
+        config: AF3Config pydantic object.
+
+    Returns:
+        A UniAF3Config.
+
+    """
+    sequences: list[Polymer | ProteinSeq | Ligand | Glycan] = []
+    if config.name:
+        warn_lossy_conversion(
+            f"AF3Config.name ('{config.name}') is not represented in UniAF3Config."
+        )
+    if config.userCCD is not None or config.userCCDPath is not None:
+        warn_lossy_conversion(
+            "AF3Config.{userCCD,userCCDPath} are not represented in UniAF3Config."
+        )
+
+    for entry in config.sequences:
+        if entry.protein is not None:
+            p = entry.protein
+            mods = (
+                [
+                    SequenceModification(ccd=m.ptmType, position=m.ptmPosition)
+                    for m in p.modifications
+                ]
+                if p.modifications
+                else None
+            )
+            if p.unpairedMsa is not None or p.pairedMsa is not None:
+                warn_lossy_conversion(
+                    "AF3Config.sequences[*].protein.{unpairedMsa,pairedMsa} are not imported; UniAF3 maps only file-based MSA paths."
+                )
+
+            templates = None
+            if p.templates:
+                if any(
+                    t.mmcif is not None and t.mmcifPath is None for t in p.templates
+                ):
+                    warn_lossy_conversion(
+                        "AF3Config.sequences[*].protein.templates[*].mmcif is not preserved; only mmcifPath maps to UniAF3 templates.path."
+                    )
+                templates = [
+                    StructuralTemplate(
+                        path=t.mmcifPath or "",
+                        query_idx=t.queryIndices,
+                        template_idx=t.templateIndices,
+                    )
+                    for t in p.templates
+                ]
+
+            seq = ProteinSeq(
+                polymer_type=PolymerType.Protein,
+                id=p.id,
+                sequence=p.sequence,
+                modifications=mods,
+                description=p.description,
+                unpaired_msa=p.unpairedMsaPath,
+                paired_msa=p.pairedMsaPath,
+                templates=templates,
+            )
+            sequences.append(seq)
+
+        elif entry.dna is not None:
+            d = entry.dna
+            mods = (
+                [
+                    SequenceModification(
+                        ccd=m.modificationType, position=m.basePosition
+                    )
+                    for m in d.modifications
+                ]
+                if d.modifications
+                else None
+            )
+            seq = Polymer(
+                polymer_type=PolymerType.DNA,
+                id=d.id,
+                sequence=d.sequence,
+                modifications=mods,
+                description=d.description,
+            )
+            sequences.append(seq)
+
+        elif entry.rna is not None:
+            r = entry.rna
+            mods = (
+                [
+                    SequenceModification(
+                        ccd=m.modificationType, position=m.basePosition
+                    )
+                    for m in r.modifications
+                ]
+                if r.modifications
+                else None
+            )
+            seq = Polymer(
+                polymer_type=PolymerType.RNA,
+                id=r.id,
+                sequence=r.sequence,
+                modifications=mods,
+                description=r.description,
+            )
+            sequences.append(seq)
+
+        elif entry.ligand is not None:
+            lg = entry.ligand
+            lig = Ligand(
+                id=lg.id,
+                ccd=lg.ccdCodes,
+                smiles=lg.smiles,
+                description=lg.description,
+            )
+            sequences.append(lig)
+
+    # Bonded atom pairs → covalent bonds
+    covalent_bonds: list[CovalentBond] = []
+    if config.bondedAtomPairs:
+        for a1, a2 in config.bondedAtomPairs:
+            covalent_bonds.append(
+                CovalentBond(
+                    atom1=Atom(
+                        chain_id=a1[0],
+                        residue_idx=a1[1],
+                        atom_name=a1[2],
+                        residue_name=None,
+                    ),
+                    atom2=Atom(
+                        chain_id=a2[0],
+                        residue_idx=a2[1],
+                        atom_name=a2[2],
+                        residue_name=None,
+                    ),
+                )
+            )
+
+    return UniAF3Config(
+        sequences=sequences,
+        covalent_bonds=covalent_bonds or None,
+        aux=AuxiliaryParams(seeds=config.modelSeeds),
+    )