> For the complete documentation index, see [llms.txt](https://jwliaomath.gitbook.io/cocofold2/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://jwliaomath.gitbook.io/cocofold2/refinement-workflows/component_parallel_tutorial.md). # Preparing component caches For the concrete two-GPU **6ZBH A + BCD (1+3)** workflow, start with [the complete example](/cocofold2/refinement-workflows/6zbh_parallel.md). It includes Contextual/Independent manifest templates, explicit BCD chain-name mapping, full-STAR refinement, automatic merged CIF output and CPU result checking. The sections below describe the general cache preparation and K-rank workflow. This tutorial describes the component-parallel particle-refinement workflow implemented in [`src/chain_parallel`](https://github.com/jwliaomath/CoCoFold2/tree/main/src/chain_parallel/README.md). It covers: 1. **Independent-component refinement**, in which Protenix caches are generated separately for user-defined component groups; 2. **Contextual-component refinement**, in which one full-complex Pairformer cache is divided into contextual component caches; and 3. a memory-conscious contextual path that saves full-complex `z_trunk`, splits it first, and then materializes component-local `pair_z` caches. The implementation is not restricted to two GPUs. With the current manifest format, a run uses one component group per process/rank. A manifest containing `K` component groups is therefore launched with `K` processes, normally on `K` GPUs. A component group may contain one chain or several chains. ## Scope and approximation boundary Each rank owns: * one frozen Protenix-v1 diffusion decoder; * one component conditioning cache; * one component-specific latent perturbation (`z_bias`); and * component-local Gaussian renderer parameters. For each particle, all ranks render their components in the same experimental coordinate frame. With the default fixed-frame projection, every rank uses the same explicitly supplied 3D origin; the rendered component images are summed. The historical legacy mode instead derives a shared per-view origin from the assembled model. Particle sign, translation, CTF and FRC loss are then applied to the assembled projection, allowing one particle-space objective to update every component. The two cache strategies differ before this shared objective: | Strategy | Pairformer calculation | Information retained in each component cache | | --------------------- | --------------------------------------- | ------------------------------------------------------------------------ | | Independent-component | Run separately for each component group | Component-only sequence context | | Contextual-component | Run once for the complete assembly | Full-complex contextual single rows and the selected pair diagonal block | Neither strategy is exact full-complex Protenix refinement. Component-local diffusion omits cross-component diffusion attention and joint coordinate generation. Contextual-component refinement also requires the full Pairformer and its full pair representation to fit once. ## 1. Installation and repository layout Follow the [installation instructions](/cocofold2/getting-started/installation.md) and run the commands below from the repository root: ```bash git clone https://github.com/jwliaomath/CoCoFold2.git cd CoCoFold2 conda env create -f environment.yml conda activate cocofold2 ``` The compatible Protenix checkpoint and common resources must be available as described in the main README. They can live outside the source checkout: ``` CoCoFold2/ ├── src/ │ ├── inference.py │ ├── get_pdb.py │ ├── train.py │ └── chain_parallel/ │ ├── train_chain_parallel_2d.py │ ├── manifest.py │ ├── distributed_gmm.py │ ├── contextual_cache.py │ ├── prepare_contextual_diffusion_caches.py │ └── materialize_local_diffusion_cache.py └── docs/ /path/to/protenix_resources/ ├── checkpoint/ └── common/ ``` Configure the Python paths once per shell: ```bash export REPO_ROOT="$(pwd)" export PROTENIX_ROOT_DIR=/absolute/path/to/protenix_resources export COCOFOLD2_ROOT="${REPO_ROOT}/src" export CHAIN_PARALLEL_DIR="${REPO_ROOT}/src/chain_parallel" export PYTHONPATH="${COCOFOLD2_ROOT}:${CHAIN_PARALLEL_DIR}:${PYTHONPATH:-}" ``` `COCOFOLD2_ROOT` must directly contain `ctf.py`, `particledataset.py`, `pts2img.py`, `utils.py` and `model/protenix.py`. `PROTENIX_ROOT_DIR` must contain the compatible `checkpoint/` and `common/` directories; it is not an output location. ## 2. Required inputs In addition to the standard particle-workflow inputs described in [Data requirements](/cocofold2/getting-started/data_requirements.md), prepare: * a chain-to-component partition; * one diffusion cache per component group; * one component topology file with exactly the same atom order as its cache; * one fitted initial structure per component group; and * one component manifest assigning component groups to ranks. All fitted component structures must be placed in the same coordinate frame defined by the experimental reconstruction and upstream particle poses. Fit the sequence-derived initial predictions, not the deposited evaluation structure. The STAR file, MRC/MRCS particle stacks, particle subset, pose/CTF metadata, box size, pixel size and loss settings must be identical across component strategies used in a controlled comparison. ## 3. Select component groups and GPU count Partition the assembly into `K` complete-chain groups. The current implementation requires: ``` number of manifest entries = torchrun world size = number of ranks ``` and each rank may occur only once in the manifest. For example, a three-GPU partition could be: ``` rank 0: chain A rank 1: chains B+C rank 2: chains D+E+F ``` Choose groups according to available GPU memory and component sizes. If there are more chains than GPUs, place multiple chains in one component group rather than assigning multiple manifest entries to one rank. The manuscript examples use two groups (6ZBH: A versus B--D; 6O77: A--B versus C--D), but these are example partitions rather than a two-GPU software limit. ## 4. Independent-component refinement ### 4.1 Generate one Protenix cache per component group Create one Protenix input JSON for each group. Each JSON should contain only the chains assigned to that group, while using the same checkpoint and the same MSA/template policy. For a three-group example: ```bash mkdir -p runs/target/independent/cache runs/target/independent/protenix python -u src/inference.py \ --resource-root "$PROTENIX_ROOT_DIR" \ --input_json_path inputs/target_A.json \ --sample_name target_A \ --output_model_dir runs/target/independent/cache/ \ --dump_dir runs/target/independent/protenix/A python -u src/inference.py \ --resource-root "$PROTENIX_ROOT_DIR" \ --input_json_path inputs/target_BC.json \ --sample_name target_BC \ --output_model_dir runs/target/independent/cache/ \ --dump_dir runs/target/independent/protenix/BC python -u src/inference.py \ --resource-root "$PROTENIX_ROOT_DIR" \ --input_json_path inputs/target_DEF.json \ --sample_name target_DEF \ --output_model_dir runs/target/independent/cache/ \ --dump_dir runs/target/independent/protenix/DEF ``` `--output_model_dir` is the cache **directory** and `--dump_dir` is the separate Protenix prediction/summary **directory**. Use new destinations: inference refuses to overwrite existing caches. For one target and one seed, the compatible filenames include: ``` runs/target/independent/cache/target_A_diffusion_data.pth runs/target/independent/cache/target_BC_diffusion_data.pth runs/target/independent/cache/target_DEF_diffusion_data.pth ``` ### 4.2 Generate and rigidly place each initial component For a compatible cache, `src/get_pdb.py` reconstructs topology directly from its saved features: ```bash python -u src/get_pdb.py \ --pdbid target_A \ --diffusion_data_dir runs/target/independent/cache/target_A_diffusion_data.pth \ --out_dir runs/target/independent/initial/A \ --output-format cif \ --device cuda:0 ``` This writes `target_A_initial_prediction.cif` and `target_A_initial_prediction_topology.json` inside `initial/A/`. If an older cache cannot supply atom identities, add `--cif_path` pointing to the **matching Protenix output** as a topology template, never to the deposited evaluation structure. The component trainer still requires the separately fitted CIFs in the manifest. Repeat for every component. Rigidly fit the generated component structures into one common experimental frame and preserve their atom ordering. The fitted files, not deposited reference structures, are supplied to the component trainer. ### 4.3 Create the component manifest Save a YAML file such as `runs/target/independent/components.yaml`: ```yaml schema_version: 1 components: - id: target_A_independent rank: 0 diffusion_data_dir: cache/target_A_diffusion_data.pth cif_path: fitted/target_A_fitted.cif - id: target_BC_independent rank: 1 diffusion_data_dir: cache/target_BC_diffusion_data.pth cif_path: fitted/target_BC_fitted.cif - id: target_DEF_independent rank: 2 diffusion_data_dir: cache/target_DEF_diffusion_data.pth cif_path: fitted/target_DEF_fitted.cif ``` Relative paths are resolved from the manifest directory. Component IDs and ranks must be unique, and ranks must cover `0` through `K-1`. ### 4.4 Launch on `K` GPUs Set the number of processes to the number of component groups: ```bash NUM_COMPONENTS=3 OUTPUT_PREFIX=runs/target/independent/results/model_ mkdir -p "$(dirname "${OUTPUT_PREFIX}")" torchrun --standalone --nproc_per_node="${NUM_COMPONENTS}" \ src/chain_parallel/train_chain_parallel_2d.py \ --component_manifest runs/target/independent/components.yaml \ --star_data_dir /path/to/particles.star \ --mrc_data_dir /path/to/particle/stacks/ \ --output_trained_model_dir "${OUTPUT_PREFIX}" \ --backend nccl \ --boxsize REPLACE_WITH_BOX_SIZE \ --apix REPLACE_WITH_PIXEL_SIZE \ --projection-frame fixed \ --projection-origin REPLACE_WITH_X_A REPLACE_WITH_Y_A REPLACE_WITH_Z_A \ --resolution REPLACE_WITH_GMM_RESOLUTION \ --map_resolution REPLACE_WITH_FRC_CUTOFF_RESOLUTION \ --batch_size REPLACE_WITH_OUTER_BATCH_SIZE \ --mini_batch_size REPLACE_WITH_MICROBATCH_SIZE \ --particle_sign -1 \ --train_deterministic \ --update_affine_mat ``` Add `--transR` only when required by the validated upstream orientation convention. Add `--update_affine_mat` only when the flip safeguard used in the registered experiment is intended. The trainer's `--output_trained_model_dir` is a **filename prefix**, not the cache directory from Section 4.1. For `.../results/model_`, epoch 1 writes component files such as `model_target_A_independent_rank0_1.cif/.pth`, plus `model_merged_1.cif` and the complete-epoch index `model_epoch_1.json` in the same `results/` directory. The component ID comes from the manifest. Choose a new prefix for a new run; existing outputs are protected. Absolute particle-stack paths in `rlnImageName` are used directly. Relative paths use `--mrc_data_dir` when supplied, otherwise the STAR file's directory. A trailing slash is not required. ## 5. Contextual-component refinement Contextual-component refinement uses one complete-assembly Pairformer cache and extracts one contextual cache for every component group. ### 5.1 Generate a full-complex cache For assemblies that can construct the standard full cache: ```bash mkdir -p runs/target/contextual/full_cache runs/target/contextual/protenix python -u src/inference.py \ --resource-root "$PROTENIX_ROOT_DIR" \ --input_json_path inputs/target_full_complex.json \ --sample_name target_full \ --output_model_dir runs/target/contextual/full_cache/ \ --dump_dir runs/target/contextual/protenix ``` The expected cache is: ``` runs/target/contextual/full_cache/target_full_diffusion_data.pth ``` If generating full-complex shared diffusion variables is impractical, use the `z_trunk` workflow in Section 6 instead. ### 5.2 Inspect numeric `asym_id` assignments ```bash python src/chain_parallel/prepare_contextual_diffusion_caches.py \ inspect \ --cache runs/target/contextual/full_cache/target_full_diffusion_data.pth \ --output-json runs/target/contextual/full_cache/inspection.json ``` Use the reported Protenix `asym_id`, token count and atom count to define the split. Do not infer `asym_id` solely from the visible CIF chain label, particularly for repeated subunits. ### 5.3 Define an arbitrary `K`-component contextual split Create `runs/target/contextual/split.yaml`: ```yaml schema_version: 1 source_cache: full_cache/target_full_diffusion_data.pth require_complete_partition: true expected_source_n_token: REPLACE_WITH_INSPECTED_INTEGER expected_source_n_atom: REPLACE_WITH_INSPECTED_INTEGER components: - id: target_A_contextual asym_ids: [REPLACE_WITH_ASYM_ID] expected_n_token: REPLACE_WITH_INSPECTED_INTEGER expected_n_atom: REPLACE_WITH_INSPECTED_INTEGER output_cache: cache/target_A_contextual.pth - id: target_BC_contextual asym_ids: [REPLACE_WITH_ASYM_IDS] expected_n_token: REPLACE_WITH_INSPECTED_INTEGER expected_n_atom: REPLACE_WITH_INSPECTED_INTEGER output_cache: cache/target_BC_contextual.pth - id: target_DEF_contextual asym_ids: [REPLACE_WITH_ASYM_IDS] expected_n_token: REPLACE_WITH_INSPECTED_INTEGER expected_n_atom: REPLACE_WITH_INSPECTED_INTEGER output_cache: cache/target_DEF_contextual.pth report_json: cache/contextual_split_report.json ``` `require_complete_partition: true` rejects missing, unknown, overlapping or duplicated `asym_id` assignments. ### 5.4 Split and rebuild component-local atom features ```bash python src/chain_parallel/prepare_contextual_diffusion_caches.py \ split \ --spec runs/target/contextual/split.yaml \ --device cuda:0 ``` For each group, the utility: * selects the corresponding `s_inputs` and `s_trunk` rows; * extracts the exact diagonal block of full-complex `z_trunk` or `pair_z`; * selects component atoms and remaps atom-to-token indices; * rebuilds `d_lm`, `v_lm` and `pad_info`; and * rebuilds the atom shared cache when the source already contains `pair_z`. The source cache is not overwritten. Output caches are saved on CPU for portability and record the source-cache SHA-256. ### 5.5 Generate, place and register contextual components Run `src/get_pdb.py` for every contextual component cache, then rigidly place the generated structures in the same experimental frame. The template-free command in Section 4.2 also applies to compatible contextual caches; add a matching Protenix topology template only if cache atom identities are insufficient. Create a training manifest with one contextual cache and fitted CIF per rank: ```yaml schema_version: 1 components: - id: target_A_contextual rank: 0 diffusion_data_dir: cache/target_A_contextual.pth cif_path: fitted/target_A_contextual_fitted.cif - id: target_BC_contextual rank: 1 diffusion_data_dir: cache/target_BC_contextual.pth cif_path: fitted/target_BC_contextual_fitted.cif - id: target_DEF_contextual rank: 2 diffusion_data_dir: cache/target_DEF_contextual.pth cif_path: fitted/target_DEF_contextual_fitted.cif ``` Do not mix independent and contextual caches in one run. The trainer also requires all contextual caches to record the same full source-cache hash. Launch the contextual run with the same `torchrun` command used in Section 4.4, changing only the manifest and output prefix. For a controlled comparison, keep the component partition, particles, checkpoint, frame placement policy, diffusion settings, batch sizes and checkpoint-selection rule matched. ## 6. Large-complex contextual cache: `z_trunk` first, local `pair_z` later The standard shared-cache path constructs full-complex `pair_z`, `p_lm` and `c_l` during inference. For a large assembly, these additional full-complex objects may be undesirable even when the full Pairformer and `z_trunk` can be computed once. The alternative workflow is: ``` full-complex input -> full Pairformer and full z_trunk -> split contextual z_trunk diagonal blocks -> materialize pair_z, p_lm and c_l separately for each smaller component -> component-parallel refinement ``` The order is important: split the full `z_trunk` before creating `pair_z`. ### 6.1 Save full-complex `z_trunk` Pass the Protenix configuration override as an explicit key-value pair: ```bash mkdir -p runs/large_target/contextual/full_cache python -u src/inference.py \ --resource-root "$PROTENIX_ROOT_DIR" \ --input_json_path inputs/large_target_full.json \ --sample_name large_target_full \ --output_model_dir runs/large_target/contextual/full_cache/ \ --dump_dir runs/large_target/contextual/protenix \ --enable_diffusion_shared_vars_cache false ``` Confirm that the log reports: ``` enable_diffusion_shared_vars_cache False shared_vars_cache=False ``` The saved cache should contain `z_trunk`, with `pair_z`, `p_lm` and `c_l` set to `None`. Inspect this before continuing: ```bash export FULL_CACHE=runs/large_target/contextual/full_cache/large_target_full_diffusion_data.pth python -c "import os, torch; d=torch.load(os.environ['FULL_CACHE'], map_location='cpu', weights_only=False); print({'z_trunk': None if d['z_trunk'] is None else tuple(d['z_trunk'].shape), 'pair_z': d['pair_z'], 'p_lm': d['p_lm'], 'c_l': d['c_l']})" ``` In the current source, the diffusion cache is written before the final full-complex coordinate-sampling call. If that later call fails, retain the error log and validate the cache explicitly; a saved cache does not imply that full-complex coordinate generation succeeded. This pathway still requires the full Pairformer and full `z_trunk` to fit. It does not solve the Pairformer memory boundary. ### 6.2 Inspect and split the `z_trunk` cache Use `prepare_contextual_diffusion_caches.py inspect` and `split` exactly as in Section 5. After splitting, every component cache should contain local `z_trunk` while `pair_z`, `p_lm` and `c_l` remain `None`. ### 6.3 Materialize a local shared diffusion cache for each component Run the following command once per component output: ```bash python src/chain_parallel/materialize_local_diffusion_cache.py \ --input-cache runs/large_target/contextual/cache/group_A_z_trunk.pth \ --output-cache runs/large_target/contextual/cache/group_A_pair_z.pth \ --device cuda:0 \ --report-json runs/large_target/contextual/cache/group_A_materialization.json ``` Repeat for all `K` components. The input files are not modified. Each output contains component-local `pair_z`, `p_lm` and `c_l`, sets `z_trunk` to `None`, and records source/output hashes, tensor shapes, finite-value checks, wall time and peak GPU memory for the materialization step. Use the materialized `*_pair_z.pth` files in the final component manifest. ### 6.4 Record the optimized representation level The trainer supports both cache forms, but their latent updates occur at different representation levels: * for a `z_trunk` cache, `z_bias` is added to `z_trunk` before diffusion conditioning is constructed; * for a materialized cache, `z_bias` is added directly to cached `pair_z`. Consequently, local `z_trunk -> pair_z` materialization is not merely a file format conversion. It selects the cached-`pair_z` optimization path. Record the active conditioning tensor in each experiment and do not describe a `z_trunk` and a `pair_z` run as optimizing identical variables. If `z_trunk`-level optimization is required, skip materialization and provide the split `z_trunk` component caches directly to the trainer. The component diffusion-conditioning cache will then be constructed during sampling. ## 7. Generic Slurm launch Request one GPU per component group and launch one `torchrun` process per GPU. A cluster-neutral skeleton is: ```bash #!/usr/bin/env bash #SBATCH --nodes=1 #SBATCH --ntasks=1 #SBATCH --gres=gpu:K #SBATCH --cpus-per-task=4 #SBATCH --output=logs/%x-%j.out #SBATCH --error=logs/%x-%j.err set -euo pipefail # Activate the environment using the commands required by your cluster. conda activate cocofold2 cd /path/to/CoCoFold2 export COCOFOLD2_ROOT="${PWD}/src" export CHAIN_PARALLEL_DIR="${PWD}/src/chain_parallel" export PYTHONPATH="${COCOFOLD2_ROOT}:${CHAIN_PARALLEL_DIR}:${PYTHONPATH:-}" NUM_COMPONENTS=K mkdir -p logs runs/target/results python -c "import torch; expected=${NUM_COMPONENTS}; assert torch.cuda.device_count() == expected, f'expected {expected} visible GPUs, got {torch.cuda.device_count()}'; assert torch.distributed.is_nccl_available()" srun --ntasks=1 \ torchrun --standalone --nproc_per_node="${NUM_COMPONENTS}" \ src/chain_parallel/train_chain_parallel_2d.py \ --component_manifest runs/target/components.yaml \ --star_data_dir /path/to/particles.star \ --mrc_data_dir /path/to/particle/stacks/ \ --output_trained_model_dir runs/target/results/model_ \ --backend nccl \ --boxsize REPLACE_WITH_BOX_SIZE \ --apix REPLACE_WITH_PIXEL_SIZE \ --projection-frame fixed \ --projection-origin REPLACE_WITH_X_A REPLACE_WITH_Y_A REPLACE_WITH_Z_A \ --resolution REPLACE_WITH_GMM_RESOLUTION \ --map_resolution REPLACE_WITH_FRC_CUTOFF_RESOLUTION \ --batch_size REPLACE_WITH_OUTER_BATCH_SIZE \ --mini_batch_size REPLACE_WITH_MICROBATCH_SIZE \ --particle_sign -1 \ --train_deterministic \ --update_affine_mat ``` Replace `K`, all paths and all data-dependent parameters. Add the partition, account, time and memory directives required by the local scheduler. The command above uses the new fixed-frame default. Replace the three origin placeholders with a point in the same experimental map frame as **every** placed component CIF. All ranks use that one origin; the recorded particle translation is applied after their rendered component images are summed. For the checked zero-origin, standard-axis 6ZBH-style 288-pixel map at 1.073 Å/pixel, use `154.512 154.512 154.512` after verifying the map and CIF coordinates. For earlier experiments, request `--projection-frame legacy` and omit the origin. See [choosing the projection origin](/cocofold2/reference/parameter_guide.md#choosing-the-projection-origin). The current implementation assumes all processes run within one `torchrun` job and all ranks load the same particle minibatches. Multi-node execution has not been documented or validated by this tutorial. ## 8. Outputs and validation For every rank, the trainer writes: * the unrefined component structure; * one component structure and `.pth` checkpoint per epoch; * `chain_parallel_2d_metrics_rank.jsonl`; and * run metadata containing component IDs, cache metadata, world size and numerical settings. Each epoch also produces a merged CIF and a complete-epoch JSON index; manual merging is unnecessary. If component CIFs reuse local chain letters, provide an explicit `chain_id_map` in each affected manifest entry. In the validated 6ZBH references the BCD component's A/B/C maps to B/C/D, while the A component keeps A. This changes merged identities, not reference coordinates. Rank 0 also writes a summary metrics file. For a distributed batch, use the maximum `batch_time_seconds` across rank-specific files as the observed step time; do not report rank-0 time alone as total distributed time. Before interpreting a run, confirm: 1. the manifest entry count equals `WORLD_SIZE`; 2. all ranks used the same frozen diffusion weights; 3. all contextual caches have the same full source-cache hash; 4. every fitted CIF matches its cache in atom count and ordering; 5. every component is in the same experimental frame; 6. all ranks received the same particle indices; 7. loss, gradients and renderer parameters remain finite; and 8. the reported checkpoint was selected without access to the deposited evaluation structure. ## 9. Public recording and restart controls `--seed` defaults to 42; `--diffusion-seed` and `--data-seed` can override the streams. Parallel retains its historical dedicated DataLoader generator and the same particle sequence on every rank. `--learn-gmm` remains on by default; `--no-learn-gmm` freezes both amplitude and width parameters. `--epochs` is configurable, with the original default of 10. `--output-format cif` is the default; requesting PDB also retains CIF for assembly. Default placement is one transform per component. `--by-chain --fit-atoms ca` fits chains independently inside each component without adding decoders or ranks; `--fit-atoms all` uses all atoms. If affine updates are enabled, chains share a trace threshold but update independently when it is triggered. These options are not enabled in the 6ZBH 1+3 main example. To resume, give the trainer `--resume /path/to/model_epoch_EPOCH.json`, the same component manifest/STAR and a new output prefix. `--epochs` is the total target, including completed epochs. Scientific settings are inherited; explicit conflicting settings, incomplete rank files, changed grouping or GPU count are rejected. Only complete epochs can resume. Old refinement files can instead be listed as component cache paths with `--warm-start`, which starts new optimization progress. Exceptions do not trigger rescue saves. Command/configuration/provenance and JSONL records are saved per rank. An explicit `--record-dir` is a shared parent containing `rank0`, `rank1`, etc. `--submission-script` records the actual submitted script. The complete-example checker does not redo diffusion decoding or the short resume experiment. --- # Agent Instructions This documentation is published with GitBook. GitBook is the documentation platform designed so that both humans and AI agents can read, navigate, and reason over technical content effectively. Learn more at gitbook.com. ## Querying This Documentation If you need additional information that is not directly available in this page, you can query the documentation dynamically by asking a question. Perform an HTTP GET request on the current page URL with the `ask` query parameter, and the optional `goal` query parameter: ``` GET https://jwliaomath.gitbook.io/cocofold2/refinement-workflows/component_parallel_tutorial.md?ask=&goal= ``` `ask` is the immediate question: it should be specific, self-contained, and written in natural language. `goal` is optional and describes the broader end goal you are ultimately trying to accomplish on behalf of the user. GitBook uses it to tailor the answer towards what is most useful for that goal. The response will contain a direct answer to the question and relevant excerpts and sources from the documentation. Use this mechanism when the answer is not explicitly present in the current page, you need clarification or additional context, or you want to retrieve related documentation sections.