feat(it-opt): per step latent space optimization - #313
Conversation
Add LatentAdaptedWrapper, a swappable transform on a model's single representation (post-trunk latent `s`) at the featurize/step boundary — the step-1 scaffold for AlphaSAXS-style latent injection / IT-optimization. - Wraps the FlowModelWrapper protocol so samplers, scalers, and eval are unchanged; identity transform (k=1, b=0) is behaviourally a no-op. - Model-agnostic via AttrLatentIO (single attr name: `s` for Boltz, `s_trunk` for Protenix/RF3) — no model packages imported. - Injected latent detached during sampling to preserve existing coordinate-guidance gradients. - Adds unit tests and design docs (roadmap + solution diagrams). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds inference-time latent optimization for cached model representations, bond geometry penalties, latent adapter plumbing, Protenix gradient preservation, CLI/runtime integration, focused tests, documentation, and local artifact ignore rules. ChangesLatent optimization
Local artifact ignore rules
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant GuidanceCLI
participant GuidanceRunner
participant LatentOptimization
participant TrajectorySampler
participant ProtenixWrapper
participant RewardFunction
GuidanceCLI->>GuidanceRunner: parse latent_opt parameters
GuidanceRunner->>LatentOptimization: construct and start sample
LatentOptimization->>ProtenixWrapper: inject optimizable latent leaves
LatentOptimization->>TrajectorySampler: run diffusion with gradient-enabled scaler
TrajectorySampler->>ProtenixWrapper: produce denoised coordinates
TrajectorySampler->>RewardFunction: evaluate denoised coordinates
RewardFunction-->>LatentOptimization: return reward loss
LatentOptimization->>LatentOptimization: update latent leaves with Adam
LatentOptimization-->>GuidanceRunner: return trajectory and optimization metadata
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
manzuoni-astera
left a comment
manzuoni-astera
left a comment
There was a problem hiding this comment.
Thanks for keeping this as a narrowly scoped step-1 scaffold. Treating the optimizer, sampler feature hook, Protenix detach handling, Boltz2 cache recomputation, and real experimental transform as intentional follow-up work, the basic read/transform/write seam looks reasonable.
Before this is marked ready, I think the scaffold contract needs a few corrections:
-
LatentAdaptedWrapperis not currently drop-in for the actual guidance entry point._run_guidancedispatches preprocessing frommodel_wrapper.__class__.__name__and will rejectLatentAdaptedWrapperas unknown (guidance_script_utils.py:444-486). Pure guidance also readsmodel_wrapper.modeldirectly at line 541, while diagnostics and MSA reporting use other concrete-wrapper attributes. The one-line factory wrapping shown in the diagram would therefore fail before sampling. Please either make orchestration decorator-aware or deliberately proxy/unwrap the required model identity and attributes, with a pipeline-level test. -
A wrong
single_attrsilently becomes a no-op (latent_adapter.py:85-88, 172-174), and the test currently enshrines that behavior. Because this string is required configuration, I would prefer a clear error, or an explicit opt-inallow_missingmode, so field drift cannot masquerade as successful adaptation. -
The mock
step()ignores conditioning, so the tests prove tensor replacement but not that rewritten conditioning affects the delegated step. A conditioning-sensitive mock should cover identity equivalence and non-identity propagation.
CI unit tests and typechecks pass across all environments, but lint currently fails with 39 introduced Ruff errors. GPU jobs are awaiting environment approval. With those scaffold-level issues addressed, the deferred model-specific work can reasonably remain in follow-ups.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (2)
src/sampleworks/core/rewards/geometry.py (1)
45-215: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftUse NumPy-style docstrings for the new Python APIs.
src/sampleworks/core/rewards/geometry.py#L45-L215: addParameters,Returns, and shape/side-effect details to public and helper APIs.src/sampleworks/core/scalers/latent_optimization.py#L71-L553: document parameters, return values, and mutation/gradient behavior in NumPy style.src/sampleworks/utils/guidance_script_arguments.py#L617-L649: documentparserand theNonereturn contract in NumPy style.As per coding guidelines, “Add NumPy-style docstrings to every function and class.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/core/rewards/geometry.py` around lines 45 - 215, Update every function and class in src/sampleworks/core/rewards/geometry.py (lines 45-215), src/sampleworks/core/scalers/latent_optimization.py (lines 71-553), and src/sampleworks/utils/guidance_script_arguments.py (lines 617-649) with NumPy-style docstrings; document parameters, return values, relevant tensor shapes, side effects, and gradient/mutation behavior, including the parser argument and None return contract for the guidance-arguments API.Source: Coding guidelines
src/sampleworks/utils/guidance_script_utils.py (1)
468-475: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winKeep Protenix diffusion caches for single-only optimization.
The cache must be disabled when optimizing
z_trunk, but--which-latent singleleavesz_trunkunchanged. Disabling it for every LATENT_OPT run forces needless recomputation each denoising step. Gate this on pair optimization instead.Suggested change
- enable_diffusion_shared_vars_cache=(guidance_type != GuidanceType.LATENT_OPT), + enable_diffusion_shared_vars_cache=( + guidance_type != GuidanceType.LATENT_OPT + or getattr(args, "which_latent", "pair") == "single" + ),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/utils/guidance_script_utils.py` around lines 468 - 475, Update the enable_diffusion_shared_vars_cache argument in _run_guidance()’s Protenix setup so caching is disabled only when LATENT_OPT is optimizing the pair representation z_trunk, while retaining the cache for single-only optimization. Use the existing pair-optimization condition or which-latent configuration to distinguish these cases.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/CODEBASE_GUIDE.md`:
- Around line 25-39: Update the documentation to reflect the shipped LATENT_OPT
implementation: in docs/CODEBASE_GUIDE.md lines 25-39 add the LATENT_OPT
execution branch, lines 54-60 include latent optimization in the scaler
inventory, lines 76-79 list LatentOptimization among trajectory scalers, and
lines 140-143 describe its construction and execution. In
docs/IT_OPTIMIZATION_PLAN.md lines 31-40 update the historical Protenix gradient
limitation, lines 73-76 state that anchoring is configurable and disabled by
default, lines 150-163 remove completed rollout items, and lines 177-183
document separate per-latent clipping.
In `@docs/IT_OPT_REFERENCE_COMPARISON.md`:
- Line 58: Specify the language for each affected fenced code block by changing
the opening fence to text in docs/IT_OPT_REFERENCE_COMPARISON.md:58-58,
docs/latent_adapter/implementation_roadmap.md:234-234, and
docs/latent_adapter/latent_space_optimization.md:72-72 and 151-151.
In `@docs/IT_OPT_TESTING_PROTENIX.md`:
- Around line 172-176: Update the LatentOptimization guidance to state that
GuidanceType.LATENT_OPT is now wired into _run_guidance and available through
the sampleworks-guidance --model protenix --guidance-type latent_opt CLI path.
Retain the existing script workflow only as an optional debugging alternative,
and remove stale Phase 2 or unavailable wording.
In `@docs/IT_OPTIMIZATION_NOTES.md`:
- Around line 53-55: Update the path references in
docs/IT_OPTIMIZATION_NOTES.md, including the sections around the listed
occurrences, to remove the developer-specific /Users/fengyu/... absolute path.
Use $IT_OPT_ROOT for the local root reference while preserving the existing
../it_opt relative path and public clone URL.
In `@docs/latent_adapter/generate_solution_diagram.py`:
- Around line 382-407: Update _emit_raw_dot() to accept an out parameter and use
it for both the DOT output filename and render instructions. Pass each builder’s
requested output name from all four __main__ fallback call sites by changing
their calls to _emit_raw_dot(out), preserving distinct fallback files instead of
overwriting latent_adapter_solution.gv.
In `@src/sampleworks/core/rewards/geometry.py`:
- Around line 58-215: Split BondGeometryReward into separate bond-length and
steric-clash reward classes, moving the corresponding topology setup and loss
logic from bond_length_loss and collision_loss into their respective classes.
Ensure each reward computes exactly one mismatch and accepts its own weight,
then compose both rewards at the optimization boundary instead of combining them
in BondGeometryReward.__call__. Preserve the existing bond and clash
calculations and make each term independently reusable and weightable.
- Around line 197-211: The clash-loss calculation around pos_row, pos_col, and
worst_overlap must avoid materializing the full [ensemble, n, n, 3] tensor and
quadratic intermediates. Process atom-pair rows in bounded blocks (or use an
equivalent sparse strategy), compute distances and overlap maxima per block
across the ensemble, apply the corresponding _scored_pairs entries, and
accumulate the same total sum while preserving the per-pair ensemble maximum.
In `@src/sampleworks/core/scalers/latent_optimization.py`:
- Around line 525-553: Update the latent optimization loop around `trajectory`,
`losses`, and `step_output.denoised` to collect each available denoised frame as
a detached CPU clone, preserving alignment with sampling steps (use an
appropriate placeholder when denoised output is absent). Return the collected
frames in `GuidanceOutput.metadata["trajectory_denoised"]` so callers can read
the denoised trajectory.
In `@src/sampleworks/models/latent_adapter.py`:
- Around line 173-181: Update DeltaInjector.__init__ and forward so delta is
registered as an eager or uninitialized trainable parameter before the first
forward, ensuring DeltaInjector().parameters() is non-empty for optimizer
construction while preserving shape-aware initialization. Add a regression test
that constructs the optimizer before calling forward and verifies the parameter
is included.
In `@src/sampleworks/utils/guidance_script_arguments.py`:
- Around line 626-648: Validate IT-opt CLI parameters during argument parsing or
scaler construction: require positive values for outer-steps, learning-rate, and
max-grad-norm, and require nonnegative values for anchor-weight and
bond-length-weight. Reject invalid inputs immediately with the parser’s standard
argument error instead of allowing optimization to be skipped or penalties
disabled.
In `@tests/eval/test_patch_output_cif_files.py`:
- Around line 21-23: Add a NumPy-style docstring to the script fixture
describing that it returns the loaded module, including a concise Returns
section. Keep the existing fixture scope and load_script behavior unchanged.
- Around line 26-30: Align the test’s _DEFAULT_REGEX with the script’s CLI
default so the real default behavior is exercised, and require the captured
four-character identifier to end at a folder boundary before matching. Update
the source default and remove the explicit override if appropriate, or add a
lookahead boundary to _DEFAULT_REGEX while preserving the expected None results
for names such as 4hhb_final.
---
Nitpick comments:
In `@src/sampleworks/core/rewards/geometry.py`:
- Around line 45-215: Update every function and class in
src/sampleworks/core/rewards/geometry.py (lines 45-215),
src/sampleworks/core/scalers/latent_optimization.py (lines 71-553), and
src/sampleworks/utils/guidance_script_arguments.py (lines 617-649) with
NumPy-style docstrings; document parameters, return values, relevant tensor
shapes, side effects, and gradient/mutation behavior, including the parser
argument and None return contract for the guidance-arguments API.
In `@src/sampleworks/utils/guidance_script_utils.py`:
- Around line 468-475: Update the enable_diffusion_shared_vars_cache argument in
_run_guidance()’s Protenix setup so caching is disabled only when LATENT_OPT is
optimizing the pair representation z_trunk, while retaining the cache for
single-only optimization. Use the existing pair-optimization condition or
which-latent configuration to distinguish these cases.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: bc57db4f-22b0-4cc5-9edd-a64f09adb0db
⛔ Files ignored due to path filters (6)
docs/latent_adapter/latent_adapter_solution.pngis excluded by!**/*.pngdocs/latent_adapter/latent_adapter_solution.svgis excluded by!**/*.svgdocs/latent_adapter/latent_adapter_solution_downstream.pngis excluded by!**/*.pngdocs/latent_adapter/latent_adapter_solution_downstream.svgis excluded by!**/*.svgdocs/latent_adapter/latent_adapter_solution_step1.pngis excluded by!**/*.pngdocs/latent_adapter/latent_adapter_solution_step1.svgis excluded by!**/*.svg
📒 Files selected for processing (21)
.gitignoredocs/CODEBASE_GUIDE.mddocs/IT_OPTIMIZATION_NOTES.mddocs/IT_OPTIMIZATION_PLAN.mddocs/IT_OPT_GIT_WORKFLOW.mddocs/IT_OPT_REFERENCE_COMPARISON.mddocs/IT_OPT_TESTING_PROTENIX.mddocs/PERSONAL_CODING_STYLE.mddocs/latent_adapter/generate_solution_diagram.pydocs/latent_adapter/implementation_roadmap.mddocs/latent_adapter/latent_space_optimization.mdsrc/sampleworks/core/rewards/geometry.pysrc/sampleworks/core/scalers/latent_optimization.pysrc/sampleworks/models/latent_adapter.pysrc/sampleworks/models/protenix/wrapper.pysrc/sampleworks/utils/guidance_constants.pysrc/sampleworks/utils/guidance_script_arguments.pysrc/sampleworks/utils/guidance_script_utils.pytests/eval/script_loader.pytests/eval/test_patch_output_cif_files.pytests/models/test_latent_adapter.py
| `LatentOptimization` is **not yet wired into the CLI** (`GuidanceType.LATENT_OPT` and | ||
| the `_run_guidance` branch are Phase 2). Until then, drive it from a script as above — | ||
| which is also the better debugging surface, since it isolates the scaler from the | ||
| grid-search/save machinery. Once validated, wiring it in makes it reachable as | ||
| `sampleworks-guidance --model protenix --guidance-type latent_opt …`. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the stale CLI-wiring guidance.
This says latent_opt is Phase 2 and unavailable, but this PR wires GuidanceType.LATENT_OPT into _run_guidance. Document the current CLI path and retain the script workflow only as a debugging alternative.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/IT_OPT_TESTING_PROTENIX.md` around lines 172 - 176, Update the
LatentOptimization guidance to state that GuidanceType.LATENT_OPT is now wired
into _run_guidance and available through the sampleworks-guidance --model
protenix --guidance-type latent_opt CLI path. Retain the existing script
workflow only as an optional debugging alternative, and remove stale Phase 2 or
unavailable wording.
| - **Root (absolute, this machine):** `/Users/fengyu/Sampleworks_0622/it_opt/` | ||
| - **Relative to the Sampleworks repo root:** `../it_opt/` | ||
| - **Public backup (clone-able anywhere):** https://github.com/sai-advaith/it_opt |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Remove developer-specific absolute paths from repository documentation.
/Users/fengyu/... exposes a local username and makes the workflow machine-specific. Replace these occurrences with $IT_OPT_ROOT, ../it_opt, or the public-clone path.
Also applies to: 79-86, 94-100
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/IT_OPTIMIZATION_NOTES.md` around lines 53 - 55, Update the path
references in docs/IT_OPTIMIZATION_NOTES.md, including the sections around the
listed occurrences, to remove the developer-specific /Users/fengyu/... absolute
path. Use $IT_OPT_ROOT for the local root reference while preserving the
existing ../it_opt relative path and public clone URL.
| def _emit_raw_dot(): | ||
| """Fallback: graphviz python package missing -> write a .gv the user can render.""" | ||
| print("graphviz python package not found; writing raw DOT instead.", file=sys.stderr) | ||
| print(f"Render with: dot -Tpng {OUT}.gv -o {OUT}.png", file=sys.stderr) | ||
| # Minimal raw-DOT mirror of the structure above. | ||
| dot = '''digraph latent_adapter_solution { | ||
| rankdir=TB; node [shape=box, style=filled, fontname=Helvetica]; | ||
| struct [label="structure: dict", fillcolor="#eeeeee"]; | ||
| reward [label="RealSpaceRewardFunction\\n(target density)", fillcolor="#eeeeee"]; | ||
| wire [label="get_model_and_device()\\nmodel = LatentAdaptedWrapper(\\n BoltzWrapper(...), AttrLatentIO(\\"s\\"), AffineInjector())", fillcolor="#fdebd0"]; | ||
| wrapper [label="LatentAdaptedWrapper[C]\\n(satisfies FlowModelWrapper)", fillcolor="#d5f5e3"]; | ||
| io [label="AttrLatentIO(\\"s\\"|\\"s_trunk\\")\\nonly model-specific knowledge", fillcolor="#d5f5e3"]; | ||
| injector[label="AffineInjector s'=k*s+b\\nk=1,b=0 -> identity", fillcolor="#d5f5e3"]; | ||
| seam [label="featurize: read s -> inject -> detach(if sampling) -> write s'", fillcolor="#fcf3cf"]; | ||
| pg [label="PureGuidance.sample() / AF3EDMSampler.step()\\nUNCHANGED", fillcolor="#eeeeee"]; | ||
| mstep [label="model.step(): cond = features.conditioning (reads s')", fillcolor="#eeeeee"]; | ||
| step2 [label="STEP 2: MLP/FiLM/Density injector + LatentGuidance training", fillcolor="#d6eaf8", style="dashed,filled"]; | ||
| struct -> wire -> wrapper -> seam -> pg -> mstep; | ||
| wrapper -> io [style=dashed,arrowhead=none]; wrapper -> injector [style=dashed,arrowhead=none]; | ||
| reward -> mstep [style=dotted,label="target density"]; | ||
| injector -> step2 [style=dashed,label="swap"]; | ||
| } | ||
| ''' | ||
| with open(f"{OUT}.gv", "w") as fh: | ||
| fh.write(dot) | ||
| print(f"Wrote {OUT}.gv") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Preserve each builder’s requested fallback output name.
When Graphviz is unavailable, every builder calls this helper without out, so the four __main__ calls overwrite latent_adapter_solution.gv. Accept out in _emit_raw_dot() and pass it from each caller.
Proposed fix
-def _emit_raw_dot():
+def _emit_raw_dot(out: str):
"""Fallback: graphviz python package missing -> write a .gv the user can render."""
- print(f"Render with: dot -Tpng {OUT}.gv -o {OUT}.png", file=sys.stderr)
+ print(f"Render with: dot -Tpng {out}.gv -o {out}.png", file=sys.stderr)
...
- with open(f"{OUT}.gv", "w") as fh:
+ with open(f"{out}.gv", "w") as fh:
fh.write(dot)
- print(f"Wrote {OUT}.gv")
+ print(f"Wrote {out}.gv")Also change each fallback call to _emit_raw_dot(out).
🧰 Tools
🪛 ast-grep (0.44.1)
[warning] 404-404: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(f"{OUT}.gv", "w")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/latent_adapter/generate_solution_diagram.py` around lines 382 - 407,
Update _emit_raw_dot() to accept an out parameter and use it for both the DOT
output filename and render instructions. Pass each builder’s requested output
name from all four __main__ fallback call sites by changing their calls to
_emit_raw_dot(out), preserving distinct fallback files instead of overwriting
latent_adapter_solution.gv.
| trajectory: list[Tensor] = [] | ||
| losses: list[float | None] = [] | ||
| steps = tqdm(range(self.num_steps), "IT-opt final sampling") | ||
| for i in steps: | ||
| context = sampler.get_context_for_step(i, schedule).with_reconciler( | ||
| reconciler=reconciler, alignment_reference=alignment_reference | ||
| ) | ||
| step_output = sampler.step( | ||
| state=coords, | ||
| model_wrapper=model, | ||
| context=context, | ||
| scaler=None, | ||
| features=frozen_features, | ||
| ) | ||
| coords = step_output.state.detach() | ||
| trajectory.append(coords.clone().cpu()) | ||
|
|
||
| if step_output.denoised is not None: | ||
| with torch.no_grad(): | ||
| loss = reward( | ||
| coordinates=step_output.denoised, | ||
| elements=reward_inputs.elements, | ||
| b_factors=reward_inputs.b_factors, | ||
| occupancies=reward_inputs.occupancies, | ||
| ) | ||
| losses.append(float(loss)) | ||
| else: | ||
| losses.append(None) | ||
| return coords, trajectory, losses |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Populate trajectory_denoised metadata for latent optimization.
This final pass records only step_output.state. The caller reads denoised frames from result.metadata["trajectory_denoised"], so LATENT_OPT silently writes an empty denoised trajectory. Collect detached denoised frames and include them in GuidanceOutput.metadata.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/sampleworks/core/scalers/latent_optimization.py` around lines 525 - 553,
Update the latent optimization loop around `trajectory`, `losses`, and
`step_output.denoised` to collect each available denoised frame as a detached
CPU clone, preserving alignment with sampling steps (use an appropriate
placeholder when denoised output is absent). Return the collected frames in
`GuidanceOutput.metadata["trajectory_denoised"]` so callers can read the
denoised trajectory.
| def __init__(self): | ||
| super().__init__() | ||
| # Reserve the parameter slot; filled on first forward once the shape is known. | ||
| self.register_parameter("delta", None) | ||
|
|
||
| def forward(self, latent: Tensor) -> Tensor: | ||
| if self.delta is None: | ||
| self.delta = nn.Parameter(torch.zeros_like(latent)) | ||
| return latent + self.delta |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
ast-grep outline src/sampleworks/models/latent_adapter.py --items all
rg -n -C3 'DeltaInjector\(|torch\.optim\.(Adam|AdamW)|optimizer\s*=' src testsRepository: diff-use/sampleworks
Length of output: 7009
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "## latent_adapter relevant lines"
sed -n '160,182p' src/sampleworks/models/latent_adapter.py
echo
echo "## latent optimization optimizer setup"
sed -n '260,282p' src/sampleworks/core/scalers/latent_optimization.py
echo
echo "## PyTorch register_parameter=None behavior"
python3 - <<'PY'
import torch
from torch import nn
class M(nn.Module):
def __init__(self):
super().__init__()
self.register_parameter("delta", None)
def forward(self, x):
if self.delta is None:
self.delta = nn.Parameter(torch.zeros_like(x))
return x + self.delta
m = M()
opt = torch.optim.Adam(m.parameters(), lr=0.01)
print("before forward parameter names:", list(p.name for p in m.parameters()))
print("before forward param group len:", len(opt.param_groups[0]["params"]))
first = next(m.parameters())
print("before forward first is delta:", first is m.delta)
out = m(torch.ones(3, 2)[:, 1:2])
print("after forward parameter names:", list(p.name for p in m.parameters()))
print("after forward param group len:", len(opt.param_groups[0]["params"]))
last = next(m.parameters())
print("after forward last is delta:", last is m.delta)
PYRepository: diff-use/sampleworks
Length of output: 2795
🌐 Web query:
PyTorch nn.Module register_parameter None returns None and is included in parameters() documentation
💡 Result:
In PyTorch, the nn.Module.register_parameter(name, param) method allows passing None as the param argument [1][2][3]. When None is registered, the following behaviors occur: 1. Parameters and state_dict: If a parameter is registered as None, it is explicitly excluded from the module's parameters() iterator and its state_dict() [1][2][3]. 2. Operations: Operations that typically run on parameters—such as moving the module to a device via .cuda() or .to(device)—will ignore the None parameter [1][2][3]. 3. Attribute access: The parameter name will be set to None as an attribute of the module (accessible via module.name) [2][4]. 4. State Loading: If a parameter is registered as None but a corresponding key exists in the state_dict during a call to load_state_dict(), PyTorch will raise a RuntimeError [3]. This behavior is documented in the official PyTorch nn.Module.register_parameter documentation, which explicitly states that if param is None, the parameter is not included in the module's state_dict and operations are ignored [3]. [1][2][3]
Citations:
- 1: https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/module.py
- 2: https://docs.pytorch.org/docs/2.3/%5Fmodules/torch/nn/modules/module.html
- 3: https://docs.pytorch.org/docs/2.3/generated/torch.nn.Module.html
- 4: https://github.com/pytorch/pytorch/blob/ebcacd5e878f3f8e84c941d6b9567c88f50a2f10/torch/nn/modules/module.py
Materialize delta before optimizer construction.
register_parameter("delta", None) makes delta absent from parameters(), so torch.optim.Adam(DeltaInjector().parameters()) sees an empty param list before the first forward pass. Register an eager/uninitialized parameter or require explicit materialization before optimizer construction; add a regression test for that order.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/sampleworks/models/latent_adapter.py` around lines 173 - 181, Update
DeltaInjector.__init__ and forward so delta is registered as an eager or
uninitialized trainable parameter before the first forward, ensuring
DeltaInjector().parameters() is non-empty for optimizer construction while
preserving shape-aware initialization. Add a regression test that constructs the
optimizer before calling forward and verifies the parameter is included.
| parser.add_argument("--learning-rate", type=float, default=0.05, help="Adam learning rate") | ||
| parser.add_argument( | ||
| "--outer-steps", | ||
| type=int, | ||
| default=2, | ||
| help="Number of optimization rounds (fresh diffusion noise per round)", | ||
| ) | ||
| parser.add_argument( | ||
| "--anchor-weight", | ||
| type=float, | ||
| default=0.0, | ||
| help="On-manifold L2-to-baseline anchor weight for the optimized latent(s)", | ||
| ) | ||
| parser.add_argument( | ||
| "--max-grad-norm", type=float, default=1.0, help="Per-latent gradient-clip threshold" | ||
| ) | ||
| parser.add_argument( | ||
| "--bond-length-weight", | ||
| type=float, | ||
| default=5e-5, | ||
| help="Weight of the coordinate-space bond-geometry penalty (bond-length + steric-clash " | ||
| "hinges) added to the IT-opt loss. Default 5e-5 = smallest weight that fixes mean clash " | ||
| "while keeping density fit and diversity; use 1e-3 for median clash 0; 0 disables it", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Validate IT-opt parameter ranges at parse time.
--outer-steps <= 0 silently skips optimization, and negative --bond-length-weight silently disables the penalty. Reject invalid round counts, learning rates, gradient norms, and nonnegative-only weights in the parser or scaler constructor.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/sampleworks/utils/guidance_script_arguments.py` around lines 626 - 648,
Validate IT-opt CLI parameters during argument parsing or scaler construction:
require positive values for outer-steps, learning-rate, and max-grad-norm, and
require nonnegative values for anchor-weight and bond-length-weight. Reject
invalid inputs immediately with the parser’s standard argument error instead of
allowing optimization to be skipped or penalties disabled.
| @pytest.fixture(scope="module") | ||
| def script(): | ||
| return load_script(_SCRIPT_PATH) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add a NumPy-style docstring to the script fixture.
This fixture is a function but is currently undocumented. Add a short Returns section describing the loaded module.
As per coding guidelines, every function and class in Python files must have a NumPy-style docstring.
Proposed fix
`@pytest.fixture`(scope="module")
def script():
+ """Load the patching script for this test module.
+
+ Returns
+ -------
+ ModuleType
+ The dynamically loaded patching-script module.
+ """
return load_script(_SCRIPT_PATH)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @pytest.fixture(scope="module") | |
| def script(): | |
| return load_script(_SCRIPT_PATH) | |
| `@pytest.fixture`(scope="module") | |
| def script(): | |
| """Load the patching script for this test module. | |
| Returns | |
| ------- | |
| ModuleType | |
| The dynamically loaded patching-script module. | |
| """ | |
| return load_script(_SCRIPT_PATH) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/eval/test_patch_output_cif_files.py` around lines 21 - 23, Add a
NumPy-style docstring to the script fixture describing that it returns the
loaded module, including a concise Returns section. Keep the existing fixture
scope and load_script behavior unchanged.
Source: Coding guidelines
| # An independent copy of the script's DEFAULT_RCSB_PATTERN, kept separate on purpose: the | ||
| # test is a spec, not a mirror of the implementation. If the script's default ever changes, | ||
| # this literal does not move with it -- the resulting behavior difference should surface as | ||
| # a failure here for a human to review, rather than being silently tracked. | ||
| _DEFAULT_REGEX = r"grid_search_results/(pdb_[A-Za-z0-9]{8}|[0-9][A-Za-z0-9]{3})" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Align the regex with the script default and require a full folder match.
The supplied script contract still declares grid_search_results/(.{4}) as the CLI default, but this test uses a different pattern explicitly, so the real default is never exercised. Also, this regex has no component boundary: under the script’s re.search(...).group(1) behavior, 4hhb_final matches 4hhb, contradicting the expected None cases at Lines 46-48. Update the source default and test it, or add a boundary such as (?=/) here.
Proposed regex fix
-_DEFAULT_REGEX = r"grid_search_results/(pdb_[A-Za-z0-9]{8}|[0-9][A-Za-z0-9]{3})"
+_DEFAULT_REGEX = (
+ r"grid_search_results/(pdb_[A-Za-z0-9]{8}|[0-9][A-Za-z0-9]{3})(?=/)"
+)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # An independent copy of the script's DEFAULT_RCSB_PATTERN, kept separate on purpose: the | |
| # test is a spec, not a mirror of the implementation. If the script's default ever changes, | |
| # this literal does not move with it -- the resulting behavior difference should surface as | |
| # a failure here for a human to review, rather than being silently tracked. | |
| _DEFAULT_REGEX = r"grid_search_results/(pdb_[A-Za-z0-9]{8}|[0-9][A-Za-z0-9]{3})" | |
| # An independent copy of the script's DEFAULT_RCSB_PATTERN, kept separate on purpose: the | |
| # test is a spec, not a mirror of the implementation. If the script's default ever changes, | |
| # this literal does not move with it -- the resulting behavior difference should surface as | |
| # a failure here for a human to review, rather than being silently tracked. | |
| _DEFAULT_REGEX = ( | |
| r"grid_search_results/(pdb_[A-Za-z0-9]{8}|[0-9][A-Za-z0-9]{3})(?=/)" | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/eval/test_patch_output_cif_files.py` around lines 26 - 30, Align the
test’s _DEFAULT_REGEX with the script’s CLI default so the real default behavior
is exercised, and require the captured four-character identifier to end at a
folder boundary before matching. Update the source default and remove the
explicit override if appropriate, or add a lookahead boundary to _DEFAULT_REGEX
while preserving the expected None results for names such as 4hhb_final.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
docs/CODEBASE_GUIDE.md (1)
17-17: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd languages to the fenced code blocks.
Specify
textfor the architecture diagrams andpythonfor the code-like diagrams so markdownlint no longer reports MD040.Also applies to: 169-169, 266-266
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/CODEBASE_GUIDE.md` at line 17, Add language identifiers to the fenced code blocks in CODEBASE_GUIDE.md: use text for architecture diagrams and python for code-like diagrams, including the additional blocks referenced in the comment. Ensure every affected fence has a language tag so markdownlint MD040 passes.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/CODEBASE_GUIDE.md`:
- Around line 85-91: Update the GenerativeModelInput documentation in
CODEBASE_GUIDE.md, including the referenced featurize() description, to remove
the stale x_init contract and state that featurize() produces cached
conditioning only. Keep the existing guidance that starting coordinates are
created by initialize_from_prior(batch_size=...) during sampling.
In `@docs/IT_OPT_TESTING.md`:
- Around line 152-182: Revise the “Verification results” section to label all
reported metrics as preliminary local results rather than verified efficacy.
Replace definitive claims such as “Verified” and “strongly improves” with
appropriately qualified wording, and explicitly state that efficacy remains
unconfirmed because the scorer is homemade, five proteins are excluded, and
repo-exact validation is still open.
---
Nitpick comments:
In `@docs/CODEBASE_GUIDE.md`:
- Line 17: Add language identifiers to the fenced code blocks in
CODEBASE_GUIDE.md: use text for architecture diagrams and python for code-like
diagrams, including the additional blocks referenced in the comment. Ensure
every affected fence has a language tag so markdownlint MD040 passes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 135bfc09-fde9-47fe-b87b-2031c967db9a
📒 Files selected for processing (7)
docs/CODEBASE_GUIDE.mddocs/IT_OPT_DESIGN.mddocs/IT_OPT_TESTING.mdsrc/sampleworks/core/scalers/latent_optimization.pysrc/sampleworks/models/latent_adapter.pysrc/sampleworks/models/protenix/wrapper.pysrc/sampleworks/utils/guidance_script_utils.py
🚧 Files skipped from review as they are similar to previous changes (4)
- src/sampleworks/models/protenix/wrapper.py
- src/sampleworks/models/latent_adapter.py
- src/sampleworks/core/scalers/latent_optimization.py
- src/sampleworks/utils/guidance_script_utils.py
marcuscollins
left a comment
marcuscollins
left a comment
There was a problem hiding this comment.
Lots of little things to do, I haven't read through absolutely everything but let's start by addressing these comments, and then I'll do another pass. This is great work, just need to make it a little easier to understand and simplify a bit.
| bond_geometry = None | ||
| if self.bond_length_weight > 0: | ||
| # Prefer the model atom array; fall back to the input atom array if it is absent. | ||
| geometry_atom_array = processed.model_atom_array or processed.atom_array |
There was a problem hiding this comment.
Does that happen? If so, when/why? I'd be concerned that they don't have the same dimensions.
GitHub Actions
Waiting
GitHub Actions
Waiting
GitHub Actions
Waiting
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/sampleworks/core/rewards/geometry.py (1)
78-98: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winValidate penalty configuration values.
Reject non-finite or negative
weight,bond_tolerance, andclash_padding. Rejectbond_power < 1. A negative weight makes distorted geometry reduce the optimization loss. A non-finite value propagatesNaNthrough the latent update.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/core/rewards/geometry.py` around lines 78 - 98, Validate the penalty configuration in the constructor before assigning fields: reject non-finite or negative weight, bond_tolerance, and clash_padding, and reject bond_power values below 1. Raise an appropriate argument-validation error identifying the invalid parameter, while preserving valid configuration behavior.
🧹 Nitpick comments (2)
src/sampleworks/core/rewards/geometry.py (2)
61-94: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse NumPy-style docstrings for the new API.
Convert the class and constructor docstrings to use
Parameters,Returns, andNotessections where applicable. Apply the same format to the new helper and loss methods in this file.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/core/rewards/geometry.py` around lines 61 - 94, Convert the BondGeometryReward class and __init__ docstrings, along with the new helper and loss-method docstrings in this module, to NumPy-style sections. Document arguments under Parameters, returned values under Returns where applicable, and behavioral or interface caveats under Notes, preserving the existing content and meaning.Source: Coding guidelines
176-238: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the geometry-loss tensor boundaries.
jaxtypingis already a dependency, and this reward class currently only carries comments about coordinate/array shapes. AnnotatecoordsasFloat[torch | Tensor, "ensemble atoms 3"]forbond_length_loss,collision_loss, and__call__, and annotate return shapes where applicable so the geometry loss boundary is machine-checkable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/core/rewards/geometry.py` around lines 176 - 238, Annotate the geometry-loss API methods bond_length_loss, collision_loss, and __call__ with jaxtyping Float types for coords using the “ensemble atoms 3” shape, and add scalar Tensor return annotations where applicable. Import or reuse the project’s established jaxtyping and torch type symbols without changing the loss behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/CODEBASE_GUIDE.md`:
- Line 79: Update the Scalers contract inventory in CODEBASE_GUIDE.md to include
the model-agnostic latent adapter interface used by LatentOptimization,
alongside the existing protocol references. Extend the wrapper integration
section around cached conditioning to document the adapter’s read/write boundary
for cached s and z latents, preserving the existing wrapper behavior
description.
In `@src/sampleworks/core/rewards/geometry.py`:
- Around line 10-14: Update the BondLengthLossFunction description to
distinguish the powered bond-length hinge from the linear non-bonded collision
hinge. State that bond_length_loss() raises the positive excess to
self.bond_power (quadratic by default), while the collision term remains an
exponent-free positive-part hinge.
---
Outside diff comments:
In `@src/sampleworks/core/rewards/geometry.py`:
- Around line 78-98: Validate the penalty configuration in the constructor
before assigning fields: reject non-finite or negative weight, bond_tolerance,
and clash_padding, and reject bond_power values below 1. Raise an appropriate
argument-validation error identifying the invalid parameter, while preserving
valid configuration behavior.
---
Nitpick comments:
In `@src/sampleworks/core/rewards/geometry.py`:
- Around line 61-94: Convert the BondGeometryReward class and __init__
docstrings, along with the new helper and loss-method docstrings in this module,
to NumPy-style sections. Document arguments under Parameters, returned values
under Returns where applicable, and behavioral or interface caveats under Notes,
preserving the existing content and meaning.
- Around line 176-238: Annotate the geometry-loss API methods bond_length_loss,
collision_loss, and __call__ with jaxtyping Float types for coords using the
“ensemble atoms 3” shape, and add scalar Tensor return annotations where
applicable. Import or reuse the project’s established jaxtyping and torch type
symbols without changing the loss behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: a5e9bca2-197f-4fdc-9d1d-ce313edb47ca
📒 Files selected for processing (8)
docs/CODEBASE_GUIDE.mddocs/IT_OPT_DESIGN.mddocs/IT_OPT_TESTING.mdsrc/sampleworks/core/rewards/geometry.pysrc/sampleworks/core/scalers/latent_optimization.pysrc/sampleworks/models/latent_adapter.pysrc/sampleworks/utils/guidance_constants.pytests/models/test_latent_adapter.py
💤 Files with no reviewable changes (2)
- tests/models/test_latent_adapter.py
- src/sampleworks/models/latent_adapter.py
🚧 Files skipped from review as they are similar to previous changes (4)
- src/sampleworks/utils/guidance_constants.py
- docs/IT_OPT_DESIGN.md
- docs/IT_OPT_TESTING.md
- src/sampleworks/core/scalers/latent_optimization.py
What
Integrates inference-time latent optimization (IT-opt) into Sampleworks as a first-class guidance type, ported from the
it_opt/reference forks (Protenix/Boltz/AF3).Treats a frozen structure predictor as a differentiable sampler and optimizes its cached post-trunk latents — the single representation
sand the pair representationz— directly against an experimental density reward on the denoised structure. No model weights are trained. Supersedes this PR's original scaffold: theLatentAdaptedWrapperaffinek*s + btransform was never wired in, and was replaced by the direct-leaf approach below.Design
LatentOptimizationscaler extracts(s, z)once, clones them into optimizable leaves, and runs one Adam over them across the diffusion trajectory. Each step is one differentiabledenoise → x̂₀, thenloss = reward(x̂₀) + anchor + bond_geometry, backward, per-latent grad-clip,adam.step(); the coordinate state is detached, so the gradient is a greedy per-step∂reward/∂latent, not backprop-through-sampling.LatentAnchor(mean-squared drift from the baseline latents) tethers the optimization;BondGeometryReward(bond-length + steric-clash hinges on the denoised coords) stopszfrom buying density fit with broken geometry.AttrLatentIO(s/zfor Boltz,s_trunk/z_trunkfor Protenix/RF3). No model packages are imported; a 4th/5th model is one more pair of strings, not new code.step()keeps an injected latent leaf attached automatically (detach_unless_leaf: detach cached latents under grad unlessrequires_grad=True), and disables the diffusion shared-vars cache forz(else thezgradient is silently zero).Included
src/sampleworks/core/scalers/latent_optimization.py—LatentOptimizationscaler +LatentAnchorsrc/sampleworks/core/rewards/geometry.py—BondGeometryReward(bond-length + collision penalties)src/sampleworks/models/latent_adapter.py—AttrLatentIO/LatentIO(+tests/models/test_latent_adapter.py,test_latent_optimization.py)src/sampleworks/models/protenix/wrapper.py—detach_unless_leafgradient routingsrc/sampleworks/utils/—GuidanceType.LATENT_OPT,add_latent_opt_args(--which-latent,--learning-rate,--outer-steps,--anchor-weight,--max-grad-norm,--bond-length-weight), and theLATENT_OPTdispatch in_run_guidanceStatus / next steps
Draft — the loop runs end to end for both latents and is reachable from the CLI (
--guidance-type latent_opt) andrun_guidance. Under the density objectivesis effectively inert whilezis the lever; unregularizedzraises the paper's altloc RSCC but drifts off-manifold and adds a few clashes (hence the anchor + bond-geometry terms). Coordinate guidance is the modest, safe baseline it must beat. No mode crosses the paper's 0.8 RSCC bar yet, so efficacy is still open.🤖 Generated with Claude Code
Summary by CodeRabbit
latent_optinference-time latent optimization with configurable latent selection, scheduling, anchoring, gradient clipping, and optional bond-geometry penalties.