Skip to content

Add shared Categories list-of-categories datastructure - #12

Merged
cortner merged 6 commits into
mainfrom
categories
Jul 6, 2026
Merged

Add shared Categories list-of-categories datastructure#12
cortner merged 6 commits into
mainfrom
categories

Conversation

@cortner

@cortner cortner commented Jul 1, 2026

Copy link
Copy Markdown
Member

Summary

Adds a single, generic, isbits list-of-categories datastructure to ACEbase,
consolidating the zlist / _i2z / _z2i machinery that ACEpotentials, ACEradials,
and AtomicOrbitalKernels each reimplement. New file src/categories.jl (plus tests);
no new dependencies; not exported (use ACEbase.Categories, ACEbase.cat2idx, …).
Version stays 0.4.7 — this is a non-breaking addition.

Design

  • Categories{N,T} stores a sorted SVector{N,T} of distinct categories; the sorted
    position is the 1-based index. Generic over the category type T (chemical species =
    T=Int atomic numbers); isbits whenever T is.
  • cat2idx(c, z) — category → index — is non-throwing and allocation-free (returns
    0 when absent), so it is safe inside GPU kernels. Compile-time branch on N: linear
    scan for small N (≤ 8), binary search (searchsortedfirst) for many species.
  • idx2cat(c, i) — index → category. Collection interface (length, iterate,
    getindex, eltype, in, ==, show). Batched lookup via broadcast:
    cat2idx.(c, zs) (a Base.broadcastable method treats a Categories as a scalar).
  • Constructor Categories(itr; allow_nonbits=false) sorts, checks distinctness, and
    requires isbitstype(T) by default (so the struct is always device-movable);
    allow_nonbits=true opts out for CPU-only / abstract-category use (e.g. Symbol).
    No AtomsBase dependency — chemical-symbol → atomic-number conversion stays in the
    consumer.
  • Related helpers in scope: make_smatrix(obj, NZ) (per-pair parameter matrices, with a
    dense-Matrix fallback for large NZ) and the pair-index utilities symidx,
    catcat2idx, catcat2idx_sym.

Notes / follow-ups

  • Sorted-order caveat: indices are defined by sorted order rather than insertion
    order; when a consumer migrates, its per-species/per-pair parameter arrays must be
    addressed consistently through idx2cat/cat2idx, and any pre-trained models would
    need their species axes permuted on load.
  • Downstream migration (ACEradials first) is out of scope here.

Testing

Pkg.test() — 634/634 pass. Covers the bijection round-trip (Int and Char), the
sorted invariant, both lookup regimes (N=3 and N=100), absent-key handling, the
isbits policy + opt-out, allocation-free lookup, broadcast lookup, and the pair
utilities / make_smatrix fallback.

🤖 Generated with Claude Code

cortner and others added 6 commits June 29, 2026 22:36
Consolidates the `zlist` / `_i2z` / `_z2i` machinery duplicated across
ACEpotentials, ACEradials and AtomicOrbitalKernels into a single generic,
sorted, isbits type in ACEbase.

- `Categories{N,T}` stores a sorted `SVector` of distinct categories; the
  sorted position is the 1-based index, so `cat2idx` is a linear scan for
  small `N` and a binary search for large `N`, over the stored list.
- `cat2idx` is non-throwing and allocation-free (returns 0 when absent) so it
  is safe to call inside GPU kernels.
- generic over the category type `T` (chemical species = `T=Int` atomic
  numbers); no chemistry-specific code or AtomsBase dependency.
- also ships `idx2cat`, the collection interface, `make_smatrix` (with a dense
  `Matrix` fallback for large `NZ`), and the pair-index utilities `symidx`,
  `catcat2idx`, `catcat2idx_sym`.
- not exported: use explicit `ACEbase.cat2idx` etc.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The constructor now errors when the category type `T` is not an `isbits`
type (e.g. `Symbol`, `String`), so a `Categories` is always safe to move to a
GPU / pass into a kernel. An `allow_nonbits = true` keyword opts out for
CPU-only / abstract-category use; the varargs constructor forwards it.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Drop the varargs constructor, the raw-list `idx2cat`/`cat2idx` methods, and the
batched `cat2idx`; construction is from a single iterable and batched lookups are
done by broadcasting. Add `Base.broadcastable(::Categories) = Ref(cats)` so
`cat2idx.(cats, zs)` treats the list as a scalar. Refresh docstrings/tests to
match; the kept methods are unchanged in behaviour.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds ACEbase's first package extension (weakdep on AtomsBase):

- `chemical_species(x)` converts an Integer (atomic number), Symbol, String, or
  ChemicalSpecies to an AtomsBase `ChemicalSpecies` (errors otherwise).
- `chemical_categories(list)` builds a `Categories{N,ChemicalSpecies}` from a list
  of chemical species (mixed input forms allowed; duplicates rejected).

Both are declared as stubs in the main package and implemented in
`ext/ACEbaseAtomsBaseExt.jl`. AtomsBase 0.5 does not define an ordering for
`ChemicalSpecies`, so the extension temporarily pirates
`Base.isless(::ChemicalSpecies, ::ChemicalSpecies)` (by atomic number) to let
`Categories` sort / binary-search; this is to be removed once the order is added
upstream in AtomsBase.

Version 0.4.7 -> 0.4.8; julia compat bumped to 1.9 (extensions).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
AtomsBase 0.5.3 defines a total order for ChemicalSpecies, so the
temporary Base.isless definition in the extension is no longer needed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@cortner
cortner merged commit 093631a into main Jul 6, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant