feat(it-opt): give each ensemble member its own trunk latent

experiments/it_opt_1gpu.toml

@@ -0,0 +1,41 @@
+description = "IT-opt s_plus_z on 1 GPU: all proteins run sequentially in one worker."
+
+# Run with:
+#   sampleworks-runs --preset it_opt_1gpu
+#
+# One job per protein, one worker, so 40 proteins is 40 jobs run back to back. The sibling
+# it_opt_4gpu.toml is this same configuration with gpu_count = 4, which splits those jobs over
+# four workers.
+#
+# The settings below are the defaults of it_opt_scratch/slurm_ensemble_and_score.sh, so a run
+# here is comparable to the recorded trees. To sample a different mode, change which-latent
+# (single = s, pair = z, both = s+z); for the unguided or coordinate-DPS references use
+# scalers = "pure_guidance" with step-scaler-type = "none" or "noisespace" instead.
+
+[defaults]
+DATA_DIR = "/data/inputs"
+RESULTS_DIR = "/data/results/it_opt_1gpu"
+PROTEINS_CSV = "${DATA_DIR}/proteins.csv"
+
+# ensemble-sizes and gradient-weights MUST stay single values. Their defaults are "1 2 4 8" and
+# "0.01 0.1 0.2", which would expand one job per protein into twelve. gradient-weights only
+# feeds the step scaler, which latent_opt does not use, so its value is arbitrary.
+# align-to-input matches the scratch runner, which set it unconditionally; the grid defaults it
+# off. bond-length-weight is named explicitly because the scratch runner defaulted it to 0.0.
+[shared_args]
+proteins = "${PROTEINS_CSV}"
+model = "protenix"
+scalers = "latent_opt"
+which-latent = "both"
+ensemble-sizes = "8"
+gradient-weights = "0.0"
+num-diffusion-steps = 200
+bond-length-weight = 5e-5
+align-to-input = true
+
+[[jobs]]
+name = "s_plus_z"
+env = "protenix"
+gpu_count = 1
+output_subdir = "s_plus_z"
+args = {}

experiments/it_opt_4gpu.toml

@@ -0,0 +1,44 @@
+description = "IT-opt s_plus_z on 4 GPUs: proteins split across four workers, one per GPU."
+
+# Run with:
+#   sampleworks-runs --preset it_opt_4gpu
+#
+# Identical to it_opt_1gpu.toml except for gpu_count and jobs-per-gpu. run_grid_search sets
+# max_workers = len(gpus) * jobs-per-gpu and gives worker i the slice jobs[i::max_workers], so
+# 4 GPUs x 2 gives 8 workers and 40 proteins become 5 jobs each rather than 40 in a row.
+#
+# jobs-per-gpu = 2 matches the recorded run (--jobs-per-gpu 2 on 4 GPUs). Two jobs share a card,
+# each holding its own copy of the weights: a measured single job peaked near 11 GiB of an 80 GiB
+# H100, so two fit with wide margin. Raise it further only after checking VRAM on the largest
+# protein in the set, since cost scales with atom count.
+#
+# Settings are the defaults of it_opt_scratch/slurm_ensemble_and_score.sh; see the sibling
+# preset for how to switch mode.
+
+[defaults]
+DATA_DIR = "/data/inputs"
+RESULTS_DIR = "/data/results/it_opt_4gpu"
+PROTEINS_CSV = "${DATA_DIR}/proteins.csv"
+
+# ensemble-sizes and gradient-weights MUST stay single values. Their defaults are "1 2 4 8" and
+# "0.01 0.1 0.2", which would expand one job per protein into twelve. gradient-weights only
+# feeds the step scaler, which latent_opt does not use, so its value is arbitrary.
+# align-to-input matches the scratch runner, which set it unconditionally; the grid defaults it
+# off. bond-length-weight is named explicitly because the scratch runner defaulted it to 0.0.
+[shared_args]
+proteins = "${PROTEINS_CSV}"
+model = "protenix"
+scalers = "latent_opt"
+which-latent = "both"
+ensemble-sizes = "8"
+gradient-weights = "0.0"
+num-diffusion-steps = 200
+bond-length-weight = 5e-5
+align-to-input = true
+
+[[jobs]]
+name = "s_plus_z"
+env = "protenix"
+gpu_count = 4
+output_subdir = "s_plus_z"
+args = { jobs-per-gpu = 2 }

run_grid_search.py

@@ -19,7 +19,12 @@
 
 from loguru import logger as log
 from sampleworks.utils.guidance_constants import GuidanceType, StructurePredictor
-from sampleworks.utils.guidance_script_arguments import GuidanceConfig, JobConfig, JobResult
+from sampleworks.utils.guidance_script_arguments import (
+    add_latent_opt_args,
+    GuidanceConfig,
+    JobConfig,
+    JobResult,
+)
 from sampleworks.utils.protein_input import ProteinInput
 
 
@@ -267,12 +272,29 @@ def run_grid_search(
     successful = 0
     failed = 0
 
-    max_workers = len(gpus)
-    log.info(f"Running {len(jobs)} jobs with {max_workers} parallel workers")
+    # Workers per GPU (--jobs-per-gpu) lets two jobs share a card, which is worthwhile because a
+    # single job leaves the GPU idle while it featurizes and writes output. Two constraints shape
+    # the worker count:
+    #   never more workers than jobs -- a surplus worker gets an empty queue, and the
+    #     worker_job_queues[worker_num][0] lookup below then raises IndexError before any model
+    #     loads. Easy to hit once jobs_per_gpu multiplies the count.
+    #   keep it a whole multiple of the GPU count, so every card carries the same load. Clamping
+    #     to the job count alone leaves e.g. 5 workers on 4 GPUs: one card runs two jobs
+    #     concurrently while the other three run one and then idle.
+    max_workers = min(len(gpus) * args.jobs_per_gpu, len(jobs))
+    if max_workers > len(gpus):
+        max_workers -= max_workers % len(gpus)
+    gpus_used = min(max_workers, len(gpus))
+    log.info(
+        f"Running {len(jobs)} jobs with {max_workers} parallel workers "
+        f"({gpus_used} GPU(s) x {max_workers // gpus_used} jobs/GPU)"
+    )
 
-    # Divide the job among the workers:
+    # Divide the job among the workers. The worker index is no longer the GPU ordinal, so the
+    # device wraps with i % len(gpus): workers 0..7 across 4 GPUs pair up as 0,1,2,3,0,1,2,3.
+    # Passing i directly would ask for cuda:4..cuda:7 on a 4-GPU box and fail immediately.
     worker_job_queues = [
-        [build_args_for_process_pool(j, args, i) for j in jobs[i::max_workers]]
+        [build_args_for_process_pool(j, args, i % len(gpus)) for j in jobs[i::max_workers]]
         for i in range(max_workers)
     ]
     # we'll pickle each job queue separately and then execute each job queue in a separate process
@@ -779,6 +801,11 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--augmentation", action="store_true", help="Enable augmentation")
     parser.add_argument("--align-to-input", action="store_true", help="Align to input structure")
 
+    # Latent optimization (IT-opt) arguments, used when --scalers includes latent_opt.
+    # Registered from the shipped adder rather than restated here so that flag names, types and
+    # defaults stay identical to the sampleworks-guidance CLI.
+    add_latent_opt_args(parser)
+
     # RF3-specific arguments
     parser.add_argument(
         "--disable-chiral-features",
@@ -803,6 +830,14 @@ def parse_args() -> argparse.Namespace:
         default="auto",
         help="Max parallel jobs (default: auto = number of GPUs)",
     )
+    parser.add_argument(
+        "--jobs-per-gpu",
+        type=int,
+        default=1,
+        choices=(1, 2),
+        help="Concurrent jobs per GPU (default: 1). Each holds its own copy of the model weights, "
+        "so 2 needs roughly twice the VRAM; 2 is the tested ceiling",
+    )
     parser.add_argument(
         "--dry-run",
         action="store_true",

src/sampleworks/core/scalers/latent_optimization.py

@@ -126,6 +126,70 @@ def __call__(self, latents: Sequence[Tensor], baselines: Sequence[Tensor]) -> Te
         return torch.stack(terms).sum()
 
 
+class _PerMemberStepper:
+    """Denoise one ensemble member at a time, each with its OWN trunk latent.
+
+    The it_opt "multiple-leaf" scheme gives every ensemble member its own latent, so ``s``/``z``
+    carry a leading ``ensemble_size`` batch dim. Stock model diffusion modules instead take an
+    un-batched conditioning and broadcast it internally, so a batched latent either crashes
+    (Protenix) or mis-broadcasts (Boltz's ``multiplicity``). This adapter sidesteps that in a
+    model-agnostic way: for each member it slices that member's latent, runs the wrapped model's
+    normal un-batched ``step`` on that member alone, and stacks the results. Members stay
+    independent, so gradients stay per-member; the summed density reward couples them only through
+    the ensemble average. Cost is N forwards instead of one batched call, with an equivalent result.
+
+    ``featurize`` / ``initialize_from_prior`` pass straight through to the wrapped model.
+    """
+
+    def __init__(
+        self,
+        model,
+        io: AttrLatentIO,
+        *,
+        optimize_single: bool,
+        optimize_pair: bool,
+        ensemble_size: int,
+    ):
+        self._model = model
+        self._io = io
+        self._optimize_single = optimize_single
+        self._optimize_pair = optimize_pair
+        self._ensemble_size = ensemble_size
+
+    def step(self, x_t: Tensor, t, *, features: GenerativeModelInput) -> Tensor:
+        """Loop the wrapped model's ``step`` over ensemble members; stack the per-member results."""
+        cond = features.conditioning
+        per_member: list[Tensor] = []
+        for i in range(self._ensemble_size):
+            # Slice only the OPTIMIZED latents (they carry the ensemble batch dim); a non-optimized
+            # latent stays the shared un-batched baseline already on ``cond``.
+            cond_i = cond
+            if self._optimize_single:
+                cond_i = self._io.write_single(cond_i, self._io.read_single(cond)[i])
+            if self._optimize_pair:
+                # read_pair returns None when the io addresses no pair rep. sample() only sets
+                # optimize_pair together with a pair_attr, so this is a misconfigured io reaching
+                # us directly; say so rather than failing on a None subscript mid-denoise.
+                pair = self._io.read_pair(cond)
+                if pair is None:
+                    raise ValueError(
+                        "optimize_pair is set, but the io addresses no pair representation."
+                    )
+                cond_i = self._io.write_pair(cond_i, pair[i])
+            t_i = t
+            if isinstance(t, Tensor) and t.ndim >= 1 and t.shape[0] == x_t.shape[0]:
+                t_i = t[i : i + 1]
+            features_i = GenerativeModelInput(conditioning=cond_i)
+            per_member.append(self._model.step(x_t[i : i + 1], t_i, features=features_i))
+        return torch.cat(per_member, dim=0)
+
+    def featurize(self, *args, **kwargs):
+        return self._model.featurize(*args, **kwargs)
+
+    def initialize_from_prior(self, *args, **kwargs):
+        return self._model.initialize_from_prior(*args, **kwargs)
+
+
 class LatentOptimization:
     """Trajectory scaler that optimizes the model's ``s``/``z`` latents (IT-opt).
 
@@ -259,6 +323,16 @@ def sample(
         schedule = sampler.compute_schedule(num_steps=self.num_steps)
         grad_enabler = _GradEnablingScaler()
 
+        # Denoise per member so each uses its own latent: stock model diffusion modules take an
+        # un-batched conditioning, so the batched per-member latents can't go through in one call.
+        stepper = _PerMemberStepper(
+            model,
+            io,
+            optimize_single=self.optimize_single,
+            optimize_pair=self.optimize_pair,
+            ensemble_size=self.ensemble_size,
+        )
+
         # --- optional coordinate-space geometry penalty -------------------------
         # BondGeometryReward penalizes stretched bonds and steric clashes in the denoised structure,
         # curbing the overshoot where an aggressive latent update trades geometry for density fit.
@@ -279,7 +353,7 @@ def sample(
         for outer in tqdm(range(self.outer_steps)):
             optimizer = torch.optim.Adam(latents, lr=self.learning_rate)  # a fresh, persistent Adam
             round_losses = self._optimize_one_round(
-                model=model,
+                model=stepper,
                 sampler=sampler,
                 reward=reward,
                 features=features,
@@ -307,7 +381,7 @@ def sample(
 
         # --- final clean sampling round with the optimized latents --------------
         final_coords, trajectory, losses = self._sample_with_frozen_latents(
-            model=model,
+            model=stepper,
             sampler=sampler,
             reward=reward,
             io=io,
@@ -341,8 +415,12 @@ def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO):
         detached baselines (anchor targets), and per-latent anchor weights. Each
         leaf is a detached clone made ``requires_grad=True`` -- a true leaf severed
         from any trunk graph, so Adam updates it directly (leaves persist and are
-        updated in place across rounds and steps). Shapes are preserved (whatever
-        the wrapper caches), so no assumption is made about a batch dimension.
+        updated in place across rounds and steps).
+
+        Each leaf gets a leading ``ensemble_size`` batch dimension -- one INDEPENDENT latent per
+        ensemble member, all cloned from the same trunk baseline (the it_opt "multiple-leaf"
+        scheme), so members can diverge rather than share one latent. The baseline kept for the
+        anchor stays un-batched and broadcasts across members.
         """
         conditioning = features.conditioning
         latents: list[Tensor] = []
@@ -367,7 +445,14 @@ def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO):
                     "conditioning does not expose it. Check the attribute names for this model."
                 )
             baseline = baseline.detach()
-            leaf = baseline.clone().requires_grad_(True)
+            # it_opt "multiple-leaf" scheme: give each ensemble member its OWN latent. We stack
+            # ensemble_size independent copies of the trunk baseline into a leading batch dim, so
+            # each member gets its own gradient and can diverge, instead of collapsing onto one
+            # shared latent. (The reference does the same via batch-expand-then-clone.) Stock
+            # diffusion modules take an un-batched conditioning, so _PerMemberStepper slices this
+            # batch dim back apart and runs one forward per member at denoise time.
+            member_copies = [baseline for _ in range(self.ensemble_size)]
+            leaf = torch.stack(member_copies).requires_grad_(True)
             # The conditioning is a frozen dataclass, so setattr would raise; replace() returns a
             # copy with this one field swapped and every sidecar field left untouched.
             conditioning = dataclasses.replace(conditioning, **{attr: leaf})

src/sampleworks/utils/cif_utils.py

@@ -12,6 +12,7 @@
 from loguru import logger
 
 from sampleworks.utils.atom_array_utils import (
+    BLANK_ALTLOC_IDS,
     find_all_altloc_ids,
     save_structure_to_cif,
     select_altloc,
@@ -247,6 +248,79 @@ def resolve_mixed_hetatm_atom_altlocs(cif_path: Path | str) -> Path:
     return tmp_path
 
 
+def remap_altlocs_to_ab(cif_path: Path | str) -> Path:
+    """Relabel a residue's non-A/B alternate altloc(s) to fill the A/B slots.
+
+    Some depositions label a residue's two conformers ``A``/``C`` (or ``A``/``D``, etc.) instead of
+    ``A``/``B``. Tools that assume the two altlocs are literally named A and B -- e.g. a
+    min-RMSD-to-altloc-A/altloc-B scorer -- then silently drop the second conformer. This is an
+    optional pre-processing step that, per residue, maps the alternate label onto the free ``B``
+    slot: an atom with altloc ``A`` stays ``A`` and the other non-blank altloc becomes ``B``; if
+    ``A`` is absent, the two labels are assigned to ``A``/``B`` in sorted order.
+
+    Only residues with **exactly two** non-blank altlocs that are **not already** ``{A, B}`` are
+    touched. Residues that are already A/B, have a single or blank altloc, or carry three or more
+    altlocs (ambiguous which two to keep) are left unchanged, so the metric's numbers do not move
+    for structures that were already A/B.
+
+    A warning is logged for every remapped ``(chain, residue)`` position.
+
+    Parameters
+    ----------
+    cif_path
+        Path to the input CIF file.
+
+    Returns
+    -------
+    Path
+        Path to a remapped temporary CIF file if any position was changed, or the original
+        ``cif_path`` unchanged if none were.
+    """
+    cif_path = Path(cif_path)
+    atom_array = load_any(cif_path, altloc="all", extra_fields=["occupancy", "b_factor"])
+    if isinstance(atom_array, AtomArrayStack):
+        atom_array = atom_array[0]
+    if not hasattr(atom_array, "altloc_id"):
+        return cif_path  # no altloc annotation -> nothing to remap
+
+    altloc = atom_array.altloc_id.copy()  # per-atom altloc characters, mutated in place below
+    chain_id = atom_array.chain_id
+    res_id = atom_array.res_id
+    remapped = 0
+
+    for chain in np.unique(chain_id):
+        for rid in np.unique(res_id[chain_id == chain]):
+            pos = (chain_id == chain) & (res_id == rid)
+            labels = sorted(set(altloc[pos].tolist()) - BLANK_ALTLOC_IDS)
+            if len(labels) != 2 or set(labels) == {"A", "B"}:
+                continue
+            if "A" in labels:
+                mapping = {next(x for x in labels if x != "A"): "B"}
+            else:
+                mapping = {labels[0]: "A", labels[1]: "B"}
+            for src, dst in mapping.items():
+                sel = pos & (altloc == src)
+                altloc[sel] = dst
+                remapped += int(sel.sum())
+            logger.warning(
+                f"Chain {chain}, residue {rid}: remapped altlocs {labels} -> A/B so the alternate "
+                "conformer is not dropped by A/B-only tooling."
+            )
+
+    if remapped == 0:
+        return cif_path
+
+    atom_array.set_annotation("altloc_id", altloc)
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".cif", prefix="sampleworks_altloc_ab_", delete=False
+    ) as tmp_file:
+        tmp_path = Path(tmp_file.name)
+
+    save_structure_to_cif(atom_array, tmp_path)
+    logger.info(f"Wrote altloc-remapped CIF to temporary file: {tmp_path}")
+    return tmp_path
+
+
 def add_category_to_cif(
     ciffile: CIFFile,
     data: dict[str, Any],