feat(rewards): add tmol plausibility and composite rewards#319
feat(rewards): add tmol plausibility and composite rewards#319etherealsunshine wants to merge 9 commits into
Conversation
Add the tmol (Rosetta beta_nov2016_cart) energy function as a pixi feature attached to all model envs, for use as a reward in guided sampling (any strategy). No prebuilt wheel matches this project's torch 2.7.1, so tmol compiles from source; the feature encodes the conda-CUDA cmake workarounds (CUDA root, NVTX3 headers, no build isolation) and a memory-bounded, sm_80;86-targeted build. See TMOL_BUILD.md for the full rationale and the Wynton build/run workflow. pixi.lock is intentionally not included; it regenerates on next `pixi install`. Co-Authored-By: Claude Opus 4.8 <[email protected]>
Add TmolPlausibilityReward, a differentiable physical-plausibility reward using tmol's beta2016 Rosetta energy. Scatters model-order coordinates into tmol's canonical per-residue grid (name-based map built once via prepare()) and scores via pose_stack_from_canonical_form + the beta2016 score function, so gradients flow to input coords (usable for gradient guidance and FK steering).
Make the tmol beta2016 plausibility reward selectable and runnable end-to-end:
- Add an optional reward.prepare(atom_array) hook in PureGuidance and FKSteering
(structural Protocol, not inheritance): rewards that need the concrete model
atom order (tmol) build their scatter map against the same array to_reward_inputs
uses; density-style rewards without prepare() are unaffected.
- Add get_tmol_reward_and_structure() and a reward_type branch in _run_guidance
(no experimental map needed for plausibility).
- Add --reward-type {density,plausibility} and --tmol-weight CLI flags plus the
matching GuidanceConfig fields; make --density/--resolution optional and required
only in density mode.
- Log the tmol scatter-map size in prepare() for run-time visibility.
Verified end-to-end (CPU): FK steering with plausibility drives beta2016 energy
from ~2684 to favorable (negative) with a geometrically sane structure
(min Ca-Ca 3.75 A, Rg 15.3 A) when combined with gradient normalization.
…bility Add CompositeReward: a weighted sum of RewardFunctionProtocol rewards that is itself a RewardFunctionProtocol, and forwards the optional prepare() hook to any sub-reward that defines it (so tmol's scatter map still gets built inside a composite). Wire it into the pipeline: get_composite_reward_and_structure builds the density and tmol sub-rewards (tmol term scaled by --tmol-weight) and is selectable via --reward-type composite.
The pyproject.toml tmol comments already document the source build and the cmake/CUDA flags next to the config, so the standalone doc was redundant.
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds tmol build configuration and integrates differentiable tmol plausibility rewards into guidance. Rewards can be selected through CLI options as density, plausibility, or composite modes, with preparation hooks and weighted composition support. Changestmol guidance reward integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant GuidanceConfig
participant GuidanceRunner
participant RewardBuilder
participant GuidanceSampler
CLI->>GuidanceConfig: parse reward-type and tmol-weight
GuidanceConfig->>GuidanceRunner: provide reward configuration
GuidanceRunner->>RewardBuilder: construct selected reward
RewardBuilder-->>GuidanceRunner: return reward and structure
GuidanceRunner->>GuidanceSampler: start sampling with reward
GuidanceSampler->>RewardBuilder: prepare structure-dependent reward
Possibly related issues
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
Apply ruff import-sorting/formatting and wrap long lines; add a runtime_checkable PreparableRewardProtocol and use isinstance() at the scaler prepare() hooks so ty narrows the optional hook to a callable (replaces the untyped hasattr check).
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
Add tests/rewards/test_composite.py covering the weighted sum, default and mismatched weights, gradient flow, and prepare()-forwarding. Replace the stray print() in build_tmol_map with logger.warning to match the codebase's logging.
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
requested review from
a team,
Abdelsalam-Abbas,
DorisMai,
denis-zaitsev,
k-chrispens,
mag-astera and
marcuscollins
as code owners
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/sampleworks/core/rewards/composite.py (1)
46-50: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the loop variable to avoid shadowing and improve readability.
The loop variable
rewardsshadows the class attributeself.rewardslogically and is semantically confusing for a single item.♻️ Proposed refactor
def prepare(self, atom_array): - for rewards in self.rewards: - if hasattr(rewards, "prepare"): - rewards.prepare(atom_array) # prepare from TmolPlausibility + for reward in self.rewards: + if hasattr(reward, "prepare"): + reward.prepare(atom_array) # prepare from TmolPlausibility return self🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/sampleworks/core/rewards/composite.py` around lines 46 - 50, Rename the singular loop variable in Composite’s prepare method to a distinct item-oriented name, and update its hasattr check and prepare call accordingly; keep self.rewards and the method’s behavior unchanged.
🤖 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/core/rewards/composite.py`:
- Around line 33-44: Update the composite reward evaluation method containing
the rewards loop so an empty self.rewards collection returns a scalar zero torch
tensor instead of None. Preserve the existing weighted accumulation behavior
when rewards are present, and ensure the zero tensor is compatible with the
computed reward dtype/device and gradient-based downstream use.
In `@src/sampleworks/core/scalers/fk_steering.py`:
- Around line 117-119: Replace the truthiness-based AtomArray fallback in the
reward preparation flow around PreparableRewardProtocol in
src/sampleworks/core/scalers/fk_steering.py lines 117-119 and
src/sampleworks/core/scalers/pure_guidance.py lines 93-95 with an explicit
model_atom_array is not None check, while preserving fallback to atom_array when
it is absent. Also update the corresponding fallback in
src/sampleworks/eval/structure_utils.py lines 33-37 the same way.
---
Nitpick comments:
In `@src/sampleworks/core/rewards/composite.py`:
- Around line 46-50: Rename the singular loop variable in Composite’s prepare
method to a distinct item-oriented name, and update its hasattr check and
prepare call accordingly; keep self.rewards and the method’s behavior unchanged.
🪄 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: d550a429-d77f-4ca8-8eb2-ed4eee5b8008
📒 Files selected for processing (9)
pyproject.tomlsrc/sampleworks/core/rewards/composite.pysrc/sampleworks/core/rewards/protocol.pysrc/sampleworks/core/rewards/tmol_plausibility.pysrc/sampleworks/core/scalers/fk_steering.pysrc/sampleworks/core/scalers/pure_guidance.pysrc/sampleworks/utils/guidance_script_arguments.pysrc/sampleworks/utils/guidance_script_utils.pytests/rewards/test_composite.py
| total = None | ||
| for reward, weight in zip(self.rewards, self.weights): | ||
| term = weight * reward( | ||
| coordinates, | ||
| elements, | ||
| b_factors, | ||
| occupancies, | ||
| unique_combinations, | ||
| inverse_indices, | ||
| ) | ||
| total = term if total is None else total + term | ||
| return total |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Return a zero tensor if self.rewards is empty.
If the composite reward is initialized with an empty list of rewards, total will remain None. This violates the Float[torch.Tensor, ""] return type annotation and will likely cause downstream crashes when gradients are computed.
🐛 Proposed fix to return a valid zero tensor
total = None
for reward, weight in zip(self.rewards, self.weights):
term = weight * reward(
coordinates,
elements,
b_factors,
occupancies,
unique_combinations,
inverse_indices,
)
total = term if total is None else total + term
- return total
+ if total is None:
+ return torch.tensor(0.0, device=coordinates.device)
+ return total📝 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.
| total = None | |
| for reward, weight in zip(self.rewards, self.weights): | |
| term = weight * reward( | |
| coordinates, | |
| elements, | |
| b_factors, | |
| occupancies, | |
| unique_combinations, | |
| inverse_indices, | |
| ) | |
| total = term if total is None else total + term | |
| return total | |
| total = None | |
| for reward, weight in zip(self.rewards, self.weights): | |
| term = weight * reward( | |
| coordinates, | |
| elements, | |
| b_factors, | |
| occupancies, | |
| unique_combinations, | |
| inverse_indices, | |
| ) | |
| total = term if total is None else total + term | |
| if total is None: | |
| return torch.tensor(0.0, device=coordinates.device) | |
| return total |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/sampleworks/core/rewards/composite.py` around lines 33 - 44, Update the
composite reward evaluation method containing the rewards loop so an empty
self.rewards collection returns a scalar zero torch tensor instead of None.
Preserve the existing weighted accumulation behavior when rewards are present,
and ensure the zero tensor is compatible with the computed reward dtype/device
and gradient-based downstream use.
There was a problem hiding this comment.
@etherealsunshine my above comment will help fix this problem too.
| if isinstance(reward, PreparableRewardProtocol): | ||
| reward.prepare(processed.model_atom_array or processed.atom_array) | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check if biotite AtomArray raises ValueError on boolean evaluation.
cat << 'EOF' > test_truth.py
import biotite.structure as struc
array = struc.AtomArray(10)
try:
bool(array)
print("Truth value OK")
except Exception as e:
print(f"Error: {e}")
EOF
python test_truth.pyRepository: diff-use/sampleworks
Length of output: 341
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== file list ==\n'
git ls-files 'src/sampleworks/core/scalers/*' 'src/sampleworks/**/protocol*' 'src/sampleworks/**/*reward*' 'src/sampleworks/**/*structure*' 'src/sampleworks/**/*atom*' | sed 's#^`#-` #'
printf '\n== search model_atom_array ==\n'
rg -n "model_atom_array|atom_array|PreparableRewardProtocol|prepare\(" src/sampleworks -g '!**/__pycache__/**'
printf '\n== outline candidate files ==\n'
for f in src/sampleworks/core/scalers/fk_steering.py src/sampleworks/core/scalers/pure_guidance.py; do
if [ -f "$f" ]; then
echo "--- $f"
ast-grep outline "$f" --view expanded || true
fi
doneRepository: diff-use/sampleworks
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
for f in \
src/sampleworks/core/scalers/fk_steering.py \
src/sampleworks/core/scalers/pure_guidance.py
do
echo "===== $f ====="
wc -l "$f"
sed -n '1,220p' "$f" | nl -ba | sed -n '1,220p'
done
printf '\n===== definitions/usages =====\n'
rg -n "class .*Processed|dataclass|model_atom_array|atom_array|PreparableRewardProtocol|Protocol" src/sampleworks -g '!**/__pycache__/**'Repository: diff-use/sampleworks
Length of output: 303
🌐 Web query:
biotite AtomArray boolean evaluation ValueError ambiguous truth value
💡 Result:
The ValueError: The truth value of an array with more than one element is ambiguous occurs because an AtomArray in Biotite is backed by NumPy ndarrays [1][2][3]. When you attempt to use an AtomArray directly in a boolean context (such as an if statement or with logical operators like and/or), Python attempts to evaluate the truthiness of the entire object, which is ambiguous for container-like structures containing multiple elements [4][5]. To resolve this error, you should specify how the boolean evaluation should be performed on the structure's data: 1. Check if the AtomArray is empty: Use the array_length method [6] or check the length of the array: if array.array_length == 0: # Handle empty array case 2. Perform element-wise evaluation: If you are trying to filter atoms or check conditions based on atom attributes, you must generate a boolean mask (an array of booleans) and then use.any or.all [6][4][5]. For example: # Incorrect: if array.res_id == "ALA": # Correct (element-wise check): if (array.res_id == "ALA").any: # At least one atom belongs to a residue named "ALA" if (array.res_id == "ALA").all: # All atoms belong to a residue named "ALA" If you are trying to combine multiple boolean conditions, ensure you use the bitwise operators (&, |, ~) instead of the Python keywords (and, or, not), and wrap each condition in parentheses [4][5]: if ((array.res_id == "ALA") & (array.atom_name == "CA")).any: # Handle case where at least one atom is a CA atom in an ALA residue Top results: [6], [4], [5], [2]
Citations:
- 1: https://www.biotite-python.org/latest/apidoc/biotite.structure.AtomArray.html
- 2: https://github.com/biotite-dev/biotite/blob/v1.7.0/src/biotite/structure/atoms.py
- 3: https://www.biotite-python.org/latest/tutorial/structure/index.html
- 4: https://stackoverflow.com/questions/10062954/valueerror-the-truth-value-of-an-array-with-more-than-one-element-is-ambiguous
- 5: https://note.nkmk.me/en/python-numpy-pandas-value-error-ambiguous/
- 6: https://www.biotite-python.org/latest/tutorial/structure/filter.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
files = [
"src/sampleworks/core/scalers/fk_steering.py",
"src/sampleworks/core/scalers/pure_guidance.py",
]
for path in files:
print(f"===== {path} =====")
lines = Path(path).read_text().splitlines()
for i, line in enumerate(lines, 1):
if "model_atom_array or" in line or "model_atom_array if" in line or "processed_structure.model_atom_array" in line or "processed.model_atom_array" in line:
start = max(1, i - 3)
end = min(len(lines), i + 3)
for j in range(start, end + 1):
print(f"{j:4d}: {lines[j-1]}")
print()
PYRepository: diff-use/sampleworks
Length of output: 2887
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '40,90p' src/sampleworks/eval/structure_utils.py | cat -nRepository: diff-use/sampleworks
Length of output: 2577
Avoid or on AtomArray fallbacks.
model_atom_array is an AtomArray | None, and the or fallback will still try to truth-test the AtomArray. Use is not None here:
src/sampleworks/core/scalers/fk_steering.py#L117-L119src/sampleworks/core/scalers/pure_guidance.py#L93-L95src/sampleworks/eval/structure_utils.py#L33-L37
📍 Affects 2 files
src/sampleworks/core/scalers/fk_steering.py#L117-L119(this comment)src/sampleworks/core/scalers/pure_guidance.py#L93-L95
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/sampleworks/core/scalers/fk_steering.py` around lines 117 - 119, Replace
the truthiness-based AtomArray fallback in the reward preparation flow around
PreparableRewardProtocol in src/sampleworks/core/scalers/fk_steering.py lines
117-119 and src/sampleworks/core/scalers/pure_guidance.py lines 93-95 with an
explicit model_atom_array is not None check, while preserving fallback to
atom_array when it is absent. Also update the corresponding fallback in
src/sampleworks/eval/structure_utils.py lines 33-37 the same way.
There was a problem hiding this comment.
@etherealsunshine I've noticed that more and more people are using the is not None form, rather than relying on the "falseness" of None. It's more specific, and it avoids the problem that the "truthiness" of different classes can be hard to define, e.g., is an empty AtomArray True or False? (I'd say False, but not everyone will agree and we don't want to rely on biotite's developers to necessarily define that for us.)
There was a problem hiding this comment.
also @etherealsunshine remind me what the difference between these two is?
There was a problem hiding this comment.
Pull request overview
This PR adds a differentiable physical-plausibility reward based on tmol’s Rosetta beta2016 energy, plus a CompositeReward that combines multiple reward terms as a weighted sum. It integrates these into the guidance CLI/config so users can select {density, plausibility, composite} reward modes, and wires an optional one-time prepare() hook into trajectory scalers for rewards that need concrete atom ordering.
Changes:
- Add
TmolPlausibilityReward(tmol beta2016) andCompositeReward(weighted sum) reward implementations. - Add
PreparableRewardProtocol+ callreward.prepare(...)once in PureGuidance/FK steering for rewards that implement it. - Extend guidance CLI/config to select reward type and make
--density/--resolutionconditional on reward mode; add tmol build/runtime configuration via pixi feature.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/rewards/test_composite.py | Adds tests for CompositeReward behavior, differentiability, and prepare forwarding. |
| src/sampleworks/utils/guidance_script_utils.py | Adds reward constructors for tmol + composite and selects reward implementation based on --reward-type. |
| src/sampleworks/utils/guidance_script_arguments.py | Adds reward_type/tmol_weight args and makes --density/--resolution optional depending on reward mode. |
| src/sampleworks/core/scalers/pure_guidance.py | Calls reward.prepare(atom_array) once when supported by the reward (protocol-based). |
| src/sampleworks/core/scalers/fk_steering.py | Same prepare() hook support as PureGuidance. |
| src/sampleworks/core/rewards/tmol_plausibility.py | Implements differentiable tmol plausibility reward and atom-order → canonical-grid mapping. |
| src/sampleworks/core/rewards/protocol.py | Introduces PreparableRewardProtocol for optional one-time reward preparation. |
| src/sampleworks/core/rewards/composite.py | Implements CompositeReward weighted-sum wrapper with a prepare() forwarder. |
| pyproject.toml | Adds pixi tmol feature + environment wiring and build env vars for compiling tmol from source. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| self._m = build_tmol_map(atom_array, self.co, self.device) | ||
| logger.info( | ||
| f"tmol plausibility map: {self._m['n_res']} residues, " | ||
| f"{len(self._m['sel'])} scored atoms" | ||
| ) | ||
| return self |
|
|
||
| if weights is None: | ||
| weights = [1.0] * len(rewards) | ||
| if len(weights) != len(rewards): |
There was a problem hiding this comment.
@etherealsunshine this is good practice: fail early with a clear explanation. Please implement.
| def prepare(self, atom_array): | ||
| for rewards in self.rewards: | ||
| if hasattr(rewards, "prepare"): | ||
| rewards.prepare(atom_array) # prepare from TmolPlausibility | ||
| return self |
There was a problem hiding this comment.
@etherealsunshine these are both good suggestions, please implement them.
| parser.add_argument( | ||
| "--density", | ||
| type=str, | ||
| default=None, | ||
| help="Input density map (required for --reward-type density)", |
There was a problem hiding this comment.
@etherealsunshine Definitely if this is now optional, you need to update the type annotation in GuidanceConfig, and either GuidanceConfig, or probably the density-based reward function itself, needs to raise a clear error if this argument is unset and --reward-type is "density". I think I prefer the error checking to be in the reward functions themselves, so that each reward function can check errors that make sense for it to check.
| import tmol | ||
| from tmol.io.canonical_ordering import ( | ||
| default_canonical_ordering, | ||
| default_packed_block_types, | ||
| ) | ||
|
|
There was a problem hiding this comment.
I think the correct solution to this has two parts:
- make sure all relevant environments have the tmol feature (or just make it default, is there a reason not to? It's possible to go too crazy paring down your environment to only what's needed for the current invocation of a command.)
- Maybe add the suggested check that tmol is importable.
| coords = _coords() | ||
| comp = CompositeReward([_SumSquares(), _Sum()], [2.0, 3.0]) | ||
| out = comp(coords, None, None, None) |
| coords = _coords() | ||
| comp = CompositeReward([_SumSquares(), _Sum()]) # weights omitted -> all 1.0 | ||
| out = comp(coords, None, None, None) |
| coords = _coords().requires_grad_(True) | ||
| comp = CompositeReward([_SumSquares(), _Sum()], [2.0, 3.0]) | ||
| comp(coords, None, None, None).backward() |
Co-authored-by: Copilot Autofix powered by AI <[email protected]>
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
etherealsunshine
had a problem deploying
to
gpu-testing
GitHub Actions
Error
marcuscollins
left a comment
marcuscollins
left a comment
There was a problem hiding this comment.
@etherealsunshine lots of comments and things to improve. The code can be more modular, and simpler at the same time. LMK if you have questions or want to discuss.
| Rewards without it (e.g. density fit) simply don't match this protocol. | ||
| """ | ||
|
|
||
| def prepare(self, atom_array: Any) -> Any: ... |
There was a problem hiding this comment.
Several issues here:
- What kind of object does
preparereturn? I think this needs to be more cleanly defined, and in fact I'd prefer if it returned None and only modified the "reward" object itself. - atom_array needs a more specific type than
Any. It is presumably atAtomArrayorAtomArrayStack.
| def prepare(self, atom_array): | ||
| for rewards in self.rewards: | ||
| if hasattr(rewards, "prepare"): | ||
| rewards.prepare(atom_array) # prepare from TmolPlausibility | ||
| return self |
There was a problem hiding this comment.
@etherealsunshine these are both good suggestions, please implement them.
|
|
||
| if weights is None: | ||
| weights = [1.0] * len(rewards) | ||
| if len(weights) != len(rewards): |
There was a problem hiding this comment.
@etherealsunshine this is good practice: fail early with a clear explanation. Please implement.
| occupancies: Float[torch.Tensor, "batch n_atoms"], | ||
| unique_combinations: torch.Tensor | None = None, | ||
| inverse_indices: torch.Tensor | None = None, | ||
| ) -> Float[torch.Tensor, ""]: |
There was a problem hiding this comment.
needs a good docstring. In particular you need to explain unique_combinations and inverse_indices.
| unique_combinations, | ||
| inverse_indices, | ||
| ) | ||
| total = term if total is None else total + term |
There was a problem hiding this comment.
What type will term be? Presumably a torch.Float of size 1? This is a technically correct but unusual way to initialize total correctly, and it doesn't provide much insight into what total is. I would initialize total to be zeros that match what's expected from the rewards, and then just add to that.
| logger.info("Creating tmol plausibility reward function") | ||
| co = default_canonical_ordering() | ||
| pbt = default_packed_block_types(device) | ||
| sfxn = tmol.beta2016_score_function(device) |
There was a problem hiding this comment.
Can you describe this function in a comment? Also, add a link to where it is documented/described. Are there others in tmol that we might eventually want to pipe through?
| "--reward-type composite needs --density and --resolution for the density term." | ||
| ) | ||
|
|
||
| density_reward, structure = get_reward_function_and_structure( |
There was a problem hiding this comment.
possibly should split up get_reward_function_and_structure and in any event rename it to get_density_reward_and_structure to distinguish it from the more general cases that now exist.
| args, | ||
| device: torch.device, | ||
| ) -> tuple[Any, dict[str, Any]]: | ||
| """Build a CompositeReward = density fit + tmol plausibility, weighted. |
There was a problem hiding this comment.
Make it more general, there's no reason to limit ourselves to just these two.
|
|
||
|
|
||
| def get_composite_reward_and_structure( | ||
| args, |
There was a problem hiding this comment.
needs type annotation (is this a GuidanceConfig?) and you need a docstring documenting all your arguments.
In general I prefer passing in specific arguments, rather than config objects, since it makes it easier to understand what's getting used and type check it.
| ) | ||
| # v1 reward selection (becomes a proper --reward flag later): "plausibility" uses the | ||
| # tmol beta2016 energy prior (no map needed); anything else falls back to density fit. | ||
| if getattr(args, "reward_type", "density") == "plausibility": |
There was a problem hiding this comment.
I think if you just make all reward functions composite, and use a configuration dictionary that maps between reward types and weights, the logic here becomes simpler and much easier to maintain (you won't have to keep adding extra conditionals here for every new possibility.
|
@etherealsunshine another big issue here is how well we understand what this code will do in actual use: what should the weights be, how much impact will it have on guidance. This is no longer purely research code, it is a product that we want others to start to use, so it needs to be very usable. I think before we merge this we need to do experiments to establish best options for using it, and whether it has measurable positive impact on the outcomes. |
|
@marcuscollins that definitely makes sense. I was talking to Karson about this yesterday about this in terms of whether he had any specific priors for how to decide weights and reward annealing for the same. Once I clean up the highlighted comments, I'm happy to run a few jobs experimenting with how strongly the density (and tmol) reward should influence the trajectory as a function of timestep and track changes |
etherealsunshine
requested a deployment
to
gpu-testing
GitHub Actions
Waiting
etherealsunshine
requested a deployment
to
gpu-testing
GitHub Actions
Waiting
etherealsunshine
requested a deployment
to
gpu-testing
GitHub Actions
Waiting
manzuoni-astera
left a comment
manzuoni-astera
left a comment
There was a problem hiding this comment.
Thanks for taking on this integration; a differentiable physical prior is valuable, and the preparation hook is a reasonable direction. I do not think the current implementation is ready to merge, however.
The main blocker is the coordinate/topology map in tmol_plausibility.py. Lines 99–102 filter the AtomArray, but lines 119–134 retain indices in that filtered space and line 246 applies them to the original coordinate tensor. A zero-occupancy atom therefore shifts every subsequent mapping. Residues are also keyed only by (chain_id, res_id), which merges insertion-code residues, while skipped nonstandard residues and other chain discontinuities are not represented through res_not_connected. tmol may consequently create peptide bonds across gaps and guide toward an incorrect topology.
There is currently no automated test of this reward itself. Before merging, I would want CPU/GPU tests covering finite score and gradients, zero-occupancy indexing, insertion codes, missing heavy atoms, unsupported-only inputs, chain breaks, batch behavior, and at least one sampler integration using actual model atom order.
The CLI/config change also needs completion. Default density mode now accepts missing density or resolution, contrary to the existing tests at test_guidance_cli.py:488-518, and GuidanceConfig.as_dict() serializes missing density as "None". Conditional validation and plausibility/composite CLI tests should be added.
Finally, please regenerate and commit pixi.lock and pin a compatible tmol version. An unbounded, source-built CUDA dependency is not reproducible, and neither standard CI nor GPU CI has run yet. The previously raised empty-composite, typing, modularity, and reward-calibration comments should also be resolved.
Summary
beta_nov2016energy) as a differentiable physical-plausibilityreward for guided sampling, plus a CompositeReward that combines rewards (e.g. density
fit + plausibility) as a weighted sum. Both satisfy the existing
RewardFunctionProtocoland are selectable via
--reward-type.What's added
TmolPlausibilityReward- differentiable beta2016 energy; maps model atom order tmol'scanonical per-residue grid via a scatter map built once in
prepare().CompositeReward- weighted sum of rewards; itself aRewardFunctionProtocol; forwardsthe optional
prepare()hook to sub-rewards that need it.--reward-type {density,plausibility,composite},--tmol-weight;--density/--resolutionnow required only in density mode.Validation
Tested FKS and DPS on Modal an A10G GPU, ensuring energy decreases monotonically and end-structures do not have any clashes.
To be discussed: CUDA architectures
tmol compiles from source on every install, so
CMAKE_CUDA_ARCHITECTURESis acommitted build setting that affects everyone. It's currently
sm_80;86(A100 / A40 /RTX 3090). Widening to tmol's own default
80;86;89;90would also cover sm_89 (L40 /RTX 4090) and sm_90 (H100), but makes every build slower and raises peak build memory (may
need a lower
CMAKE_BUILD_PARALLEL_LEVELto avoid OOM). Let me know how to resolve thisSummary by CodeRabbit
CompositeRewardto combine multiple rewards via weighted sums.