Add shared Categories list-of-categories datastructure - #12
Merged
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a single, generic,
isbitslist-of-categories datastructure to ACEbase,consolidating the
zlist/_i2z/_z2imachinery 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 sortedSVector{N,T}of distinct categories; the sortedposition is the 1-based index. Generic over the category type
T(chemical species =T=Intatomic numbers);isbitswheneverTis.cat2idx(c, z)— category → index — is non-throwing and allocation-free (returns0when absent), so it is safe inside GPU kernels. Compile-time branch onN: linearscan 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)(aBase.broadcastablemethod treats aCategoriesas a scalar).Categories(itr; allow_nonbits=false)sorts, checks distinctness, andrequires
isbitstype(T)by default (so the struct is always device-movable);allow_nonbits=trueopts out for CPU-only / abstract-category use (e.g.Symbol).No AtomsBase dependency — chemical-symbol → atomic-number conversion stays in the
consumer.
make_smatrix(obj, NZ)(per-pair parameter matrices, with adense-
Matrixfallback for largeNZ) and the pair-index utilitiessymidx,catcat2idx,catcat2idx_sym.Notes / follow-ups
order; when a consumer migrates, its per-species/per-pair parameter arrays must be
addressed consistently through
idx2cat/cat2idx, and any pre-trained models wouldneed their species axes permuted on load.
Testing
Pkg.test()— 634/634 pass. Covers the bijection round-trip (IntandChar), thesorted invariant, both lookup regimes (
N=3andN=100), absent-key handling, theisbits policy + opt-out, allocation-free lookup, broadcast lookup, and the pair
utilities /
make_smatrixfallback.🤖 Generated with Claude Code