Julia Implementation (#79)

---------

Co-authored-by: Christoph Ortner <[email protected]>
Co-authored-by: frostedoyster <[email protected]>

.github/workflows/tests.yml

@@ -23,6 +23,11 @@ jobs:
         with:
           python-version: "3.10"
 
+      - name: setup Julia
+        uses: julia-actions/setup-julia@v1
+        with:
+          version: '1.9.2'
+
       - name: install tests dependencies
         run: |
           python -m pip install --upgrade pip
@@ -42,6 +47,16 @@ jobs:
           # Use the CPU only version of torch when building/running the code
           PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
 
+      - name: run Julia tests
+        run: |
+          cd julia/test/
+          julia --project=.. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile();'
+          julia --project=.. runtests.jl
+          cd python/
+          julia -e 'using Pkg; Pkg.add("PyCall"); Pkg.instantiate();'
+          python -m pip install numpy pytest sphericart julia
+          python -m pytest .
+
   # check that we can build Python wheels on any Python version
   python-build:
     runs-on: ubuntu-20.04

.gitignore

@@ -3,3 +3,5 @@ dist/
 *.egg-info
 .tox/
 *.pyc
+.vscode
+Manifest.toml

LICENSE-MIT

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 sphericart contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -25,6 +25,8 @@ If you are using sphericart for your academic work, you can cite it as
 }
 ```
 
+This library is dual-licensed under the Apache License 2.0 and the MIT license. You can use to use it under either of the two licenses.
+
 ## Installation
 
 ### Python
@@ -55,6 +57,14 @@ pip install --extra-index-url https://download.pytorch.org/whl/cpu .[torch]
 Building from source is also necessary to use sphericart's PyTorch GPU 
 functionalities, and it requires a CUDA compiler.
 
+### Julia 
+
+A native Julia implementation of `sphericart` is provided, called `SpheriCart`. 
+Install the package by opening a REPL, switch to the package manager by 
+typing `]` and then `add SpheriCart`. 
+See [julia/README.md](julia/README.md) for usage.
+
+
 ### C and C++
 
 From source

julia/Project.toml

@@ -0,0 +1,19 @@
+name = "SpheriCart"
+uuid = "5caf2b29-02d9-47a3-9434-5931c85ba645"
+authors = ["Christoph Ortner <[email protected]> and contributors"]
+version = "0.0.2-dev"
+
+[deps]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+
+[extras]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test", "LinearAlgebra", "Random"]

julia/README.md

@@ -0,0 +1,65 @@
+# SpheriCart 
+
+A Julia implementation of real solid and spherical harmonics, following
+```quote
+Fast evaluation of spherical harmonics with sphericart, 
+Filippo Bigi, Guillaume Fraux, Nicholas J. Browning and Michele Ceriotti; 
+J. Chem. Phys. 159, 064802 (2023); arXiv:2302.08381
+```
+
+
+`SpheriCart.jl` is released under MIT license and under Apache 2.0 license.
+
+## Installation 
+
+Install the package by opening a REPL, switch to the package manager by typing `]` and then `add SpheriCart`.
+
+## Basic Usage
+
+There are two implementations of real solid harmonics and real spherical harmonics
+- a generated  implementation for a single `𝐫::SVector{3, T}` input, returning the spherical harmonics as an `SVector{T}`. 
+- a generic implementation that is optimized for evaluating over batches of inputs, exploiting SIMD vectorization. 
+
+For large enough batches (system dependent) the second implementation is comparable to or faster than broadcasting over the generated implementation. For single inputs, the generated implementation is far superior in performance. 
+
+
+```julia
+using SpheriCart, StaticArrays 
+
+# generate the basis object 
+L = 5
+basis = SolidHarmonics(L)
+# Replace this with 
+#  basis = SphericalHarmonics(L) 
+# to evaluate the spherical instead of solid harmonics 
+
+# evaluate for a single input 
+𝐫 = @SVector randn(3) 
+# Z : SVector of length (L+1)²
+Z = basis(𝐫)  
+Z = compute(basis, 𝐫)
+# ∇Z : SVector of length (L+1)², each ∇Z[i] is an SVector{3, T}
+Z, ∇Z = compute_with_gradients(basis, 𝐫)
+
+# evaluate for many inputs 
+nX = 32
+Rs = [ @SVector randn(3)  for _ = 1:nX ]
+# Z : Matrix of size nX × (L+1)² of scalar 
+# dZ : Matrix of size nX × (L+1)² of SVector{3, T}
+Z = basis(Rs)  
+Z = compute(basis, Rs)
+Z, ∇Z = compute_with_gradients(basis, Rs)
+
+# in-place evaluation to avoid the allocation 
+compute!(Z, basis, Rs)
+compute_with_gradients!(Z, ∇Z, basis, Rs)
+```
+
+Note that Julia uses column-major indexing, which means that for batched output the loop over inputs is contiguous in memory. 
+
+<!-- ## Advanced Usage
+
+TODO:  
+- different normalizations
+- enforce static versus dynamic 
+- wrapping outputs into zvec for easier indexing  -->

julia/benchmarks/benchmark.jl

@@ -0,0 +1,104 @@
+# to run this benchmark add `BenchmarkTools` to the standard julia 
+# environment (e.g. `julia` and `] add BenchmarkTools`)
+# then run this scrict via from the current folder via 
+# `julia --project=.. -O3 benchmark.jl`
+
+# This is just a preliminary benchmark for initial testing. 
+# We need to write a proper benchmark suite to check performance 
+# regressions and compare against the C++ implementation on the same 
+# system. 
+
+using StaticArrays, BenchmarkTools, SpheriCart
+using SpheriCart: SolidHarmonics, compute, compute!, 
+                  static_solid_harmonics, 
+                  compute_with_gradients, compute_with_gradients!
+
+##
+
+@info("static_solid_harmonics")
+𝐫 = @SVector randn(3)
+for L = 1:12
+   @show L
+   basis = SolidHarmonics(L; static=true)
+   basis(𝐫) # warmup 
+   @btime ($basis)($𝐫)
+   # these two are equivalent - just keeping them here for testing 
+   # since there is an odd effect that in some environments there there 
+   # is an unexplained overhead in the `compute` interface. 
+   # @btime static_solid_harmonics($(Val(L)), $𝐫)
+   # @btime compute($basis, $𝐫)
+end
+
+
+##
+
+@info("batched evaluation vs broadcast")
+@info("nX = 32 (try a nice number)")
+@info("broadcast! cost is almost exactly single * nX")
+
+Rs = [ (@SVector randn(3)) for _ = 1:32 ]
+
+for L = 3:3:15
+   @show L
+   basis = SolidHarmonics(L)
+   Zs = static_solid_harmonics.(Val(L), Rs)
+   print("    single: "); @btime static_solid_harmonics($(Val(L)), $(Rs[1]))
+   print("broadcast!: "); @btime broadcast!(𝐫 -> static_solid_harmonics($(Val(L)), 𝐫), $Zs, $Rs)
+   Zb = compute(basis, Rs)
+   print("   batched: "); (@btime compute!($Zb, $basis, $Rs));
+end
+
+##
+
+@info("static vs generic basis for single input (nX = 1)")
+@info("  this shouws that the generic single-input implementation needs work")
+
+for L = 3:3:30 
+   local 𝐫
+   @show L 
+   𝐫 = @SVector randn(3)
+   basis_st = SolidHarmonics(L; static=true)
+   basis_dy = SolidHarmonics(L; static=false)
+   print("  static: "); @btime compute($basis_st, $𝐫)
+   print(" generic: "); @btime compute($basis_dy, $𝐫)
+end 
+
+##
+
+
+@info("compute! vs compute_with_gradients! (using 32 inputs)")
+
+for L = 1:2:15
+   @show L 
+   nX = 32
+   basis = SolidHarmonics(L)
+   Rs = [ (@SVector randn(3)) for _ = 1:nX ]
+   Z, dZ = compute_with_gradients(basis, Rs)
+   print("        compute!: "); @btime compute!($Z, $basis, $Rs)
+   print(" _with_gradient!: "); @btime compute_with_gradients!($Z, $dZ, $basis, $Rs)
+end 
+
+
+##
+
+
+@info("compute vs compute_with_gradients for code-generated basis")
+for L = 1:10
+   @show L 
+   basis = SolidHarmonics(L; static=true)
+   𝐫 = @SVector randn(3)
+   Z1 = compute(basis, 𝐫)
+   Z, dZ = compute_with_gradients(basis, 𝐫)
+   print("        compute: "); @btime compute($basis, $𝐫)
+   print(" _with_gradient: "); @btime compute_with_gradients($basis, $𝐫)
+end 
+
+
+
+## -------------- 
+
+# @profview let Z = Z, uf_Zlm = uf_Zlm, XX = XX
+#    for n = 1:3_000_000
+#       evaluate!(Z, uf_Zlm, XX)
+#    end
+# end

julia/src/SpheriCart.jl

@@ -0,0 +1,19 @@
+module SpheriCart
+
+using StaticArrays, OffsetArrays, ObjectPools
+
+export SolidHarmonics, 
+       compute, 
+       compute!, 
+       compute_with_gradients,
+       compute_with_gradients!
+
+include("indexing.jl")
+include("normalisations.jl")
+include("api.jl")
+include("generated_kernels.jl")
+include("batched_kernels.jl")
+include("spherical.jl")
+
+end
+