Complexa Sweep

complexa-sweep

Run cartesian-product parameter sweeps over Proteina-Complexa design pipelines. Pick or author a sweeper YAML in configs/sweeps/, expand it to N inference configs with script_utils/generate_inference_configs.py, loop complexa design over those configs, then aggregate per-config success metrics into a ranked summary CSV plus a manifest.

Important: The complexa design CLI does NOT accept --sweeper directly. Sweeps are driven by script_utils/generate_inference_configs.py, which writes one configs/inference_configs/inf_{idx}_{run_name}.yaml per sweep combination. You then loop complexa design over those generated configs.

Instructions

Sweeper authoring, Cartesian enumeration, cost estimates, and analysis of supplied CSVs need only the relevant configs and data. Complete those artifacts without setting up inference. Read the relevant axis section in sweep_axes.md, then check the matching keys in the selected source config for current defaults. Summary CSVs with supplied counts and objective columns do not require the full inference or metric guides. For a partial snapshot, enumerate combinations locally and provide the future generator command; running that command requires the full project and its dependencies.

For a requested launch, check readiness and the budget below. A single existing job or a SLURM template does not require a parameter sweep.

Step 1: Pre-flight

Before GPU execution, or when the request conditions its workflow on hardware, set SKILL_DIR to the directory containing this manifest and run preflight once. Use a separate shell assignment; an inline environment assignment cannot set $SKILL_DIR's expansion in the same command. Reuse an existing probe from the same request if the environment is unchanged:

bash "$SKILL_DIR"/scripts/preflight.sh

Read ./complexa_setup/preflight.json. A sweep multiplies GPU time by the number of configs and targets. Compute the budget before launch and obtain confirmation if the user has not already authorized that scope and cost:

"This sweep produces N configs × ~M minutes per config ≈ TOTAL GPU-hours. OK to proceed? (y / reduce / cancel)"

If gpu.available=false, stop before GPU execution. Sweeper authoring, Cartesian-product enumeration, cost estimates, and analysis of supplied CSVs can run on CPU. For those requests, finish the requested artifacts and record the launch prerequisites. Label synthetic input data as synthetic and keep estimated GPU time separate from measured runtime.

For a hardware-conditional request, verify the allocated hardware before selecting the CPU fallback. Missing weights, dependencies, or sufficient VRAM for an allocated GPU are launch blockers, not evidence of a CPU-only environment. Report probe or execution failures without reclassifying the selected branch. The shared preflight's gpu.available=false also covers failed or missing nvidia-smi; check device/allocation evidence before treating it as GPU absence. In containers, /proc/driver/nvidia/gpus and loaded kernel modules may describe the host, not this task's allocation. Check container device exposure and any available scheduler allocation information together. A CPU allocation can share a GPU host; an allocated GPU with inaccessible devices is blocked. If allocation remains unclear, report it as undetermined.

Step 2: Pick the pipeline + target

Use the user's supplied pipeline, targets, and run name. Protein, ligand, and AME runs use search_binder_local_pipeline, search_ligand_binder_local_pipeline, and search_ame_local_pipeline, respectively. Ask only for missing values:

• pipeline_config_name — e.g. search_binder_local_pipeline (default), search_ligand_binder_local_pipeline, search_ame_local_pipeline.

• task_name — e.g. 02_PDL1, 22_DerF21, 39_7V11_LIGAND. Passed as --override generation.task_name=<task>.

• run_name — short tag for output dir naming.

Examples

Step 3: Pick or author the sweeper YAML

Sweeper YAMLs live in configs/sweeps/. Each key is a dot-notation Hydra path; each value is a list. The cartesian product becomes N configs.

Canned sweepers

Use a matching sweeper, then check that its axes affect the selected algorithm. The binder config defaults to best-of-n: a beam-width sweep must also pin generation.search.algorithm: [beam-search]. A replicas sweep pins generation.search.algorithm: [best-of-n]. For AME, reward-guided search also requires a configured reward model; its default is null.

Authoring a new sweeper

Minimal multi-axis example (saved to configs/sweeps/my_sweep.yaml):

# 3 beam widths × 2 nsteps = 6 configs
generation.search.algorithm: [beam-search]

generation.search.beam_search.beam_width:
  - 2
  - 4
  - 8

generation.args.nsteps:
  - 200
  - 400

Rules (from script_utils/generate_inference_configs.py:load_sweeper_file):

• Top-level mapping only. Keys are dot-notation Hydra paths.

• Keep budgets, launch status, and other planning metadata in the launch plan, outside the sweeper. Every top-level entry becomes an override axis; a separate targets list would multiply the grid again.

• Values must be lists. A scalar is auto-wrapped into a single-element list (which pins a value without adding a dimension).

• For a list-valued parameter, nest each candidate list: metric.sequence_types: [[self, mpnn]] pins one list. The flat form sweeps two strings; --override parses scalars only.

• Cartesian product: total configs = product of list lengths. Multiply by the number of targets unless target names are already one of the axes; do not count them twice.

• If a key appears in both the sweeper file and an --override, the override wins and that axis collapses.

• Reject empty axes before launch: an axis of [] produces zero configurations.

See references/sweep_axes.md for the full catalogue of swept keys (typical ranges, cost multipliers, what improves/regresses).

Dry-run preview before generating

With the full project available, confirm the config count first:

python script_utils/generate_inference_configs.py \
    --config_name search_binder_local_pipeline \
    --sweeper configs/sweeps/my_sweep.yaml \
    --override generation.task_name=22_DerF21 \
    --run_name my_sweep \
    --dryrun

The output lists every axis + value list and prints DRY RUN — would generate N config pair(s). This previews enumeration only; it returns before Hydra composition and does not validate the pipeline, parameter types, or launch prerequisites.

When the full project is available, inspect the materialized inference configs before launching. Check the resolved target, selected search algorithm, and requested parameter values across the generated grid. The generator permits unknown keys during merging, so a successful write alone does not establish that a misspelled or misplaced axis changes the intended parameter. With only a partial snapshot, check each dot-path against its source config and report that materialization remains unverified.

Step 4: Generate configs + loop complexa design

Once the dry-run looks right, drop --dryrun to materialize inf_{idx}_{run_name}.yaml + eval_{idx}_{run_name}.yaml pairs under configs/inference_configs/ and configs/eval_configs/. Then loop complexa design over the inference configs.

For multiple targets, put generation.task_name in the sweeper as an axis, or give each target's generator invocation a distinct --run_name. Indices restart at zero on each invocation and filenames omit the target; reusing the same run name and output directories overwrites earlier configs and reuses result paths.

# 1. Generate one inf_*.yaml + one eval_*.yaml per combination.
python script_utils/generate_inference_configs.py \
    --config_name search_binder_local_pipeline \
    --sweeper configs/sweeps/my_sweep.yaml \
    --override generation.task_name=22_DerF21 \
    --run_name my_sweep

# 2. Loop complexa design over each generated inference config.
for cfg in configs/inference_configs/inf_*_my_sweep.yaml; do
    complexa design "$cfg" || echo "FAILED: $cfg"
done

This serialises the sweep on a single GPU — be honest with the user that wall-clock = N × per-run.

Multi-GPU host (optional speed-up)

On a host with K GPUs, shard the config list K ways and launch one loop per GPU in parallel. Each complexa design call uses exactly one GPU at default gen_njobs=1 / eval_njobs=1, so pinning via CUDA_VISIBLE_DEVICES keeps them from colliding:

CONFIGS=(configs/inference_configs/inf_*_my_sweep.yaml)
N=${#CONFIGS[@]}; K=4   # 4 GPUs
for gpu in $(seq 0 $((K-1))); do
    (
        for i in $(seq $gpu $K $((N-1))); do
            CUDA_VISIBLE_DEVICES=$gpu complexa design "${CONFIGS[$i]}" \
                || echo "FAILED on GPU $gpu: ${CONFIGS[$i]}"
        done
    ) &
done
wait

Drop gen_njobs / eval_njobs overrides into the loop only if you intentionally want each complexa design invocation to consume multiple GPUs (and you have a way to keep them out of each other's way).

Step 5: Collect results

Each swept config writes to its own ./inference/inf_{idx}_{run_name}/ directory; the analyze step writes ./evaluation_results/eval_{idx}_{run_name}/results_*.csv. After the sweep finishes:

ls -d ./inference/inf_*_my_sweep/
ls ./evaluation_results/eval_*_my_sweep/results_*.csv

Inspect the supplied CSV schema before aggregating. If it already contains one row per configuration with n_passed and n_designs, compute success_rate = n_passed / n_designs; flag zero denominators and invalid counts. Do not treat those summary rows as individual designs. For per-design rows, use existing pass flags, parsing string booleans explicitly. If flags are absent, use the run's analysis thresholds, result_type, column names, and units from EVALUATION_METRICS.md. PAE may be normalized by 31 and pLDDT may use different scales; do not substitute generic thresholds.

Step 6: Rank configs

Emit sweep_summary.csv to the run directory. One row per config:

Populate only metrics present in the inputs or computable from them. Missing diversity or timing data stays unavailable.

Then report:

• Best config = argmax of success_rate. Tie-break on mean_i_pae (lower).

• Pareto frontier: choose the requested objectives and directions. For time versus success, minimize time and maximize success; for time versus iPAE, minimize both. A point is dominated when another is no worse on every objective and strictly better on at least one. Keep distinct config IDs at equal nondominated coordinates; report invalid/missing measurements separately.

Save the full table to sweep_summary.csv, then reopen the saved results and print a compact summary: row count, best config and metric values, and frontier IDs with their objective values. Report the CSV paths and keep the hardware probe and configuration dumps separate from this result summary.

Step 7: Emit manifest

Capture the resolved invocation + outputs for replay. The shared helper takes a single --output-dir (it walks for CSVs and pulls the Hydra config from the run's .hydra/), so call it against the best-config's evaluation directory and pass the parent complexa design command for that run:

BEST_ID=4   # from Step 6 ranking
python3 "$SKILL_DIR"/scripts/write_manifest.py \
    --output-dir ./evaluation_results/eval_${BEST_ID}_my_sweep \
    --command "python script_utils/generate_inference_configs.py --config_name search_binder_local_pipeline --sweeper configs/sweeps/my_sweep.yaml --override generation.task_name=22_DerF21 --run_name my_sweep && for cfg in configs/inference_configs/inf_*_my_sweep.yaml; do complexa design \"\$cfg\"; done" \
    --skill complexa-sweep \
    --out ./sweep_runs/my_sweep/sweep_manifest.json

Alongside the manifest, save the ranked sweep_summary.csv from Step 6 to ./sweep_runs/my_sweep/sweep_summary.csv and surface both paths to the user.

Recommended sweep recipes

Hardware

Total GPU-time for a sweep = N_configs × per_run_GPU_time. A single search_binder_local_pipeline run on one A100 is roughly 30–90 min (binder length + nsteps dependent). A 4-axis × 4-value sweep = 256 configs × ~60 min = ~256 GPU-hours — at that scale, plan on a multi-GPU host with the Step-4 sharding pattern.

Refer to references/hardware.md for the per-run baseline + VRAM minima.

Troubleshooting

For per-axis reference (typical ranges, cost, what gets better/worse), see references/sweep_axes.md.

For the user-facing sweep system overview (config generation, output layout), see references/SWEEP.md.