EquivariantTensors Backend #305

cortner · 2025-11-18T09:43:15Z

dev-0.10 uses ET only superficially, this PR is to integrate ET more deeply, ~~remove duplicated functionality~~, with the end-goal to enable GPU evaluation.

Removing duplicated functionality will be done in a future version.

jameskermode · 2025-12-12T23:20:10Z

See #312 which targets this PR and implements most of the remaining steps (but still needs to be cleaned up)

cortner · 2025-12-19T19:24:39Z

UPDATE: I think I now have all functionality that we can expect ET to provide. My plan is to finalize the unit tests and then I think we can make a plan how to proceed. Here is a tentative list:

Backend

splines (still missing)
gpu evaluation (I broke this again)

Low model level

pair potentials, E0s => I'd like to treat these as stacked models instead of integrated, to be discussed.
wrap the new ACE model into a calculator that takes care of energy, forces, virial, etc
provide a wrapper that generates an old ACE model, then converts to new
extensive correctness tests and unit tests to compare the two models for correctness and performance

User-facing

Throuhgout 0.10 keep both frameworks alive
For 0.11 retire the entirety of ACE 0.8, 0.9 infrastructure which is still surviving in 0.10

Speculative

bring back purification

@ntuple

- Extract core helpers (_core_site_energies, _core_site_grads) for shared evaluation logic between ETACEPotential and WrappedETACE - Refactor WrappedETACE and ETACEPotential to use core helpers - Simplify stackedcalc.jl: replace manual AST building (_gen_sum, _gen_broadcast_sum) with idiomatic @nexprs/@ntuple from Base.Cartesian - Net reduction of ~50 lines while maintaining identical behavior 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Removes duplicate AtomsCalculators interface and evaluation logic by: - Making WrappedETACE mutable with co_ps field for training - Defining ETACEPotential as const alias for WrappedSiteCalculator{WrappedETACE} - Removing duplicate _evaluate_* functions and AtomsCalculators methods - Adding accessor helpers (_etace, _ps, _st) for training functions The evaluation now flows through WrappedSiteCalculator's generic methods which call site_energies/site_energy_grads on the WrappedETACE model. This reduces ~66 lines of duplicated code. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Remove duplicate E0Model in favor of upstream ETOneBody - Unify WrappedSiteCalculator to work with all ETACE-pattern models directly - Document that ETACE, ETPairModel, ETOneBody share identical interface - Plan Phase 6 refactoring to eliminate WrappedETACE indirection - Update architecture diagrams showing target unified structure 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Refactor WrappedSiteCalculator to store ps, st, rcut directly - Remove E0Model (use upstream ETOneBody instead) - Remove WrappedETACE (functionality merged into WrappedSiteCalculator) - Remove old SiteEnergyModel interface (site_energies, site_energy_grads) - Update ETACEPotential to be type alias for WrappedSiteCalculator{ETACE} - Update training assembly accessors for new flat structure All ETACE-pattern models (ETACE, ETPairModel, ETOneBody) now work directly with WrappedSiteCalculator via their common interface: - model(G, ps, st) -> (site_energies, st) - site_grads(model, G, ps, st) -> edge gradients 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Add convert2et_full() to convert complete ACE model to StackedCalculator - Combines ETOneBody (E0), ETPairModel (pair), and ETACE (many-body) - Returns StackedCalculator compatible with AtomsCalculators - Add _copy_ace_params!() for many-body parameter copying - Copies radial basis Wnlq parameters - Copies readout WB parameters - Add _copy_pair_params!() for pair potential parameter copying - Based on mapping from test/etmodels/test_etpair.jl - Copies pair radial basis and readout parameters 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Replace E0Model tests with ETOneBody (upstream) tests - Remove WrappedETACE tests (no longer exists) - Update WrappedSiteCalculator tests for new (model, ps, st, rcut) signature - Update ETACEPotential construction test for direct model access - Update silicon integration test to use ETOneBody and unified wrapper Tests now use upstream models directly: - ETOneBody instead of E0Model - WrappedSiteCalculator(model, ps, st, rcut) instead of nested wrappers 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Return NamedTuple with empty edge_data matching ETACE/ETPairModel interface - Remove unnecessary Zygote import (hand-coded since gradient is trivially zero) - Update test to check isempty(∂G.edge_data) instead of zero norms The calling code in et_calculators.jl checks isempty(∂G.edge_data) and returns zero forces/virial when true, which is the correct behavior for ETOneBody (energy depends only on atom types, not positions). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Comment out Pkg.activate in test_committee.jl that was switching away from the test project environment - Update test_etonebody.jl gradient test to check for NamedTuple return type with .edge_data field (matching the updated ETOneBody interface that returns consistent structure with ETACE/ETPairModel) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

ET ACE (site_basis_jacobian): - Remove ps.basis and st.basis from _jacobian_X call - The upstream ET._jacobian_X for SparseACEbasis only takes 5 args: (basis, Rnl, Ylm, dRnl, dYlm) ET Pair (site_grads): - Implement hand-coded gradient using evaluate_ed instead of Zygote - Avoids Zygote InplaceableThunk issue with upstream EdgeEmbed rrule - Matches the pattern used in site_basis_jacobian Also inline _apply_etpairmodel to avoid calling site_basis (cleaner). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Fix z.number → z.atomic_number in E0_dict creation - Fix _copy_ace_params! path: rembed.basis.linl.W → rembed.post.W - Fix _copy_pair_params! path: rembed.basis.rbasis.linl.W → rembed.rbasis.post.W - Add benchmark comparing ACE vs ETACE StackedCalculator for full model 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Address moderator concern about commit 50ed668: - Avoid forming O(nnodes * nbasis) dense intermediate matrix - Compute edge gradients directly using loops - Same numerical results, better memory characteristics 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Bump EquivariantTensors compat to 0.4.2 in main and test Project.toml - Simplify site_basis_jacobian to use 5-arg _jacobian_X API (requires ET >= 0.4.2) - Improve ETPairModel site_grads memory efficiency: - Avoid O(nnodes * nbasis) intermediate matrix allocation - Compute edge gradients directly using loops 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Now that test project uses ET 0.4.2 (which fixed InplaceableThunk bug in EdgeEmbed rrule), we can use the simpler Zygote-based gradient computation for ETPairModel. Also fix _jacobian_X call in ETACE to use 7-arg API (requires ps, st). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

- Add benchmark/gpu_benchmark.jl for GPU energy/forces benchmarks - Test both many-body only (ETACE) and full model (E0 + Pair + Many-Body) - Add LuxCUDA to test/Project.toml for GPU testing support - GPU forces now work with Polynomials4ML v0.5.8+ (bug fix Dec 29, 2024) Results show significant GPU speedups: - Many-body energy: 6x-48x speedup (64-800 atoms) - Many-body forces: 3x-18x speedup (64-800 atoms) - Full model energy: 3x-36x speedup (64-800 atoms) - Full model forces: 1x-14x speedup (64-800 atoms) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

- Mark all core phases as complete - Add GPU benchmark results (energy and forces) - Document outstanding work: pair training assembly, ACEfit integration - Note basis index design discussion needed with maintainer - Update dependencies: ET 0.4.2, P4ML 0.5.8+ 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

- Add ETPairPotential and ETOneBodyPotential type aliases - Implement length_basis, energy_forces_virial_basis, potential_energy_basis for ETPairPotential, ETOneBodyPotential, and StackedCalculator - Add get/set_linear_parameters for all calculator types - Add ACEfit.basis_size dispatch for all calculator types - Import and extend length_basis, energy_forces_virial_basis from Models - ACEfit.assemble now works with full ETACE StackedCalculator 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

…neBodyPotential, StackedCalculator Tests cover: - ETOneBodyPotential returns empty arrays (0 learnable parameters) - ETPairPotential training assembly with learnable pair basis - StackedCalculator concatenation of basis from all components - Linear combinations reproduce energy/forces/virial - get/set_linear_parameters round-trip - ACEfit.basis_size dispatch 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Update EquivariantTensors compat to 0.4.3 in Project.toml and test/Project.toml - Fix ETOneBody site_grads to return fill(VState(), length(X.edge_data)) instead of similar(X.edge_data, 0) for type stability - Empty VState() acts as additive identity when summed with other VStates - Update test to verify new behavior 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

This commit addresses all review comments from PR #313 and fixes pre-existing test failures in ETOneBody. **Changes from PR #313 review:** 1. **test/et_models/test_et_calculators.jl**: - Removed 44 println status messages that provided misleading output - Deleted embedded performance benchmarks (CPU and GPU, ~100 lines) - Replaced manual parameter copying with utility function calls - Total: 169 lines removed for cleaner, more maintainable tests 2. **src/et_models/et_calculators.jl**: - Made parameter copying functions public: * _copy_ace_params! → copy_ace_params! * _copy_pair_params! → copy_pair_params! - These are now part of the public API since used by tests 3. **test/runtests.jl**: - Added ET Calculators test to CI suite (line 23) - Test now runs automatically with full test suite **Additional fix - ETOneBody site_grads:** 4. **src/et_models/onebody.jl**: - Fixed site_grads to return empty array instead of array of empty VStates - Resolves test failure at test_et_calculators.jl:234 - Resolves FieldError at test_et_calculators.jl:254 - ETOneBody energy depends only on atom types, not positions, so there are no position-dependent gradients **Test Results:** - ET Calculators: 182 tests passed (previously 94 passed, 2 failed/errored) - Overall: 1127 tests passed (up from 1040) - All PR #313 review items addressed ✓ 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Demonstrates two approaches for working with ET backend: 1. Converting from existing ACE model (recommended) - convert2et_full: Full model to StackedCalculator - convert2et: Many-body only to ETACE - convertpair: Pair potential to ETPairModel 2. Creating ETACE from scratch (advanced) - Direct EquivariantTensors component construction - ETOneBody, ETACE, StackedCalculator assembly Also shows training assembly interface for ACEfit integration. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Integrate the ETACE tutorial from examples/etmodels/ into the Documenter.jl-based documentation using Literate.jl. - Point Literate to examples/etmodels/etace_tutorial.jl (no duplication) - Add to Tutorials section in navigation - Add entry in tutorials/index.md 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Add StaticArrays, Lux, EquivariantTensors, and Polynomials4ML which are required by the ETACE tutorial examples. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Use ## instead of # for inline comment inside code block to prevent Literate.jl from splitting the code block at that line. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

ETACE Calculator Interface and Training Assembly

cortner · 2026-01-05T18:44:38Z

@jameskermode - I will merge once tests pass, unless I hear otherwise.

jameskermode · 2026-01-05T19:10:06Z

Go ahead. I suggest we tag this as v0.10.1. LAMMPS export can then come in v0.10.2 in case it causes any problems.

prototype ET Rnl

caeb7f0

cortner changed the base branch from main to dev-0.10 November 18, 2025 09:43

Merge branch 'dev-0.10' into co/etback

c468196

jameskermode mentioned this pull request Nov 20, 2025

EquivariantTensors + J1.12 Port #304

Merged

2 tasks

cortner added 9 commits November 23, 2025 12:07

start ET ace model prototype

291f384

simple test improvements

76e27b7

integrate Rnl -> ET + improved tests

e5e8a6a

fixed the energy test in new_backend

350001b

first complete Lux-style linear ace model draft

5bf14b7

converted the transforms to new ET implementation

c29a063

towards finding the Rnl bug

bc3a76f

Rnl now evaluates on the GPU

46b4175

first fully working GPU evaluation of E

ffe87d1

cortner changed the base branch from dev-0.10 to main December 10, 2025 04:23

Merge branch 'main' into co/etback

70e16ad

cortner added 8 commits December 18, 2025 13:02

update version bounds

109efc8

checkpoint new rbasis construction

86adc37

new backend now mostly up to date again

73e1109

Zygote.gradient evaluates (correctness tests missing)

f973935

prototype ETACE living inside ACEpotentials

a9a70ef

some test cleanup

7b23831

site gradients, basis and jacobians draft

74ec7e8

remove ET path

fc78b79

cortner added 5 commits December 19, 2025 20:38

derivative correctness tests

5233ba0

some gpu bugfixes

8cf42c6

steps towards GPU gradients

c794258

gpu gradient tests pass

1429040

update versioning

ce24dcb

jameskermode and others added 25 commits January 5, 2026 13:58

update plan

1e22042

Add missing dependencies to docs/Project.toml

4971a56

Add StaticArrays, Lux, EquivariantTensors, and Polynomials4ML which are required by the ETACE tutorial examples. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Merge pull request #313 from jameskermode/jrk/etcalculators

d7de233

ETACE Calculator Interface and Training Assembly

cortner changed the title ~~WIP: EquivariantTensors Backend~~ EquivariantTensors Backend Jan 5, 2026

cortner merged commit be22044 into main Jan 5, 2026
3 checks passed

cortner deleted the co/etback branch January 5, 2026 20:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

EquivariantTensors Backend #305

EquivariantTensors Backend #305

Uh oh!

cortner commented Nov 18, 2025 •

edited

Loading

Uh oh!

jameskermode commented Dec 12, 2025

Uh oh!

cortner commented Dec 19, 2025

Uh oh!

cortner commented Jan 5, 2026

Uh oh!

jameskermode commented Jan 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

EquivariantTensors Backend #305

EquivariantTensors Backend #305

Uh oh!

Conversation

cortner commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jameskermode commented Dec 12, 2025

Uh oh!

cortner commented Dec 19, 2025

Backend

Low model level

User-facing

Speculative

Uh oh!

cortner commented Jan 5, 2026

Uh oh!

jameskermode commented Jan 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cortner commented Nov 18, 2025 •

edited

Loading