Skip to content

AllScAIP and UMA integration#496

Open
allaffa wants to merge 4 commits into
ORNL:mainfrom
allaffa:allscaip_uma_integration
Open

AllScAIP and UMA integration#496
allaffa wants to merge 4 commits into
ORNL:mainfrom
allaffa:allscaip_uma_integration

Conversation

@allaffa
Copy link
Copy Markdown
Collaborator

@allaffa allaffa commented Apr 30, 2026

No description provided.

allaffa added 2 commits April 30, 2026 15:37
- Vendor minimal AllScAIP code under hydragnn/utils/model/allscaip/
  (plus shared EScAIP utilities under hydragnn/utils/model/escaip/utils/).
- Add hydragnn/models/AllScAIPStack.py wrapper:
  * runs the full transformer backbone in _embedding(), exposes the
    result as the invariant node feature for HydraGNN's standard
    multi-head decoders;
  * uses standard num_conv_layers as the AllScAIP depth and replaces
    Base's per-layer loop with a single Identity no-op iteration;
  * inherits HydraGNN's activation_function (no separate
    allscaip_activation knob);
  * documents that AllScAIP is NOT e3nn-equivariant (scalar
    transformer; SH used as input features only).
- Wire allscaip_* hyperparameters as first-class Architecture keys via
  create_model / create_model_config / update_config (matches PNA,
  MACE, EGNN, ... patterns).
- setup.py: switch to find_namespace_packages so the vendored
  subpackages are installed.
- CI: add 'AllScAIP' to the basic mpnn parametrization in
  tests/test_graphs.py::pytest_train_model and
  tests/test_graphs_graphattr.py::pytest_train_model_graphattr
  (8 new cases pass). Excluded from equivariant / lengths /
  vectoroutput / global-attention families since AllScAIP is
  is_edge_model=False and not equivariant.
@allaffa allaffa requested a review from messerlyra April 30, 2026 20:48
@allaffa allaffa self-assigned this Apr 30, 2026
@allaffa allaffa added the enhancement New feature or request label Apr 30, 2026
- Add UMAStack thin wrapper around fairchem.core eSCNMDBackbone,
  with lazy import + clear error if fairchem-core is missing.
- Wire uma_* config keys through create_model and update_config
  (sphere_channels<-hidden_dim, lmax<-max_ell, etc.); uma_hidden_channels
  defaults to hidden_dim. Clamp uma_mmax<=max_ell to avoid zero-channel
  SO(2) convs. Robust cell/pbc handling for batched PyG graphs.
- Add fairchem-core to requirements-optional.txt.
- Add UMA to CI parametrization in tests/test_graphs.py and
  tests/test_graphs_graphattr.py (basic + equivariant suites).
- Add UMA to LennardJones training CI in
  tests/test_forces_equivariant_training.py and to the energy-only
  equivariance benchmark in tests/test_forces_equivariant.py.
- Apply black==21.5b1 formatting.
@allaffa allaffa changed the title Allscaip uma integration AllScAIP and UMA integration Apr 30, 2026
Adds the standard 10-line BSD-3-clause header to every Python file
that lacked one. The year reflects the file's last git-commit year
(falls back to mtime). Vendored FairChem trees under
examples/*/fairchem/core/** are intentionally skipped; they retain
Meta's MIT header.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant