ES/01 core - #174
Conversation
…riable Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Greptile SummaryThe PR introduces the core enhanced-sampling API, including conservative autograd-derived bias outputs, result aggregation, and a differentiable pair-distance collective variable, while deprecating the legacy bias hook.
Important Files Changed
Reviews (1): Last reviewed commit: "redesigned to provide batteries as compo..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if ( | ||
| wants_stress | ||
| and original_cell is not None | ||
| and (pbc is None or bool(pbc.any())) | ||
| ): |
There was a problem hiding this comment.
When a batch contains both a periodic graph and a non-periodic graph with a zero-volume placeholder cell, the batch-wide pbc.any() enables stress calculation for every graph. The shared stress helper then divides the non-periodic graph's strain derivative by zero, producing non-finite output that BiasResult rejects.
Knowledge Base Used: Models
| any_periodic = ( | ||
| has_cell | ||
| and has_pbc | ||
| and (torch.compiler.is_compiling() or bool(batch.pbc.any())) | ||
| ) | ||
|
|
||
| if any_periodic: | ||
| if not torch.compiler.is_compiling(): | ||
| _check_minkowski_reduced(batch.cell, batch.pbc) | ||
| dr = _apply_mic(dr, batch.cell, batch.pbc) |
There was a problem hiding this comment.
When a mixed batch contains one periodic graph and a non-periodic graph with a permitted degenerate placeholder cell, the batch-wide periodicity check sends every graph through _apply_mic. torch.linalg.inv then attempts to invert the non-periodic graph's singular cell before PBC masking occurs, raising LinAlgError instead of returning its distance.
| return BiasResult( | ||
| energy=summed.get("energy"), | ||
| forces=summed.get("forces"), | ||
| stress=summed.get("stress"), | ||
| virial=summed.get("virial"), | ||
| observables=observables_total, |
There was a problem hiding this comment.
aggregate_bias_results constructs the combined BiasResult without preserving or validating any input state_version. This silently discards the version IDs documented for replica-state coherence and forces future consumers to bypass or supplement the aggregation API.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
So I'm not convinced that pair distance is really a collective variable. It should be more like a geometric primitive (mic_displacement, pair_distance, et cetera). Second, this should really be a function that lives in nvalchemi-toolkit-ops and recieves a proper warp kernel for forwards, backward, double-backward kernels with a torch.custom op and autograd wrappers registered. We can expose them here for utility.
| class BiasResult: | ||
| """Immutable, fully-detached output of a single bias evaluation. | ||
|
|
||
| All tensor fields must be detached (``requires_grad=False`` and | ||
| ``grad_fn is None``). Energy, forces, stress, and virial are | ||
| independently optional. Provide **either** ``stress`` or ``virial``, | ||
| not both; the runner converts stress to virial or vice-versa as needed. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| energy: | ||
| Per-graph bias energy, shape ``[B, 1]``, unit eV. | ||
| forces: | ||
| Per-atom bias forces, shape ``[N_atoms, 3]``, unit eV/Å. | ||
| stress: | ||
| Tensile-positive Cauchy stress ``σ = −W/V``, shape ``[B, 3, 3]``, | ||
| unit eV/ų. This is the toolkit-wide convention and what | ||
| :class:`ConservativeBias` produces. Mutually exclusive with | ||
| ``virial``. | ||
| virial: | ||
| Virial ``W = −dE/dε`` with ``ε`` the symmetric infinitesimal | ||
| strain tensor, shape ``[B, 3, 3]``, unit eV. Provided for biases | ||
| that compute a virial directly. Mutually exclusive with | ||
| ``stress``. | ||
| state_version: | ||
| Integer version IDs used by ``ReplicaExchange`` to validate that | ||
| accepted state assignments are coherent, shape ``[B]``. | ||
| observables: | ||
| Named diagnostic tensors exposed as ``bias/<name>/<key>`` in the | ||
| runner's output dict. All tensors must be detached. | ||
| """ |
There was a problem hiding this comment.
So BiasResult is basically an unnecessary class that attempts to be ModelOutputs specialized for the enhanced_sampling class. I would rather you do the following: Use ModelOutputs, modify to include observables or state_version if necessary (some way to make it more generic, to serve other workflows down the road would be smart). If you need validation then write a single function in nvalchemi/models/_utils.py called validate_contribution that essentially encodes _validate_bias_result.
| package beyond :class:`BiasResult`. | ||
|
|
||
| Batteries are supplied as **composable mixins that satisfy this protocol, | ||
| not as parallel hierarchies**. A bias mixes in only what applies to it: |
There was a problem hiding this comment.
Since we presently use BaseModelMixin for other additive potentials with no learned parameters (dftd3, lennardjones). I don't know if this should survive as a separate class. The diagnostics and versioning should move into ModelOutputs and the update / comit_epoch should move onto hooks.
| def evaluate(self, current: Batch) -> BiasResult: | ||
| """Evaluate the bias on the current batch. | ||
|
|
||
| Must be **read-only**: it must not mutate bias internal state, | ||
| deposit hills, write any storage, or communicate across workers. | ||
| It is safe to call ``evaluate`` multiple times on the same batch | ||
| without side effects. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| current: | ||
| The live ``Batch`` from the dynamics step. Treat as | ||
| read-only; do not modify any field. | ||
|
|
||
| Returns | ||
| ------- | ||
| BiasResult | ||
| Fully detached outputs. All tensor fields must satisfy | ||
| ``requires_grad=False`` and ``grad_fn is None``. | ||
| """ | ||
| ... |
There was a problem hiding this comment.
Separating a read-only evaluation path from a state-mutating one is the best
idea in this PR.
Hook already carries frequency and stage, and DynamicsStage already has AFTER_COMPUTE and AFTER_STEP. Training reached the same conclusion already: TrainingUpdateHook claims DO_BACKWARD
and DO_OPTIMIZER_STEP rather than inventing a parallel lifecycle.
# nvalchemi/hooks/_protocol.py
class Hook(Protocol):
frequency: int
stage: Enum | None
class StatefulHook(Hook, Protocol):
"""A hook whose state evolves with the trajectory."""
read_only = False
def commit(self) -> None: ... # sync boundary; optional
def state_dict(self) -> Mapping[str, Any]: ...
def load_state_dict(self, state: Mapping[str, Any]) -> None: ...A bias then becomes a StatefulHook that is also a BaseModelMixin. The same
lifecycle covers NEB's climbing-image promotion, an adaptive thermostat, and an
adaptive neighbour skin — none of which are biases, and all of which need
"read-only during compute, mutate after the step, sync occasionally".
| class ConservativeBias(nn.Module, BaseModelMixin): | ||
| """Autograd helper that derives atomic forces and stress from energy. | ||
|
|
||
| Composed as ``nn.Module, BaseModelMixin`` — the house multiple-inheritance | ||
| idiom (``LennardJonesModelWrapper(nn.Module, BaseModelMixin)``, | ||
| ``TrainingStrategy(BaseModel, HookRegistryMixin)``). | ||
| ``BaseModelMixin.__init_subclass__`` is cooperative (it calls | ||
| ``super().__init_subclass__(**kwargs)``), so it composes correctly with | ||
| further mixins added later for adaptive or checkpointable biases. | ||
|
|
||
| A conservative bias is an additive potential with no learned parameters, | ||
| which is exactly the shape of ``DFTD3ModelWrapper`` and | ||
| ``LennardJonesModelWrapper``. The cost of the abstraction is the two | ||
| ``BaseModelMixin`` abstract methods (:attr:`embedding_shapes` and | ||
| :meth:`compute_embeddings`), stubbed here the same way those two wrappers | ||
| stub them. What it buys: | ||
|
|
||
| * ``model_config.active_outputs`` declares which outputs this bias | ||
| produces, replacing an ad-hoc private flag. | ||
| * :meth:`distribution_spec` gives domain decomposition a defined answer | ||
| instead of an undefined one — see that method for why the default is | ||
| deliberately ``None``. | ||
| * ``+`` composition with a model via ``PipelineModelWrapper``, and | ||
| ``state_dict``/``load_state_dict`` from ``nn.Module`` for checkpointing. | ||
|
|
||
| .. note:: | ||
|
|
||
| Subclasses **must** call ``super().__init__(name=...)``. This is the | ||
| ``nn.Module`` requirement (attribute assignment before | ||
| ``Module.__init__`` raises), and ``BaseModelMixin.__init_subclass__`` | ||
| additionally verifies ``self.model_config`` is set after construction. | ||
|
|
||
| Subclass ``ConservativeBias`` and override :meth:`energy` to return a | ||
| differentiable per-graph bias energy ``[B, 1]``. The base class | ||
| provides :meth:`evaluate`, which: | ||
|
|
||
| 1. Enters a local ``torch.enable_grad()`` region (safe inside | ||
| ``torch.no_grad()`` outer contexts). | ||
| 2. Creates a detached positions leaf ``pos_leaf`` (for forces) and, | ||
| for periodic batches, a per-graph strain leaf via | ||
| :func:`~nvalchemi.models._utils.prepare_strain` (for stress). | ||
| 3. Substitutes the strained positions and cell into the batch, calls | ||
| :meth:`energy`, and restores the original tensors unconditionally | ||
| in a ``finally`` block. | ||
| 4. Derives forces and stress in one ``autograd.grad`` call through | ||
| :func:`~nvalchemi.models._utils.autograd_forces_and_stresses`. | ||
| 5. Constructs a ``BiasResult`` from fully detached output tensors. | ||
|
|
||
| The framework must never place a tensor with ``requires_grad=True`` or | ||
| a non-null ``grad_fn`` into the live ``Batch``, ``BiasResult``, | ||
| retained history, bias state, observables, or a checkpoint. |
There was a problem hiding this comment.
This class is 415 lines to turn an energy function into forces and stress by
autograd. models/_utils.py already does the numerical part:
autograd_forces_and_stresses takes exactly these arguments and is what every
model wrapper already calls.
What the class adds that is genuinely valuable is isolation: it evaluates the
energy on a detached copy of the batch, restores every field it touched in a
finally block, and guarantees the tensors it hands back carry no grad graph.
That protection is useful to anything that computes derivatives against a live
batch, not just to biases — which means it belongs in the shared helper rather
than in enhanced_sampling.
# nvalchemi/models/_utils.py (~60 lines)
def isolated_energy_derivatives(
energy_fn: Callable[[Batch], Energy], batch: Batch, *,
want_forces: bool = True, want_stress: bool = True, allow_unused: bool = False,
) -> ModelOutputs:
"""Evaluate energy_fn on a detached view of batch and return detached
derivatives. Every mutated field is restored in a finally block, so the
caller's Batch never carries a grad_fn after this returns."""
# nvalchemi/enhanced_sampling/_bias.py (~40 lines, was 415)
class ConservativeBias(nn.Module, BaseModelMixin):
def energy(self, current: Batch) -> Energy: ...
def forward(self, data: Batch, **kwargs) -> ModelOutputs:
return isolated_energy_derivatives(self.energy, data, ...)Once the helper is shared, someone writing an NEB spring term or a custom wall
potential can call isolated_energy_derivatives directly and get the same
safety guarantees, without depending on the enhanced-sampling package at all.
laserkelvin
left a comment
laserkelvin
left a comment
There was a problem hiding this comment.
I reproduced five additional blocking issues on the current head and left each one inline.
|
|
||
| **BiasedPotentialHook** — add an external bias potential for enhanced sampling. | ||
|
|
||
| > **Deprecated.** Use `nvalchemi.enhanced_sampling` (`ConservativeBias`) for new |
There was a problem hiding this comment.
You do not need to deprecate things within skills. Just remove them entirely.
| sweep the target position along a reaction coordinate and post-process the | ||
| windowed histograms with WHAM or MBAR. | ||
|
|
||
| .. note:: |
There was a problem hiding this comment.
Just refactor the example, instead of just adding a deprecation note
| "the caller's responsibility before aggregation." | ||
| ) | ||
|
|
||
| summed = sum_outputs( |
There was a problem hiding this comment.
blocking (P1): sum_outputs permits broadcasting, so an energy shaped [1, 1] silently applies to every graph in a [B, 1] result; [1, 3] forces and [1, 3, 3] stress broadcast the same way. All results here describe the same batch, so repeated fields should have exact matching shapes. Validate shape equality before summing and add mismatch regressions.
| ``energy``, ``forces``, and — when active and the batch is | ||
| periodic — ``stress``. | ||
| """ | ||
| result = self.evaluate(data) # type: ignore[arg-type] |
There was a problem hiding this comment.
blocking (P1): forward() always enters evaluate(), so an autograd pipeline narrows this step to {"energy"} but receives detached energy and then fails with element 0 of tensors does not require grad. Energy-only calls also compute forces and stress before dropping them. Preserve the energy graph when only energy is active, and reserve evaluate() for direct derivative outputs.
|
|
||
| if strain_cell is not None: | ||
| pos_for_energy, cell_for_energy, displacement = prepare_strain( | ||
| pos_leaf, strain_cell, current.batch_idx |
There was a problem hiding this comment.
blocking (P1): forward() accepts AtomicData, but periodic evaluation reaches current.batch_idx and raises AttributeError because only Batch provides it. The BaseModelMixin forward contract supports either input. Convert AtomicData to a one-system Batch at the boundary, as PipelineModelWrapper.forward() does, or narrow the public contract to Batch.
| for name, t in tensor_fields.items(): | ||
| if t is None or not t.is_floating_point(): | ||
| continue | ||
| if not t.isfinite().all(): |
There was a problem hiding this comment.
blocking (P1): each Python if not t.isfinite().all() reads a CUDA scalar on the host, so a result with energy, forces, and stress performs three full tensor scans and three stream synchronizations every step; aggregation scans them again. This hot path should stay asynchronous. Make full finiteness checks opt-in or lower-frequency at the runner boundary, while keeping cheap shape and graph-state checks here.
| pbc_mask = pbc.to(dtype=cell.dtype) # [B, 3] | ||
|
|
||
| # Fractional displacement | ||
| cell_inv = torch.linalg.inv(cell) # [B, 3, 3] |
There was a problem hiding this comment.
blocking (P1): the full-cell inverse mixes non-periodic lattice vectors into fractional coordinates. For cell=[[1,0,0],[100,1,0],[0,0,1]], pbc=[T,F,F], and Cartesian displacement [0,0.4,0], eager and compiled paths return 39.002 instead of 0.4. Partial-axis MIC must solve and wrap only in the periodic lattice subspace; add a skewed mixed-axis regression.
| # limitations under the License. | ||
| """Collective-variable functions for enhanced sampling. | ||
|
|
||
| CVs are plain callables — no class hierarchy, no registration. Any |
There was a problem hiding this comment.
Why is this docstring here?
| """Collective-variable functions for enhanced sampling. | ||
|
|
||
| CVs are plain callables — no class hierarchy, no registration. Any | ||
| differentiable function ``cv(batch: Batch) -> Tensor[B, D]`` satisfies |
There was a problem hiding this comment.
If this was the case, then we should have a CollectiveVariable protocol. You can't guarantee differentiability in the Python API, but it would at least allow type hinting
| differentiable function ``cv(batch: Batch) -> Tensor[B, D]`` satisfies | ||
| the CV interface. | ||
|
|
||
| Available CVs: :func:`pair_distance`. |
There was a problem hiding this comment.
This is not the way to document modules
| * Zarr checkpoint support | ||
| * General triclinic MIC for unreduced cells | ||
|
|
||
| Relationship to ``BiasedPotentialHook`` |
There was a problem hiding this comment.
Why is this commentary left in here?
There was a problem hiding this comment.
Why is this a private module? We have public API declared in here
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class BiasResult: |
There was a problem hiding this comment.
I'm not sure if I'm in favor of having this declared as a dedicated class on its own
We already have Batch, and even with ModelOutputs we are intentionally unopinionated. The problem is declaring these variables again elsewhere outside of Batch makes it difficult and confusing to keep track.
| runner's output dict. All tensors must be detached. | ||
| """ | ||
|
|
||
| energy: Tensor | None = None |
There was a problem hiding this comment.
You have declared shapes in the docstring but not in the dataclass. You should be documenting them in the code with jaxtyping
3 participants
ALCHEMI Toolkit Pull Request
Description
Compile spike (CPU + CUDA): ConservativeBias autograd helper skeleton; pair_distance with nonperiodic and MIC; BiasPotential protocol and BiasResult; bias aggregation; validates no graph breaks or memory growth across 10 repeated calls; documents fallback
Type of Change
Related Issues
Changes Made
Testing
make pytest)make lint)Checklist
Additional Notes
Tip
This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.