molforge.folding¶

folding ¶

Data types for multi-component structure prediction.

molforge's single-sequence folding interface (:meth:FoldingEngine.predict) covers the bread-and-butter case of "predict the structure of one protein chain." For the AlphaFold-3- class engines :class:Boltz and :class:Chai1, that's a small fraction of what the underlying models can do — both engines fold complexes of multiple proteins, nucleic acids, and small-molecule ligands in a single forward pass. Modelling those complexes is the headline drug-discovery and structural-biology use case.

This module provides the shared input shape for multi-component prediction:

:class:Entity — one component of a complex (a protein chain, a ligand, a DNA strand, etc.).
:class:ComplexSpec — an ordered list of entities defining the complete system to fold.
:class:TemplatePolicy — whether, and how, an engine may lean on structural templates. Ablating templates is how you find out what a model can predict without being shown the answer.

Engines that support multi-component prediction expose a predict_complex(spec) method (see :class:Boltz.predict_complex and :class:Chai1.predict_complex). The method returns a multi-chain :class:Protein whose atom_array has one logical chain per entity (or per copy, for homo-oligomers).

Engines that don't support complexes (:class:ESMFold, :class:AlphaFold in single-chain mode) simply don't expose predict_complex. Check with hasattr(engine, "predict_complex") or by importing the engine and reading its documented capabilities.

Design notes¶

The :class:Entity shape is a common denominator of Boltz YAML and Chai-1 FASTA. Both engines support:

Polymer entities (protein / DNA / RNA) specified by one-letter sequence.
Ligand entities specified by SMILES or CCD code (mutually exclusive).
Multiple copies of an entity (homo-oligomers).
Per-entity chain IDs (auto-assigned A, B, C, ... if omitted).

What this v1 deliberately doesn't model:

Modified residues (Boltz uses modifications lists; Chai needs a separate mechanism). Carry the unmodified sequence here; add engine-specific modification kwargs as a follow-up.
Custom MSAs per entity (use use_msa_server for v1). MSA depth is controllable — see the engines' msa_depth parameter.
Per-entity templates. :class:TemplatePolicy is set on the engine and applies to the whole prediction; targeting one chain means dropping down to the engine's own input format.
Restraints (covalent bonds, pocket constraints). Boltz supports these via top-level YAML keys; deferred.
Ions specified as a separate entity type — Boltz models them as CCD ligand entries with codes like ZN / MG; Chai accepts them as ligand|name= SMILES like [Zn+2]. We treat them as ligands.

For these advanced features, drop down to the engine's raw input format using the engine instance's documented advanced API.

Entity `dataclass` ¶

Entity(
    kind: EntityKind,
    sequence: str | None = None,
    smiles: str | None = None,
    ccd: str | None = None,
    chain_id: str | None = None,
    copies: int = 1,
    name: str | None = None,
    msa: MultipleSequenceAlignment
    | Literal["empty"]
    | None = None,
)

One component of a biomolecular complex.

An Entity describes a single chain (for polymers) or a single small molecule (for ligands). It does not specify a 3D structure; that's what the folding engine produces.

For polymer entities (protein, dna, rna), supply sequence as a one-letter string. For ligand entities, supply exactly one of smiles (a SMILES string) or ccd (a 3-letter PDB Chemical Component Dictionary code, e.g. "ATP", "NAD", "ZN").

Parameters:

Name	Type	Description	Default
`kind`	`EntityKind`	One of `"protein"`, `"dna"`, `"rna"`, `"ligand"`.	required
`sequence`	`str \| None`	One-letter sequence string. Required for polymer kinds; must be `None` for ligand.	`None`
`smiles`	`str \| None`	SMILES string. Only valid for `kind="ligand"`. Mutually exclusive with `ccd`.	`None`
`ccd`	`str \| None`	3-letter CCD code. Only valid for `kind="ligand"`. Mutually exclusive with `smiles`.	`None`
`chain_id`	`str \| None`	Logical chain identifier (single letter A-Z, or two-letter for large complexes). `None` means "auto-assign at serialization time" in the order entities appear in the :class:`ComplexSpec`.	`None`
`copies`	`int`	Number of identical copies for homo-oligomers. `copies=2` declares a homodimer of this entity; the engine produces two chains with the same sequence but distinct chain IDs.	`1`
`name`	`str \| None`	Human-readable identifier for logs and reports. `None` defaults to `"{kind}_{chain_id}"` at serialization time.	`None`
`msa`	`MultipleSequenceAlignment \| Literal['empty'] \| None`	A precomputed alignment for this chain. `None` (default) leaves the engine to its own behaviour — usually a public MSA server. A :class:`~molforge.sequence.MultipleSequenceAlignment` is handed to the engine instead, which is how you fold against an alignment you computed and can inspect. `"empty"` forces single-sequence mode for this chain specifically. Protein chains only: nucleic acids and ligands have no alignment. To use an alignment from disk, read it first — `Entity(..., msa=read_a3m("hits.a3m"))` — rather than passing a path, so a typo surfaces as a missing file instead of being mistaken for a mode.	`None`

Examples:

Single protein chain::

Entity(kind="protein", sequence="MKQHKAMIVAL...")

Aspirin (small molecule via SMILES)::

Entity(kind="ligand", smiles="CC(=O)OC1=CC=CC=C1C(=O)O")

ATP cofactor (via CCD code)::

Entity(kind="ligand", ccd="ATP")

Homodimer (two identical protein chains)::

Entity(kind="protein", sequence="MKQH...", copies=2)

DNA double helix (two complementary strands)::

Entity(kind="dna", sequence="ATCGATCG", chain_id="A")
Entity(kind="dna", sequence="CGATCGAT", chain_id="B")

is_polymer `property` ¶

is_polymer: bool

True for protein / DNA / RNA entities (anything with a sequence).

is_ligand `property` ¶

is_ligand: bool

True for ligand entities (small molecules, ions, cofactors).

normalized_sequence ¶

normalized_sequence() -> str

The polymer sequence with whitespace stripped and uppercased.

Raises:

Type	Description
`ValueError`	If called on a non-polymer entity.

ComplexSpec `dataclass` ¶

ComplexSpec(entities: tuple[Entity, ...])

An ordered specification of a biomolecular complex to fold.

A :class:ComplexSpec is the input to multi-component prediction via :meth:Boltz.predict_complex and :meth:Chai1.predict_complex. It's the engine-agnostic shape; serializers in each wrapper convert it to the engine's native format (Boltz YAML or Chai typed FASTA).

The spec is ordered: the order of entities maps deterministically to chain IDs A, B, C, ... when none are explicitly assigned. For multi-copy entities (Entity.copies > 1), each copy gets its own chain ID in sequence.

Parameters:

Name	Type	Description	Default
`entities`	`tuple[Entity, ...]`	Tuple of :class:`Entity` defining the complex. Must be non-empty.	required

Validation

At least one entity required.
All explicitly-set chain_ids must be unique across the spec.
At most one ligand entity may carry an affinity_binder marker (deferred; not exposed in v1).

Examples:

Single-protein "complex" (degenerate; usually go through :meth:FoldingEngine.predict instead)::

spec = ComplexSpec.from_protein("MKQH...")

Protein + small-molecule drug::

spec = ComplexSpec.protein_ligand(
    protein_sequence="MKQHKAMIVAL...",
    ligand_smiles="CC(=O)OC1=CC=CC=C1C(=O)O",
)

Antibody-antigen (3-chain protein complex: heavy, light, antigen)::

spec = ComplexSpec(entities=(
    Entity(kind="protein", sequence=heavy_chain_seq,
           chain_id="H", name="heavy"),
    Entity(kind="protein", sequence=light_chain_seq,
           chain_id="L", name="light"),
    Entity(kind="protein", sequence=antigen_seq,
           chain_id="A", name="antigen"),
))

Transcription factor on DNA::

spec = ComplexSpec(entities=(
    Entity(kind="protein", sequence=tf_seq, copies=1),
    Entity(kind="dna", sequence=forward_strand),
    Entity(kind="dna", sequence=reverse_strand),
))

from_protein `classmethod` ¶

from_protein(
    sequence: str, *, chain_id: str = "A"
) -> ComplexSpec

Single-protein convenience constructor.

Equivalent to ComplexSpec(entities=(Entity(kind="protein", sequence=sequence, chain_id=chain_id),)). Used internally so :meth:Boltz.predict(sequence) can delegate to the same code path as :meth:Boltz.predict_complex(...).

protein_ligand `classmethod` ¶

protein_ligand(
    *,
    protein_sequence: str,
    ligand_smiles: str | None = None,
    ligand_ccd: str | None = None,
    protein_chain_id: str = "A",
    ligand_chain_id: str = "B",
) -> ComplexSpec

Headline protein + small-molecule constructor.

The standard drug-discovery shape: one protein chain plus one small molecule. Supply the ligand as ligand_smiles or ligand_ccd (mutually exclusive).

Example::

spec = ComplexSpec.protein_ligand(
    protein_sequence="MVTPEG...",
    ligand_smiles="N[C@@H](Cc1ccc(O)cc1)C(=O)O",
)

to_provenance ¶

to_provenance() -> dict[str, Any]

Flatten to the JSON-safe shape recorded in :class:Provenance.

Provenance.inputs must be JSON-serialisable and a :class:ComplexSpec is not — it's a frozen dataclass of :class:Entity tuples. This flattens it to a list of dicts with engine-agnostic keys, so the same structure appears in provenance whichever engine produced the prediction, and a Boltz run can be compared with a Chai-1 one side by side.

Chain IDs are recorded as assigned, not as declared, so the record says which chains actually came back rather than leaving the reader to re-run the allocator.

Returns:

Type	Description
`dict[str, Any]`	`{"entities": [{...}, ...]}`, one dict per entity.

from_provenance `classmethod` ¶

from_provenance(payload: Mapping[str, Any]) -> ComplexSpec

Rebuild a spec from :meth:to_provenance output.

The inverse exists so a recorded complex prediction can be replayed: without it a manifest can describe a multi-component fold perfectly well and still be unable to re-run it.

Each entity's first assigned chain ID is pinned explicitly and the remaining copies are counted, which reproduces the original assignment exactly — including when some entities had explicit IDs and others were auto-assigned around them.

Parameters:

Name	Type	Description	Default
`payload`	`Mapping[str, Any]`	A dict shaped like :meth:`to_provenance`'s output.	required

Returns:

Name	Type	Description
`A`	`ComplexSpec`	class:`ComplexSpec` that re-serialises to `payload`.

Raises:

Type	Description
`ValueError`	If `payload` has no `entities` list, or an entity is missing the fields its kind requires.

assigned_chain_ids ¶

assigned_chain_ids() -> list[list[str]]

Return chain IDs per entity after auto-assignment.

Returns a list of length len(self.entities). Each element is itself a list of length entity.copies containing the chain IDs that will appear in the final structure for that entity's copies.

Auto-assignment uses A, B, C, ..., Z, AA, AB, ... in spec order, skipping any IDs already claimed by explicit Entity.chain_id values.

Example::

spec = ComplexSpec(entities=(
    Entity(kind="protein", sequence="MKQ", copies=2),     # auto
    Entity(kind="protein", sequence="HIS", chain_id="X"), # explicit
    Entity(kind="ligand", smiles="O"),                    # auto
))
spec.assigned_chain_ids()
# → [["A", "B"], ["X"], ["C"]]

TemplatePolicy `dataclass` ¶

TemplatePolicy(
    mode: Literal["none", "structures", "hits", "server"],
    structures: tuple[str, ...] = (),
    hits_file: str | None = None,
)

Whether — and how — a folding engine may use structural templates.

A template lets the model copy a known fold rather than derive one. That is usually what you want, and occasionally exactly what you don't: to see whether a model can produce a conformation on its own, you have to be able to take its templates away. Comparing runs across template settings is an ablation, and an ablation needs the setting to be a parameter rather than an engine default.

The engines disagree about what a template is, so the policy names the mechanism rather than pretending they're interchangeable:

================= ====================================== =========== Mode Means Supported by ================= ====================================== =========== "none" No templates. Predict from sequence. Boltz, Chai-1 "structures" These specific structure files. Boltz "hits" A precomputed template-hits file. Chai-1 "server" Let the engine search for templates. Chai-1 ================= ====================================== ===========

An engine raises :class:ValueError for a mode it cannot honour, naming what it does support — rather than quietly folding without the templates you asked for.

Note that molforge does not decide which templates are appropriate. Restricting a run to, say, only unbound structures is a judgement about your targets that molforge has no way to make; you make it by choosing what to pass to :meth:structures or by filtering the hits file. What molforge guarantees is that whatever you chose is applied, recorded in provenance, and therefore part of the cache key — so two policies never share a cached result.

Attributes:

Name	Type	Description
`mode`	`Literal['none', 'structures', 'hits', 'server']`	One of `"none"`, `"structures"`, `"hits"`, `"server"`.
`structures`	`tuple[str, ...]`	Paths to template structure files, for `mode="structures"`. Empty otherwise.
`hits_file`	`str \| None`	Path to a template-hits file, for `mode="hits"`. `None` otherwise.

Example

from molforge.folding import TemplatePolicy TemplatePolicy.none().mode 'none' TemplatePolicy.from_structures("4hhb.cif", "1ubq.cif").structures ('4hhb.cif', '1ubq.cif')

none `classmethod` ¶

none() -> TemplatePolicy

No templates: make the model predict from sequence alone.

from_structures `classmethod` ¶

from_structures(
    *paths: str | PathLike[str],
) -> TemplatePolicy

Use these structure files as templates.

Parameters:

Name	Type	Description	Default
`*paths`	`str \| PathLike[str]`	Template structure files. Boltz reads `.cif` and `.pdb`.	`()`

Raises:

Type	Description
`ValueError`	If no paths are given.

from_hits `classmethod` ¶

from_hits(path: str | PathLike[str]) -> TemplatePolicy

Use a precomputed template-hits file (Chai-1's .m8).

from_server `classmethod` ¶

from_server() -> TemplatePolicy

Let the engine search a template server for itself.

to_provenance ¶

to_provenance() -> dict[str, object]

A JSON-safe summary for :class:~molforge.core.Provenance.

Only the keys that carry information are emitted, so a policy's contribution to a cache key stays as small as what it actually says.

molforge.folding¶

folding ¶

Design notes¶

Entity dataclass ¶

is_polymer property ¶

is_ligand property ¶

normalized_sequence ¶

ComplexSpec dataclass ¶

from_protein classmethod ¶

protein_ligand classmethod ¶

to_provenance ¶

from_provenance classmethod ¶

assigned_chain_ids ¶

TemplatePolicy dataclass ¶

none classmethod ¶

from_structures classmethod ¶

from_hits classmethod ¶

from_server classmethod ¶

to_provenance ¶

Entity `dataclass` ¶

is_polymer `property` ¶

is_ligand `property` ¶

ComplexSpec `dataclass` ¶

from_protein `classmethod` ¶

protein_ligand `classmethod` ¶

from_provenance `classmethod` ¶

TemplatePolicy `dataclass` ¶

none `classmethod` ¶

from_structures `classmethod` ¶

from_hits `classmethod` ¶

from_server `classmethod` ¶