feat(add protpardelle): Add the Protpardelle-1c cc89 protein generative model.

[marcuscollins]

Adding Protpardelle, attempt 2 (see PR #274)

Closes #258 — protpardelle sampling options were piped through the guidance CLI in #270 (merged into mdc/add-protpardelle); lands in main with this PR.
Closes #286 — protpardelle test temp-dir cleanup was done in #297 (merged into mdc/add-protpardelle); lands in main with this PR.

Original description:

This is a working version of Protpardelle-1c in SampleWorks. It may still require parameter tuning and other updates--in particular it isn't clear how to use self-conditioning. However we can generate structures with this version and so I'm making this PR as a milestone.

Only a couple significant changes have been made outside the Protpardelle wrapper code itself. One is to make a features class that is unfrozen, so that we can pass self-conditioning input forward to the next Euler step during sampling. Another is that we define different sampling parameters for Protpardelle when instantiating the sampler. There is no CLI control for either change in this PR.

Summary by CodeRabbit

• New Features
  • Added support for the Protpardelle-1c structure prediction model, including sequence-conditioned generation and dedicated command-line options.
  • Added a packaged default Protpardelle configuration and new installation environments.
  • Added model availability checks and clearer setup guidance.
• Improvements
  • Renamed model configuration fields from model to model_name across guidance and grid-search workflows.
  • Added a new sample configuration for training and evaluation.
• Documentation
  • Updated grid-search troubleshooting examples to reflect the revised configuration format.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a Protpardelle model wrapper, guidance and CLI integration, Pixi environments and dependencies, a packaged training configuration, availability checks, and tests. Guidance configuration and grid-search job metadata consistently use model_name instead of model.

Changes

Protpardelle model integration

Layer / File(s)	Summary
Wrapper implementation and coordinate handling src/sampleworks/models/protpardelle/*	Adds sequence conditioning, atom37 mapping, featurization, diffusion stepping, self-conditioning, and Gaussian prior initialization.
Wrapper fixtures and behavioral tests tests/models/protpardelle/*	Adds optional-dependency fixtures and tests for sequence extraction, structure annotation, tensor shapes, atom mapping, gradients, initialization, and sampling.

Guidance configuration and execution

Layer / File(s)	Summary
Model configuration contract src/sampleworks/utils/guidance_constants.py, src/sampleworks/utils/guidance_script_arguments.py, run_grid_search.py	Adds the Protpardelle model key and CLI options, and renames guidance, job, result, and grid-search fields from model to model_name.
Runtime model loading and sampling src/sampleworks/cli/guidance.py, src/sampleworks/utils/guidance_script_utils.py	Passes complete guidance configuration into model loading, resolves the packaged or supplied Protpardelle config, annotates structures, and applies Protpardelle EDM sampler settings.
Availability handling src/sampleworks/utils/imports.py	Adds Protpardelle availability detection, decorators, checks, and aggregate installation guidance while supporting bare decorator usage for existing checks.

Packaging and regression coverage

Layer / File(s)	Summary
Packaged configuration and environments pyproject.toml, src/sampleworks/data/cc89_epoch415.yaml, .actlignore	Adds Protpardelle Pixi features and environments, dependency pins, macOS dependencies, a default YAML configuration, and package-data ignore exceptions.
CLI, metadata, and runner validation tests/cli/*, tests/utils/*, tests/integration/*, tests/runs/*, GRID_SEARCH.md	Updates tests and documentation for model_name, adds Protpardelle CLI parsing coverage, and makes runner assertions and GPU settings environment-stable.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related issues

• diff-use/sampleworks issue 289 — Covers the same GuidanceConfig.model to model_name migration across configuration, job results, callers, and tests.

Possibly related PRs

• diff-use/sampleworks#207 — Introduces the guidance CLI stack extended here with Protpardelle options and model-name plumbing.
• diff-use/sampleworks#255 — Modifies the Pixi configuration areas updated here for macOS and feature environments.
• diff-use/sampleworks#261 — Modifies the .actlignore package-data rules updated here.

Suggested reviewers: denis-zaitsev

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 49.65% 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 is concise and accurately summarizes the addition of the Protpardelle-1c model.
Linked Issues check	✅ Passed	The PR passes Protpardelle options through the guidance flow and cleans up test temp dirs via atexit/fixture cleanup.
Out of Scope Changes check	✅ Passed	The changes stay focused on Protpardelle integration, related config plumbing, tests, and docs, with no clear unrelated additions.

✨ 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 mdc/protpardelle

---

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

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (8)

src/sampleworks/utils/guidance_script_arguments.py (2)

619-619: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Docstring is not NumPy-style.

Only a one-line summary is provided; no Parameters section documenting parser. As per coding guidelines, "**/*: Always include NumPy-style docstrings for 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/utils/guidance_script_arguments.py` at line 619, The
docstring on the Protpardelle guidance argument helper is missing the required
NumPy-style format, so update the function docstring in
guidance_script_arguments.py to document the parser parameter with a proper
Parameters section and keep the one-line summary concise. Use the nearby
function that adds CLI arguments for Protpardelle guidance runs as the target,
and make the docstring consistent with the rest of the module’s NumPy-style
documentation.

Source: Coding guidelines

---

648-651: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Inconsistent boolean flag conventions.

--protpardelle-sidechain-mode and --protpardelle-jump-steps use plain store_true (no way to explicitly disable via CLI), while --protpardelle-uniform-steps uses BooleanOptionalAction. Consider standardizing on BooleanOptionalAction for all protpardelle boolean flags for a consistent, explicit CLI surface.

Also applies to: 658-662, 663-668

🤖 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 648 - 651,
The Protpardelle CLI boolean flags are inconsistent:
`--protpardelle-sidechain-mode` and `--protpardelle-jump-steps` currently use
`store_true`, while `--protpardelle-uniform-steps` uses `BooleanOptionalAction`.
Update the argument definitions in the guidance script parser so all
Protpardelle boolean options use the same boolean optional pattern, giving users
explicit enable/disable forms. Keep the existing flag names and help text
behavior aligned across the related parser entries.

src/sampleworks/cli/guidance.py (1)

10-17: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a NumPy-style docstring to main.

This changed function currently has no docstring.

Proposed docstring

 def main(argv: list[str] | None = None) -> int:
+    """Run the guidance CLI.
+
+    Parameters
+    ----------
+    argv
+        Optional command-line arguments. When ``None``, arguments are read from
+        the process command line.
+
+    Returns
+    -------
+    int
+        Process exit code from the guidance run.
+    """
     config = GuidanceConfig.from_cli(argv)

As per coding guidelines, **/*: Always include NumPy-style docstrings for 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/cli/guidance.py` around lines 10 - 17, Add a NumPy-style
docstring to the main function in guidance.py. Document the argv parameter and
the integer return value, and keep the docstring aligned with the existing
main/GuidanceConfig.from_cli/run_guidance flow so the function follows the
project’s required docstring standard.

Source: Coding guidelines

src/sampleworks/utils/guidance_script_utils.py (1)

184-197: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use a NumPy-style docstring for get_model_and_device.

This changed docstring currently uses Arguments: and omits a NumPy-style Returns section.

Proposed docstring shape

-    Arguments:
-        device_str: The device to use, e.g. "cuda:0" or "cpu".
-        model_checkpoint_path: The path to the model checkpoint.
-        model_type: The type of model to use.
-        config: The configuration object, usually GuidanceConfig, from which extra
-           model-specific settings are read, e.g. a path to a YAML config file.
-        model: The model to use, if provided, helps to prevent re-loading the actual weights.
+    Parameters
+    ----------
+    device_str
+        Device to use, e.g. ``"cuda:0"`` or ``"cpu"``.
+    model_checkpoint_path
+        Path to the model checkpoint.
+    model_type
+        Type of model to use.
+    config
+        Configuration object with model-specific settings.
+    model
+        Pre-loaded model to reuse, when supported by the wrapper.
+
+    Returns
+    -------
+    tuple[torch.device, Any]
+        Selected device and constructed model wrapper.

As per coding guidelines, **/*: Always include NumPy-style docstrings for 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/utils/guidance_script_utils.py` around lines 184 - 197, The
`get_model_and_device` docstring is not following the required NumPy style
because it uses `Arguments:` and is missing a `Returns` section. Update the
docstring in `guidance_script_utils` for `get_model_and_device` to use standard
NumPy-style sections (`Parameters`, `Returns`, and any needed `Raises`), keeping
the existing parameter descriptions but reformatted to match the project’s
docstring conventions.

Source: Coding guidelines

pyproject.toml (1)

64-65: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Redundant pandas dependency-override.

pandas = "==2.3.1" is pinned both under [tool.pixi.feature.protpardelle.pypi-options.dependency-overrides] (line 65) and the workspace-level [tool.pixi.pypi-options.dependency-overrides] (line 187). Having the same override declared in two places invites drift if one is updated but not the other.

Also applies to: 182-188

🤖 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 `@pyproject.toml` around lines 64 - 65, The pandas override is duplicated in
both the feature-level dependency-overrides block and the workspace-level
dependency-overrides block, which can drift over time. Remove the redundant
`pandas = "==2.3.1"` entry from one place and keep the override defined in a
single authoritative location, checking the
`[tool.pixi.feature.protpardelle.pypi-options.dependency-overrides]` and
`[tool.pixi.pypi-options.dependency-overrides]` sections for consistency.

src/sampleworks/models/protpardelle/wrapper.py (3)

134-139: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Missing NumPy-style docstrings.

__setattr__ and the device property have no docstrings, unlike the rest of the file.
As per coding guidelines, **/*: "Always include NumPy-style docstrings for every function and class."

📝 Proposed docstrings

     def __setattr__(self, key, value):
+        """Raise if a frozen conditioning field is reassigned after construction."""
         if key in self._FROZEN and getattr(self, "_initialized", False):

     `@property`
     def device(self) -> torch.device:
+        """torch.device: Device the wrapped model and its tensors live on."""
         return self._device

Also applies to: 347-349

🤖 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/protpardelle/wrapper.py` around lines 134 - 139, Add
NumPy-style docstrings to the wrapper’s __setattr__ method and the device
property so they match the rest of the module’s documentation style. Update the
__setattr__ definition to include a short summary plus Parameters/Returns/Raises
as appropriate, and add the same for the device property accessor (and setter if
present) using the unique symbols __setattr__ and device to locate the spots.

Source: Coding guidelines

---

55-62: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Module-level file I/O at import time.

Reading and parsing the default sampling YAML happens unconditionally when this module is imported. This means every availability check in imports.py (which imports ProtpardelleWrapper just to test importability) also pays this I/O cost, and it makes the module harder to unit-test/mock in isolation (any missing/corrupt file fails the whole import rather than a specific call).

Consider lazily loading and caching these defaults (e.g., via functools.lru_cache on a small loader function) instead of doing it at module scope.

♻️ Proposed lazy-loading refactor

-# get default model configurations from protpardelle-1c/cc89
-sampling_partial_diffusion_all_atom = files("protpardelle").joinpath(
-    "configs/running/sampling_partial_diffusion_allatom.yaml"
-)
-with open(str(sampling_partial_diffusion_all_atom), "r") as f:
-    DEFAULT_CONFIG = yaml.safe_load(f)
-
-DEFAULT_CONDITIONAL_CFG: dict[str, Any] = DEFAULT_CONFIG["sampling"]["conditional_cfg"]
-DEFAULT_PARTIAL_DIFFUSION: dict[str, Any] = DEFAULT_CONFIG["sampling"]["partial_diffusion"]
+@functools.lru_cache(maxsize=1)
+def _load_default_sampling_config() -> dict[str, Any]:
+    """Lazily load and cache the default sampling config from the protpardelle package."""
+    path = files("protpardelle").joinpath("configs/running/sampling_partial_diffusion_allatom.yaml")
+    with open(str(path), "r") as f:
+        return yaml.safe_load(f)["sampling"]

🤖 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/protpardelle/wrapper.py` around lines 55 - 62, The
module-level YAML read in ProtpardelleWrapper is happening at import time, which
makes import checks and tests pay unnecessary I/O and fail on missing files.
Move the sampling config load out of the module scope into a small helper in
wrapper.py, and cache it with a lazy mechanism such as a loader function used by
the DEFAULT_CONFIG/DEFAULT_CONDITIONAL_CFG/DEFAULT_PARTIAL_DIFFUSION access path
so ProtpardelleWrapper only reads the file when those defaults are actually
needed.

---

692-714: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Self-conditioning gate should not use a magic threshold. self_cond_train_prob is the only training-time signal here, so > 0.5 will skip struct_self_cond for any model trained with self-conditioning below that rate. Use an explicit config flag if available, or gate on the intended probability instead.

🤖 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/protpardelle/wrapper.py` around lines 692 - 714, The
self-conditioning decision in the model.forward call is using a hardcoded > 0.5
threshold on self_cond_train_prob, which can incorrectly disable
struct_self_cond for valid models. Update the gating logic in wrapper.py around
self.model.forward to use an explicit config flag if one exists, or otherwise
base the decision directly on the intended self-conditioning setting from
self.model.config.train.self_cond_train_prob rather than a magic threshold.

🤖 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 `@src/sampleworks/data/cc89_epoch415.yaml`:
- Around line 10-11: The bundled Protpardelle config includes developer-local
absolute paths that will ship into installs and break portability. Update the
cc89_epoch415 YAML to remove the `/scratch/users/...` style values and replace
`pdb_paths`, `full_mpnn_model_path`, `ckpt_path`, and `shapes_path` with
package-relative resource references or safe empty defaults. Use the existing
config keys in cc89_epoch415 to locate and change these defaults so the package
no longer depends on local filesystem paths.

In `@src/sampleworks/models/protpardelle/wrapper.py`:
- Line 664: The docstring/type reference for the features parameter is copied
from the wrong model type and should match the wrapper’s actual conditioning
type. Update the documentation around the relevant parameter in wrapper.py to
reference ProtpardelleConditioning instead of BoltzConditioning, keeping it
consistent with the rest of the wrapper and the associated parameter
annotations.

In `@src/sampleworks/utils/guidance_script_utils.py`:
- Around line 252-256: The `get_model_and_device(..., model=...)` reuse path is
being dropped in the `ProtpardelleWrapper` construction, so the preloaded model
is never forwarded and weights are reloaded. Update the `model_wrapper`
instantiation in the `guidance_script_utils` branch to pass through the existing
`model` argument to `ProtpardelleWrapper.__init__`, alongside the current
config, checkpoint, and device values, so the wrapper can reuse the
already-loaded model when provided.
- Around line 69-72: The Protpardelle import warning text in the guidance script
is missing a space before the environment variable name, causing the sentence to
run together and reducing readability. Update the string assembly in the
Protpardelle import warning message so the mention of the environment variable
is properly separated from the preceding word, keeping the message clear in the
Protpardelle-related import path.

---

Nitpick comments:
In `@pyproject.toml`:
- Around line 64-65: The pandas override is duplicated in both the feature-level
dependency-overrides block and the workspace-level dependency-overrides block,
which can drift over time. Remove the redundant `pandas = "==2.3.1"` entry from
one place and keep the override defined in a single authoritative location,
checking the
`[tool.pixi.feature.protpardelle.pypi-options.dependency-overrides]` and
`[tool.pixi.pypi-options.dependency-overrides]` sections for consistency.

In `@src/sampleworks/cli/guidance.py`:
- Around line 10-17: Add a NumPy-style docstring to the main function in
guidance.py. Document the argv parameter and the integer return value, and keep
the docstring aligned with the existing
main/GuidanceConfig.from_cli/run_guidance flow so the function follows the
project’s required docstring standard.

In `@src/sampleworks/models/protpardelle/wrapper.py`:
- Around line 134-139: Add NumPy-style docstrings to the wrapper’s __setattr__
method and the device property so they match the rest of the module’s
documentation style. Update the __setattr__ definition to include a short
summary plus Parameters/Returns/Raises as appropriate, and add the same for the
device property accessor (and setter if present) using the unique symbols
__setattr__ and device to locate the spots.
- Around line 55-62: The module-level YAML read in ProtpardelleWrapper is
happening at import time, which makes import checks and tests pay unnecessary
I/O and fail on missing files. Move the sampling config load out of the module
scope into a small helper in wrapper.py, and cache it with a lazy mechanism such
as a loader function used by the
DEFAULT_CONFIG/DEFAULT_CONDITIONAL_CFG/DEFAULT_PARTIAL_DIFFUSION access path so
ProtpardelleWrapper only reads the file when those defaults are actually needed.
- Around line 692-714: The self-conditioning decision in the model.forward call
is using a hardcoded > 0.5 threshold on self_cond_train_prob, which can
incorrectly disable struct_self_cond for valid models. Update the gating logic
in wrapper.py around self.model.forward to use an explicit config flag if one
exists, or otherwise base the decision directly on the intended
self-conditioning setting from self.model.config.train.self_cond_train_prob
rather than a magic threshold.

In `@src/sampleworks/utils/guidance_script_arguments.py`:
- Line 619: The docstring on the Protpardelle guidance argument helper is
missing the required NumPy-style format, so update the function docstring in
guidance_script_arguments.py to document the parser parameter with a proper
Parameters section and keep the one-line summary concise. Use the nearby
function that adds CLI arguments for Protpardelle guidance runs as the target,
and make the docstring consistent with the rest of the module’s NumPy-style
documentation.
- Around line 648-651: The Protpardelle CLI boolean flags are inconsistent:
`--protpardelle-sidechain-mode` and `--protpardelle-jump-steps` currently use
`store_true`, while `--protpardelle-uniform-steps` uses `BooleanOptionalAction`.
Update the argument definitions in the guidance script parser so all
Protpardelle boolean options use the same boolean optional pattern, giving users
explicit enable/disable forms. Keep the existing flag names and help text
behavior aligned across the related parser entries.

In `@src/sampleworks/utils/guidance_script_utils.py`:
- Around line 184-197: The `get_model_and_device` docstring is not following the
required NumPy style because it uses `Arguments:` and is missing a `Returns`
section. Update the docstring in `guidance_script_utils` for
`get_model_and_device` to use standard NumPy-style sections (`Parameters`,
`Returns`, and any needed `Raises`), keeping the existing parameter descriptions
but reformatted to match the project’s docstring conventions.

🪄 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

Run ID: c9f36ab8-9000-4ad9-bda2-59022088c5f3

📥 Commits

Reviewing files that changed from the base of the PR and between 15b0063 and 100a456.

⛔ Files ignored due to path filters (1)

• pixi.lock is excluded by !**/*.lock

📒 Files selected for processing (16)

• .actlignore
• pyproject.toml
• src/sampleworks/cli/guidance.py
• src/sampleworks/core/samplers/edm.py
• src/sampleworks/data/cc89_epoch415.yaml
• src/sampleworks/models/protpardelle/__init__.py
• src/sampleworks/models/protpardelle/wrapper.py
• src/sampleworks/utils/guidance_constants.py
• src/sampleworks/utils/guidance_script_arguments.py
• src/sampleworks/utils/guidance_script_utils.py
• src/sampleworks/utils/imports.py
• tests/cli/test_guidance_cli.py
• tests/models/protpardelle/__init__.py
• tests/models/protpardelle/conftest.py
• tests/models/protpardelle/test_protpardelle_wrapper.py
• tests/runs/test_runner.py

[k-chrispens]

Looks good!! A few downstream issues we should track. Another main concern is that we don't have a GPU test for this, which will be important for integration and regressions down the line. We should also add this to the list of modelwrappers in tests/conftest.py so it gets passed through the current integration tests

[smallfishabc]

A quick question, why should we import it here not at the starting of the file?

[marcuscollins]

@smallfishabc Good question, I think something must have gotten messed up with the commit history, like a rebase was missed somewhere. I shouldn't need to change this file at all.

In general, imports inside functions are made to speed things up--e.g. if I only needed to import something when one particular function was called, and it was a slow import, I might guard it like this.

[marcuscollins]

Oh, I see what's going on.... Hmmm. This is a change I should have made separately, but it had to do with reconciling arguments across the different models (protenix, protpardelle, rf3, boltz1/2)

[DorisMai]

I have some confusion points about ProtpardelleWrapper, the main one is probably regarding the topology, as other things might get sorted out in future prs.

[DorisMai]

What happens if the input reference structure has missing atoms? it seems that the num_atom count from this atom_mask would be different from that using other models?

[marcuscollins]

@DorisMai do you mean what happens if we try to align this structure to the reference? We only convert to (and from) the atom37 layout inside ProtpardelleWrapper.step, the alignment is done inside AF3EDMSampler.step, the same for all models, and the reconciler Karson created handles this by aligning only on atoms the two have in common.

It turns out atom_mask's only use is to convert back from atom37 to "normal" coordinates.

[DorisMai]

same concern as before (line 418-423), if there are missing residues (flexible loop), would this numbering be off then? might be good to add a real structure test in addition to the existing toy/minimal ones.

[marcuscollins]

This is only used to map from an AtomArray{Stack} to the atom37 format, so it's a contiguous tensor index, not an actual residue index. Now, there is a question about whether those missing residues would be in our input to begin with. The models we use generally require a complete sequence, so for the coordinates the model is operating on, there shouldn't be any missing residues.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

pyproject.toml (1)

27-28: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Remaining dependency cleanup

• sfcalculator-torch is still >=0.3.2 in [tool.pixi.feature.boltz-osx]; bump it to >=0.3.3 if that was the intended floor.
• [tool.pixi.feature.protpardelle.pypi-dependencies] still uses an unpinned Git ref for protpardelle; pin it to a tag if you want reproducible installs.
• [tool.pixi.feature.protpardelle.pypi-options.dependency-overrides] still duplicates the workspace-level pandas==2.3.1 override; keep only one source of truth.

🤖 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 `@pyproject.toml` around lines 27 - 28, Update the boltz-osx pypi-dependencies
entry to require sfcalculator-torch >=0.3.3, pin protpardelle’s Git dependency
in its pypi-dependencies to a specific tag, and remove the duplicate
pandas==2.3.1 override from protpardelle’s dependency-overrides so the
workspace-level override remains the sole source of truth.

🧹 Nitpick comments (2)

src/sampleworks/models/protpardelle/wrapper.py (1)

101-110: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add atom37_residue_index and atom37_atom_index to _FROZEN.

These index maps are set once during featurize and never modified, but they're absent from _FROZEN, so they can be accidentally mutated after initialization. Adding them keeps the immutability contract consistent with the other structural fields.

♻️ Proposed fix

     _FROZEN = frozenset(
         {
             "aatype",
             "seq_mask",
             "residue_index",
             "chain_index",
             "atom_mask",
+            "atom37_residue_index",
+            "atom37_atom_index",
             "sequences",
         }
     )

🤖 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/protpardelle/wrapper.py` around lines 101 - 110, Add
"atom37_residue_index" and "atom37_atom_index" to the _FROZEN frozenset in the
relevant model wrapper, ensuring these featurize-initialized index maps are
protected from mutation like the other structural fields.

src/sampleworks/utils/msa.py (1)

25-29: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add full NumPy-style docstring to _msa_data_key_sort_key.

The docstring is a single summary line without Parameters/Returns sections.

As per coding guidelines, **/*: "Always include NumPy-style docstrings for every function and class."

📝 Proposed docstring fix

 def _msa_data_key_sort_key(key: str | int) -> tuple[int, int | str]:
-    """Return a deterministic sort key for supported MSA sequence keys."""
+    """Return a deterministic sort key for supported MSA sequence keys.
+
+    Parameters
+    ----------
+    key : str | int
+        A target/chain key from an ``MSAData`` mapping.
+
+    Returns
+    -------
+    tuple[int, int | str]
+        A tuple placing all ``int`` keys (sorted numerically) before all
+        ``str`` keys (sorted lexicographically).
+    """

🤖 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/msa.py` around lines 25 - 29, Expand the
`_msa_data_key_sort_key` docstring to NumPy style by documenting the `key`
parameter, its supported `int`/`str` types, and the returned `(int, int | str)`
sort-key tuple in `Parameters` and `Returns` sections.

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.

Outside diff comments:
In `@pyproject.toml`:
- Around line 27-28: Update the boltz-osx pypi-dependencies entry to require
sfcalculator-torch >=0.3.3, pin protpardelle’s Git dependency in its
pypi-dependencies to a specific tag, and remove the duplicate pandas==2.3.1
override from protpardelle’s dependency-overrides so the workspace-level
override remains the sole source of truth.

---

Nitpick comments:
In `@src/sampleworks/models/protpardelle/wrapper.py`:
- Around line 101-110: Add "atom37_residue_index" and "atom37_atom_index" to the
_FROZEN frozenset in the relevant model wrapper, ensuring these
featurize-initialized index maps are protected from mutation like the other
structural fields.

In `@src/sampleworks/utils/msa.py`:
- Around line 25-29: Expand the `_msa_data_key_sort_key` docstring to NumPy
style by documenting the `key` parameter, its supported `int`/`str` types, and
the returned `(int, int | str)` sort-key tuple in `Parameters` and `Returns`
sections.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: dec962de-033f-474a-bce6-841fb7901371

📥 Commits

Reviewing files that changed from the base of the PR and between ef4e143 and a9ef868.

📒 Files selected for processing (25)

• Dockerfile.astera
• docker/astera/install-ext.sh
• pyproject.toml
• src/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/sf.py
• src/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/transformer.py
• src/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/unitcell.py
• src/sampleworks/core/samplers/edm.py
• src/sampleworks/core/scalers/fk_steering.py
• src/sampleworks/metrics/lddt.py
• src/sampleworks/metrics/rmsd.py
• src/sampleworks/models/boltz/wrapper.py
• src/sampleworks/models/protenix/wrapper.py
• src/sampleworks/models/protpardelle/wrapper.py
• src/sampleworks/models/rf3/wrapper.py
• src/sampleworks/utils/framework_utils.py
• src/sampleworks/utils/guidance_script_arguments.py
• src/sampleworks/utils/guidance_script_utils.py
• src/sampleworks/utils/imports.py
• src/sampleworks/utils/msa.py
• tests/cli/test_guidance_cli.py
• tests/conftest.py
• tests/models/protpardelle/test_protpardelle_wrapper.py
• tests/utils/test_atom_array_utils.py
• tests/utils/test_framework_utils.py
• tests/utils/test_msa.py

💤 Files with no reviewable changes (5)

• src/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/transformer.py
• src/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/unitcell.py
• src/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/sf.py
• tests/cli/test_guidance_cli.py
• src/sampleworks/utils/guidance_script_arguments.py

✅ Files skipped from review due to trivial changes (8)

• docker/astera/install-ext.sh
• src/sampleworks/metrics/rmsd.py
• src/sampleworks/metrics/lddt.py
• src/sampleworks/utils/framework_utils.py
• tests/utils/test_framework_utils.py
• tests/utils/test_atom_array_utils.py
• src/sampleworks/models/boltz/wrapper.py
• src/sampleworks/models/protenix/wrapper.py

🚧 Files skipped from review as they are similar to previous changes (2)

• src/sampleworks/utils/imports.py
• src/sampleworks/utils/guidance_script_utils.py

[manzuoni-astera]

Blocking issues remain in the Protpardelle integration at current head a88eb4d.

1. Atom37 output is flattened in canonical mask order instead of being mapped back through explicit residue/atom indices to the input AtomArray order. Common PDB atom ordering therefore associates generated coordinates with the wrong atoms.
2. Residue positions are derived from only residues present in the AtomArray while conditioning uses full canonical sequences. Missing loops shift every subsequent mapping.
3. Non-protein chains are skipped in sequence conditioning but remain in atom37 conversion, causing unknown atoms or out-of-range/misaligned residue indices.
4. The published/Astera image does not install the Protpardelle environment or provide checkpoint/model-parameter resolution, and CI never runs protpardelle-dev; the new tests therefore skip in existing jobs.

Please use explicit atom index maps/reconciliation for realistic incomplete and heterogeneous structures, provision the image runtime, and add a Protpardelle CI smoke test with realistically ordered atoms.

[coderabbitai]

Actionable comments posted: 5

🤖 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 `@pyproject.toml`:
- Around line 27-28: Update the sfcalculator-torch dependency in the
pypi-dependencies configuration from >=0.3.2 to >=0.3.3, leaving the other
dependency constraints unchanged.
- Around line 50-55: Move the PyTorch dependency from the Pixi dependencies
table into the appropriate pypi-dependencies table, renaming it to torch while
preserving the >=2.6.0,<2.8 constraint so the GPU-enabled wheel is selected.
Update the protpardelle entry in
[tool.pixi.feature.protpardelle.pypi-dependencies] to include a fixed git
revision, using the reviewed commit pin for reproducible builds.

In `@run_grid_search.py`:
- Line 590: Update the existing-run key construction around
existing_run.get("model_name") to fall back to the legacy "model" field when
"model_name" is absent, then normalize the migrated entry to use "model_name" in
the merged output. Preserve current deduplication behavior for new-schema
results while ensuring legacy results match new_run_keys.

In `@src/sampleworks/models/protpardelle/wrapper.py`:
- Around line 116-121: Add NumPy-style docstrings to the __setattr__ method and
the device property, documenting their purpose and relevant parameters or return
value. Keep the existing frozen-attribute enforcement and device behavior
unchanged.
- Around line 193-204: Update the chain-processing flow around the sequence
collection and asym_unit mapping to retain only L-polypeptide chain IDs, then
filter the atom array to those IDs before applying sequence-derived layout
logic. Ensure mixed inputs cannot pass ligand or non-protein atoms into the
mapping; apply the same filtering at the corresponding logic around the
additional referenced section.

🪄 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

Run ID: 23207694-ac9a-413d-9c35-cfd97770dc94

📥 Commits

Reviewing files that changed from the base of the PR and between a9ef868 and 543c63a.

⛔ Files ignored due to path filters (1)

• pixi.lock is excluded by !**/*.lock

📒 Files selected for processing (21)

• .actlignore
• GRID_SEARCH.md
• pyproject.toml
• run_grid_search.py
• src/sampleworks/cli/guidance.py
• src/sampleworks/data/cc89_epoch415.yaml
• src/sampleworks/models/protpardelle/__init__.py
• src/sampleworks/models/protpardelle/wrapper.py
• src/sampleworks/utils/guidance_constants.py
• src/sampleworks/utils/guidance_script_arguments.py
• src/sampleworks/utils/guidance_script_utils.py
• src/sampleworks/utils/imports.py
• tests/cli/test_guidance_cli.py
• tests/integration/test_mismatch_integration.py
• tests/models/protpardelle/__init__.py
• tests/models/protpardelle/conftest.py
• tests/models/protpardelle/test_protpardelle_wrapper.py
• tests/runs/test_runner.py
• tests/utils/conftest.py
• tests/utils/test_guidance_script_arguments.py
• tests/utils/test_guidance_script_utils.py

🚧 Files skipped from review as they are similar to previous changes (10)

• src/sampleworks/models/protpardelle/init.py
• .actlignore
• src/sampleworks/utils/guidance_constants.py
• tests/models/protpardelle/conftest.py
• src/sampleworks/data/cc89_epoch415.yaml
• src/sampleworks/cli/guidance.py
• tests/runs/test_runner.py
• tests/models/protpardelle/test_protpardelle_wrapper.py
• src/sampleworks/utils/guidance_script_utils.py
• src/sampleworks/utils/imports.py

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Update sfcalculator-torch minimum version.

A previous review noted that sfcalculator-torch should be pinned to >=0.3.3, but it remains >=0.3.2.

🐛 Proposed fix

-pypi-dependencies = {"boltz" = "*", "cuequivariance-torch" = "*", "rdkit" = ">=2025.3.6", "joblib" = "*", "sfcalculator-torch" = ">=0.3.2"}
+pypi-dependencies = {"boltz" = "*", "cuequivariance-torch" = "*", "rdkit" = ">=2025.3.6", "joblib" = "*", "sfcalculator-torch" = ">=0.3.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 `@pyproject.toml` around lines 27 - 28, Update the sfcalculator-torch
dependency in the pypi-dependencies configuration from >=0.3.2 to >=0.3.3,
leaving the other dependency constraints unchanged.

[coderabbitai]

🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win

Fix PyTorch dependency and pin protpardelle git revision.

Two issues identified in previous reviews remain unaddressed in this block:

1. pytorch under the Pixi dependencies table resolves to the CPU version. It should be moved to pypi-dependencies as torch = ">=2.6.0,<2.8" to ensure the correct GPU-enabled wheel is installed.
2. The protpardelle git dependency is missing a specific commit pin (e.g., rev = "ee378400f25b801fa481028000f9060183d7fb4c") to ensure reproducible builds.

🐛 Proposed fix

-pytorch = ">=2.6.0,<2.8"
 tqdm = "*"
 typer = "*"
 
 [tool.pixi.feature.protpardelle.pypi-dependencies]
+torch = ">=2.6.0,<2.8"
-protpardelle = {git = "https://github.com/ProteinDesignLab/protpardelle-1c.git"}
+protpardelle = {git = "https://github.com/ProteinDesignLab/protpardelle-1c.git", rev = "ee378400f25b801fa481028000f9060183d7fb4c"}

🤖 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 `@pyproject.toml` around lines 50 - 55, Move the PyTorch dependency from the
Pixi dependencies table into the appropriate pypi-dependencies table, renaming
it to torch while preserving the >=2.6.0,<2.8 constraint so the GPU-enabled
wheel is selected. Update the protpardelle entry in
[tool.pixi.feature.protpardelle.pypi-dependencies] to include a fixed git
revision, using the reviewed commit pin for reproducible builds.

[coderabbitai]

🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

Migrate legacy job results to prevent duplication and schema mismatch.

Because the model field was renamed to model_name, reading a results.json from a previous grid search will yield None for existing_run.get("model_name"). This causes the old run's key to mismatch new_run_keys, resulting in duplicated entries in the merged results rather than proper deduplication when a run is re-executed.

Migrate the key on the fly to maintain backward compatibility, properly deduplicate runs, and standardize the resulting output schema.

🐛 Proposed fix to migrate legacy results

     merged_runs = [asdict(r) for r in results]
     for existing_run in existing_runs:
+        if "model_name" not in existing_run and "model" in existing_run:
+            existing_run["model_name"] = existing_run.pop("model")
         key = (
             existing_run.get("protein"),
             existing_run.get("model_name"),

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

Suggested change

	existing_run.get("model_name"),
	if "model_name" not in existing_run and "model" in existing_run:
	existing_run["model_name"] = existing_run.pop("model")
	existing_run.get("model_name"),

🤖 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 `@run_grid_search.py` at line 590, Update the existing-run key construction
around existing_run.get("model_name") to fall back to the legacy "model" field
when "model_name" is absent, then normalize the migrated entry to use
"model_name" in the merged output. Preserve current deduplication behavior for
new-schema results while ensuring legacy results match new_run_keys.

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Add docstrings for __setattr__ and device.

Proposed fix

 def __setattr__(self, key, value):
+    """Set an attribute while preserving immutable conditioning fields."""
     if key in self._FROZEN and getattr(self, "_initialized", False):

 `@property`
 def device(self) -> torch.device:
+    """Return the device on which the wrapped model runs."""
     return self._device

As per coding guidelines, every function and class must include a NumPy-style docstring.

Also applies to: 259-261

🤖 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/protpardelle/wrapper.py` around lines 116 - 121, Add
NumPy-style docstrings to the __setattr__ method and the device property,
documenting their purpose and relevant parameters or return value. Keep the
existing frozen-attribute enforcement and device behavior unchanged.

Source: Coding guidelines

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Filter the atom array to the retained protein chains.

Non-protein chains are reported as skipped, but the full asym_unit is still mapped. Ligand atoms can trigger ValueError, while extra residues can exceed the sequence-derived layout. Filter by the retained protein chain IDs—or reject mixed inputs consistently.

Also applies to: 321-325

🤖 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/protpardelle/wrapper.py` around lines 193 - 204,
Update the chain-processing flow around the sequence collection and asym_unit
mapping to retain only L-polypeptide chain IDs, then filter the atom array to
those IDs before applying sequence-derived layout logic. Ensure mixed inputs
cannot pass ligand or non-protein atoms into the mapping; apply the same
filtering at the corresponding logic around the additional referenced section.