Added GPU support for both CUDA and Metal, re-enabled 2D support, and updated tests

.github/workflows/CI.yml

@@ -1,6 +1,6 @@
 name: CI
-
 on:
+  workflow_dispatch:
   push:
     branches:
       - master
@@ -14,13 +14,13 @@ concurrency:
 
 env:
   JULIA_PKG_SERVER: ""
-  # Fix for Windows: Forces GR/Plots to be headless
   JULIA_NUM_THREADS: 1
   JULIA_PKG_PRECOMPILE_AUTO: 0
+  # headless GR/Plots
   GKSwstype: "100"
-  # Fix for macOS: Forces Matplotlib to use non-interactive backend
+  # non-interactive Matplotlib for PyPlot
   MPLBACKEND: "Agg"
-  # Force Julia to use Conda.jl
+  # force Conda.jl rather than a system Python
   PYTHON: ""
 
 jobs:
@@ -31,25 +31,41 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.10' # LTS
-          - '1'    # Latest Stable
+          - '1.12'
         os:
           - ubuntu-latest
-          - windows-latest
-          - macos-latest   # Apple Silicon (ARM64)
+          - macOS-latest   # Apple Silicon (ARM64)
         arch:
           - x64
-          - aarch64
         exclude:
-          - os: ubuntu-latest
-            arch: aarch64
-          - os: windows-latest
+          - os: macOS-latest
+            arch: x64
+        include:
+          - os: macOS-latest
             arch: aarch64
-          - os: macos-latest
-            arch: x64     # ARM only
+            version: '1.12'
 
     steps:
       - uses: actions/checkout@v4
+
+      - name: Check out ContinuousWavelets.jl
+        uses: actions/checkout@v4
+        with:
+          repository: JaredW40/ContinuousWavelets.jl
+          ref: master
+          path: deps/ContinuousWavelets.jl
+      - name: Check out FourierFilterFlux.jl
+        uses: actions/checkout@v4
+        with:
+          repository: JaredW40/FourierFilterFlux.jl
+          ref: master
+          path: deps/FourierFilterFlux.jl
+      - name: Check out MonogenicFilterFlux.jl
+        uses: actions/checkout@v4
+        with:
+          repository: JaredW40/MonogenicFilterFlux.jl
+          ref: main
+          path: deps/MonogenicFilterFlux.jl
 
       - uses: julia-actions/setup-julia@v2
         with:
@@ -58,27 +74,67 @@ jobs:
 
       - uses: julia-actions/cache@v2
 
-      # --- Custom Registry Configuration ---
       - name: Configure Git to use HTTPS
-        run: |
-          git config --global url."https://github.com/".insteadOf "git@github.com:"
+        run: git config --global url."https://github.com/".insteadOf "git@github.com:"
 
       - name: Add Registries
-        shell: julia --project={0} {0}
+        shell: julia --color=yes {0}
         run: |
           using Pkg
-          Pkg.Registry.add(RegistrySpec(name="General"))
+          Pkg.Registry.add("General")
           Pkg.Registry.add(RegistrySpec(url="https://github.com/UCD4IDS/TetrapodsRegistry"))
-      # -------------------------------------
 
-      - uses: julia-actions/julia-buildpkg@v1
+      - name: Generate the CI environment
+        run: |
+          julia --project=ci -e '
+            using Pkg
+            Pkg.develop([
+              PackageSpec(path = "deps/ContinuousWavelets.jl"),
+              PackageSpec(path = "deps/FourierFilterFlux.jl"),
+              PackageSpec(path = "deps/MonogenicFilterFlux.jl"),
+              PackageSpec(path = "."),
+            ])
+            Pkg.add([
+              PackageSpec(name = "Test"),
+              PackageSpec(name = "Wavelets"),
+              PackageSpec(name = "Zygote"),
+              PackageSpec(name = "Flux"),
+              PackageSpec(name = "FFTW"),
+              PackageSpec(name = "AbstractFFTs"),
+              PackageSpec(name = "Functors"),
+              PackageSpec(name = "Adapt"),
+              PackageSpec(name = "MLDataDevices"),
+              PackageSpec(name = "LinearAlgebra"),
+              PackageSpec(name = "Logging"),
+              PackageSpec(name = "Random"),
+            ])'
 
-      - uses: julia-actions/julia-runtest@v1
+      - name: Add the CUDA stack (Linux only)
+        if: runner.os == 'Linux'
+        run: |
+          julia --project=ci -e '
+            using Pkg
+            Pkg.add([
+              PackageSpec(name = "CUDA"),
+              PackageSpec(name = "cuDNN"),
+              PackageSpec(name = "cuFFT"),
+              PackageSpec(name = "CUDNN_jll", version = "9.20"),
+            ])'
+      - name: Add Metal (macOS only)
+        if: runner.os == 'macOS'
+        run: julia --project=ci -e 'using Pkg; Pkg.add("Metal")'
 
-      - uses: julia-actions/julia-processcoverage@v1
+      - name: Precompile
+        run: julia --project=ci -e 'using Pkg; Pkg.precompile()'
 
-      - uses: codecov/codecov-action@v5
+      - name: Run tests
+        env:
+          GROUP: ${{ runner.os == 'Linux' && 'CUDA' || runner.os == 'macOS' && 'Metal' || 'CPU' }}
+        run: julia --project=ci --code-coverage=user -e 'include("test/runtests.jl")'
+
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v4
         with:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: false
+          fail_ci_if_error: false

Project.toml

@@ -1,80 +1,60 @@
 name = "ScatteringTransform"
 uuid = "eadaac29-8b6b-5395-80fd-9ce36bb46293"
-version = "0.9.0"
-author = ["David Weber <david.weber2@gmail.com>", "Naoki Saito <nsaito@ucdavis.edu>", "Jared White <jgwhite@ucdavis.edu>", "Brian Knight <bcknight@ucdavis.edu>"]
+version = "0.8.3"
+authors = ["David Weber <david.weber2@gmail.com>", "Naoki Saito <nsaito@ucdavis.edu>", "Jared White <jgwhite@ucdavis.edu>", "Brian Knight <bcknight@ucdavis.edu>"]
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-BlackBoxOptim = "a134a8b2-14d6-55f6-9291-3336d3ab0209"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ContinuousWavelets = "96eb917e-2868-4417-9cb6-27e7ff17528f"
-DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2"
+CUDNN_jll = "62b44479-cb7b-5706-934f-f13b2eb2e645"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
-Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
-DistributedArrays = "aaf54ef3-cdf8-58ed-94cc-d582ad619b94"
-Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
-FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
-FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 FourierFilterFlux = "3d7dfd45-6c90-4c9b-b697-194a05757159"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-GLMNet = "8d5ece8b-de18-5317-b113-243142960cc6"
-HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
-Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
-JLD = "4138dd39-2aa7-5051-a626-17a0bb65d9c8"
-JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
-LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
-Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
-MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 MonogenicFilterFlux = "27e6407a-9159-4acb-a68a-dfae66d408cf"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
-Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
-Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
-Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-SharedArrays = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
-SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
-WaveletsExt = "8f464e1e-25db-479f-b0a5-b7680379e03f"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CUDA_Runtime_Discovery = "1af6417a-86b4-443c-805f-a4643ffb695f"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 cuFFT = "533571aa-0936-420e-b4be-9c66f5f626ca"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [extensions]
+CUDAExt = ["CUDA", "cuDNN", "cuFFT"]
+MetalExt = "Metal"
 ScatteringPlotsExt = "Plots"
 
 [compat]
 AbstractFFTs = "1"
-Adapt = "3, 4"
-CUDA = "4 - 6"
+Adapt = "4"
 ContinuousWavelets = "1"
-Distributions = "0.25"
+CUDA = "6.1"
+cuDNN = "6.1"
+CUDNN_jll = "~9.20"
+cuFFT = "6.1"
 FFTW = "1"
-FileIO = "1"
-FiniteDifferences = "0.12.34"
 Flux = "0.13, 0.14, 0.15, 0.16"
 Functors = "0.5.2"
-JLD2 = "0.4, 0.5, 0.6"
+julia = "1.12"
 Plots = "1.41.6"
-Reexport = "1"
 Wavelets = "0.9, 0.10"
-WaveletsExt = "0.2.3"
 Zygote = "0.6, 0.7"
-cuFFT = "6.2.0"
-julia = "1.10"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "CUDA", "BenchmarkTools"]
+test = ["Test", "BenchmarkTools"]

docs/src/plots.md

@@ -3,6 +3,7 @@
 
 ## 1D Plotting Functions
 ```@docs
+ScatteringPlotsExt.plotOriginalSignal1D
 ScatteringPlotsExt.plotZerothLayer1D
 ScatteringPlotsExt.plotFirstLayer1DSingleWavelet
 ScatteringPlotsExt.gifFirstLayer1D

ext/CUDAExt.jl

@@ -0,0 +1,20 @@
+module CUDAExt
+
+#=  This extension is intentionally minimal. `stFlux`'s device transfer is
+    handled generically by `Functors.@functor stFlux (mainChain,)` plus 
+    Flux's own `gpu`/`cpu` - no CUDA-specific code is needed here for that. 
+    `ConvFFT`'s actual CUDA-specific behavior already lives entirely in 
+    FourierFilterFlux's own CUDAExt, and gets inherited automatically once 
+    FourierFilterFlux is loaded.
+
+    This file's real job is triggering activation: `[extensions] CUDAExt =
+    ["CUDA", "cuDNN", "cuFFT"]` in Project.toml means all three must be
+    loaded before this extension (and therefore this package's CUDA support)
+    activates at all, so this `using` line is what actually makes that
+    happen, not a formality. It also guarantees FourierFilterFlux's and
+    ContinuousWavelets' own CUDA extensions are active by the time anything
+    in ScatteringTransform runs on the GPU, since their trigger sets
+    (CUDA+cuFFT, and CUDA+cuDNN+cuFFT, respectively) are subsets of this one. =#
+using ScatteringTransform, CUDA, cuDNN, cuFFT
+
+end

ext/MetalExt.jl

@@ -0,0 +1,5 @@
+module MetalExt
+
+using ScatteringTransform, Metal
+
+end

src/ScatteringTransform.jl

@@ -2,24 +2,23 @@ module ScatteringTransform
 
 using Core: @__doc__
 using ChainRules
-using CUDA
 using Wavelets, ContinuousWavelets
 using Zygote, Flux, LinearAlgebra, AbstractFFTs
-using Flux
 using FourierFilterFlux
 using MonogenicFilterFlux
 using Adapt
+using Functors
 using RecipesBase
 using Base: tail
 using ChainRulesCore
-# using Plots # who cares about weight really?
 using Statistics
 using Dates
 
 import Adapt: adapt
 import ChainRules: rrule
 import Zygote: has_chain_rrule, rrule
 import Wavelets: eltypes
+import Flux: gpu, cpu
 
 
 include("shared.jl")
@@ -31,7 +30,7 @@ export Scattered, ScatteredFull, ScatteredOut, nonZeroPaths, cat, addNextPath
 include("pool.jl")
 export RationPool, nPoolDims, outputsize, poolSize
 include("transform.jl")
-export cu
+export gpu, cpu
 
 @doc """
      scatteringTransform(inputSize, m=2, backend::UnionAll=stFlux; kwargs...)
@@ -63,4 +62,4 @@ for f in [:plotOriginalSignal1D, :plotZerothLayer1D, :plotFirstLayer1DSingleWave
     @eval export $f
 end
 
-end # end Module
+end # end Module