feat(add protpardelle): Add the Protpardelle-1c cc89 protein generative model.#300
feat(add protpardelle): Add the Protpardelle-1c cc89 protein generative model.#300marcuscollins wants to merge 24 commits into
Conversation
|
Warning Review limit reached
Next review available in: 3 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (15)
📝 WalkthroughWalkthroughThis PR adds Protpardelle-1c support across model wrapping, guidance configuration, availability checks, packaging, and tests. It also adds a packaged model configuration and updates unrelated typing, sampler, MSA, runner, and Astera support code. ChangesProtpardelle model integration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant Guidance
participant ProtpardelleWrapper
participant ProtpardelleModel
CLI->>Guidance: parse model and configuration options
Guidance->>ProtpardelleWrapper: initialize with checkpoint and config
Guidance->>ProtpardelleWrapper: featurize structure
loop diffusion steps
Guidance->>ProtpardelleWrapper: step coordinates and noise level
ProtpardelleWrapper->>ProtpardelleModel: forward conditioning
ProtpardelleModel-->>ProtpardelleWrapper: atom37 coordinates
ProtpardelleWrapper-->>Guidance: flattened coordinates
end
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 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 |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (8)
src/sampleworks/utils/guidance_script_arguments.py (2)
619-619: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring is not NumPy-style.
Only a one-line summary is provided; no
Parameterssection documentingparser. 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 valueInconsistent boolean flag conventions.
--protpardelle-sidechain-modeand--protpardelle-jump-stepsuse plainstore_true(no way to explicitly disable via CLI), while--protpardelle-uniform-stepsusesBooleanOptionalAction. Consider standardizing onBooleanOptionalActionfor 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 winAdd 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 winUse a NumPy-style docstring for
get_model_and_device.This changed docstring currently uses
Arguments:and omits a NumPy-styleReturnssection.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 valueRedundant
pandasdependency-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 winMissing NumPy-style docstrings.
__setattr__and thedeviceproperty 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._deviceAlso 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 winModule-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 importsProtpardelleWrapperjust 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_cacheon 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 winSelf-conditioning gate should not use a magic threshold.
self_cond_train_probis the only training-time signal here, so> 0.5will skipstruct_self_condfor 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
⛔ Files ignored due to path filters (1)
pixi.lockis excluded by!**/*.lock
📒 Files selected for processing (16)
.actlignorepyproject.tomlsrc/sampleworks/cli/guidance.pysrc/sampleworks/core/samplers/edm.pysrc/sampleworks/data/cc89_epoch415.yamlsrc/sampleworks/models/protpardelle/__init__.pysrc/sampleworks/models/protpardelle/wrapper.pysrc/sampleworks/utils/guidance_constants.pysrc/sampleworks/utils/guidance_script_arguments.pysrc/sampleworks/utils/guidance_script_utils.pysrc/sampleworks/utils/imports.pytests/cli/test_guidance_cli.pytests/models/protpardelle/__init__.pytests/models/protpardelle/conftest.pytests/models/protpardelle/test_protpardelle_wrapper.pytests/runs/test_runner.py
k-chrispens
left a comment
There was a problem hiding this comment.
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
|
|
||
| from loguru import logger | ||
|
|
||
| from sampleworks.utils.guidance_script_utils import get_model_and_device, run_guidance |
There was a problem hiding this comment.
A quick question, why should we import it here not at the starting of the file?
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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.
| # This should be the way to make the atom mask: | ||
| # atom_mask = atom37_mask_from_aatype(aatype, seq_mask) but it doesn't handle OXT, so: | ||
| atom_mask = torch.zeros( | ||
| (padded_len, ATOM37_NUM_ATOMS), dtype=torch.float, device=self.device | ||
| ) | ||
| atom_mask[atom37_residue_index, atom37_atom_index] = 1 |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
@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.
| # Contiguous residue ordinal (0..L-1) per atom, in atom order across chains. | ||
| residue_starts = struc.get_residue_starts(atom_array) | ||
| is_start = np.zeros(len(atom_names), dtype=np.int64) | ||
| is_start[residue_starts] = 1 | ||
| residue_ordinals = np.cumsum(is_start) - 1 |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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 winRemaining dependency cleanup
sfcalculator-torchis still>=0.3.2in[tool.pixi.feature.boltz-osx]; bump it to>=0.3.3if that was the intended floor.[tool.pixi.feature.protpardelle.pypi-dependencies]still uses an unpinned Git ref forprotpardelle; pin it to a tag if you want reproducible installs.[tool.pixi.feature.protpardelle.pypi-options.dependency-overrides]still duplicates the workspace-levelpandas==2.3.1override; 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 valueAdd
atom37_residue_indexandatom37_atom_indexto_FROZEN.These index maps are set once during
featurizeand 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 valueAdd full NumPy-style docstring to
_msa_data_key_sort_key.The docstring is a single summary line without
Parameters/Returnssections.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
📒 Files selected for processing (25)
Dockerfile.asteradocker/astera/install-ext.shpyproject.tomlsrc/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/sf.pysrc/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/transformer.pysrc/sampleworks/core/forward_models/xray/real_space_density_deps/qfit/unitcell.pysrc/sampleworks/core/samplers/edm.pysrc/sampleworks/core/scalers/fk_steering.pysrc/sampleworks/metrics/lddt.pysrc/sampleworks/metrics/rmsd.pysrc/sampleworks/models/boltz/wrapper.pysrc/sampleworks/models/protenix/wrapper.pysrc/sampleworks/models/protpardelle/wrapper.pysrc/sampleworks/models/rf3/wrapper.pysrc/sampleworks/utils/framework_utils.pysrc/sampleworks/utils/guidance_script_arguments.pysrc/sampleworks/utils/guidance_script_utils.pysrc/sampleworks/utils/imports.pysrc/sampleworks/utils/msa.pytests/cli/test_guidance_cli.pytests/conftest.pytests/models/protpardelle/test_protpardelle_wrapper.pytests/utils/test_atom_array_utils.pytests/utils/test_framework_utils.pytests/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
## Summary Fixes #283 and unbreaks CI on #274 (`mdc/add-protpardelle`). The EDM sampler called `model_wrapper.step(noisy_state, t_hat, eps, features=features)`, passing a **third positional `eps`** that is not part of the `FlowModelWrapper.step(x_t, t, *, features=None)` protocol. Because boltz/protenix/rf3 (and the test mocks) implement the protocol exactly — no extra positional, no `**kwargs` — every sampler-path test failed with: ``` TypeError: <Wrapper>.step() takes 3 positional arguments but 4 positional arguments (and 1 keyword-only argument) were given src/sampleworks/core/samplers/edm.py:428 ``` That single call site accounts for **all 164 failing tests** across the `boltz-dev` / `protenix-dev` / `rf3-dev` jobs (110 `MismatchCaseWrapper` + 54 `MockFlowModelWrapper`). The extra argument was only ever absorbed by Protpardelle's `step()`, where it landed in a `sigma_float` parameter that is **never referenced** in the method body — so removing it changes no behavior. This matches @k-chrispens's note on #283 (already addressed by #267; `eps` is the noise tensor, not a noise level) and restores the resolution that was lost when the branch was force-updated. ## Changes **Protocol reconciliation (#283)** - `core/samplers/edm.py`: drop the extra `eps` positional from the `model_wrapper.step()` call. `eps` is still computed locally and used to build the noisy state and the working-frame guidance math — only the (dead) hand-off to the wrapper is removed. Deleted the temporary "I need to modify the Protocol itself" TODO. - `models/protpardelle/wrapper.py`: remove the unused `sigma_float` parameter so `step()` matches `FlowModelWrapper`. **Lint job (was failing, 8 errors)** - `cli/guidance.py`: remove unused top-level `from loguru import logger` (it's re-imported inside `main()`) — fixes F401 / F811 / I001. - `models/protpardelle/wrapper.py`: use the repo's existing leading-space single-axis jaxtyping convention (`Int[Tensor, " atoms"]`, as in `edm.py`/`step_scalers.py`) to avoid UP037 / F821. - `tests/runs/test_runner.py`: wrap an over-length inline comment (E501). ## Validation - `ruff 0.15.8 check` passes on all touched files locally (same version as CI). - Test suites can't run on this macOS checkout (pixi envs are linux-64); CI on this PR exercises them. Every failing test traced to the removed `edm.py:428` positional and the tests already call `step(x, t, features=...)`, so they should go green. Closes #283. Co-authored-by: xraymemory <[email protected]>
…ardelle temp dirs (#297) ## Summary Two P1 quick wins that live on this branch (both flagged by CodeRabbit on #274). ### #285 — `require_*` decorators don't support bare `@decorator` usage The `require_protpardelle` factory (and its siblings `require_boltz` / `require_protenix` / `require_rf3` / `require_any_model`) only worked as `@require_x("msg")`. Used bare (`@require_x`), the decorated function was passed in as `message`, so the function was never wrapped/guarded — silently unprotected. Added a `callable(message)` check that re-dispatches through the factory, so **all three** forms now work: ```python @require_protpardelle # bare @require_protpardelle() # empty parens @require_protpardelle("custom") # with message ``` The docstrings' bare examples are now accurate. Every existing call site uses `@require_x()`, so no behavior change there. ### #286 — temp dirs leaked in protpardelle tests `tests/models/protpardelle/conftest.py` and `test_protpardelle_wrapper.py` set the model-params dir with `os.environ.setdefault("PROTPARDELLE_MODEL_PARAMS", tempfile.mkdtemp(...))`. `setdefault` evaluates its default eagerly, so `mkdtemp()` ran — and leaked a dir — on **every** session even when the var was already set (e.g. real weights configured, or the other module already set it). Now guard on the var being unset, then register `atexit.register(shutil.rmtree, ..., ignore_errors=True)` so the throwaway dir is removed at process exit. Whichever module runs first owns creation+cleanup; the other no-ops. Targets `mdc/add-protpardelle` since the affected code isn't on `main` yet. Fixes #285. Fixes #286. Co-authored-by: xraymemory <[email protected]>
## Summary Fixes the Protenix MSA typing/runtime issues from ENG-76 / #247. - Introduces an `MSAData` mapping type so MSA helpers can accept string-keyed, int-keyed, or mixed-key sequence mappings without tripping `ty`. - Adds an explicit output type for `_compute_msa()` so returned MSA paths remain `dict[str | int, Path]`. - Makes Protenix MSA sequence ordering deterministic for mixed key types by sorting items by their stringified key before splitting keys/sequences. - Passes `str(out_dir)` to Protenix `msa_search`, matching the external API expectation. ## Validation - `uvx ty check scripts/eval/classify_altloc_regions.py src/sampleworks/utils/msa.py tests/utils/test_msa.py` - `uvx ruff check src/sampleworks/utils/msa.py tests/utils/test_msa.py scripts/eval/classify_altloc_regions.py` - `PYTHONPATH=src uvx --with pytest --with loguru --with requests --with tqdm pytest --noconftest tests/utils/test_msa.py` - `7 passed` Note: the targeted local pytest invocation uses `--noconftest` because the repository-level test conftest imports heavier model/runtime dependencies such as `torch` that are not installed in the lightweight local `uvx` environment. Fixes #247. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Improved MSA input handling to accept broader mapping-based formats beyond plain dictionaries. * Made Protenix-specific MSA generation order deterministic for mixed numeric and string keys. * **Bug Fixes** * Improved consistency in how mixed-key MSA inputs are processed. * **Tests** * Added coverage to verify the mixed numeric/string key sorting behavior used during MSA generation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: xraymemory <[email protected]>
Completes the Protpardelle debugging work from ENG-75 / #257 on top of `mdc/add-protpardelle`. This PR intentionally targets `mdc/add-protpardelle` so the review contains only the final Protpardelle fixes, not the whole in-progress integration branch. Key changes: - Restores the EDM sampler/model-wrapper call shape to the protocol-compatible `step(noisy_state, t_hat, features=features)` after removing the temporary extra `eps` argument. - Implements Protpardelle step-time noise handling via `_expand_noise_level()`, broadcasting scalar or per-batch EDM timesteps to Protpardelle's expected `B x L` tensor. - Moves Protpardelle step inputs/conditioning tensors onto the wrapper device before the model forward pass. - Keeps `prot_lens_per_chain` on CPU when calling Protpardelle's sampling helper to avoid CPU/GPU device mismatch in helper-created residue indices. - Maps Atomworks selenomethionine atom name `SE` into Protpardelle atom37's methionine sulfur slot `SD`. - Fixes `ProtpardelleConditioning` immutability so dataclass construction can complete before selected conditioning fields become frozen. - Allows packaged Protpardelle config/data under `src/sampleworks/data/**` through `.actlignore` so ACTL sync includes the runtime YAML config. - Re-enables the slow `step()` behavior test and adds coverage for the MSE selenium atom mapping. Local static checks: - `uvx ty check src/sampleworks/models/protpardelle/wrapper.py src/sampleworks/core/samplers/edm.py tests/models/protpardelle/test_protpardelle_wrapper.py` - `uvx ruff check src/sampleworks/models/protpardelle/wrapper.py src/sampleworks/core/samplers/edm.py tests/models/protpardelle/test_protpardelle_wrapper.py` Remote ACTL / Protpardelle environment checks: - `pixi run -e protpardelle-dev pytest tests/models/protpardelle -m "not slow"` - `27 passed, 1 deselected` - `pixi run -e protpardelle-dev pytest tests/models/protpardelle/test_protpardelle_wrapper.py::TestStep::test_step_returns_coords` - `1 passed` - Exact reduced Protpardelle guidance smoke command from #257 using the shared checkpoint and 1VME inputs - completed successfully with `Guidance run successfully!` - wrote results to `output/protpardelle-smoke-guided/` Non-blocking warnings observed during the smoke run were expected environment warnings for unavailable optional model/tool paths and missing mirror environment variables; they did not prevent the Protpardelle run from completing. Refs #257. Co-authored-by: xraymemory <[email protected]>
…flags ProtpardelleWrapper.step() calls Protpardelle.forward() directly, and the sampleworks EDM sampler owns the schedule, stochasticity, and step scaling, so the sampling kwargs assembled for Protpardelle.sample() never took effect. Remove the dead _build_sampling_kwargs method, the ProtpardelleConditioning sampling_kwargs field, and every ProtpardelleConfig field except ensemble_size (num_steps, s_churn, step_scale, sidechain_mode, skip_mpnn_proportion, jump_steps, uniform_steps, temperature, top_p, extra_sampling_kwargs), plus the now-orphaned copy/yaml/files/Any imports and DEFAULT_CONFIG loading. Drop the matching --protpardelle-* CLI flags, their wiring in guidance_script_utils, and the associated tests. BREAKING CHANGE: the --protpardelle-s-churn, --protpardelle-step-scale, --protpardelle-sidechain-mode, --protpardelle-skip-mpnn-proportion, --protpardelle-jump-steps, --protpardelle-uniform-steps, --protpardelle-temperature, and --protpardelle-top-p CLI flags are removed. They were inert (ignored during generation); scripts passing them must drop them. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
## Summary Fixes the Protenix MSA typing/runtime issues from ENG-76 / #247. - Introduces an `MSAData` mapping type so MSA helpers can accept string-keyed, int-keyed, or mixed-key sequence mappings without tripping `ty`. - Adds an explicit output type for `_compute_msa()` so returned MSA paths remain `dict[str | int, Path]`. - Makes Protenix MSA sequence ordering deterministic for mixed key types by sorting items by their stringified key before splitting keys/sequences. - Passes `str(out_dir)` to Protenix `msa_search`, matching the external API expectation. ## Validation - `uvx ty check scripts/eval/classify_altloc_regions.py src/sampleworks/utils/msa.py tests/utils/test_msa.py` - `uvx ruff check src/sampleworks/utils/msa.py tests/utils/test_msa.py scripts/eval/classify_altloc_regions.py` - `PYTHONPATH=src uvx --with pytest --with loguru --with requests --with tqdm pytest --noconftest tests/utils/test_msa.py` - `7 passed` Note: the targeted local pytest invocation uses `--noconftest` because the repository-level test conftest imports heavier model/runtime dependencies such as `torch` that are not installed in the lightweight local `uvx` environment. Fixes #247. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Improved MSA input handling to accept broader mapping-based formats beyond plain dictionaries. * Made Protenix-specific MSA generation order deterministic for mixed numeric and string keys. * **Bug Fixes** * Improved consistency in how mixed-key MSA inputs are processed. * **Tests** * Added coverage to verify the mixed numeric/string key sorting behavior used during MSA generation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: xraymemory <[email protected]>
Completes the Protpardelle debugging work from ENG-75 / #257 on top of `mdc/add-protpardelle`. This PR intentionally targets `mdc/add-protpardelle` so the review contains only the final Protpardelle fixes, not the whole in-progress integration branch. Key changes: - Restores the EDM sampler/model-wrapper call shape to the protocol-compatible `step(noisy_state, t_hat, features=features)` after removing the temporary extra `eps` argument. - Implements Protpardelle step-time noise handling via `_expand_noise_level()`, broadcasting scalar or per-batch EDM timesteps to Protpardelle's expected `B x L` tensor. - Moves Protpardelle step inputs/conditioning tensors onto the wrapper device before the model forward pass. - Keeps `prot_lens_per_chain` on CPU when calling Protpardelle's sampling helper to avoid CPU/GPU device mismatch in helper-created residue indices. - Maps Atomworks selenomethionine atom name `SE` into Protpardelle atom37's methionine sulfur slot `SD`. - Fixes `ProtpardelleConditioning` immutability so dataclass construction can complete before selected conditioning fields become frozen. - Allows packaged Protpardelle config/data under `src/sampleworks/data/**` through `.actlignore` so ACTL sync includes the runtime YAML config. - Re-enables the slow `step()` behavior test and adds coverage for the MSE selenium atom mapping. Local static checks: - `uvx ty check src/sampleworks/models/protpardelle/wrapper.py src/sampleworks/core/samplers/edm.py tests/models/protpardelle/test_protpardelle_wrapper.py` - `uvx ruff check src/sampleworks/models/protpardelle/wrapper.py src/sampleworks/core/samplers/edm.py tests/models/protpardelle/test_protpardelle_wrapper.py` Remote ACTL / Protpardelle environment checks: - `pixi run -e protpardelle-dev pytest tests/models/protpardelle -m "not slow"` - `27 passed, 1 deselected` - `pixi run -e protpardelle-dev pytest tests/models/protpardelle/test_protpardelle_wrapper.py::TestStep::test_step_returns_coords` - `1 passed` - Exact reduced Protpardelle guidance smoke command from #257 using the shared checkpoint and 1VME inputs - completed successfully with `Guidance run successfully!` - wrote results to `output/protpardelle-smoke-guided/` Non-blocking warnings observed during the smoke run were expected environment warnings for unavailable optional model/tool paths and missing mirror environment variables; they did not prevent the Protpardelle run from completing. Refs #257. Co-authored-by: xraymemory <[email protected]>
## Summary Fixes the Protenix MSA typing/runtime issues from ENG-76 / #247. - Introduces an `MSAData` mapping type so MSA helpers can accept string-keyed, int-keyed, or mixed-key sequence mappings without tripping `ty`. - Adds an explicit output type for `_compute_msa()` so returned MSA paths remain `dict[str | int, Path]`. - Makes Protenix MSA sequence ordering deterministic for mixed key types by sorting items by their stringified key before splitting keys/sequences. - Passes `str(out_dir)` to Protenix `msa_search`, matching the external API expectation. ## Validation - `uvx ty check scripts/eval/classify_altloc_regions.py src/sampleworks/utils/msa.py tests/utils/test_msa.py` - `uvx ruff check src/sampleworks/utils/msa.py tests/utils/test_msa.py scripts/eval/classify_altloc_regions.py` - `PYTHONPATH=src uvx --with pytest --with loguru --with requests --with tqdm pytest --noconftest tests/utils/test_msa.py` - `7 passed` Note: the targeted local pytest invocation uses `--noconftest` because the repository-level test conftest imports heavier model/runtime dependencies such as `torch` that are not installed in the lightweight local `uvx` environment. Fixes #247. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Improved MSA input handling to accept broader mapping-based formats beyond plain dictionaries. * Made Protenix-specific MSA generation order deterministic for mixed numeric and string keys. * **Bug Fixes** * Improved consistency in how mixed-key MSA inputs are processed. * **Tests** * Added coverage to verify the mixed numeric/string key sorting behavior used during MSA generation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: xraymemory <[email protected]>
fbf9c8e to
9a82574
Compare
3aa618f to
a88eb4d
Compare
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 inmainwith this PR.Closes #286 — protpardelle test temp-dir cleanup was done in #297 (merged into
mdc/add-protpardelle); lands inmainwith 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