Crystal-mates + cache meta-data - #88
Conversation
📝 WalkthroughWalkthroughThe dataset now supports filtered protein and ligand symmetry mates, deduplicates mate entities, maps mate atoms to ASU embedding rows, validates cache filter metadata, and prevents mate atoms from anchoring generated waters. Filtering defaults and documentation were updated. ChangesSymmetry-aware dataset preprocessing
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant DatasetPreprocessor
participant GeometryCache
participant EmbeddingAnnotator
participant FlowMatcher
DatasetPreprocessor->>GeometryCache: write mate flags and embedding indices
GeometryCache->>EmbeddingAnnotator: load cached graph metadata
EmbeddingAnnotator->>EmbeddingAnnotator: inherit ASU rows for mates
FlowMatcher->>FlowMatcher: exclude mate atoms from water anchors
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 Generate docstrings
🧪 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.
Pull request overview
This PR improves how crystal symmetry mates are handled in dataset preprocessing and sampling, adds cache provenance metadata to prevent mixing incompatible caches, and updates defaults/tests/docs to reflect the new behavior.
Changes:
- Refactors crystal-mate collection to exclude mate waters, optionally include ligand mates, and deduplicate coincident mate atoms/ligands; adds
is_mateandemb_res_idxprovenance fields. - Anchors the uniform-ball water prior on ASU (non-mate) atoms when mates exist.
- Writes
_filter_meta.jsoninto each geometry cache directory and refuses to read/extend caches built with different filter/graph settings; updates defaults and expands test coverage.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/test_flow.py | Adds unit tests for anchor_mask behavior and per-graph anchoring semantics in uniform-ball sampling. |
| tests/test_dataset.py | Adds extensive unit/integration tests for mate deduplication, residue-id parsing, mate/ligand separation, embedding inheritance, and cache filter sidecar behavior. |
| tests/conftest.py | Updates PDB fixture documentation to note special-position water scenario for 4h0b. |
| src/flow.py | Adds anchor_mask support to uniform-ball sampling and uses is_mate to anchor sampling on ASU atoms. |
| src/dataset.py | Implements mate/ligand selection changes, deduplication utilities, is_mate/emb_res_idx fields, and _filter_meta.json cache provenance enforcement. |
| src/constants.py | Removes outdated comment text about embedding dims (constants unchanged). |
| scripts/train.py | Updates default thresholds and clarifies help text about cache-write-time filtering implications. |
| scripts/inference.py | Updates default filter thresholds to match training/cache expectations. |
| scripts/generate_slae_embeddings.py | Updates script note to reflect current usage status more accurately. |
| README.md | Documents new mate-handling rules, node ordering/masks, embedding indexing, and cache provenance sidecar behavior. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| is_mate = getattr(batch_data["protein"], "is_mate", None) | ||
| anchor_mask = ( | ||
| ~is_mate.bool() if is_mate is not None and bool(is_mate.any()) else None | ||
| ) |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
src/flow.py (1)
105-111: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueThe starvation fallback is global, not per graph.
If one graph in a batch has no eligible anchor, the mask is dropped for the whole batch. Every other graph then anchors on mate atoms too, which is the behaviour the mask exists to prevent. With
is_matefrom the dataset a starved graph means all-mate protein nodes, which is unlikely, so this is a batch-quality concern rather than a defect.Consider falling back per graph: keep the mask where
counts > 0and re-enable all atoms only for the starved graphs. The docstring at lines 77-80 already states the current global behaviour, so update it if you change this.🤖 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/flow.py` around lines 105 - 111, Update the anchor filtering logic around eligible and counts to apply starvation fallback per graph: retain eligible anchors for graphs with counts greater than zero, while allowing all protein atoms only for graphs with zero eligible anchors. Update the nearby docstring describing global fallback to document the per-graph behavior, preserving correct protein_pos and batch_p alignment.src/dataset.py (1)
1267-1275: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider including mate-ligand coordinates in the distance-filter reference.
The reference adds
crystal_data["mate_coords"]only. Mate ligand atoms become protein-type nodes later at lines 1464-1467, so a water that contacts only a mate ligand surface still failsmax_protein_distand is dropped, although its context atom is in the graph. That is the same asymmetry the comment says the mate coordinates exist to remove.Note that the dedup pass has not run yet at this point, so the raw
crystal_data["mate_ligand_coords"]would be the value to concatenate.🤖 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/dataset.py` around lines 1267 - 1275, The quality-filter reference in the include_mates branch must include both mate protein coordinates and raw crystal_data["mate_ligand_coords"], so waters contacting mate ligands satisfy max_protein_dist. Update the concatenation that builds filter_protein_coords while preserving the existing non-mate path.
🤖 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/dataset.py`:
- Around line 1117-1129: Update the write branch around the sidecar creation so
it emits the existing missing-metadata warning before writing when geometry_dir
already contains .pt entries. Ensure this warning also occurs for preprocessing
runs that create FILTER_META_FILENAME, while preserving the atomic
temporary-file replacement and existing warning text.
- Around line 1405-1411: The mate-ligand grouping keys currently omit the
symmetry-object identity, allowing atoms from distinct operators to collapse.
Update dedup_mate_ligands_by_residue at src/dataset.py:312-322 and
mate_residue_keys at src/dataset.py:1405-1411 to include atom.model alongside
chain and resi, preserving the existing ordering and index-mapping behavior.
In `@tests/test_dataset.py`:
- Around line 2917-2919: Update the assertion near the mate-ordering test to
match the documented layout: verify that every node after the ASU boundary is
not necessarily a mate, while asserting the mate protein block is contiguous in
its expected range. Revise the misleading comment accordingly, and avoid using
the total mate count to slice across ASU ligand nodes.
---
Nitpick comments:
In `@src/dataset.py`:
- Around line 1267-1275: The quality-filter reference in the include_mates
branch must include both mate protein coordinates and raw
crystal_data["mate_ligand_coords"], so waters contacting mate ligands satisfy
max_protein_dist. Update the concatenation that builds filter_protein_coords
while preserving the existing non-mate path.
In `@src/flow.py`:
- Around line 105-111: Update the anchor filtering logic around eligible and
counts to apply starvation fallback per graph: retain eligible anchors for
graphs with counts greater than zero, while allowing all protein atoms only for
graphs with zero eligible anchors. Update the nearby docstring describing global
fallback to document the per-graph behavior, preserving correct protein_pos and
batch_p alignment.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 5df00e45-9dfb-45a4-8de1-2cc84a420150
📒 Files selected for processing (10)
README.mdscripts/generate_slae_embeddings.pyscripts/inference.pyscripts/train.pysrc/constants.pysrc/dataset.pysrc/flow.pytests/conftest.pytests/test_dataset.pytests/test_flow.py
💤 Files with no reviewable changes (1)
- src/constants.py
| if write: | ||
| self.geometry_dir.mkdir(parents=True, exist_ok=True) | ||
| # Written through a temp file: cache builds fan out over processes, | ||
| # and a reader must never catch a half-written sidecar. | ||
| tmp_path = meta_path.with_suffix(f".{os.getpid()}.tmp") | ||
| with open(tmp_path, "w") as f: | ||
| json.dump(current, f, indent=2) | ||
| tmp_path.replace(meta_path) | ||
| elif any(self.geometry_dir.glob("*.pt")): | ||
| logger.warning( | ||
| f"{self.geometry_dir} has no {FILTER_META_FILENAME}; the settings " | ||
| "its entries were built with cannot be verified." | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Warn when a preprocessing run stamps a directory that already holds entries.
The write branch runs whenever preprocess=True and no sidecar exists. If the directory already holds .pt files from a run made before the sidecar existed, this run writes its own settings over that unverified population. The sidecar then asserts a provenance nobody checked, and every later run trusts it. The warning branch is elif, so it never fires in this case.
Emit the same warning before claiming a non-empty unlabelled directory.
🛠️ Proposed fix
if write:
+ if any(self.geometry_dir.glob("*.pt")):
+ logger.warning(
+ f"{self.geometry_dir} already holds entries but no "
+ f"{FILTER_META_FILENAME}; stamping it with this run's settings "
+ "does not verify the entries already there."
+ )
self.geometry_dir.mkdir(parents=True, exist_ok=True)📝 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.
| if write: | |
| self.geometry_dir.mkdir(parents=True, exist_ok=True) | |
| # Written through a temp file: cache builds fan out over processes, | |
| # and a reader must never catch a half-written sidecar. | |
| tmp_path = meta_path.with_suffix(f".{os.getpid()}.tmp") | |
| with open(tmp_path, "w") as f: | |
| json.dump(current, f, indent=2) | |
| tmp_path.replace(meta_path) | |
| elif any(self.geometry_dir.glob("*.pt")): | |
| logger.warning( | |
| f"{self.geometry_dir} has no {FILTER_META_FILENAME}; the settings " | |
| "its entries were built with cannot be verified." | |
| ) | |
| if write: | |
| if any(self.geometry_dir.glob("*.pt")): | |
| logger.warning( | |
| f"{self.geometry_dir} already holds entries but no " | |
| f"{FILTER_META_FILENAME}; stamping it with this run's settings " | |
| "does not verify the entries already there." | |
| ) | |
| self.geometry_dir.mkdir(parents=True, exist_ok=True) | |
| # Written through a temp file: cache builds fan out over processes, | |
| # and a reader must never catch a half-written sidecar. | |
| tmp_path = meta_path.with_suffix(f".{os.getpid()}.tmp") | |
| with open(tmp_path, "w") as f: | |
| json.dump(current, f, indent=2) | |
| tmp_path.replace(meta_path) | |
| elif any(self.geometry_dir.glob("*.pt")): | |
| logger.warning( | |
| f"{self.geometry_dir} has no {FILTER_META_FILENAME}; the settings " | |
| "its entries were built with cannot be verified." | |
| ) |
🧰 Tools
🪛 ast-grep (0.45.0)
[warning] 1121-1121: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(tmp_path, "w")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/dataset.py` around lines 1117 - 1129, Update the write branch around the
sidecar creation so it emits the existing missing-metadata warning before
writing when geometry_dir already contains .pt entries. Ensure this warning also
occurs for preprocessing runs that create FILTER_META_FILENAME, while preserving
the atomic temporary-file replacement and existing warning text.
| # compute mate residue indices (group atoms by actual residue) | ||
| mate_residue_keys = [(a.chain, a.resi) for a in mate_atoms] | ||
| unique_mate_res = list(dict.fromkeys(mate_residue_keys)) # keeps order | ||
| mate_res_map = {k: i for i, k in enumerate(unique_mate_res)} | ||
| mate_res_idx = torch.tensor( | ||
| [mate_res_map[k] for k in mate_residue_keys], dtype=torch.long | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the attribute PyMOL exposes for an atom's source object, and list every
# place the repo keys mate entities so both sites are covered.
rg -nP --type=py -C4 'atom\.(chain|resi|segi|model)|\.resi\b' src/ | head -60Repository: diff-use/WaterFlow
Length of output: 2262
🌐 Web query:
PyMOL cmd.get_model atom object attribute model symexp copy chain resi
💡 Result:
The PyMOL commands cmd.get_model and symexp serve distinct purposes and do not share the specific parameter list or attribute structure you described. cmd.get_model The cmd.get_model command returns a ChemPy 'Indexed' model object containing atom data from a selection [1][2]. It is used to access atom information directly in Python rather than through the iterate command [3][4]. Usage: cmd.get_model(selection="(all)", state=1) The resulting object contains a list of atom objects (e.g., atoms = cmd.get_model("all")). Each atom in the list (e.g., at = atoms.atom) possesses attributes that can be accessed via dot notation, including: - model: The object name [3][5][6] - chain: The chain identifier [3][4][6] - resi: The residue number (as a string) [3][4][6] - name: The atom name [3][4][6] - resn: The residue name [3][6] - coord: Coordinates (as a list/tuple) [3][4] - index: The internal atom index [3][4][6] - b: The B-factor [3][4] symexp The symexp command is a separate function used to generate crystallographic symmetry mates for a structure [7][8]. It does not return a model object with these specific attributes; instead, it creates new objects in the PyMOL scene based on symmetry operations within a defined cutoff [7][9][8]. Usage: symexp(prefix, object, selection, cutoff, segi=0, quiet=1) [9][10] In this context, 'copy' is not a parameter of symexp or get_model. If you are looking to copy atoms or objects, you might use the create command or specific selection-based operations [11][12]. The attributes you listed (chain, resi, etc.) are standard atom properties accessible via the model object returned by cmd.get_model, but they are not parameters of the symexp command [3][4][8].
Citations:
- 1: https://pymol.org/dokuwiki/doku.php?id=api%3Acmd%3Aget_model
- 2: https://pymol.sourceforge.net/newman/ref/S1000comref.html
- 3: https://pymolwiki.org/Get_Model
- 4: https://wiki.pymol.org/index.php/Get_Model
- 5: https://www.pymolwiki.org/Get_Model
- 6: https://pymolwiki.org/Iterate
- 7: https://pymol.org/dokuwiki/doku.php?id=command%3Asymexp
- 8: https://wiki.pymol.org/index.php/Symexp
- 9: https://pymol.org/dokuwiki/doku.php?id=api%3Acmd%3Asymexp
- 10: https://github.com/schrodinger/pymol-open-source/blob/46530027/modules/pymol/creating.py
- 11: https://pymol.org/pymol-command-ref.html
- 12: https://github.com/schrodinger/pymol-open-source/blob/46530027/modules/pymol/api.py
🌐 Web query:
Sybyx PyMOL symexp model attribute chain resi atom object name source
💡 Result:
In PyMOL, the symexp command is used to generate symmetry-related objects (crystal lattice partners) for a given structure [1][2]. The basic syntax is: symexp prefix, object, selection, cutoff [, segi] Arguments: - prefix: A string used as a name prefix for all newly created symmetry-related objects [2][3]. - object: The name of the existing PyMOL object to be replicated [2][4]. - selection: An atom selection used to determine the boundary for creating symmetry mates; only mates within the specified cutoff distance of this selection will be generated [2][5]. - cutoff: A float specifying the distance (in Angstroms) around the selection within which symmetry mates are created [2][3]. - segi (optional): If set to 1, assigns a unique 4-character segment identifier to each generated symmetry mate [2][3]. Naming Convention: Created objects are named using the pattern AAXXYYZZ [4]: - AA: A zero-based index for the symmetry operator. - XX, YY, ZZ: The unit cell offsets (-1, 0, or 1) along the X, Y, and Z axes, respectively [4]. Regarding your mention of "chain, resi, atom, object name, source": In PyMOL, these terms typically refer to properties used in atom selection expressions, rather than arguments for the symexp command itself [6][7]. When referencing atoms in PyMOL, you can use a selection macro syntax, such as /object_name/segi/chain/resi/name (where name corresponds to the atom name) [6]. The symexp command relies on an object name as its primary input (the source of the symmetry information) and an atom selection expression to define the spatial extent of the expansion [2][5].
Citations:
- 1: https://pymol.org/dokuwiki/doku.php?id=command%3Asymexp
- 2: https://wiki.pymol.org/index.php/Symexp
- 3: https://pymolwiki.org/index.php/Symexp
- 4: https://pymolwiki.org/Symexp
- 5: https://pymol.sourceforge.net/newman/user/S0400xtal.html
- 6: https://dasher.wustl.edu/bio5357/software/pymol/user-guide.pdf
- 7: https://wiki.pymol.org/index.php/Property_Selectors
🏁 Script executed:
#!/bin/bash
# Inspect the surrounding implementation and any tests/fixtures related to mate handling.
set -euo pipefail
echo "== src/dataset.py relevant ranges =="
sed -n '260,350p' src/dataset.py
printf '\n'
sed -n '1350,1470p' src/dataset.py
echo
echo "== pytest references =="
rg -n --type=py -C3 'symexp|sym|mate_res_idx|dedup_mate_ligands_by_residue|ligand|image_frac' tests src/dataset.py | head -200Repository: diff-use/WaterFlow
Length of output: 20714
🏁 Script executed:
#!/bin/bash
# Inspect relevant implementation and tests/fixtures related to mate handling.
set -euo pipefail
echo "== src/dataset.py relevant ranges =="
sed -n '260,350p' src/dataset.py
printf '\n'
sed -n '1350,1470p' src/dataset.py
echo
echo "== references =="
rg -n --type=py -C3 'symexp|sym|mate_res_idx|dedup_mate_ligands_by_residue|ligand|image_frac' tests src/dataset.py | head -220Repository: diff-use/WaterFlow
Length of output: 22117
🏁 Script executed:
#!/bin/bash
# Check whether tests or helper code expose PyMOL symmetry-object models directly.
set -euo pipefail
printf '%s\n' '== files =='
git ls-files | sed -n '1,160p'
printf '\n== model/atom access in repo =='
rg -n --type=py -C3 'get_model|\.atom|model|symexp|cmd\.select|sym\*' .Repository: diff-use/WaterFlow
Length of output: 50374
Use the PyMOL symmetry-object name as part of mate-entity IDs. get_model("sym* ...") returns atoms from all symmetry-prefixed objects, and atom chain/resi are not unique across those objects. Add atom.model to the mate ligand grouping key in dedup_mate_ligands_by_residue and to mate_residue_keys, so distinct operators do not collapse into one entity.
📍 Affects 1 file
src/dataset.py#L1405-L1411(this comment)src/dataset.py#L312-L322
🤖 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/dataset.py` around lines 1405 - 1411, The mate-ligand grouping keys
currently omit the symmetry-object identity, allowing atoms from distinct
operators to collapse. Update dedup_mate_ligands_by_residue at
src/dataset.py:312-322 and mate_residue_keys at src/dataset.py:1405-1411 to
include atom.model alongside chain and resi, preserving the existing ordering
and index-mapping behavior.
| # Ligands ride behind the mates, so the mate block is contiguous but need | ||
| # not run to the end; every marked node is past the ASU. | ||
| assert is_mate[num_asu:][: is_mate.sum().item()].all() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The mate block is not contiguous, so this assertion can fail.
The node order is [ASU protein | mate protein | ASU ligand | mate ligand], as src/dataset.py lines 1450-1491 and README.md lines 124-129 state. is_mate.sum() counts mate protein plus mate ligand atoms. When the structure has ASU ligands and include_ligands stays at its default True, the slice is_mate[num_asu:][: is_mate.sum()] spans the ASU-ligand nodes, which are False. The comment above the assertion also states the opposite of the documented layout.
Assert the property that actually holds: every mate node is past the ASU block, and the mate protein block is contiguous.
💚 Proposed test fix
- # Ligands ride behind the mates, so the mate block is contiguous but need
- # not run to the end; every marked node is past the ASU.
- assert is_mate[num_asu:][: is_mate.sum().item()].all()
-
mate_protein = is_mate & ~cached["is_ligand"]
+ # Node order is [ASU protein | mate protein | ASU ligand | mate ligand],
+ # so only the mate *protein* block is contiguous, right after the ASU.
+ assert mate_protein[:num_asu].sum().item() == 0
+ assert mate_protein[num_asu : num_asu + mate_protein.sum().item()].all()
assert (emb_res_idx[mate_protein] >= 0).all()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_dataset.py` around lines 2917 - 2919, Update the assertion near
the mate-ordering test to match the documented layout: verify that every node
after the ASU boundary is not necessarily a mate, while asserting the mate
protein block is contiguous in its expected range. Revise the misleading comment
accordingly, and avoid using the total mate count to slice across ASU ligand
nodes.
DorisMai
left a comment
DorisMai
left a comment
There was a problem hiding this comment.
I have some questions on the mates/ligands implementation. The connection between pymol handled mates/ligands and biotite handled ASU is in my opinion a brittle / risky point that can benefit from more real test cases and integration tests. In addition, before this PR, there is already code assumes or uses mates, and it is outside the diff for review. I would encourage double checking the correctness of what you added given the existing code (see the two major issues raised by coderabbitai).
There was a problem hiding this comment.
why removing the comments?
There was a problem hiding this comment.
I thought it was over-explanatory in this case, since we do mention this later in the code.
| │ - protein_pos: centered protein coordinates (N, 3) | ||
| │ - protein_x: element one-hot encoding (N, 16) | ||
| │ - protein_res_idx: residue indices for grouping | ||
| │ - is_ligand: bool mask marking the appended ASU ligand atoms (N,) | ||
| │ - is_ligand: bool mask marking the ligand atoms (N,) | ||
| │ - is_mate: bool mask marking the symmetry-mate atoms (N,) | ||
| │ - emb_res_idx: embedding row per atom; -1 means no row (N,) |
There was a problem hiding this comment.
you might want to rename protein_pos/x/res_idx to be more generic, maybe like node_pos/x/res_idx (referencing your line 124: Node order is [ASU protein | mate protein | ASU ligand | mate ligand]), since at this point they are not just protein atoms any more, and the array shape annotation N would get confusing.
| with open(tmp_path, "w") as f: | ||
| json.dump(current, f, indent=2) | ||
| tmp_path.replace(meta_path) | ||
| elif any(self.geometry_dir.glob("*.pt")): |
There was a problem hiding this comment.
related to coderabbitai's comment. you probably want to move this elif branch up to be if the meta_path file does NOT exist, but *.pt exists, then log warning, regardless of whether write or not. This would change your test_unlabelled_cache_warns.
| key = ( | ||
| str(sanitized_for_idx.chain_id[start]).strip(), | ||
| int(sanitized_for_idx.res_id[start]), | ||
| normalize_ins_code(sanitized_for_idx.ins_code[start]), |
There was a problem hiding this comment.
seems not needed, already normalized in line 1329
| asu_reskey_to_residx.get((str(atom.chain).strip(), *parsed), -1) | ||
| if parsed is not None | ||
| else -1 |
There was a problem hiding this comment.
in what scenario would you expect to have a non-matched mate protein reskey? Should here be a warning or raise an error? And what does a zero-embedding of those extra atoms functionally do to the water prediction?
| batch_w: Tensor, | ||
| cutoff: float = 8.0, | ||
| device: torch.device | None = None, | ||
| anchor_mask: Tensor | None = None, |
There was a problem hiding this comment.
should you add anchor_mask for sample_waters_scaled_gaussian too? would mate protein affect the sigma_per_graph computation?
| return mate_coords[keep_idx], [mate_atoms[i] for i in keep_idx] | ||
|
|
||
|
|
||
| def dedup_mate_ligands_by_residue( |
There was a problem hiding this comment.
I am quite confused by the different treatment of mate protein vs ligand in the two dedup_* methods here. (1) Why are you not treating the mate protein in whole blocks (symmetry copy 1, copy 2, ...) like the whole ligand block? (2) Why does protein dedup needs first sweep + second round of neighbor look up, but ligand dedup only needs first sweep?
|
|
||
|
|
||
| @pytest.mark.unit | ||
| class TestDedupMateAtoms: |
There was a problem hiding this comment.
I think it's important for TestDedupMateAtoms and TestDedupMateLigandsByResidue to have a real test case where you see dedup is really needed.
There was a problem hiding this comment.
Currently you don't have integration tests where both ligands and mates are turned on. The ligand integration tests all have include_mates=False. The mates tests uses 6eey which doesn't contain ligands. You should add a test class/case. Also, I think most integration tests here uses encoder_type="gvp". If the paper uses esm, it's probably good to have at least one test here.
| assert (emb_res_idx[mate_protein] >= 0).all() | ||
| assert (emb_res_idx[mate_protein] < data["protein"].num_protein_residues).all() |
There was a problem hiding this comment.
worth addressing the coderabbitai's comment above.
Also, here, asserting just the range is quite weak. Since getting the emb_res_idx wrong was what prevented symmetry mates from being useful in the past, I think a stronger test on exact equivalence (for residues that you know are the same in ASU vs in mate) would be helpful.
is_mateemb_res_idxfields. Closes the special-position label leak, a water on a rotation axis is its own symmetry copy, i.e. the target sitting in the input._filter_meta.jsonper cache dir; a run with different filter settings is refused instead of silently extending the cache.Default changes to match current checkpoint runs:
max_bfactor_zscore1.5 -> 2.0,min_water_residue_ratio0.6 -> 0.1. Every cache used for training was built at the these values.Summary by CodeRabbit
New Features
Bug Fixes
Documentation