rc-foundry 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

foundry/utils/ddp.py

@@ -2,7 +2,7 @@ import logging
 
 import torch
 from beartype.typing import Any
-from lightning_fabric.utilities import rank_zero_only
+from lightning.fabric.utilities import rank_zero_only
 from lightning_utilities.core.rank_zero import rank_prefixed_message
 from omegaconf import DictConfig
 

foundry/utils/logging.py

@@ -4,7 +4,7 @@ from contextlib import contextmanager
 
 import pandas as pd
 from beartype.typing import Any
-from lightning_fabric.utilities import rank_zero_only
+from lightning.fabric.utilities import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
 from rich.console import Console
 from rich.syntax import Syntax

foundry/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.7'
-__version_tuple__ = version_tuple = (0, 1, 7)
+__version__ = version = '0.1.9'
+__version_tuple__ = version_tuple = (0, 1, 9)
 
 __commit_id__ = commit_id = None

{rc_foundry-0.1.7.dist-info → rc_foundry-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rc-foundry
-Version: 0.1.7
+Version: 0.1.9
 Summary: Shared utilities and training infrastructure for biomolecular structure prediction models.
 Author-email: Institute for Protein Design <contact@ipd.uw.edu>
 License: BSD 3-Clause License
@@ -97,11 +97,15 @@ Foundry provides tooling and infrastructure for using and training all classes o
 
 All models within Foundry rely on [AtomWorks](https://github.com/RosettaCommons/atomworks) - a unified framework for manipulating and processing biomolecular structures - for both training and inference. 
 
+
+> [!NOTE]
+> We have a slack now! Join for updates and to get your questions answered [here](https://join.slack.com/t/proteinmodelfoundry/shared_invite/zt-3kpwru8c6-nrmTW6LNHnSE7h16GNnfLA).
+
 ## Getting Started
 ### Quickstart guide
 **Installation**
 ```bash
-pip install rc-foundry[all]
+pip install "rc-foundry[all]"
 ```
 
 **Downloading weights** Models can be downloaded to a target folder with:

{rc_foundry-0.1.7.dist-info → rc_foundry-0.1.9.dist-info}/RECORD

@@ -1,7 +1,7 @@
 foundry/__init__.py,sha256=H8S1nl5v6YeW8ggn1jKy4GdtH7c-FGS-j7CqUCAEnAU,1926
 foundry/common.py,sha256=Aur8mH-CNmcUqSsw7VgaCQSW5sH1Bqf8Da91jzxPV1Y,3035
 foundry/constants.py,sha256=0n1wBKCvNuw3QaQehSbmsHYkIdaGn3tLeRFItBrdeHY,913
-foundry/version.py,sha256=szvPIs2C82UunpzuvVg3MbF4QhzbBYTsVJ8DmPfq6_E,704
+foundry/version.py,sha256=ib8ckvf-NNDfacXd8unW0p5cf-gl57XyQvjoEMc_pvc,704
 foundry/callbacks/__init__.py,sha256=VsRT1e4sqlJHPcTCsfupMEx82Iz-LoOAGPpwvf_OJeE,126
 foundry/callbacks/callback.py,sha256=xZBo_suP4bLrP6gl5uJPbaXm00DXigePa6dMeDxucgg,3890
 foundry/callbacks/health_logging.py,sha256=tEtkByOlaAA7nnelxb7PbM9_dcIgOsdbxCdQY3K5pMc,16664
@@ -25,9 +25,9 @@ foundry/training/schedulers.py,sha256=StmXegPfIdLAv31FreCTrDh9dsOvNUfzG4YGa61Y4o
 foundry/utils/alignment.py,sha256=2anqy0mn9zeFEiVWS_EG7zHiyPk1C_gbUu-SRvQ5mAM,2502
 foundry/utils/components.py,sha256=Piw2TfQF26uuxC3hXG3iv_4rgud1lVO-cv6N-p05EDY,15200
 foundry/utils/datasets.py,sha256=pLBxVezm-TSrYuC5gFnJZdGnNWV7aPH2QiWIVE2hkdQ,16629
-foundry/utils/ddp.py,sha256=ydHrO6peGbRnWAwgH5rmpHuQd55g2gFzzoZJYypn7GU,3970
+foundry/utils/ddp.py,sha256=202_7qqm4ihPjpB5Q9NhUjDl4u22pu5JvY0ui0UkRUQ,3970
 foundry/utils/instantiators.py,sha256=oGCp6hrmY-QPPPEjxKxe5uVFL125fH1RaLxjMKWCD_8,2169
-foundry/utils/logging.py,sha256=jrDgiB_56q_hWDc0jkBFekvqnNWcowJBt4B-S-ipJmM,9312
+foundry/utils/logging.py,sha256=ywV75MBlQsothV0IBvqoAQTNg6pjo2-Cib7Uo080nzQ,9312
 foundry/utils/rigid.py,sha256=_Z1pmitb6xgxyguLj_TukKscUBJjQsU4bsBD24GVS84,44444
 foundry/utils/rotation_augmentation.py,sha256=7q1WEX2iJ0i7-2aV-M97nEaEdpqexDTaZn5JquYpkUk,1927
 foundry/utils/squashfs.py,sha256=QlcwuJyVe-QVfIOS7o1QfLhaCQPNzzox7ln4n8dcYEg,5234
@@ -63,8 +63,8 @@ rf3/__init__.py,sha256=XBb5hF2RqBPHODGRmjiRbfTXgOGfOzdY91GbS4Vex00,70
 rf3/_version.py,sha256=fCfpbI5aeA6yHqjo3tK78-l2dPGxhp-AyKSoCXp34Nc,739
 rf3/alignment.py,sha256=BvvwMqQGCVxV20xIsTighD1kXMadXXL2SkckLjTerx0,2102
 rf3/chemical.py,sha256=VECnRPgVm-icXbZeUG4svcENzdUiIupP6dhka_8zCrg,26572
-rf3/cli.py,sha256=dPKJFRHYoV2XS6xc_ZmdLTz6frqa6HZg4qgZU5oJcXU,2356
-rf3/inference.py,sha256=_AAJ07AfSeU3xTM2_KH9n_H12EK4qZ23IJuyauOrMaQ,2466
+rf3/cli.py,sha256=jxjq8u77J8itIK4aNTfIpnMsNgg8brW1A3NfVjlgE0s,2743
+rf3/inference.py,sha256=yUjSqCjIqsGErUezq42-8K9570Opm_2eYa-bNnwltwA,2494
 rf3/kinematics.py,sha256=V3yjalPupu1X2FEp7l3XZR-qzLKrhWLZyECk6RgIkcs,10901
 rf3/scoring.py,sha256=dTllswE-6Fgli2eLiNzLFc2Rhz4ouDT4WL-sVbvLTGU,41541
 rf3/train.py,sha256=V4nqCC_1JKLI3WQ-nErNa8sqFpvb1mFhXSe6ZPpEheM,7945
@@ -119,18 +119,18 @@ rfd3/__init__.py,sha256=2Wto2IsUIj2lGag9m_gqgdCwBNl5p21-Xnr7W_RpU3c,348
 rfd3/callbacks.py,sha256=Zjt8RiaYWquoKOwRmC_wCUbRbov-V4zd2_73zjhgDHE,2783
 rfd3/cli.py,sha256=ka3K5H117fzDYIDXFpOpJV21w_XBrHYJZdFE0thsGBI,1644
 rfd3/constants.py,sha256=wLvDzrThpOrK8T3wGFNQeGrhAXOJQze8l3v_7pjIdMM,13141
-rfd3/engine.py,sha256=NwATrhYFyqT7C9Bie8mWtUiqqzXgs9x6nOCkmZYPiT4,21224
+rfd3/engine.py,sha256=viXzVGYkHPyzv5d0Ifg8zFZ6Wqg-U4I8Y4ldQLPe9x4,21536
 rfd3/run_inference.py,sha256=HfRMQ30_SAHfc-VFzBV52F-aLaNdG6PW8VkdMyB__wE,1264
 rfd3/train.py,sha256=rHswffIUhOae3_iYyvAiQ3jALoFuzrcRUgMlbJLinlI,7947
-rfd3/inference/datasets.py,sha256=u-2U7deHXu-iOs7doiKKynewP-NEyJfdORSTDzUSaQI,6538
-rfd3/inference/input_parsing.py,sha256=TyEzCzeCaNhuNi0RjMcq9fF2j3Sp36KbuZ1FUjlBTZ8,45442
+rfd3/inference/datasets.py,sha256=9VLbzl7dpG8mk_pjs0R5C2wFYUoRIgXXoZcS9IohSy0,6510
+rfd3/inference/input_parsing.py,sha256=pocqnnhE3-szeBbCL9gy9E3kZJSP_CGHXd6FFRxfv0c,46563
 rfd3/inference/legacy_input_parsing.py,sha256=G2XxkrjdIpL6i1YY7xEmkFitVv__Pc45ow6IKKPHw64,28855
 rfd3/inference/parsing.py,sha256=ktAMUuZE3Pe4bKAjjV3zjqcEDmGlMZ-cotIUhJsEQQA,5402
-rfd3/inference/symmetry/atom_array.py,sha256=HfFagFUB5yB-Y4IfUM5nuVGWHC5AEkyHqt0JcIqTQ_E,10922
-rfd3/inference/symmetry/checks.py,sha256=y-Kq0l5OhEmmxsPBBsMMB0qaAt18FeEicD3-jSMQFa0,9900
+rfd3/inference/symmetry/atom_array.py,sha256=j8yhtZAjRNm4d06KyS4gk2XCquHPCwR0k9WiNmxz7WA,12941
+rfd3/inference/symmetry/checks.py,sha256=ZWpC1JjrAjXY__xDt8EFYb5WUdSF0kobZyxxmacFU7U,10076
 rfd3/inference/symmetry/contigs.py,sha256=6OvbZ2dJg-a0mvvKAC0VkzUH5HpUDxOJvkByIst_roU,2127
-rfd3/inference/symmetry/frames.py,sha256=aEwkmlUsYexERX9hu09JMhisC8QTpHPVhfITbL80-EE,10819
-rfd3/inference/symmetry/symmetry_utils.py,sha256=p_PkxU3sw6gYGO2EmZTrbNQdLjz1mdTWEIl5MjQdIuY,14664
+rfd3/inference/symmetry/frames.py,sha256=kwog5jU_wgv6ACcUER2iU5qz-mdfdOe0cpi2OTEDKMU,18894
+rfd3/inference/symmetry/symmetry_utils.py,sha256=CGUzMI5CKVIcNi5_l2-YRu-ExroZX54GndxLb5P7RtY,14680
 rfd3/metrics/design_metrics.py,sha256=O1RqZdjQPNlAWYRg6UJTERYg_gUI1_hVleKsm9xbWBY,16836
 rfd3/metrics/hbonds_hbplus_metrics.py,sha256=Sewy9KzmrA1OnfkasN-fmWrQ9IRx9G7Yyhe2ua0mk28,11518
 rfd3/metrics/hbonds_metrics.py,sha256=SIR4BnDhYdpVSqwXXRYpQ_tB-M0_fVyugGl08WivCmE,15257
@@ -140,11 +140,11 @@ rfd3/metrics/sidechain_metrics.py,sha256=EGZuFuWQ0cCe83EVPAf4eysN8vP9ifNjfnmE0o5
 rfd3/model/RFD3.py,sha256=95aKzye-XzuDyLGgost-Wsfu8eT635zHIRky-pNoHSA,3569
 rfd3/model/RFD3_diffusion_module.py,sha256=BPjKGyQpbnqdzii3gXMKLhhijNqV8Xh4bSosmfDBt8w,12094
 rfd3/model/cfg_utils.py,sha256=XPBLyoB_bQRLmdrJ1Z0hCjcVvgUMGIPuw4rxTlHjB_s,2575
-rfd3/model/inference_sampler.py,sha256=qge7BNJttW0NXgerg3msPY3izxQ-6FsvWSTAMhZ4GJs,24696
+rfd3/model/inference_sampler.py,sha256=5k_UIkJCL6QO4BEqAxkR9GKxNAOz-NRq5Jh51Wm5MU8,25152
 rfd3/model/layers/attention.py,sha256=XuNA7WyFlRfLnAgky1PtGvXFCnDGv7GeEcXz8hodTBo,19472
-rfd3/model/layers/block_utils.py,sha256=EZq2qYUeO6_VCLKDVC60cxfBE_EPwvp84FPmqLr28ZQ,21197
+rfd3/model/layers/block_utils.py,sha256=oN0aD-vZiH4JbIFs2CzDmb2B74GNPKzdFurmGd-dirE,21244
 rfd3/model/layers/blocks.py,sha256=MOjJ53THxM2MMM27Ap7xiIXRCdI_SHzqKzLLQVX6FEc,24888
-rfd3/model/layers/chunked_pairwise.py,sha256=de5Qc3P7GEfZlX-QLaKfJxr6Ky5vgLcWWogatCw2UnY,14582
+rfd3/model/layers/chunked_pairwise.py,sha256=B8oUxXbdm9akwdbrmjA-7HtcNQBTppMB_gAMAz9VbvY,14712
 rfd3/model/layers/encoders.py,sha256=CqByjHNSbtMIaNP_h2iEJZdTbm-N8SGo1bZgvRNpMJ8,15207
 rfd3/model/layers/layer_utils.py,sha256=UPYo-DYa__93KONSEj2YZWLtBqvYNSA9_wHDDPhVrIc,5710
 rfd3/model/layers/pairformer_layers.py,sha256=uimskhN-Ec0apEXAU6JqomyKX5-6ormrEsCFJotkBtM,3991
@@ -166,11 +166,11 @@ rfd3/transforms/ncaa_transforms.py,sha256=Lz4L8OGuOOG53sKJHcLSdV7WPQ3YzOzwd5tJG4
 rfd3/transforms/pipelines.py,sha256=FGH-XH3taTWQ6k1zpDO_d-097EQdXmL6uqXZXw4HIMs,22086
 rfd3/transforms/ppi_transforms.py,sha256=7rXyf-tn2TLz6ybYR_YVDtSDG7hOgqhYY4shNviA_Sw,23493
 rfd3/transforms/rasa.py,sha256=a4IPFvVMMxldoGLyJQiSlGg7IyUkcBASbRZLWmguAKk,4156
-rfd3/transforms/symmetry.py,sha256=GSnMF7oAnUxPozfafsRuHEv0yKXW0BpLTI6wsKGZrbc,2658
+rfd3/transforms/symmetry.py,sha256=9I9gzAZkk5vMUJm7x8XCDSHtNPYYLAHt4meXxOczGT0,2970
 rfd3/transforms/training_conditions.py,sha256=UXiUPjDwrNKM95tRe0eXrMeRN8XlTPc_MXUvo6UpePo,19510
 rfd3/transforms/util_transforms.py,sha256=2AcLkzx-73ZFgcWD1cIHv7NyniRPI4_zThHK8azyQaY,18119
 rfd3/transforms/virtual_atoms.py,sha256=UpmxzPPd5FaJigcRoxgLSHHrLLOqsCvZ5PPZfQSGqII,12547
-rfd3/utils/inference.py,sha256=-8IKzkB9ulhLEJgapvnZSdIaIPQDPMpyPpHTQlFS7r0,27317
+rfd3/utils/inference.py,sha256=Yf3aAUk_YZi58uIJr5Y2wfVnQ-2bh3S5GHLBPzCRjUs,26448
 rfd3/utils/io.py,sha256=wbdjUTQkDc3RCSM7gdogA-XOKR68HeQ-cfvyN4pP90w,9849
 rfd3/utils/vizualize.py,sha256=HPlczrA3zkOuxV5X05eOvy_Oga9e3cPnFUXOEP4RR_g,11046
 rf3/configs/inference.yaml,sha256=JmEZdkAnbnOrX79lGS5xrYYho9aBFfVxfUp-8KjJV5I,309
@@ -304,8 +304,8 @@ rfd3/configs/trainer/rfd3_base.yaml,sha256=R3lZxdyjUirjlLU31qWlnZgHaz4GcWTGGIz4f
 rfd3/configs/trainer/loss/losses/diffusion_loss.yaml,sha256=FE4FCEfurE0ekwZ4YfS6wCvPSNqxClwg_kc73cPql5Y,323
 rfd3/configs/trainer/loss/losses/sequence_loss.yaml,sha256=kezbQcqwAZ0VKQPUBr2MsNr9DcDL3ENIP1i-j7h-6Co,64
 rfd3/configs/trainer/metrics/design_metrics.yaml,sha256=xVDpClhHqSHvsf-8StL26z51Vn-iuWMDG9KMB-kqOI0,719
-rc_foundry-0.1.7.dist-info/METADATA,sha256=zlvCxfZ5-Ow7WuGKskfW6P1DGhZB9OfLIIBUBGncFeQ,11309
-rc_foundry-0.1.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-rc_foundry-0.1.7.dist-info/entry_points.txt,sha256=BmiWCbWGtrd_lSOFMuCLBXyo84B7Nco-alj7hB0Yw9A,130
-rc_foundry-0.1.7.dist-info/licenses/LICENSE.md,sha256=NKtPCJ7QMysFmzeDg56ZfUStvgzbq5sOvRQv7_ddZOs,1533
-rc_foundry-0.1.7.dist-info/RECORD,,
+rc_foundry-0.1.9.dist-info/METADATA,sha256=SBHGrkr8RsLTCTOhJnbusBs8G7C8HxJUPajCbmU6OyE,11502
+rc_foundry-0.1.9.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rc_foundry-0.1.9.dist-info/entry_points.txt,sha256=BmiWCbWGtrd_lSOFMuCLBXyo84B7Nco-alj7hB0Yw9A,130
+rc_foundry-0.1.9.dist-info/licenses/LICENSE.md,sha256=NKtPCJ7QMysFmzeDg56ZfUStvgzbq5sOvRQv7_ddZOs,1533
+rc_foundry-0.1.9.dist-info/RECORD,,

rf3/cli.py

@@ -23,10 +23,19 @@ def fold(
         configure_minimal_inference_logging()
 
     # Find the RF3 configs directory relative to this file
-    # This file is at: models/rf3/src/rf3/cli.py
-    # Configs are at: models/rf3/configs/
-    rf3_package_dir = Path(__file__).parent.parent.parent  # Go up to models/rf3/
-    config_path = str(rf3_package_dir / "configs")
+    # In development: models/rf3/src/rf3/cli.py -> models/rf3/configs/
+    # When installed: site-packages/rf3/cli.py -> site-packages/rf3/configs/
+    rf3_file_dir = Path(__file__).parent
+
+    # Check if we're in installed mode (configs are sibling to this file)
+    # or development mode (configs are ../../../configs)
+    if (rf3_file_dir / "configs").exists():
+        # Installed mode
+        config_path = str(rf3_file_dir / "configs")
+    else:
+        # Development mode
+        rf3_package_dir = rf3_file_dir.parent.parent  # Go up to models/rf3/
+        config_path = str(rf3_package_dir / "configs")
 
     # Get all arguments
     args = ctx.params.get("args", []) + ctx.args

rf3/inference.py

@@ -16,7 +16,9 @@ rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
 
 load_dotenv(override=True)
 
-_config_path = os.path.join(os.environ["PROJECT_ROOT"], "models/rf3/configs")
+_config_path = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "configs"
+)
 
 
 @hydra.main(

rfd3/engine.py

@@ -21,11 +21,13 @@ from rfd3.constants import SAVED_CONDITIONING_ANNOTATIONS
 from rfd3.inference.datasets import (
     assemble_distributed_inference_loader_from_json,
 )
-from rfd3.inference.input_parsing import DesignInputSpecification
+from rfd3.inference.input_parsing import (
+    DesignInputSpecification,
+    ensure_input_is_abspath,
+)
 from rfd3.model.inference_sampler import SampleDiffusionConfig
 from rfd3.utils.inference import (
     ensure_inference_sampler_matches_design_spec,
-    ensure_input_is_abspath,
 )
 from rfd3.utils.io import (
     CIF_LIKE_EXTENSIONS,
@@ -391,7 +393,13 @@ class RFD3InferenceEngine(BaseInferenceEngine):
         design_specifications = {}
         for prefix, example_spec in inputs.items():
             # Record task name in the specification
-            example_spec["extra"]["task_name"] = prefix
+            if isinstance(example_spec, DesignInputSpecification):
+                example_spec.extra = example_spec.extra or {}
+                example_spec.extra["task_name"] = prefix
+            else:
+                if "extra" not in example_spec:
+                    example_spec["extra"] = {}
+                example_spec["extra"]["task_name"] = prefix
 
             # ... Create n_batches for example
             for batch_id in range((n_batches) if exists(n_batches) else 1):

rfd3/inference/datasets.py

@@ -14,8 +14,8 @@ from atomworks.ml.transforms.base import Compose, Transform
 from omegaconf import DictConfig, OmegaConf
 from rfd3.inference.input_parsing import (
     DesignInputSpecification,
+    ensure_input_is_abspath,
 )
-from rfd3.utils.inference import ensure_input_is_abspath
 from torch.utils.data import (
     DataLoader,
     SequentialSampler,

rfd3/inference/input_parsing.py

@@ -5,6 +5,7 @@ import os
 import time
 import warnings
 from contextlib import contextmanager
+from os import PathLike
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
@@ -1121,3 +1122,33 @@ def accumulate_components(
     if atom_array_accum.bonds is None:
         atom_array_accum.bonds = BondList(atom_array_accum.array_length())
     return atom_array_accum
+
+
+def ensure_input_is_abspath(args: Dict[str, Any], path: PathLike | None):
+    """
+    Ensures the input source is an absolute path if exists, if not it will convert
+
+    args:
+        args: Inference specification for atom array
+        path: None or file to which the input is relative to.
+    """
+    if isinstance(args, str):
+        raise ValueError(
+            "Expected args to be a dictionary, got a string: {}. If you are using an input JSON ensure it contains dictionaries of arguments".format(
+                args
+            )
+        )
+    if "input" not in args or not exists(args["input"]):
+        return args
+    input = str(args["input"])
+    if not os.path.isabs(input):
+        if path is None:
+            raise ValueError(
+                "Input path is relative, but no base path was provided to resolve it against."
+            )
+        input = os.path.abspath(os.path.join(os.path.dirname(str(path)), input))
+        logger.info(
+            f"Input source path is relative, converted to absolute path: {input}"
+        )
+        args["input"] = input
+    return args

rfd3/inference/symmetry/atom_array.py

@@ -1,3 +1,5 @@
+import string
+
 import numpy as np
 from rfd3.inference.symmetry.frames import (
     decompose_symmetry_frame,
@@ -7,6 +9,68 @@ from rfd3.inference.symmetry.frames import (
 FIXED_TRANSFORM_ID = -1
 FIXED_ENTITY_ID = -1
 
+# Alphabet for chain ID generation (uppercase letters only, per wwPDB convention)
+_CHAIN_ALPHABET = string.ascii_uppercase
+
+
+def index_to_chain_id(index: int) -> str:
+    """
+    Convert a zero-based index to a chain ID following wwPDB convention.
+
+    The naming follows the wwPDB-assigned chain ID system:
+    - 0-25: A-Z (single letter)
+    - 26-701: AA-ZZ (double letter)
+    - 702-18277: AAA-ZZZ (triple letter)
+    - And so on...
+
+    This is similar to Excel column naming (A, B, ..., Z, AA, AB, ...).
+
+    Arguments:
+        index: zero-based index (0 -> 'A', 25 -> 'Z', 26 -> 'AA', etc.)
+    Returns:
+        chain_id: string chain identifier
+    """
+    if index < 0:
+        raise ValueError(f"Chain index must be non-negative, got {index}")
+
+    result = ""
+    remaining = index
+
+    # Convert to bijective base-26 (like Excel columns)
+    while True:
+        result = _CHAIN_ALPHABET[remaining % 26] + result
+        remaining = remaining // 26 - 1
+        if remaining < 0:
+            break
+
+    return result
+
+
+def chain_id_to_index(chain_id: str) -> int:
+    """
+    Convert a chain ID back to a zero-based index.
+
+    Inverse of index_to_chain_id.
+
+    Arguments:
+        chain_id: string chain identifier (e.g., 'A', 'Z', 'AA', 'AB')
+    Returns:
+        index: zero-based index
+    """
+    if not chain_id or not all(c in _CHAIN_ALPHABET for c in chain_id):
+        raise ValueError(f"Invalid chain ID: {chain_id}")
+
+    # Offset for all shorter chain IDs (26 + 26^2 + ... + 26^(len-1))
+    offset = sum(26**k for k in range(1, len(chain_id)))
+
+    # Value within the current length group (standard base-26)
+    value = 0
+    for char in chain_id:
+        value = value * 26 + _CHAIN_ALPHABET.index(char)
+
+    return offset + value
+
+
 ########################################################
 # Symmetry annotations
 ########################################################
@@ -247,11 +311,13 @@ def reset_chain_ids(atom_array, start_id):
     Reset the chain ids and pn_unit_iids of an atom array to start from the given id.
     Arguments:
         atom_array: atom array with chain_ids and pn_unit_iids annotated
+        start_id: starting chain ID (e.g., 'A')
     """
     chain_ids = np.unique(atom_array.chain_id)
-    new_chain_range = range(ord(start_id), ord(start_id) + len(chain_ids))
-    for new_id, old_id in zip(new_chain_range, chain_ids):
-        atom_array.chain_id[atom_array.chain_id == old_id] = chr(new_id)
+    start_index = chain_id_to_index(start_id)
+    for i, old_id in enumerate(chain_ids):
+        new_id = index_to_chain_id(start_index + i)
+        atom_array.chain_id[atom_array.chain_id == old_id] = new_id
     atom_array.pn_unit_iid = atom_array.chain_id
     return atom_array
 
@@ -259,15 +325,18 @@ def reset_chain_ids(atom_array, start_id):
 def reannotate_chain_ids(atom_array, offset, multiplier=0):
     """
     Reannotate the chain ids and pn_unit_iids of an atom array.
+
+    Uses wwPDB-style chain IDs (A-Z, AA-ZZ, AAA-ZZZ, ...) to support
+    any number of chains.
+
     Arguments:
         atom_array: protein atom array with chain_ids and pn_unit_iids annotated
-        offset: offset to add to the chain ids
-        multiplier: multiplier to add to the chain ids
+        offset: offset to add to the chain ids (typically num_chains in ASU)
+        multiplier: multiplier for the offset (typically transform index)
     """
-    chain_ids_int = (
-        np.array([ord(c) for c in atom_array.chain_id]) + offset * multiplier
-    )
-    chain_ids = np.array([chr(id) for id in chain_ids_int], dtype=str)
+    chain_ids_indices = np.array([chain_id_to_index(c) for c in atom_array.chain_id])
+    new_indices = chain_ids_indices + offset * multiplier
+    chain_ids = np.array([index_to_chain_id(idx) for idx in new_indices], dtype="U4")
     atom_array.chain_id = chain_ids
     atom_array.pn_unit_iid = chain_ids
     return atom_array

rfd3/inference/symmetry/checks.py

@@ -24,7 +24,16 @@ def check_symmetry_config(
     assert sym_conf.id, "symmetry_id is required. e.g. {'id': 'C2'}"
     # if unsym motif is provided, check that each motif name is in the atom array
 
+    is_motif_atom = get_motif_features(atom_array)["is_motif_atom"]
     is_unsym_motif = np.zeros(atom_array.shape[0], dtype=bool)
+
+    if not is_motif_atom.any():
+        sym_conf.is_symmetric_motif = None
+        ranked_logger.warning(
+            "No motifs found in atom array. Setting is_symmetric_motif to None."
+        )
+        return sym_conf
+
     if sym_conf.is_unsym_motif:
         assert (
             src_atom_array is not None
@@ -36,21 +45,20 @@ def check_symmetry_config(
             if (sm and n not in sm.split(",")) and (n not in atom_array.src_component):
                 raise ValueError(f"Unsym motif {n} not found in atom_array")
 
-    is_motif_token = get_motif_features(atom_array)["is_motif_token"]
     if (
-        is_motif_token[~is_unsym_motif].any()
+        is_motif_atom[~is_unsym_motif].any()
         and not sym_conf.is_symmetric_motif
         and not has_dist_cond
     ):
         raise ValueError(
-            "Asymmetric motif inputs should be distance constrained."
-            "Use atomwise_fixed_dist to constrain the distance between the motif atoms."
+            "Asymmetric motif inputs are not supported yet. Please provide a symmetric motif."
         )
 
     if partial and not sym_conf.is_symmetric_motif:
         raise ValueError(
             "Partial diffusion with symmetry is only supported for symmetric inputs."
         )
+    return sym_conf
 
 
 def check_atom_array_is_symmetric(atom_array):