refactor(models): remove unused x_init

[xraymemory]

Summary

• Remove the unused GenerativeModelInput.x_init field and wrapper-specific coordinate initialization from Boltz, Protenix, and RF3.
• Remove the now-dead wrapper ensemble_size plumbing while retaining ensemble ownership in trajectory scalers.
• Update tests and wrapper documentation so sampling starts through initialize_from_prior().

Protpardelle is not yet on main and is intentionally outside this change.

Testing

• pixi run -e boltz-osx pytest tests/models/test_model_wrapper_protocol.py -k generative_model_input_carries_only_conditioning
• pixi run -e boltz-osx pytest tests/integration/test_pipeline_integration.py::TestTrajectoryScalerMatrixMock::test_trajectory_scaler_initializes_from_prior_and_returns_guidance_output
• pixi run -e boltz-osx pytest tests/models/boltz/test_boltz_wrapper.py -k 'AnnotateStructureForBoltz or BoltzConfig'
• pixi run -e boltz-osx pytest tests/integration/test_mismatch_integration.py::TestTrajectoryScalers
• pixi run -e boltz-osx ruff check ...
• pixi run -e boltz-osx ruff format --check ...

Closes #309

Summary by CodeRabbit

• Changes

  • Model featurization now returns conditioning data only; initial coordinates are created separately when sampling begins.
  • Updated Boltz, Protenix, and RF3 configuration and annotation interfaces to remove ensemble-size settings.
  • Initial-state handling now derives atom counts from conditioning information.
  • Updated guidance workflows and sampling behavior to use the revised interface.

• Documentation

  • Updated model-wrapper examples and descriptions to reflect conditioning-only inputs and separate prior initialization.

[coderabbitai]

📝 Walkthrough

Walkthrough

Model wrappers now return GenerativeModelInput with conditioning only. Diffusion coordinates are initialized through initialize_from_prior(), while ensemble_size is removed from model configuration and annotation APIs. Tests and examples are updated for the revised contract.

Changes

Conditioning-only model input flow

Layer / File(s)	Summary
Conditioning-only input contract and wrapper outputs src/sampleworks/models/protocol.py, src/sampleworks/models/{boltz,protenix,rf3}/wrapper.py	GenerativeModelInput no longer exposes x_init; Boltz, Protenix, and RF3 featurization returns conditioning-only inputs, with prior initialization retained for diffusion state creation.
Configuration and preprocessing API updates src/sampleworks/models/{boltz,protenix,rf3}/wrapper.py, src/sampleworks/utils/guidance_script_utils.py, AGENTS.md	Model configuration and annotation calls no longer accept or store ensemble_size; examples use conditioning.num_atoms for prior initialization.
Mock wrappers and protocol pipeline validation tests/mocks/model_wrappers.py, tests/integration/test_pipeline_integration.py, tests/models/test_model_wrapper_protocol.py	Mocks and integration tests validate conditioning-only features and separate prior-state initialization before stepping.
Wrapper and annotation regression coverage tests/models/boltz/test_boltz_wrapper.py, tests/models/test_rf3_*.py, tests/integration/test_no_guidance_geometry.py	Boltz tests cover updated configuration and explicit state stepping; RF3 and integration tests stop passing explicit ensemble sizes while preserving existing behavior checks.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Annotator
  participant ModelWrapper
  participant GenerativeModelInput
  participant PriorInitializer
  participant DiffusionStep
  Annotator->>ModelWrapper: annotated structure
  ModelWrapper->>GenerativeModelInput: conditioning
  PriorInitializer->>GenerativeModelInput: conditioning.num_atoms
  PriorInitializer-->>DiffusionStep: initial state
  DiffusionStep->>ModelWrapper: state, time, conditioning
  ModelWrapper-->>DiffusionStep: denoised state

Loading

Possibly related PRs

• diff-use/sampleworks#73: Introduced the GenerativeModelInput and model-wrapper protocol that this PR changes to remove x_init.
• diff-use/sampleworks#182: Also modified RF3 wrapper configuration, annotation, and featurization behavior.

Suggested reviewers: marcuscollins, k-chrispens, abdelsalam-abbas

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 54.29% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly reflects the main change: removing unused x_init from the model protocol and wrappers.
Linked Issues check	✅ Passed	The PR removes x_init from the protocol and updated Boltz, Protenix, and RF3 wrappers/tests to use initialize_from_prior().
Out of Scope Changes check	✅ Passed	The edits stay within x_init removal and related wrapper/test cleanup; no unrelated feature work is evident.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch xraymemory/issue-309-remove-x-init

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Copilot]

Pull request overview

This PR refactors the model-wrapper protocol so sampling state initialization is handled exclusively via FlowModelWrapper.initialize_from_prior(), removing the unused GenerativeModelInput.x_init field and associated wrapper-specific “initial coordinate” plumbing. This aligns wrapper responsibilities with the framework’s design (trajectory scalers own ensemble semantics; wrappers own prior sampling and stepping).

Changes:

• Remove GenerativeModelInput.x_init from src/sampleworks/models/protocol.py and update wrappers (Boltz, Protenix, RF3) to return conditioning-only features.
• Remove wrapper-side ensemble_size configuration/plumbing in Boltz/Protenix/RF3 annotation/config helpers.
• Update tests and docs to initialize sampling state through initialize_from_prior() and assert x_init is absent.

Confidence: ~90% (verified no remaining .x_init references in src/, and no remaining calls passing ensemble_size into the removed wrapper annotation parameters).

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
src/sampleworks/models/protocol.py	Drops x_init from GenerativeModelInput, making conditioning the sole wrapper feature payload.
src/sampleworks/models/boltz/wrapper.py	Removes ensemble_size config and x_init creation; featurize() returns conditioning only.
src/sampleworks/models/protenix/wrapper.py	Removes ensemble_size config and x_init creation; featurize() returns conditioning only.
src/sampleworks/models/rf3/wrapper.py	Removes ensemble_size config and x_init creation; featurize() returns conditioning only.
src/sampleworks/utils/guidance_script_utils.py	Stops passing ensemble_size into wrapper-specific annotation helpers.
tests/models/test_model_wrapper_protocol.py	Updates protocol tests to use initialize_from_prior() and adds a regression test ensuring no x_init.
tests/models/boltz/test_boltz_wrapper.py	Updates wrapper tests to initialize state via initialize_from_prior() and removes ensemble_size expectations.
tests/models/test_rf3_chiral_features.py	Removes now-invalid ensemble_size argument from RF3 structure annotation calls.
tests/models/test_rf3_atom_ordering.py	Removes now-invalid ensemble_size argument from RF3 structure annotation calls.
tests/mocks/model_wrappers.py	Updates mock wrappers to construct conditioning-only GenerativeModelInput and derive prior shape from conditioning.
tests/integration/test_pipeline_integration.py	Strengthens integration test to assert trajectory scalers initialize via prior using conditioning-only features.
tests/integration/test_no_guidance_geometry.py	Removes now-invalid ensemble_size plumbing from wrapper annotation call.
AGENTS.md	Updates the model-wrapper example to omit x_init and derive atom count from conditioning for prior initialization.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[marcuscollins]

Looks good to me. I wonder if we should test all the step methods with a starting set of coordinates that have batch > 1.

I'd like @k-chrispens to take a look as well but I think this is safe so if he doesn't get to it before end of Wednesday let's go ahead and merge, assuming GPU tests pass.

[k-chrispens]

Looks good to me - this x_init was a vestigial structure from back at the beginning, and I don't think there's anything it'd be useful for now.

[xraymemory]

Added Marcus/Karson's requested multi-batch step coverage in 0f53cca. The shared wrapper protocol test now runs step() from initialize_from_prior() at batch sizes 1 and 2, asserting output shape, batch size, and finite coordinates.

Focused protocol test, Ruff, and Ty pass locally; model-specific CI has been re-triggered.

[coderabbitai]

🧹 Nitpick comments (1)

AGENTS.md (1)

412-421: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Reference wrapper skeleton doesn't fail fast on missing features/shape.

n_atoms = features.conditioning.num_atoms if features else shape[-2] will raise an opaque AttributeError/TypeError if both features and shape are None, or if features.conditioning is None. This contradicts the "Fail fast with clear messages" principle stated later in this same document, and diverges from every real wrapper (Boltz/Protenix/RF3), which explicitly raise ValueError for these cases. Since this is the canonical skeleton new wrapper authors will copy, consider aligning it with the real pattern.

♻️ Suggested alignment with real wrapper error handling

     def initialize_from_prior(
         self,
         batch_size: int,
         features: GenerativeModelInput | None = None,
         *,
         shape: tuple[int, ...] | None = None,
     ) -> Tensor:
         """Sample from the prior (typically Gaussian noise)."""
-        n_atoms = features.conditioning.num_atoms if features else shape[-2]
+        if shape is not None:
+            n_atoms = shape[-2]
+        elif features is not None and features.conditioning is not None:
+            n_atoms = features.conditioning.num_atoms
+        else:
+            raise ValueError("Either features or shape must be provided to initialize_from_prior()")
         return torch.randn(batch_size, n_atoms, 3, device=self.device)

🤖 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 `@AGENTS.md` around lines 412 - 421, Update initialize_from_prior to fail fast
with clear ValueError messages when both features and shape are missing, or when
features.conditioning is unavailable. Preserve the existing atom-count
derivation for valid inputs and continue generating the prior tensor through
torch.randn.

🤖 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.

Nitpick comments:
In `@AGENTS.md`:
- Around line 412-421: Update initialize_from_prior to fail fast with clear
ValueError messages when both features and shape are missing, or when
features.conditioning is unavailable. Preserve the existing atom-count
derivation for valid inputs and continue generating the prior tensor through
torch.randn.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 856cb075-fb13-4c7b-ac59-c6199ce1ef49

📥 Commits

Reviewing files that changed from the base of the PR and between 97203e3 and 0f53cca.

📒 Files selected for processing (13)

• AGENTS.md
• src/sampleworks/models/boltz/wrapper.py
• src/sampleworks/models/protenix/wrapper.py
• src/sampleworks/models/protocol.py
• src/sampleworks/models/rf3/wrapper.py
• src/sampleworks/utils/guidance_script_utils.py
• tests/integration/test_no_guidance_geometry.py
• tests/integration/test_pipeline_integration.py
• tests/mocks/model_wrappers.py
• tests/models/boltz/test_boltz_wrapper.py
• tests/models/test_model_wrapper_protocol.py
• tests/models/test_rf3_atom_ordering.py
• tests/models/test_rf3_chiral_features.py

💤 Files with no reviewable changes (2)

• src/sampleworks/models/protocol.py
• tests/models/test_rf3_chiral_features.py