Decouple SpheriCart from Polynomials4ML

Project.toml

@@ -1,7 +1,7 @@
 name = "Polynomials4ML"
 uuid = "03c4bcba-a943-47e9-bfa1-b1661fc2974f"
 authors = ["Christoph Ortner <christophortner@gmail.com> and contributors"]
-version = "0.5.10"
+version = "0.6.0"
 
 [deps]
 ACEbase = "14bae519-eb20-449c-a949-9c58ed33163e"
@@ -21,13 +21,18 @@ QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
-SpheriCart = "5caf2b29-02d9-47a3-9434-5931c85ba645"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 WithAlloc = "fb1aa66a-603c-4c1d-9bc4-66947c7b08dd"
 
+[weakdeps]
+SpheriCart = "5caf2b29-02d9-47a3-9434-5931c85ba645"
+
+[extensions]
+SpheriCartExt = "SpheriCart"
+
 [compat]
-ACEbase = "0.4.5"
+ACEbase = "0.4.7"
 BenchmarkTools = "1"
 Bumper = "0.7.0"
 ChainRulesCore = "1"
@@ -44,7 +49,7 @@ QuadGK = "2"
 Random = "1.10"
 SparseArrays = "1.10"
 SpecialFunctions = "2.2"
-SpheriCart = "0.2.3"
+SpheriCart = "0.2.4"
 StaticArrays = "1.5"
 Test = "1.10"
 WithAlloc = "0.1.0"
@@ -53,8 +58,9 @@ julia = "1.10"
 [extras]
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+SpheriCart = "5caf2b29-02d9-47a3-9434-5931c85ba645"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Printf", "Optimisers", "Zygote"]
+test = ["Test", "Printf", "Optimisers", "Zygote", "SpheriCart"]

docs/Project.toml

@@ -3,4 +3,5 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f"
+SpheriCart = "5caf2b29-02d9-47a3-9434-5931c85ba645"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

docs/src/polynomials.md

@@ -12,11 +12,13 @@ This page documents the public API for polynomial bases: the list of bases and f
 * 2D harmonics: 
    - Complex trigonometric polynomials [`CTrigBasis`](@ref)
    - Real trigonometric polynomials [`RTrigBasis`](@ref)
-* 3D harmonics: 
-   - Complex spherical harmonics [`complex_sphericalharmonics`](@ref)
-   - Real spherical harmonics [`real_sphericalharmonics`](@ref)
-   - Complex solid harmonics [`complex_solidharmonics`](@ref)
-   - Real solid harmonics [`real_solidharmonics`](@ref)
+* 3D harmonics: real/complex spherical and solid harmonics are provided by
+   [SpheriCart.jl](https://github.com/lab-cosmo/sphericart); with `ACEbase`
+   loaded they expose the same `evaluate` / `evaluate_ed` / `natural_indices`
+   interface used here.
+* Atomic orbitals: [`AtomicOrbitals`](@ref) — quantum-chemistry `Pn * Dn * Ylm`
+   product bases. The angular `Ylm` is a SpheriCart harmonics basis, so using
+   these requires `import SpheriCart` (see the docstring).
 
 ### In-place Evaluation  
 
@@ -72,7 +74,7 @@ The arrays `P1, P2, dP2` are Bumper-allocated i.e. are not allowed to leave the
 
 ### Utility functions
 
-The basis specification can be obtained using [`Polynomials4ML.natural_indices`](@ref)
+The basis specification can be obtained using `natural_indices`
 ```julia
 spec = natural_indices(basis)
 ```

ext/SpheriCartExt.jl

@@ -0,0 +1,21 @@
+module SpheriCartExt
+
+# SpheriCart-specific glue for the `AtomicOrbitals` basis: the precise complex
+# value type and the default angular basis (`SolidHarmonics`) used by the
+# `_rand_*` example constructors. Everything else about `AtomicOrbitals` is
+# generic over `basis.Ylm` (used via the ACEbase `evaluate` interface) and lives
+# in Polynomials4ML core.
+
+import Polynomials4ML: _ylm_valtype, _default_ylm
+using SpheriCart: SolidHarmonics, ComplexSolidHarmonics, ComplexSphericalHarmonics
+using StaticArrays
+
+# precise value type for the complex harmonics; the real harmonics use the
+# generic `_ylm_valtype` fallback in P4ML core (which returns `S`).
+_ylm_valtype(::Union{ComplexSolidHarmonics, ComplexSphericalHarmonics},
+             ::Type{<: SVector{3, S}}) where {S} = Complex{S}
+
+# default angular basis for the `_rand_*` example constructors in P4ML core
+_default_ylm(L) = SolidHarmonics(L)
+
+end

src/Polynomials4ML.jl

@@ -4,9 +4,10 @@ module Polynomials4ML
 
 import ACEbase 
 
-import ACEbase: evaluate, evaluate_d, evaluate_ed, 
-                evaluate!, evaluate_d!, evaluate_ed!, 
-                pullback, pullback!, pushforward, pushforward!
+import ACEbase: evaluate, evaluate_d, evaluate_ed,
+                evaluate!, evaluate_d!, evaluate_ed!,
+                pullback, pullback!, pushforward, pushforward!,
+                natural_indices
 
 import ChainRulesCore: rrule, frule, NoTangent, ZeroTangent
 import LuxCore: AbstractLuxLayer, initialparameters, initialstates                 
@@ -44,20 +45,9 @@ _generate_batch(basis::AbstractP4MLBasis; nbatch = rand(7:16)) =
          [ _generate_input(basis) for _ = 1:nbatch ]
 
 
-""" 
-   natural_indices(basis) -> AbstractVector
-
-Returns an abstract vector of "natural" descriptions of the basis functions in 
-the order that they are stored in the computed vector of basis function values.
-For example, for Chebyshev polynomials, `natural_indices(basis)` returns 
-`0:N`, where `N+1` is the length of the basis. For Spherical Harmmonics, 
-a natural description requires two indices `(l, m)`, so the output will be a 
-vector of tuples.
-
-At the moment, this function is used only for inspection and testing so no 
-strict format is enforced.
-"""
-function natural_indices end   # could rename this get_spec or similar ... 
+# `natural_indices(basis)` (owned by ACEbase, imported above) returns a vector
+# of "natural" descriptions of the basis functions, in storage order. For
+# Chebyshev this is `0:N`; for spherical harmonics it is a vector of `(l, m)`.
 
 """
    index(basis, k) -> Integer
@@ -101,11 +91,13 @@ include("splinify.jl")
 include("ctrig.jl")
 include("rtrig.jl")
 
-# 3d harmonics 
-include("sphericart.jl")
-
-# quantum chemistry 
+# 3d spherical harmonics (real + complex solid/spherical) are owned by SpheriCart
+# (with an ACEbase extension for the evaluate interface). The quantum-chemistry
+# atomic-orbital basis `AtomicOrbitals = Pn * Dn * Ylm` lives here; the
+# SpheriCart-specific glue (complex value types, default-Ylm constructors) is in
+# ext/SpheriCartExt.jl, loaded when SpheriCart is available.
 include("atomicorbitals/atomicorbitals.jl")
+export AtomicOrbitals, RadialDecay, GaussianDecay, SlaterDecay
 
 # generating product bases (generalisation of tensor products)
 # RETIRE - to be discussed?