openprotein-python 0.10.1__tar.gz → 0.11.1__tar.gz

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/.gitignore

@@ -367,3 +367,6 @@ __marimo__/
 
 .envrc
 /.direnv/
+
+scratch.org
+gptel

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openprotein-python
-Version: 0.10.1
+Version: 0.11.1
 Summary: OpenProtein Python interface.
 Author-email: Mark Gee <markgee@ne47.bio>, "Timothy Truong Jr." <ttruong@ne47.bio>, Tristan Bepler <tbepler@ne47.bio>
 License-Expression: MIT

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/openprotein/embeddings/poet.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from openprotein.base import APISession
-from openprotein.common import ModelMetadata, ReductionType
+from openprotein.common import ModelMetadata, Reduction, ReductionType
 from openprotein.data import AssayDataset, AssayMetadata
 from openprotein.prompt import Prompt
 
@@ -51,9 +51,9 @@ class PoETModel(EmbeddingModel):
 
     def embed(
         self,
-        sequences: list[bytes],
-        prompt: str | Prompt | None = None,
+        sequences: list[bytes] | list[str],
         reduction: ReductionType | None = ReductionType.MEAN,
+        prompt: str | Prompt | None = None,
         **kwargs,
     ) -> EmbeddingsResultFuture:
         """
@@ -74,6 +74,12 @@ class PoETModel(EmbeddingModel):
         -------
         EmbeddingsResultFuture
             Future object that returns the embeddings of the submitted sequences.
+
+            Note: The embeddings for PoET can have an extra first dimension if using ensemble
+            prompts, where the first dimension is the number of replicates in the ensemble
+            prompt. i.e. the shape is ``(N, L, D)`` if ``N`` > 1 else ``(L, D)`` where ``N`` is
+            the number of replicates in the prompt, ``L`` is the length of the sequence, ``D`` is
+            the dimensions of the ensemble.
         """
         if prompt is None:
             prompt_id = None
@@ -88,7 +94,7 @@ class PoETModel(EmbeddingModel):
 
     def logits(
         self,
-        sequences: list[bytes],
+        sequences: list[bytes] | list[str],
         prompt: str | Prompt | None = None,
         **kwargs,
     ) -> EmbeddingsResultFuture:
@@ -108,6 +114,12 @@ class PoETModel(EmbeddingModel):
         -------
         EmbeddingsResultFuture
             Future object that returns the logits of the submitted sequences.
+
+            Note: The logits for PoET can have an extra first dimension if using ensemble
+            prompts, where the first dimension is the number of replicates in the ensemble
+            prompt. i.e. the shape is ``(N, L, D)`` if ``N`` > 1 else ``(L, D)`` where ``N`` is
+            the number of replicates in the prompt, ``L`` is the length of the sequence, ``D`` is
+            the size of the vocabulary.
         """
         if prompt is None:
             prompt_id = None
@@ -317,11 +329,11 @@ class PoETModel(EmbeddingModel):
 
     def fit_svd(
         self,
-        prompt: str | Prompt | None = None,
         sequences: list[bytes] | list[str] | None = None,
-        assay: AssayDataset | None = None,
+        assay: AssayDataset | AssayMetadata | None = None,
         n_components: int = 1024,
-        reduction: ReductionType | None = None,
+        reduction: Reduction | ReductionType | None = None,
+        prompt: str | Prompt | None = None,
         **kwargs,
     ) -> "SVDModel":
         """
@@ -365,11 +377,11 @@ class PoETModel(EmbeddingModel):
 
     def fit_umap(
         self,
-        prompt: str | Prompt | None = None,
         sequences: list[bytes] | list[str] | None = None,
-        assay: AssayDataset | None = None,
+        assay: AssayDataset | AssayMetadata | None = None,
         n_components: int = 2,
-        reduction: ReductionType = ReductionType.MEAN,
+        reduction: Reduction | ReductionType = ReductionType.MEAN,
+        prompt: str | Prompt | None = None,
         **kwargs,
     ) -> "UMAPModel":
         """
@@ -413,8 +425,11 @@ class PoETModel(EmbeddingModel):
 
     def fit_gp(
         self,
-        assay: AssayMetadata | AssayDataset | str,
+        assay: AssayDataset | AssayMetadata | str,
         properties: list[str],
+        reduction: ReductionType,
+        name: str | None = None,
+        description: str | None = None,
         prompt: str | Prompt | None = None,
         **kwargs,
     ) -> "PredictorModel":
@@ -444,6 +459,9 @@ class PoETModel(EmbeddingModel):
         return super().fit_gp(
             assay=assay,
             properties=properties,
+            reduction=reduction,
+            name=name,
+            description=description,
             prompt_id=prompt_id,
             **kwargs,
         )

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/openprotein/embeddings/poet2.py

@@ -86,6 +86,12 @@ class PoET2Model(PoETModel, EmbeddingModel):
         -------
         EmbeddingsResultFuture
             A future object that returns the embeddings of the submitted sequences.
+
+            Note: The embeddings for PoET can have an extra first dimension if using ensemble
+            prompts, where the first dimension is the number of replicates in the ensemble
+            prompt. i.e. the shape is ``(N, L, D)`` if ``N`` > 1 else ``(L, D)`` where ``N`` is
+            the number of replicates in the prompt, ``L`` is the length of the sequence, ``D`` is
+            the dimensions of the ensemble.
         """
         prompt_api = getattr(self.session, "prompt", None)
         assert isinstance(prompt_api, PromptAPI)
@@ -127,6 +133,12 @@ class PoET2Model(PoETModel, EmbeddingModel):
         -------
         EmbeddingsResultFuture
             A future object that returns the logits of the submitted sequences.
+
+            Note: The logits for PoET can have an extra first dimension if using ensemble
+            prompts, where the first dimension is the number of replicates in the ensemble
+            prompt. i.e. the shape is ``(N, L, D)`` if ``N`` > 1 else ``(L, D)`` where ``N`` is
+            the number of replicates in the prompt, ``L`` is the length of the sequence, ``D`` is
+            the size of the vocabulary.
         """
         prompt_api = getattr(self.session, "prompt", None)
         assert isinstance(prompt_api, PromptAPI)

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/openprotein/fold/boltz.py

@@ -1,7 +1,7 @@
 """Community-based Boltz models for complex structure prediction with ligands/dna/rna."""
 
 import warnings
-from typing import Sequence
+from typing import Mapping, Sequence, cast
 
 from pydantic import BaseModel, Field, TypeAdapter, model_validator
 
@@ -10,6 +10,8 @@ from openprotein.base import APISession
 from openprotein.common import ModelMetadata
 from openprotein.fold.common import normalize_inputs, serialize_input
 from openprotein.molecules import Complex, Ligand, Protein
+from openprotein.molecules.template import Template
+from openprotein.prompt import PromptAPI
 
 from . import api
 from .complex import id_generator
@@ -40,7 +42,7 @@ class BoltzModel(FoldModel):
         num_steps: int = 200,
         step_scale: float = 1.638,
         use_potentials: bool = False,
-        constraints: list[dict] | None = None,
+        constraints: Sequence[Mapping] | None = None,
         **kwargs,
     ) -> FoldResultFuture:
         """
@@ -83,19 +85,9 @@ class BoltzModel(FoldModel):
 
         # build the normalized_models from msa
         if isinstance(sequences, MSAFuture):
-            id_gen = id_generator()
-            align_api = getattr(self.session, "align", None)
-            assert isinstance(align_api, AlignAPI)
-            msa = sequences  # rename
-            seed = align_api.get_seed(job_id=msa.job.job_id)
-            _proteins: dict[str, Protein] = {}
-            for seq in seed.split(":"):
-                protein = Protein(sequence=seq)
-                id = next(id_gen)
-                protein.msa = msa.id
-                _proteins[id] = protein
-            normalized_complexes = [Complex(chains=_proteins)]
-
+            normalized_complexes = [
+                _msa_future_to_complex(session=self.session, msa=sequences)
+            ]
         else:
             normalized_complexes = normalize_inputs(sequences)
 
@@ -139,9 +131,9 @@ class Boltz2Model(BoltzModel, FoldModel):
         num_steps: int = 200,
         step_scale: float = 1.638,
         use_potentials: bool = False,
-        constraints: list[dict] | None = None,
-        templates: list[dict] | None = None,
-        properties: list[dict] | None = None,
+        constraints: Sequence[Mapping] | None = None,
+        templates: Sequence[Protein | Complex | Template] | None = None,
+        properties: Sequence[Mapping] | None = None,
         method: str | None = None,
     ) -> FoldResultFuture:
         """
@@ -163,7 +155,7 @@ class Boltz2Model(BoltzModel, FoldModel):
             Whether or not to use potentials.
         constraints : list[dict] | None = None
             List of constraints.
-        templates: list[dict] | None = None
+        templates: list[Protein | Complex | Template] | None = None
             List of templates to use for structure prediction.
         properties: list[dict] | None = None
             List of additional properties to predict. Should match the `BoltzProperties`
@@ -180,24 +172,98 @@ class Boltz2Model(BoltzModel, FoldModel):
         Returns
         -------
         FoldResultFuture
-            Future for the folding result.
+             Future for the folding result.
         """
-
+        prompt_api = getattr(self.session, "prompt", None)
+        assert isinstance(prompt_api, PromptAPI)
+
+        # validate templates
+        # mapping chain_id (to predict) to template
+        # needs to be consistent
+        templates_: list[Template] = []
+        if not isinstance(sequences, MSAFuture):
+            first_chain_id_to_template = {}
+            for batch_idx, seq in enumerate(sequences):
+                # validate templates and normalize to complex
+                if isinstance(seq, str) or isinstance(seq, bytes):
+                    seq = Protein(seq)
+                seq._assert_valid_templates()
+                if isinstance(seq, Protein):
+                    complex = Complex({"A": seq})
+                else:
+                    complex = seq
+                # resolve chain-level templates
+                for chain_id, protein in complex.get_proteins().items():
+                    # Verify same chain_id should have same templates
+                    if batch_idx == 0:
+                        first_chain_id_to_template[chain_id] = protein.templates
+                        for template in protein.templates:
+                            templates_.append(_to_template(template, chain_id=chain_id))
+                    elif first_chain_id_to_template[chain_id] != protein.templates:
+                        raise ValueError(
+                            "Expected same chain across batches to have the same templates"
+                        )
+                # resolve complex-level templates
+                if batch_idx == 0:
+                    first_templates = complex.templates
+                    for template in complex.templates:
+                        templates_.append(_to_template(template))
+                elif first_templates != complex.templates:
+                    raise ValueError(
+                        "Expected templates across complexes in batch to be the same"
+                    )
+        # method level argument
         if templates is not None:
-            raise ValueError("`templates` not yet supported!")
+            if isinstance(sequences, MSAFuture):
+                # need to convert to complex for template validation
+                sequences = [
+                    _msa_future_to_complex(session=self.session, msa=sequences)
+                ]
+            for template in templates:
+                template = _to_template(template)
+                # validate the template for all sequences before accepting it
+                for seq in sequences:
+                    if isinstance(seq, str) or isinstance(seq, bytes):
+                        seq = Protein(seq)
+                    template.validate_for_target(seq)
+                templates_.append(template)
+
+        # resolve list of Templates into expected dict arg
+        template_dicts: list[dict] = []
+        # track resolved queries to reduce network calls - use id() for identity-based caching
+        struct_id_to_query_id = {}
+
+        for template in templates_:
+            # Use id() for caching - only resolve each unique structure once
+            struct_id = id(template.template)
+            if struct_id not in struct_id_to_query_id:
+                struct_id_to_query_id[struct_id] = prompt_api._resolve_query(
+                    query=template.template
+                )
+
+            template_dict = {"query_id": struct_id_to_query_id[struct_id]}
+
+            if template.mapping is not None:
+                if isinstance(template.mapping, str):
+                    template_dict["chain_id"] = template.mapping
+                else:
+                    template_dict["chain_id"] = list(template.mapping.values())
+                    template_dict["template_id"] = list(template.mapping.keys())
+
+            template_dicts.append(template_dict)
 
         # validate properties
         if properties is not None:
             props = TypeAdapter(list[BoltzProperty]).validate_python(properties)
             # Only allow affinity for ligands, and check binder refers to a ligand chain_id (str, not list)
             ligand_chain_ids = set()
-            if isinstance(sequences, list):
+            if not isinstance(sequences, MSAFuture):
                 for protein in sequences:
                     if isinstance(protein, Complex):
                         complex = protein
-                        for id, chain in complex.get_chains().items():
+                        for chain_id, chain in complex.get_chains().items():
                             if isinstance(chain, Ligand):
-                                ligand_chain_ids.add(id)
+                                ligand_chain_ids.add(chain_id)
             for prop in props:
                 if hasattr(prop, "affinity") and prop.affinity is not None:
                     binder_id = prop.affinity.binder
@@ -214,7 +280,7 @@ class Boltz2Model(BoltzModel, FoldModel):
             step_scale=step_scale,
             use_potentials=use_potentials,
             constraints=constraints,
-            templates=templates,
+            templates=template_dicts or None,
             properties=properties,
             method=method,
         )
@@ -235,7 +301,7 @@ class Boltz1Model(BoltzModel, FoldModel):
         num_steps: int = 200,
         step_scale: float = 1.638,
         use_potentials: bool = False,
-        constraints: list[dict] | None = None,
+        constraints: Sequence[Mapping] | None = None,
     ) -> FoldResultFuture:
         """
         Request structure prediction with Boltz-1 model.
@@ -305,7 +371,7 @@ class Boltz1xModel(Boltz1Model, BoltzModel, FoldModel):
         num_recycles: int = 3,
         num_steps: int = 200,
         step_scale: float = 1.638,
-        constraints: list[dict] | None = None,
+        constraints: Sequence[Mapping] | None = None,
     ) -> FoldResultFuture:
         """
         Request structure prediction with Boltz-1x model. Uses potentials with Boltz-1 model.
@@ -516,3 +582,21 @@ class BoltzAffinity(BaseModel):
 
     class Config:
         extra = "allow"  # Allow extra fields
+
+
+def _msa_future_to_complex(session: APISession, msa: MSAFuture) -> Complex:
+    align_api = getattr(session, "align", None)
+    assert isinstance(align_api, AlignAPI)
+    seed = align_api.get_seed(job_id=msa.job.job_id)
+    proteins: dict[str, Protein] = {}
+    for chain_id, seq in zip(id_generator(), seed.split(":")):
+        protein = Protein(sequence=seq)
+        protein.msa = msa.id
+        proteins[chain_id] = protein
+    return Complex(chains=proteins)
+
+
+def _to_template(obj, chain_id: str | None = None):
+    if not isinstance(obj, Template):
+        obj = Template(template=obj, mapping=chain_id)
+    return obj

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/openprotein/fold/common.py

@@ -16,21 +16,20 @@ def normalize_inputs(
     normalized_complexes: list[Complex] = []
     remaining_proteins: list[Protein] = []
     remaining_protein_strings: list[str] = []
-    if isinstance(proteins, list):
-        for protein in proteins:
-            if isinstance(protein, Protein):
+    for protein in proteins:
+        if isinstance(protein, Protein):
+            # collate these to init with id_gen
+            remaining_proteins.append(protein)
+        elif isinstance(protein, Complex):
+            used_ids.extend(list(protein.get_chains().keys()))
+            normalized_complexes.append(protein)
+        else:
+            if isinstance(protein, bytes):
+                protein = protein.decode()
+            # handle ':'-delimited
+            for seq in protein.split(":"):
                 # collate these to init with id_gen
-                remaining_proteins.append(protein)
-            elif isinstance(protein, Complex):
-                used_ids.extend(list(protein.get_chains().keys()))
-                normalized_complexes.append(protein)
-            else:
-                if isinstance(protein, bytes):
-                    protein = protein.decode()
-                # handle ':'-delimited
-                for seq in protein.split(":"):
-                    # collate these to init with id_gen
-                    remaining_protein_strings.append(seq)
+                remaining_protein_strings.append(seq)
 
     # auto generate the chain ids
     id_gen = id_generator(used_ids)

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/openprotein/molecules/__init__.py

@@ -1,4 +1,5 @@
 from .chains import DNA, RNA, Ligand
 from .complex import Complex
-from .protein import Protein, Binding
+from .protein import Binding, Protein
 from .structure import Structure
+from .template import Template

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/openprotein/molecules/complex.py

@@ -4,19 +4,21 @@ from collections.abc import Mapping, Sequence
 from functools import reduce
 from pathlib import Path
 from types import MappingProxyType
-from typing import Literal, overload
+from typing import TYPE_CHECKING, Literal, overload
 
+import gemmi
 import numpy as np
 import numpy.typing as npt
 
-import gemmi
-
 import openprotein.utils.chain_id as _chain_id_utils
 import openprotein.utils.cif as _cif_utils
 
 from .chains import DNA, RNA, Ligand
 from .protein import Protein
 
+if TYPE_CHECKING:
+    from .template import Template
+
 
 # TODO: deserialization note about plddt parsed per residue
 class Complex:
@@ -26,6 +28,7 @@ class Complex:
         name: bytes | str | None = None,
     ):
         self._chains = dict(sorted(chains.items())) if chains is not None else {}
+        self._templates: "Sequence[Protein | Complex | Template]" = ()
         self.name = name
 
     @property
@@ -43,6 +46,24 @@ class Complex:
         self.name = x
         return self
 
+    @property
+    def templates(self) -> "Sequence[Protein | Complex | Template]":
+        """A list of templates for guiding the structure prediction of this molecular complex."""
+        return self._templates
+
+    @templates.setter
+    def templates(self, templates: "Sequence[Protein | Complex | Template]") -> None:
+        self._templates = tuple(templates)
+
+    def get_templates(self) -> "Sequence[Protein | Complex | Template]":
+        return self.templates
+
+    def set_templates(
+        self, templates: "Sequence[Protein | Complex | Template]"
+    ) -> "Complex":
+        self.templates = templates
+        return self
+
     def get_chains(self) -> Mapping[str, Protein | DNA | RNA | Ligand]:
         return MappingProxyType(self._chains)
 
@@ -246,6 +267,21 @@ class Complex:
             chains={k: v.copy() for k, v in self._chains.items()}, name=self._name
         )
 
+    def _assert_valid_templates(self):
+        from .template import Template
+
+        for template in self.templates:
+            (
+                template if isinstance(template, Template) else Template(template)
+            ).validate_for_target(self)
+        for chain_id, protein in self.get_proteins().items():
+            for template in protein.templates:
+                (
+                    template
+                    if isinstance(template, Template)
+                    else Template(template, mapping=chain_id)
+                ).validate_for_target(Complex({chain_id: protein}))
+
     @staticmethod
     def _from_structure_block(
         structure_block: _cif_utils.StructureCIFBlock,

{openprotein_python-0.10.1 → openprotein_python-0.11.1}/openprotein/molecules/protein.py

@@ -20,6 +20,7 @@ from .. import fasta
 if TYPE_CHECKING:
     from ..align.msa import MSAFuture
     from .complex import Complex
+    from .template import Template
 
 V = TypeVar("V")
 
@@ -70,6 +71,7 @@ class Protein:
         # sequence-level properties
         self._cyclic: bool = False
         self._msa: "str | MSAFuture | None | Type[Protein.NullMSA]" = None
+        self._templates: "Sequence[Protein | Complex | Template]" = ()
         # per-residue arrays
         self._data: dict[str, npt.NDArray] = {}
 
@@ -146,6 +148,24 @@ class Protein:
         self._msa = x
         return self
 
+    @property
+    def templates(self) -> "Sequence[Protein | Complex | Template]":
+        """A list of templates for guiding the structure prediction of this protein."""
+        return self._templates
+
+    @templates.setter
+    def templates(self, templates: "Sequence[Protein | Complex | Template]") -> None:
+        self._templates = tuple(templates)
+
+    def get_templates(self) -> "Sequence[Protein | Complex | Template]":
+        return self.templates
+
+    def set_templates(
+        self, templates: "Sequence[Protein | Complex | Template]"
+    ) -> "Protein":
+        self.templates = templates
+        return self
+
     def __len__(self):
         return len(self.sequence)
 
@@ -641,6 +661,14 @@ class Protein:
     def copy(self) -> "Protein":
         return self[:]
 
+    def _assert_valid_templates(self):
+        from .template import Template
+
+        for template in self.templates:
+            (
+                template if isinstance(template, Template) else Template(template)
+            ).validate_for_target(self)
+
     @staticmethod
     def _from_structure_block(
         structure_block: _cif_utils.StructureCIFBlock,

openprotein_python-0.11.1/openprotein/molecules/template.py

@@ -0,0 +1,122 @@
+import dataclasses
+from collections.abc import Mapping
+from types import NoneType
+
+from .complex import Complex
+from .protein import Protein
+
+TemplateSource = Protein | Complex
+TargetMolecule = Protein | Complex
+ChainMapping = Mapping[str, str] | str | None
+
+
+@dataclasses.dataclass(frozen=True)
+class Template:
+    """
+    A structural template used to guide the folding of a target chain or complex.
+
+    This class wraps a structural source (Protein or Complex) and defines how it
+    should map to the target(s).
+
+    Attributes:
+        template (Protein | Complex): The structural object to be used as a template.
+            Must contain structural data (coordinates).
+        mapping (Mapping[str, str] | str | None): The rule for assigning this template
+            to the target.
+            - Mapping[str, str]: Explicitly maps {template_chain_id: target_chain_id}.
+            - str: Apply this template to a specific target_chain_id. (If template is
+              a Complex, a selection algorithm is used to pick the best source chain).
+            - None: Automatic assignment. The folding algorithm will determine which
+              chain(s) this template applies to.
+    """
+
+    template: TemplateSource
+    mapping: ChainMapping = None
+
+    def __post_init__(self) -> None:
+        """Validates the template upon initialization."""
+        self._validate_self()
+
+    def _validate_self(self) -> None:
+        """Checks internal consistency of the Template."""
+        if isinstance(self.template, Protein):
+            if not self.template.has_structure:
+                raise ValueError("Provided template Protein has no structural data.")
+            # A single Protein object is treated as an atomic unit (anonymous chain).
+            # It cannot support a dictionary mapping because it has no internal Chain IDs
+            # to map *from*.
+            if not isinstance(self.mapping, (str, NoneType)):
+                raise ValueError(
+                    f"Invalid mapping type '{type(self.mapping)}' for Protein template. "
+                    "Expected 'str' (target ID) or 'None'. A dict mapping is only valid "
+                    "if the template is a Complex with named chains."
+                )
+        elif isinstance(self.template, Complex):
+            # Ensure all parts of the complex have structure
+            for chain_id, protein in self.template.get_proteins().items():
+                if not protein.has_structure:
+                    raise ValueError(
+                        f"Template Chain '{chain_id}' has no structural data."
+                    )
+            # If mapping is explicit (dict), ensure source keys exist in the template
+            if not isinstance(self.mapping, (str, NoneType)):
+                template_chains = set(self.template.get_chains().keys())
+                mapping_keys = set(self.mapping.keys())
+                if not mapping_keys.issubset(template_chains):
+                    missing = mapping_keys - template_chains
+                    raise ValueError(
+                        f"Mapping contains source chain IDs {missing} that do not "
+                        f"exist in the template complex (available: {template_chains})."
+                    )
+        else:
+            raise TypeError(
+                f"Template source must be Protein or Complex, got {type(self.template)}"
+            )
+
+    def validate_for_target(self, target: TargetMolecule) -> None:
+        """
+        Ensures this Template is compatible with a specific target Molecule.
+
+        Args:
+            target: The Protein or Complex that is being folded.
+
+        Raises:
+            ValueError: If this Template is invalid, or if chain IDs referenced in
+                        mapping do not exist in the target.
+            TypeError: If the template/target combination is structurally incompatible.
+        """
+        self._validate_self()
+        if isinstance(target, Protein):
+            # Target is a single Protein (implies anonymous/single context).
+            # We cannot map to a specific chain ID because the target Protein object
+            # doesn't have a chain ID.
+            if self.mapping is not None:
+                raise ValueError(
+                    "Cannot use a specific chain mapping when the target is a standalone Protein. "
+                    "Mapping must be None."
+                )
+        elif isinstance(target, Complex):
+            target_chains = {
+                chain_id
+                for chain_id, chain in target.get_chains().items()
+                if isinstance(self.template, Complex)
+                or isinstance(chain, type(self.template))
+            }
+            if isinstance(self.mapping, str):
+                # Mapping points to a specific target chain ID
+                if self.mapping not in target_chains:
+                    raise ValueError(
+                        f"Template maps to target chain '{self.mapping}', but this chain "
+                        f"does not exist in the target Complex (available: {target_chains})."
+                    )
+            elif self.mapping is not None:
+                # Mapping points from Source -> specific target chain IDs
+                target_values = set(self.mapping.values())
+                if not target_values.issubset(target_chains):
+                    missing = target_values - target_chains
+                    raise ValueError(
+                        f"Template maps to target chains {missing} which do not exist "
+                        f"in the target Complex (available: {target_chains})."
+                    )
+        else:
+            raise TypeError(f"Target must be Protein or Complex, got {type(target)}")