Overview
In /rewards/real_space_density.py, RealSpaceRewardFunction.structure_to_reward_input() duplicates much of the preprocessing in extract_density_inputs_from_atomarray().
Both paths convert an AtomArray/AtomArrayStack into batched Torch tensors for real-space density evaluation:
- coordinates
- scattering element indices
- B-factors
- occupancies
However, the latter also removes invalid coordinates and replaces NaN B-factors with a default value, while structure_to_reward_input() only filters zero-occupancy atoms.
Proposed change
Refactor RealSpaceRewardFunction.structure_to_reward_input() to reuse extract_density_inputs_from_atomarray()
Overview
In /rewards/real_space_density.py,
RealSpaceRewardFunction.structure_to_reward_input()duplicates much of the preprocessing inextract_density_inputs_from_atomarray().Both paths convert an AtomArray/AtomArrayStack into batched Torch tensors for real-space density evaluation:
However, the latter also removes invalid coordinates and replaces NaN B-factors with a default value, while
structure_to_reward_input()only filters zero-occupancy atoms.Proposed change
Refactor
RealSpaceRewardFunction.structure_to_reward_input()to reuseextract_density_inputs_from_atomarray()