diff --git a/.github/agents/julia.agent.md b/.github/agents/julia.agent.md new file mode 100644 index 0000000..b06708f --- /dev/null +++ b/.github/agents/julia.agent.md @@ -0,0 +1,116 @@ +--- +description: "Use when improving Julia code quality with very long test suites, slow CI, TestItemRunner tagging/filtering, iterative fix-and-rerun loops, or flaky tests. Keywords: Julia, TestItemRunner, @testitem, tags, filter, long-running Julia process, code quality, assertions, source fixes." +name: "Julia Long-Test Quality" +tools: [read, search, edit, execute, todo] +user-invocable: true +--- +You are a specialist for improving Julia code quality in repositories with long-running test suites. + +## Mission +- Make tests reliable and informative without weakening test intent. +- Use TestItemRunner capabilities to speed iteration and triage by tags and filters. +- Iterate until the targeted test scope passes, then validate broader scopes. + +## Hard Constraints +- Never remove assertions to make tests pass. +- If a failure reflects a real implementation bug, fix source code instead of loosening tests. +- Preserve operator names in tags exactly as implemented (CamelCase, no renamed variants). +- Keep changes minimal and localized; avoid unrelated refactors. + +## Repository-Specific Engineering Rules +- Respect package structure and boundaries: + - `src/linearoperators/` for concrete linear operators. + - `src/nonlinearoperators/` for nonlinear operators. + - `src/calculus/` for operator calculus/composition. + - `src/batching/` for batch operators. +- For new or changed operators, ensure implementation completeness: + - Struct with concrete, inference-friendly field types. + - Constructors for dimension tuple and/or data-driven construction. + - Forward path `mul!(y, op, x)` and adjoint path dispatch via `AdjointOperator`. + - Trait and property behavior remains consistent (`is_linear`, `is_diagonal`, rank/invertibility traits). + - Storage traits stay valid (`domain_array_type`, `codomain_array_type`) for CPU/GPU paths. +- Prefer `copy_operator(op; array_type=nothing, threaded=nothing)` behavior when changing copy semantics: + - Deep-copy mutable working buffers only. + - Share immutable and read-only references. +- Keep test files standalone-capable and aligned with TestItems setup modules. +- Preserve quality gates: JET, Aqua, and doctests should remain passing together. +- Use Runic formatting checks when editing Julia source or tests. + +## Julia Performance Playbook +- Put performance-critical code in functions, not top-level scope. +- Avoid untyped globals in hot paths; use function arguments and `const` globals where appropriate. +- Prefer concrete field/container types; avoid abstract fields like `Function`, `AbstractArray`, or `Integer` in performance-sensitive structs. +- Maintain type stability: + - Avoid variable type changes within loops. + - Use `zero(x)`, `oneunit(T)`, and stable return types. + - Use function barriers for setup-vs-kernel separation. +- Measure, do not guess: + - Use `BenchmarkTools` for benchmarks. + - Track allocations (`@time`, `@allocated`) and treat unexpected allocations as defects to investigate. + - Use `@code_warntype` and JET to diagnose inference issues. +- Minimize allocations in inner loops: + - Preallocate outputs and favor `mul!`/in-place APIs. + - Use broadcast fusion (`@.` / dotted ops) when beneficial. + - Unfuse broadcasts when repeated subexpressions are recomputed unnecessarily. + - Use `@views` for slicing when copy cost dominates. +- Iterate arrays in memory-friendly order (column-major access patterns). +- For threaded Julia code that also calls BLAS, avoid oversubscription (often `OPENBLAS_NUM_THREADS=1` is best with multithreaded Julia; validate on workload). +- Use performance annotations (`@inbounds`, `@simd`, `@fastmath`) only when correctness assumptions are explicitly validated. + +## Test Architecture Rules +- Prefer `@testitem` with explicit `tags` and optional `setup` modules. +- Use tags that encode both operator and test type. +- Mixed tests may include multiple operator tags when behavior genuinely spans operators. +- Test type tags should come from: `:linearoperator`, `:nonlinearoperator`, `:batching`, `:calculus`, `:jet`, `:quality`, `:misc`. +- Operator tags should use exact CamelCase names, for example: `:MatrixOp`, `:FiniteDiff`, `:Compose`, `:SpreadingBatchOp`. +- Use `@run_package_tests filter=ti->...` to run focused slices. +- For grouped runs, prefer strict type-tag exclusion filters (for example, `ti -> !(:jet in ti.tags)`). + +## JET.jl Requirements +- Treat JET coverage as mandatory for all public API. +- Ensure JET test coverage includes all three modes: + - `JET.test_package(...)` for package-level inference/type diagnostics on exported/public API paths. + - `@test_opt ...` for representative public operations and constructors. + - `@test_call ...` for key public call signatures and runtime-like call paths. +- Do not accept partial JET migration: missing any of the three test modes is incomplete. +- When adding or changing public API, update JET tests in the same change. + +## Fast Iteration Workflow +1. Start one long-running Julia REPL in the package test environment. +2. Load TestItemRunner once. +3. Run filtered test slices repeatedly (by operator/type tags). +4. Fix failures immediately; rerun the same filtered slice until green. +5. Expand to adjacent slices, then run full suite. +6. Capture outputs from each run into `.temp/` files for traceability. + +Recommended REPL pattern: +```julia +using TestItemRunner +run_tests("test"; filter = ti -> (:MatrixOp in ti.tags) && (:linearoperator in ti.tags)) +``` + +Recommended shell pattern for captured logs: +```sh +mkdir -p .temp +julia --project=test test/runtests.jl > .temp/test_runtests.log 2>&1 +julia --project=test test/jet/test_package.jl > .temp/test_jet_package.log 2>&1 +``` + +## Failure Triage +1. Read the exact failing assertion and stacktrace first. +2. Classify failure: + - Test setup/import/tagging issue + - Real source bug + - Environment/performance instability +3. For real bugs, patch source and keep/assert expected behavior in tests. +4. For flaky perf tests, stabilize methodology (workload, sampling, thresholds) without dropping coverage. +5. Re-run the smallest relevant filtered subset before broad reruns. + +## Output Requirements +- Report what was changed and why. +- List files touched. +- Provide exact filtered test commands used. +- State pass/fail counts for the final run. +- Call out remaining risks or follow-up items. +- IMPORTANT! Store all temporary run outputs only under `.temp/` inside the repository (no temp scripts and logs elsewhere). +- When performance work is included, report allocation deltas and the exact benchmark commands used. diff --git a/.github/instructions/julia-operator-engineering.instructions.md b/.github/instructions/julia-operator-engineering.instructions.md index 68cf0c3..498297e 100644 --- a/.github/instructions/julia-operator-engineering.instructions.md +++ b/.github/instructions/julia-operator-engineering.instructions.md @@ -18,7 +18,12 @@ applyTo: "src/**/*.jl" - size/domain/codomain/storage traits, - property traits such as linearity, diagonal structure, and rank-related predicates. - `check` utility function must be called in all effective `mul!` paths to ensure consistent argument validation and error messages. -- Preserve `domain_storage_type` and `codomain_storage_type` semantics and dispatch compatibility. +- Preserve `domain_array_type` and `codomain_array_type` semantics and dispatch compatibility. +- Constructors should expose an `array_type` keyword where storage backend selection is meaningful. +- `domain_array_type`/`codomain_array_type` must remain consistent with constructor-selected storage. +- When storage checks become stricter, fix operator traits and tests instead of relaxing `check`. - Prefer behavior-preserving refactors: extract helpers, separate setup from kernels, reduce method size, but do not weaken checks. - If modifying copy semantics, preserve the package convention that immutable/read-only arrays are shared while mutable working buffers are copied deliberately. - Keep source formatted with Runic-compatible Julia style. +- GPU extensions live under `ext/GpuExt/` (triggered by `GPUArrays`). Override `mul!` there for any operator whose base implementation uses scalar indexing loops (`@nloops`, `@nref`, `@inbounds y[i] = b[j]`); replace with broadcast-over-view (`y .= view(b, idx...)`). +- When overriding a threaded operator (e.g. `Variation{..., true}`) for GPU, delegate to the non-threaded variant (`Variation{..., false}`) — the threading strategy is CPU-only. diff --git a/.github/instructions/julia-performance.instructions.md b/.github/instructions/julia-performance.instructions.md index 0277822..3f0beb4 100644 --- a/.github/instructions/julia-performance.instructions.md +++ b/.github/instructions/julia-performance.instructions.md @@ -24,3 +24,5 @@ applyTo: "src/**/*.jl,benchmark/**/*.jl" - benchmark representative workloads, - inspect allocations, - use JET and `@code_warntype` for inference issues. +- For benchmark harnesses, derive element types robustly when operator type traits may return wrapped array types. +- Keep benchmark setup deterministic (`Random.seed!(0)`) and validate key benchmark states with one smoke `mul!` path before full runs. diff --git a/.github/instructions/julia-testing-and-jet.instructions.md b/.github/instructions/julia-testing-and-jet.instructions.md index 19a8dbf..f90f31f 100644 --- a/.github/instructions/julia-testing-and-jet.instructions.md +++ b/.github/instructions/julia-testing-and-jet.instructions.md @@ -19,3 +19,13 @@ applyTo: "test/**/*.jl,docs/**/*.md" - Keep Aqua and doctests passing alongside functional tests. - Never remove assertions to force green tests. - All temporary test and benchmark outputs must go under `.temp/` only. +- If GPU tests are backend-specific, keep them in separate `@testitem`s and use `:gpu` tag. +- When `VERB` is enabled, print each running testitem name at test-runner filter time. +- For local coverage, mirror CI with `julia --project=test --code-coverage=user test/runtests.jl`, then process `*.cov` / `*.info` artifacts into `lcov.info` if needed. +- Subpackages (DSPOperators, FFTWOperators, NFFTOperators, WaveletOperators) have no standalone `test/` directory; they are tested and their coverage is gathered exclusively through the parent package's `test/` project. Do not attempt a separate subpackage coverage run. +- Extension coverage should be gathered through the parent-package tests that load the relevant trigger packages; do not assume a separate extension-only coverage run exists. +- JET `@test_opt` flags `array_type::Type` (unparameterized keyword) as a source of runtime dispatch. Use `array_type::Type{<:AbstractArray}` and avoid kwarg-to-kwarg forwarding; use a typed positional-arg helper (e.g., `_make_eye(T, dims, S)`) so JET can resolve dispatch statically. +- When Aqua reports "Unexpected Pass" on a `@test_broken`/`broken=true` check, the underlying issue is now fixed — remove the workaround and use `Aqua.test_all(pkg)` unconditionally. +- Agent sub-tasks frequently generate `Eye(T, dims, array_type)` (3 positional args) instead of `Eye(T, dims; array_type=...)` (keyword). Always verify agent output for this pattern. +- Stochastic test assertions `op * randn(n) ≈ other_op * (op * randn(n))` are wrong when the two `randn` calls produce different vectors; always capture into a variable first. +- When testing GPU storage-type propagation, add `@test domain_array_type(op) <: CUDA.CuArray` / `<: AMDGPU.ROCArray` assertions directly in the per-operator CUDA/AMDGPU `@testitem`. diff --git a/.github/skills/julia-gpu-implementation/SKILL.md b/.github/skills/julia-gpu-implementation/SKILL.md new file mode 100644 index 0000000..cbaff6b --- /dev/null +++ b/.github/skills/julia-gpu-implementation/SKILL.md @@ -0,0 +1,50 @@ +--- +name: julia-gpu-implementation +description: 'Use for GPU operator implementations, GPU extension fixes, backend-specific testitems, and GPU benchmark validation in AbstractOperators.jl.' +argument-hint: 'Describe the operator, GPU backend, or benchmark you want to implement or validate' +user-invocable: true +--- + +# Julia GPU Implementation + +## When To Use + +- Implementing or fixing GPU overrides under `ext/GpuExt/`. +- Adding or updating CUDA/AMDGPU testitems. +- Debugging backend-specific dispatch, storage traits, or array conversion issues. +- Extending benchmark coverage for GPU behavior. +- Checking whether a CPU operator should get a GPU path or stay CPU-only. + +## Implementation Rules + +- Julia package extensions can only `import` the parent package, trigger package(s), and stdlib; if extension code needs a parent dependency API, expose it from the parent module first. +- For FFT plans, prefer `inv(plan)` (AbstractFFTs-generic) over backend-specific `FFTW.plan_inv(...)` to keep CUDA/AMDGPU compatibility. +- With JLArrays/GPUArrays, avoid `copyto!(gpu, cpu_view)` where the source is a `SubArray`; materialize first, for example with `src[1:n]`, or copy from a plain array. +- Preserve backend storage semantics and trait dispatch when adding GPU methods. +- Keep CPU-only implementation details out of GPU overrides unless the backend truly supports them. +- For GPU `GetIndex` overrides, keep boolean-mask and integer-vector fancy indexing in CPU paths unless the backend support is verified. +- When overriding a threaded operator for GPU, delegate to the non-threaded variant; threading strategy is CPU-only. +- Prefer direct `CuArray(arr)` / `CUDA.zeros(...)` / `AMDGPU.ROCArray(arr)` / `AMDGPU.zeros(...)` calls over intermediate conversion variables. +- Benchmark setup code should normalize wrapped domain and codomain type traits to scalar element types before calling `randn` or `zeros`. + +## Testing Rules + +- For honest GPU coverage, keep JLArray checks separate from real device checks and add backend-specific tags such as `:cuda` and `:amdgpu` plus runtime skip guards. +- In `test/runtests.jl`, filter backend-tagged testitems when the runtime is unavailable, but keep per-test safety checks too. +- Add explicit tests for `domain_array_type` and `codomain_array_type`, and verify that `op * x` allocates on the active backend. +- When adding CUDA/AMDGPU companion tests, prefer direct backend array construction instead of temporary conversion variables. +- For GPU `GetIndex` tests, restrict indices to ranges, colons, and scalar integers; bool-mask and integer-vector `view` forms are not universally supported across GPU backends. +- Migrate GPU-backend storage-type assertions from central quality files into each operator's own CUDA/AMDGPU `@testitem` so they run with the functional tests. +- Use direct `import CUDA` / `import AMDGPU` plus `functional()` guards in testitems; avoid try/catch gating. + +## Benchmarking Rules + +- Benchmark scripts under `benchmark/` must prefer local workspace package paths over registry-installed copies, otherwise GPU fixes in sibling packages can be silently skipped. +- Use representative large inputs for GPU crossover studies and keep the measurement setup deterministic. +- Capture benchmark logs and generated reports under `.temp/`. + +## Tooling Reminders + +- Agent sub-tasks frequently generate `Eye(T, dims, array_type)` with three positional arguments instead of `Eye(T, dims; array_type=...)` with a keyword; verify this pattern. +- JET `@test_opt` catches runtime dispatch from `array_type::Type` when it is unparameterized; use `array_type::Type{<:AbstractArray}` and avoid kwarg-to-kwarg forwarding by routing through an internal helper. +- When fixing an "unexpected pass" Aqua error, remove the workaround and use `Aqua.test_all(pkg)` once the underlying issue is fixed. diff --git a/.github/skills/julia-long-test-workflow/SKILL.md b/.github/skills/julia-long-test-workflow/SKILL.md index c61a788..5a185f3 100644 --- a/.github/skills/julia-long-test-workflow/SKILL.md +++ b/.github/skills/julia-long-test-workflow/SKILL.md @@ -29,6 +29,24 @@ user-invocable: true ## Common Commands +Main package coverage: + +```sh +julia --project=test --code-coverage=user test/runtests.jl +``` + +Subpackage coverage (DSPOperators, FFTWOperators, NFFTOperators, WaveletOperators have **no** standalone `test/` directory): + +> All subpackage code and their GPU extensions are exercised by the parent package's +> `test/` project. Run the same coverage command above; the `.cov` files under each +> subpackage's `src/` will be populated automatically. + +Process coverage after a local run: + +```sh +julia -e 'using Coverage; Coverage.LCOV.writefile("lcov.info", Coverage.process_folder())' +``` + Filtered test run: ```julia @@ -37,9 +55,13 @@ TestItemRunner.run_tests(pwd(); filter = ti -> :MatrixOp in ti.tags) # example o TestItemRunner.run_tests(pwd(); filter = ti -> ti.name == "DCT") # example of filtering by test name instead of tags ``` -AirSpeedVelocity comparison: +### Local benchmark comparison with AirspeedVelocity + +AirspeedVelocity works well for local branch-vs-master comparisons and is the +recommended tool for interactive performance investigation: ```sh +mkdir -p .temp/asv benchpkg \ --path . \ --rev master,dirty \ @@ -48,6 +70,20 @@ benchpkg \ --exeflags="--threads=4" ``` +Filtered AirSpeedVelocity comparison for a single benchmark family: + +```sh +mkdir -p .temp/asv +benchpkg \ + --path . \ + --rev master,dirty \ + --script benchmark/benchmarks.jl \ + --output-dir .temp/asv \ + --exeflags="--threads=4" \ + --add RecursiveArrayTools \ + --filter MIMOFilt +``` + Render a comparison table: ```sh @@ -59,6 +95,44 @@ benchpkgtable \ --mode time,memory ``` +> **Note:** Use AirspeedVelocity with an explicit `--script` path when comparing +> against revisions that do not yet contain the benchmark file. + +### CI benchmark comparison (GitHub Actions) + +The GitHub Actions benchmark CI does **not** use the AirspeedVelocity action +because the root-level Julia workspace (`[workspace]` in `Project.toml`) causes +that action's revision-management to mis-resolve the monorepo subprojects. +Instead, two workflows implement a fork-safe two-stage approach: + +- **`benchmark.yml`** – unprivileged `pull_request` job that checks out both + the base and head revisions, runs `benchmark/compare.jl` against explicit + worktree paths, and uploads `body.md`, `pr_number.txt`, and + `julia_version.txt` as an artifact. +- **`post_benchmark_comment.yml`** – privileged `workflow_run` job that + downloads the artifact and creates or updates the PR comment. + +The comparison table mirrors AirspeedVelocity output with separate Time and +Memory sections, base/head columns, a ratio column, and emoji indicators: +- 🚀 significant speedup: `ratio − ratio_err > 1.2` (time) or `ratio < 0.5` (memory) +- 🐢 significant slowdown: `ratio + ratio_err < 0.8` (time) or `ratio > 1.5` (memory) + +To run the comparison locally with the same script used by CI: + +```sh +# Check out base separately, e.g. in a worktree: +git worktree add .temp/base master + +julia --project=benchmark benchmark/compare.jl \ + --base-dir .temp/base \ + --head-dir . \ + --output-dir .temp/bench-compare \ + --pr 0 \ + --julia-version "$(julia -e 'print(VERSION)')" + +cat .temp/bench-compare/body.md +``` + ## Done Criteria - Targeted tests pass. diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 69cdb1c..a2d8daa 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -36,6 +36,7 @@ jobs: run: | julia --project=test -e ' using Pkg + Pkg.rm(["DSPOperators", "FFTWOperators", "NFFTOperators", "WaveletOperators"]) Pkg.develop(path = pwd()) for pkg in ("DSPOperators", "FFTWOperators", "NFFTOperators", "WaveletOperators") Pkg.develop(path = joinpath(pwd(), pkg)) diff --git a/.gitignore b/.gitignore index c33897e..bcc7eb8 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,5 @@ docs/Manifest.toml Manifest.toml Manifest-v*.toml .temp/ +test/gpu_env/ +benchmark/gpu_env/ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..8b50d61 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,25 @@ +# AGENTS.md + +This repository uses layered guidance. Follow it in this order: + +1. Read this file first. +2. Read any applicable files under `.github/instructions/` whose `applyTo` pattern matches the files you will edit. +3. Read the matching skill under `.github/skills/` when the task clearly matches a skill's scope. +4. Then inspect the target source files before editing. + +## How to choose guidance + +- Use `julia-operator-engineering.instructions.md` for changes under `src/**/*.jl`. +- Use `julia-performance.instructions.md` for code under `src/**/*.jl` and `benchmark/**/*.jl` when performance is relevant. +- Use `julia-testing-and-jet.instructions.md` for `test/**/*.jl` and docs-backed test guidance. +- Use `.github/skills/julia-long-test-workflow/SKILL.md` for long Julia test runs, filtered `TestItemRunner` work, JET triage, and AirspeedVelocity comparisons. +- Use `.github/skills/julia-gpu-implementation/SKILL.md` for GPU operator implementations, GPU extensions, GPU-specific tests, and GPU benchmark validation. + +## Working rules + +- Prefer the smallest skill and instruction set that fully covers the task. +- Do not ignore a matching instruction file because a skill also exists; use both when they apply. +- If multiple instruction files match, combine them rather than choosing only one. +- If a task touches both implementation and tests, read both the source and test instruction files before editing. +- Keep temporary artifacts under `.temp/`. +- When in doubt, inspect the relevant files before making changes. diff --git a/DSPOperators/Project.toml b/DSPOperators/Project.toml index 003c38b..614f0d0 100644 --- a/DSPOperators/Project.toml +++ b/DSPOperators/Project.toml @@ -3,14 +3,21 @@ uuid = "d5a72628-6e2f-430e-82f5-561df0bb8116" version = "0.1.0" [deps] +AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" AbstractOperators = "d9c5613a-d543-52d8-9afd-8f241a8c3f1c" -DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +[weakdeps] +GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" + +[extensions] +GpuExt = "GPUArrays" + [compat] +AbstractFFTs = "1.5.0" AbstractOperators = "0.4" -DSP = "0.7, 0.8" FFTW = "1.9.0" +GPUArrays = "10, 11" LinearAlgebra = "1.10.0" julia = "1.10" diff --git a/DSPOperators/README.md b/DSPOperators/README.md index af97e5b..37b5b35 100644 --- a/DSPOperators/README.md +++ b/DSPOperators/README.md @@ -19,6 +19,10 @@ DSPOperators.jl is a **subpackage** of the AbstractOperators.jl ecosystem. While pkg> add DSPOperators ``` +## GPU Support + +DSPOperators.jl follows the generic GPU-array support provided by AbstractOperators.jl. In practice, its operators are intended to work with CUDA.jl, AMDGPU.jl, oneAPI.jl, and OpenCL.jl when the underlying array operations are GPU-compatible. + ## Usage Example ```julia diff --git a/DSPOperators/ext/GpuExt/GpuExt.jl b/DSPOperators/ext/GpuExt/GpuExt.jl new file mode 100644 index 0000000..f6daeee --- /dev/null +++ b/DSPOperators/ext/GpuExt/GpuExt.jl @@ -0,0 +1,214 @@ +module GpuExt + +using GPUArrays +using DSPOperators +import LinearAlgebra: mul! +import DSPOperators: Filt, MIMOFilt, AbstractFilt, AbstractMIMOFilt +import AbstractFFTs: plan_rfft, plan_irfft +import AbstractOperators: AdjointOperator, check + +_wrapper_type(x::AbstractArray) = Base.typename(typeof(x)).wrapper + +struct GpuFilt{T, N, S <: AbstractArray{T}, B <: AbstractGPUArray{T}, C <: AbstractGPUArray, P1, P2} <: + AbstractFilt{T, N, S} + dim_in::NTuple{N, Int} + b::Vector{T} + a::Vector{T} + si::Vector{T} + h_fft::C + h_fft_conj::C + buf::B + buf_fft::C + buf_out::B + plan_fwd::P1 + plan_inv::P2 +end + +struct GpuMIMOFilt{T, S <: AbstractArray{T}, F <: GpuFilt} <: AbstractMIMOFilt{T, S} + dim_out::Tuple{Int, Int} + dim_in::Tuple{Int, Int} + filters::Vector{F} +end + +function GpuMIMOFilt(dim_out, dim_in, filters::Vector{F}) where {F <: GpuFilt{T, N, S}} where {T, N, S} + return GpuMIMOFilt{T, S, F}(dim_out, dim_in, filters) +end + +function _gpu_complex_buffer(ref::AbstractGPUArray{T}, fftlen::Int) where {T} + return similar(ref, Complex{T}, fftlen ÷ 2 + 1) +end + +function _make_gpu_filt(ref::AbstractGPUArray{T}, cpu_op::AbstractFilt{T, N}) where {T <: Real, N} + fftlen = nextpow(2, cpu_op.dim_in[1] + length(cpu_op.b) - 1) + array_type = _wrapper_type(ref){T} + buf = similar(ref, T, fftlen) + buf_fft = _gpu_complex_buffer(ref, fftlen) + buf_out = similar(ref, T, fftlen) + plan_fwd = plan_rfft(buf) + plan_inv = plan_irfft(buf_fft, fftlen) + + h_pad_cpu = zeros(T, fftlen) + h_pad_cpu[1:length(cpu_op.b)] .= cpu_op.b + h_pad = similar(ref, T, fftlen) + copyto!(h_pad, h_pad_cpu) + + h_fft = similar(buf_fft) + mul!(h_fft, plan_fwd, h_pad) + h_fft_conj = similar(h_fft) + h_fft_conj .= conj.(h_fft) + + return GpuFilt{T, N, array_type, typeof(buf), typeof(buf_fft), typeof(plan_fwd), typeof(plan_inv)}( + cpu_op.dim_in, + collect(cpu_op.b), + collect(cpu_op.a), + collect(cpu_op.si), + h_fft, + h_fft_conj, + buf, + buf_fft, + buf_out, + plan_fwd, + plan_inv, + ) +end + +function Filt(x::AbstractGPUArray{T}, b::AbstractVector{T}, a::AbstractVector{T}) where {T <: Real} + cpu_op = Filt(T, size(x), copy(b), copy(a)) + return _make_gpu_filt(x, cpu_op) +end + +Filt(x::AbstractGPUArray{T}, b::AbstractVector{T}) where {T <: Real} = Filt(x, b, T[one(T)]) + +function _load_signal!(buf::AbstractGPUArray{T}, x::AbstractGPUArray{T}) where {T} + fill!(buf, zero(T)) + copyto!(view(buf, 1:length(x)), x) + return buf +end + +function _copy_head!(y::AbstractGPUArray, buf_out::AbstractGPUArray, n::Int) + copyto!(y, view(buf_out, 1:n)) + return y +end + +function _mul_fir_forward!(y::AbstractGPUArray{T}, A::GpuFilt{T}, b::AbstractGPUArray{T}) where {T} + _load_signal!(A.buf, b) + mul!(A.buf_fft, A.plan_fwd, A.buf) + A.buf_fft .*= A.h_fft + mul!(A.buf_out, A.plan_inv, A.buf_fft) + return _copy_head!(y, A.buf_out, length(y)) +end + +function _mul_fir_adjoint!(y::AbstractGPUArray{T}, A::GpuFilt{T}, b::AbstractGPUArray{T}) where {T} + _load_signal!(A.buf, b) + mul!(A.buf_fft, A.plan_fwd, A.buf) + A.buf_fft .*= A.h_fft_conj + mul!(A.buf_out, A.plan_inv, A.buf_fft) + return _copy_head!(y, A.buf_out, length(y)) +end + +function _add_fir_forward!(y::AbstractGPUArray{T}, A::GpuFilt{T}, b::AbstractGPUArray{T}) where {T} + _load_signal!(A.buf, b) + mul!(A.buf_fft, A.plan_fwd, A.buf) + A.buf_fft .*= A.h_fft + mul!(A.buf_out, A.plan_inv, A.buf_fft) + y .+= view(A.buf_out, 1:length(y)) + return y +end + +function _add_fir_adjoint!(y::AbstractGPUArray{T}, A::GpuFilt{T}, b::AbstractGPUArray{T}) where {T} + _load_signal!(A.buf, b) + mul!(A.buf_fft, A.plan_fwd, A.buf) + A.buf_fft .*= A.h_fft_conj + mul!(A.buf_out, A.plan_inv, A.buf_fft) + y .+= view(A.buf_out, 1:length(y)) + return y +end + +# ─── Filt GPU mul! ──────────────────────────────────────────────────────────── + +function mul!(y::AbstractGPUArray{T}, A::GpuFilt{T}, b::AbstractGPUArray{T}) where {T <: Real} + check(y, A, b) + length(A.a) == 1 || + throw(ArgumentError("IIR Filt is not supported on GPU arrays. Use CpuOperatorWrapper to run on CPU.")) + for col in 1:size(b, 2) + _mul_fir_forward!(view(y, :, col), A, view(b, :, col)) + end + return y +end + +function mul!( + y::AbstractGPUArray{T}, L::AdjointOperator{<:GpuFilt{T}}, b::AbstractGPUArray{T} + ) where {T <: Real} + check(y, L, b) + A = L.A + length(A.a) == 1 || + throw(ArgumentError("IIR Filt adjoint is not supported on GPU arrays. Use CpuOperatorWrapper.")) + for col in 1:size(b, 2) + _mul_fir_adjoint!(view(y, :, col), A, view(b, :, col)) + end + return y +end + +# ─── MIMOFilt GPU mul! ──────────────────────────────────────────────────────── + +function _make_gpu_mimofilt(ref::AbstractGPUArray{T}, cpu_op::AbstractMIMOFilt{T}) where {T <: Real} + filters = [ + _make_gpu_filt(ref, Filt(T, (cpu_op.dim_in[1],), copy(cpu_op.B[i]), copy(cpu_op.A[i]))) for + i in eachindex(cpu_op.B) + ] + return GpuMIMOFilt(cpu_op.dim_out, cpu_op.dim_in, filters) +end + +function MIMOFilt( + x::AbstractGPUArray{T}, b::Vector{<:AbstractVector{T}}, a::Vector{<:AbstractVector{T}} + ) where {T <: Real} + cpu_op = MIMOFilt(T, size(x), copy(b), copy(a)) + return _make_gpu_mimofilt(x, cpu_op) +end + +function MIMOFilt(x::AbstractGPUArray{T}, b::Vector{<:AbstractVector{T}}) where {T <: Real} + return MIMOFilt(x, b, [T[one(T)] for _ in eachindex(b)]) +end + +function mul!(y::AbstractGPUArray{T}, L::GpuMIMOFilt{T}, b::AbstractGPUArray{T}) where {T <: Real} + check(y, L, b) + for filter in L.filters + length(filter.a) == 1 || throw( + ArgumentError("IIR MIMOFilt is not supported on GPU arrays. Use CpuOperatorWrapper."), + ) + end + fill!(y, zero(T)) + cnt = 0 + for cy in 1:L.dim_out[2] + for cx in 1:L.dim_in[2] + cnt += 1 + _add_fir_forward!(view(y, :, cy), L.filters[cnt], view(b, :, cx)) + end + end + return y +end + +function mul!( + y::AbstractGPUArray{T}, M::AdjointOperator{<:GpuMIMOFilt{T}}, b::AbstractGPUArray{T} + ) where {T <: Real} + check(y, M, b) + L = M.A + for filter in L.filters + length(filter.a) == 1 || throw( + ArgumentError( + "IIR MIMOFilt adjoint is not supported on GPU arrays. Use CpuOperatorWrapper." + ), + ) + end + fill!(y, zero(T)) + cnt = 0 + for cy in 1:L.dim_out[2] + for cx in 1:L.dim_in[2] + cnt += 1 + _add_fir_adjoint!(view(y, :, cx), L.filters[cnt], view(b, :, cy)) + end + end + return y +end + +end # module GpuExt diff --git a/DSPOperators/src/Conv.jl b/DSPOperators/src/Conv.jl index c70eade..2c27f43 100644 --- a/DSPOperators/src/Conv.jl +++ b/DSPOperators/src/Conv.jl @@ -45,7 +45,7 @@ function Conv(domain_type::Type, dim_in::NTuple{N, Int}, h::H) where {N, H <: Ab R = plan_fft(buf) buf_c1 = similar(buf) complex_type = domain_type - I = FFTW.plan_inv(R) + I = inv(R) end buf_c2 = similar(buf_c1) return Conv{domain_type, N, H, typeof(buf_c1), typeof(R), typeof(I)}(dim_in, h, buf, buf_c1, buf_c2, R, I) @@ -60,12 +60,12 @@ Conv(x::H, h::H) where {H <: AbstractArray} = Conv(eltype(x), size(x), h) function mul!( y::AbstractArray{T, N}, A::AbstractConv{T, N}, b::AbstractArray{T, N} ) where {T, N} - #y .= conv(A.h,b) #naive implementation + check(y, A, b) fill!(A.buf, zero(T)) - A.buf[CartesianIndices(A.h)] .= A.h + view(A.buf, axes(A.h)...) .= A.h mul!(A.buf_c1, A.R, A.buf) fill!(A.buf, zero(T)) - A.buf[CartesianIndices(b)] .= b + view(A.buf, axes(b)...) .= b mul!(A.buf_c2, A.R, A.buf) A.buf_c2 .*= A.buf_c1 return mul!(y, A.I, A.buf_c2) @@ -74,24 +74,25 @@ end function mul!( y::AbstractArray{T, N}, L::AdjointOperator{C}, b::AbstractArray{T, N} ) where {T, N, C <: AbstractConv{T, N}} - #y .= xcorr(b,L.A.h)[size(L.A.h,1)[1]:end-length(L.A.h)+1] #naive implementation + check(y, L, b) fill!(L.A.buf, zero(T)) - L.A.buf[CartesianIndices(L.A.h)] .= L.A.h + view(L.A.buf, axes(L.A.h)...) .= L.A.h mul!(L.A.buf_c1, L.A.R, L.A.buf) fill!(L.A.buf, zero(T)) - L.A.buf[CartesianIndices(b)] .= b + view(L.A.buf, axes(b)...) .= b mul!(L.A.buf_c2, L.A.R, L.A.buf) L.A.buf_c2 .*= conj.(L.A.buf_c1) mul!(L.A.buf, L.A.I, L.A.buf_c2) - return y .= L.A.buf[CartesianIndices(y)] + y .= view(L.A.buf, axes(y)...) + return y end # Properties domain_type(::AbstractConv{T}) where {T} = T codomain_type(::AbstractConv{T}) where {T} = T -domain_storage_type(::AbstractConv{T, N, H}) where {T, N, H} = H -codomain_storage_type(::AbstractConv{T, N, H}) where {T, N, H} = H +domain_array_type(::AbstractConv{T, N, H}) where {T, N, H} = H +codomain_array_type(::AbstractConv{T, N, H}) where {T, N, H} = H is_thread_safe(::Conv) = false #TODO find out a way to verify this, diff --git a/DSPOperators/src/DSPOperators.jl b/DSPOperators/src/DSPOperators.jl index d0e6de3..cc129ea 100644 --- a/DSPOperators/src/DSPOperators.jl +++ b/DSPOperators/src/DSPOperators.jl @@ -1,9 +1,9 @@ module DSPOperators using AbstractOperators, FFTW +using AbstractFFTs: AbstractFFTs import LinearAlgebra: mul! import Base: size, ndims -using DSP: xcorr, conv import AbstractOperators: domain_type, @@ -13,13 +13,12 @@ import AbstractOperators: get_normal_op, allocate_in_domain, allocate_in_codomain, - domain_storage_type, - codomain_storage_type, + domain_array_type, + codomain_array_type, is_full_column_rank, is_full_row_rank, is_thread_safe - include("Conv.jl") include("Filt.jl") include("MIMOFilt.jl") diff --git a/DSPOperators/src/Filt.jl b/DSPOperators/src/Filt.jl index 27fc850..c17e971 100644 --- a/DSPOperators/src/Filt.jl +++ b/DSPOperators/src/Filt.jl @@ -1,5 +1,7 @@ export Filt +abstract type AbstractFilt{T, N, S <: AbstractArray} <: LinearOperator end + """ Filt([domain_type=Float64::Type,] dim_in::Tuple, b::AbstractVector, [a::AbstractVector,]) Filt(x::AbstractVector, b::AbstractVector, [a::AbstractVector,]) @@ -14,7 +16,7 @@ IIR ℝ^10 -> ℝ^10 ``` """ -struct Filt{T, N, B <: AbstractVector{T}} <: LinearOperator +struct Filt{T, N, B <: AbstractVector{T}} <: AbstractFilt{T, N, Array{T}} dim_in::NTuple{N, Int} b::B a::B @@ -84,7 +86,7 @@ function mul!(y::AbstractArray, L::Filt, x::AbstractArray) fir!(y, L.b, x, L.si, col, col) end end - return + return y end function mul!(y::AbstractArray, L::AdjointOperator{<:Filt}, x::AbstractArray) @@ -97,7 +99,7 @@ function mul!(y::AbstractArray, L::AdjointOperator{<:Filt}, x::AbstractArray) fir_rev!(y, A.b, x, A.si, col, col) end end - return + return y end function iir!(y, b, a, x, si, coly, colx) @@ -112,7 +114,7 @@ function iir!(y, b, a, x, si, coly, colx) y[i, coly] = val end si .= 0.0 #reset state - return + return nothing end # Utilities @@ -129,7 +131,7 @@ function iir_rev!(y, b, a, x, si, coly, colx) y[i, coly] = val end si .= 0.0 - return + return nothing end function fir!(y, b, x, si, coly, colx) @@ -144,7 +146,7 @@ function fir!(y, b, x, si, coly, colx) y[i, coly] = val end si .= 0.0 - return + return nothing end function fir_rev!(y, b, x, si, coly, colx) @@ -159,20 +161,22 @@ function fir_rev!(y, b, x, si, coly, colx) y[i, coly] = val end si .= 0.0 - return + return nothing end # Properties -domain_type(::Filt{T}) where {T} = T -codomain_type(::Filt{T}) where {T} = T -is_thread_safe(::Filt) = false +domain_type(::AbstractFilt{T}) where {T} = T +codomain_type(::AbstractFilt{T}) where {T} = T +domain_array_type(::AbstractFilt{T, N, S}) where {T, N, S} = S +codomain_array_type(::AbstractFilt{T, N, S}) where {T, N, S} = S +is_thread_safe(::AbstractFilt) = false -size(L::Filt) = L.dim_in, L.dim_in +size(L::AbstractFilt) = L.dim_in, L.dim_in -fun_name(L::Filt) = size(L.a, 1) != 1 ? "IIR" : "FIR" +fun_name(L::AbstractFilt) = size(L.a, 1) != 1 ? "IIR" : "FIR" #TODO find out a way to verify this, # probably for IIR it means zeros inside unit circle -is_full_row_rank(L::Filt) = true -is_full_column_rank(L::Filt) = true +is_full_row_rank(::AbstractFilt) = true +is_full_column_rank(::AbstractFilt) = true diff --git a/DSPOperators/src/MIMOFilt.jl b/DSPOperators/src/MIMOFilt.jl index ec74fd5..28f9ca2 100644 --- a/DSPOperators/src/MIMOFilt.jl +++ b/DSPOperators/src/MIMOFilt.jl @@ -1,5 +1,7 @@ export MIMOFilt +abstract type AbstractMIMOFilt{T, S <: AbstractArray} <: LinearOperator end + """ MIMOFilt([domain_type=Float64::Type,] dim_in::Tuple, B::Vector{AbstractVector}, [A::Vector{AbstractVector},]) @@ -49,10 +51,10 @@ true julia> Y[:,2] ≈ filt(B[4],A[4],X[:,1])+filt(B[5],A[5],X[:,2])+filt(B[6],A[6],X[:,3]) true - + ``` """ -struct MIMOFilt{T, A <: AbstractVector{T}} <: LinearOperator +struct MIMOFilt{T, A <: AbstractVector{T}} <: AbstractMIMOFilt{T, Array{T}} dim_out::Tuple{Int, Int} dim_in::Tuple{Int, Int} B::Vector{A} @@ -148,7 +150,7 @@ function mul!(y::AbstractArray, L::MIMOFilt, x::AbstractArray) end cx = 0 end - return + return y end function mul!(y::AbstractArray, M::AdjointOperator{<:MIMOFilt}, x::AbstractArray) @@ -177,23 +179,25 @@ function mul!(y::AbstractArray, M::AdjointOperator{<:MIMOFilt}, x::AbstractArray end cx = 0 end - return + return y end # Properties -domain_type(::MIMOFilt{T}) where {T} = T -codomain_type(::MIMOFilt{T}) where {T} = T -is_thread_safe(::MIMOFilt) = false +domain_type(::AbstractMIMOFilt{T}) where {T} = T +codomain_type(::AbstractMIMOFilt{T}) where {T} = T +domain_array_type(::AbstractMIMOFilt{T, S}) where {T, S} = S +codomain_array_type(::AbstractMIMOFilt{T, S}) where {T, S} = S +is_thread_safe(::AbstractMIMOFilt) = false -size(L::MIMOFilt) = L.dim_out, L.dim_in +size(L::AbstractMIMOFilt) = L.dim_out, L.dim_in #TODO find out a way to verify this, # probably for IIR it means zeros inside unit circle -is_full_row_rank(L::MIMOFilt) = true -is_full_column_rank(L::MIMOFilt) = true +is_full_row_rank(::AbstractMIMOFilt) = true +is_full_column_rank(::AbstractMIMOFilt) = true -fun_name(L::MIMOFilt) = "※" +fun_name(::AbstractMIMOFilt) = "※" # Utilities @@ -209,7 +213,7 @@ function add_iir!(y, b, a, x, si, coly, colx) y[i, coly] += val end si .= 0.0 #reset state - return + return nothing end function add_iir_rev!(y, b, a, x, si, coly, colx) @@ -224,7 +228,7 @@ function add_iir_rev!(y, b, a, x, si, coly, colx) y[i, coly] += val end si .= 0.0 - return + return nothing end function add_fir!(y, b, x, si, coly, colx) @@ -239,7 +243,7 @@ function add_fir!(y, b, x, si, coly, colx) y[i, coly] += val end si .= 0.0 - return + return nothing end function add_fir_rev!(y, b, x, si, coly, colx) @@ -254,5 +258,5 @@ function add_fir_rev!(y, b, x, si, coly, colx) y[i, coly] += val end si .= 0.0 - return + return nothing end diff --git a/DSPOperators/src/Xcorr.jl b/DSPOperators/src/Xcorr.jl index 8db2e93..c1e1714 100644 --- a/DSPOperators/src/Xcorr.jl +++ b/DSPOperators/src/Xcorr.jl @@ -1,11 +1,27 @@ export Xcorr -#TODO make more efficient + +# Adjoint FFT state for non-CPU array backends. +# CPU uses a tiled FIR loop instead; adj_fft is Nothing for CPU arrays. +# Hfft and Hc are separate type params because plan_rfft on a CPU-backed GPU +# mock array (e.g. JLArrays) may return a plain Vector from `R * buf`, while +# similar(h, Complex{T}, ...) returns the array's own complex type. +struct XcorrAdjFFT{ + Hfft <: AbstractVector, H <: AbstractVector, Hc <: AbstractVector, + P3 <: AbstractFFTs.Plan, P4 <: AbstractFFTs.Plan, + } + fftlen::Int + h_fft::Hfft # rfft(h padded); type may differ from buf_c + buf::H # scratch buffer size fftlen + buf_c::Hc # complex scratch buffer + R::P3 # rfft/fft plan + I::P4 # irfft/ifft plan +end """ Xcorr([domain_type=Float64::Type,] dim_in::Tuple, h::AbstractVector) Xcorr(x::AbstractVector, h::AbstractVector) -Creates a `LinearOperator` which, when multiplied with an array `x::AbstractVector`, returns the cross correlation between `x` and `h`. Uses `xcorr` from `DSP.jl`. +Creates a `LinearOperator` which, when multiplied with an array `x::AbstractVector`, returns the cross correlation between `x` and `h`. Uses FFT-based implementation. Examples ```jldoctest @@ -15,41 +31,182 @@ julia> Xcorr(Float64, (10,), [1.0, 0.5, 0.2]) ◎ ℝ^10 -> ℝ^19 ``` """ -struct Xcorr{T, H <: AbstractVector{T}} <: LinearOperator +struct Xcorr{ + T, H <: AbstractVector{T}, Hc <: AbstractVector, + P1 <: AbstractFFTs.Plan, P2 <: AbstractFFTs.Plan, Adj, + } <: LinearOperator dim_in::Tuple{Int} h::H + # Forward pass (xcorr) + fftlen_fwd::Int + padlen::Int + h_fft_conj::Hc # conj(rfft(h padded to fftlen_fwd)) + buf_fwd::H # scratch buffer size fftlen_fwd + buf_fwd_c::Hc # complex scratch buffer + R_fwd::P1 # rfft plan, fftlen_fwd + I_fwd::P2 # irfft plan, fftlen_fwd + # Adjoint: XcorrAdjFFT{...} for GPU backends, Nothing for CPU + adj_fft::Adj end +# FFT planning flags: FFTW.MEASURE only for CPU Arrays; no flags for GPU backends. +_xcorr_plan_kwargs(::Type{<:Array}) = (flags = FFTW.MEASURE,) +_xcorr_plan_kwargs(::Type) = (;) + # Constructors function Xcorr(domain_type::Type, DomainDim::NTuple{N, Int}, h::H) where {H <: AbstractVector, N} eltype(h) != domain_type && error("eltype(h) is $(eltype(h)), should be $(domain_type)") N != 1 && error("Xcorr treats only SISO, check Filt and MIMOFilt for MIMO") - return Xcorr{domain_type, H}(DomainDim, h) + + n = DomainDim[1] + m = length(h) + padlen = max(n, m) + outlen = 2 * padlen - 1 + plan_kw = _xcorr_plan_kwargs(H) + + # Forward pass plans + fftlen_fwd = nextpow(2, outlen) + buf_fwd = similar(h, fftlen_fwd) + if domain_type <: Real + R_fwd = plan_rfft(buf_fwd; plan_kw...) + complex_type = Complex{domain_type} + buf_fwd_c = similar(h, complex_type, fftlen_fwd ÷ 2 + 1) + I_fwd = plan_irfft(buf_fwd_c, fftlen_fwd; plan_kw...) + else + R_fwd = plan_fft(buf_fwd; plan_kw...) + buf_fwd_c = similar(buf_fwd) + I_fwd = inv(R_fwd) + end + fill!(buf_fwd, zero(domain_type)) + copyto!(view(buf_fwd, 1:m), h) + h_fft_conj = conj.(R_fwd * buf_fwd) + fill!(buf_fwd, zero(domain_type)) + + # Adjoint: CPU uses tiled FIR — no FFT state needed. + # GPU backends allocate FFT plans; same fftlen as forward pass is correct + # (wrap-around from h only affects positions < m ≤ padlen, outside the + # extracted range padlen..padlen+n-1). + if H <: Array + adj_fft = nothing + else + fftlen_adj = fftlen_fwd + buf_adj = similar(h, fftlen_adj) + if domain_type <: Real + R_adj = plan_rfft(buf_adj; plan_kw...) + buf_adj_c = similar(h, Complex{domain_type}, fftlen_adj ÷ 2 + 1) + I_adj = plan_irfft(buf_adj_c, fftlen_adj; plan_kw...) + else + R_adj = plan_fft(buf_adj; plan_kw...) + buf_adj_c = similar(buf_adj) + I_adj = inv(R_adj) + end + fill!(buf_adj, zero(domain_type)) + copyto!(view(buf_adj, 1:m), h) + h_fft_adj = R_adj * buf_adj + fill!(buf_adj, zero(domain_type)) + adj_fft = XcorrAdjFFT(fftlen_adj, h_fft_adj, buf_adj, buf_adj_c, R_adj, I_adj) + end + + return Xcorr{ + domain_type, typeof(h), typeof(buf_fwd_c), + typeof(R_fwd), typeof(I_fwd), typeof(adj_fft), + }( + DomainDim, h, + fftlen_fwd, padlen, h_fft_conj, buf_fwd, buf_fwd_c, R_fwd, I_fwd, + adj_fft, + ) end + Xcorr(x::H, h::H) where {H} = Xcorr(eltype(x), size(x), h) # Mappings -function mul!(y::AbstractArray, A::Xcorr, b::AbstractArray) +function mul!(y, A::Xcorr{T}, b) where {T} check(y, A, b) - return y .= xcorr(b, A.h; padmode = :longest) + n = length(b) + # Forward: xcorr(b, h; padmode=:longest) + # = irfft(rfft(b_padded) .* conj(rfft(h_padded)), fftlen)[fftlen-padlen+2:fftlen, 1:padlen] + fill!(A.buf_fwd, zero(T)) + copyto!(view(A.buf_fwd, 1:n), b) + mul!(A.buf_fwd_c, A.R_fwd, A.buf_fwd) + A.buf_fwd_c .*= A.h_fft_conj + mul!(A.buf_fwd, A.I_fwd, A.buf_fwd_c) + # Gather: DSP.xcorr format = [neg lags ascending, non-neg lags ascending] + # neg lags -(padlen-1) to -1 are at positions fftlen-padlen+2 to fftlen + # pos lags 0 to padlen-1 are at positions 1 to padlen + fftlen, padlen = A.fftlen_fwd, A.padlen + neg_start = fftlen - padlen + 2 + copyto!(view(y, 1:(padlen - 1)), view(A.buf_fwd, neg_start:fftlen)) + copyto!(view(y, padlen:length(y)), view(A.buf_fwd, 1:padlen)) + return y end -function mul!(y::AbstractArray, L::AdjointOperator{<:Xcorr}, b::AbstractArray) +# CPU adjoint: tiled 8-wide FIR (fast for cache-resident accumulators) +function mul!(y, L::AdjointOperator{<:Xcorr{T, <:Array{T}}}, b) where {T} check(y, L, b) A = L.A - l = floor(Int64, size(A, 1)[1] / 2) - idx = (l + 1):(l + length(y)) - return y .= conv(b, A.h)[idx] + _xcorr_fir_adj!(y, b, A.h, A.padlen) + return y +end + +# GPU adjoint: FFT-based conv via XcorrAdjFFT +function mul!(y, L::AdjointOperator{<:Xcorr{T, <:Any, <:Any, <:Any, <:Any, <:XcorrAdjFFT}}, b) where {T} + check(y, L, b) + A = L.A + adj = A.adj_fft + n = length(y) + outlen = length(b) + fill!(adj.buf, zero(T)) + copyto!(view(adj.buf, 1:outlen), b) + mul!(adj.buf_c, adj.R, adj.buf) + adj.buf_c .*= adj.h_fft + mul!(adj.buf, adj.I, adj.buf_c) + y .= @view(adj.buf[A.padlen:(A.padlen + n - 1)]) + return y +end + +# Tiled FIR adjoint: y[j] = Σ_k h[k] * b[padlen+j-k] (k = 1..m) +# Processes 8 output samples per outer iteration to keep accumulators in +# registers and avoid repeated reads/writes of y. +function _xcorr_fir_adj!(y, b, h, padlen) + m = length(h); n = length(y); T = eltype(y) + j = 1 + @inbounds while j ≤ n - 7 + a0 = a1 = a2 = a3 = a4 = a5 = a6 = a7 = zero(T) + for k in 1:m + hk = h[k] + base = padlen + j - k + a0 = muladd(hk, b[base], a0) + a1 = muladd(hk, b[base + 1], a1) + a2 = muladd(hk, b[base + 2], a2) + a3 = muladd(hk, b[base + 3], a3) + a4 = muladd(hk, b[base + 4], a4) + a5 = muladd(hk, b[base + 5], a5) + a6 = muladd(hk, b[base + 6], a6) + a7 = muladd(hk, b[base + 7], a7) + end + y[j] = a0; y[j + 1] = a1; y[j + 2] = a2; y[j + 3] = a3 + y[j + 4] = a4; y[j + 5] = a5; y[j + 6] = a6; y[j + 7] = a7 + j += 8 + end + @inbounds while j ≤ n + acc = zero(T) + for k in 1:m + acc = muladd(h[k], b[padlen + j - k], acc) + end + y[j] = acc + j += 1 + end end # Properties domain_type(::Xcorr{T}) where {T} = T codomain_type(::Xcorr{T}) where {T} = T +domain_array_type(::Xcorr{T, H}) where {T, H} = H +codomain_array_type(::Xcorr{T, H}) where {T, H} = H is_thread_safe(::Xcorr) = false -#TODO find out a way to verify this, is_full_row_rank(L::Xcorr) = true is_full_column_rank(L::Xcorr) = true diff --git a/FFTWOperators/Project.toml b/FFTWOperators/Project.toml index 342b407..fc91051 100644 --- a/FFTWOperators/Project.toml +++ b/FFTWOperators/Project.toml @@ -3,14 +3,26 @@ uuid = "c59a084b-ba08-4f3f-af9e-f4298d6caa94" version = "0.1.0" [deps] +AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" AbstractOperators = "d9c5613a-d543-52d8-9afd-8f241a8c3f1c" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588" +[weakdeps] +AcceleratedDCTs = "16e42a1c-b1d1-5566-807a-fe3b9ee7f6e1" +GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" + +[extensions] +AcceleratedDCTsExt = ["GPUArrays", "AcceleratedDCTs"] +GpuExt = "GPUArrays" + [compat] +AcceleratedDCTs = "0.4" +AbstractFFTs = "1.5.0" AbstractOperators = "0.4" FFTW = "1.9.0" +GPUArrays = "10, 11" LinearAlgebra = "1.10.0" Polyester = "0.7.18" julia = "1.10" diff --git a/FFTWOperators/README.md b/FFTWOperators/README.md index 5940ea2..c455ff0 100644 --- a/FFTWOperators/README.md +++ b/FFTWOperators/README.md @@ -19,6 +19,13 @@ FFTWOperators.jl is a **subpackage** of the AbstractOperators.jl ecosystem. Whil pkg> add FFTWOperators ``` +## GPU Support + +FFTWOperators.jl has mixed GPU support: + +- All operators except `DCT` and `IDCT` work with GPU arrays on CUDA.jl, AMDGPU.jl, oneAPI.jl, and OpenCL.jl. +- `DCT` and `IDCT` use FFTW CPU plans by default, but loading `AcceleratedDCTs` activates GPU support for those operators. If `AcceleratedDCTs` is not imported, these operators will not work with GPU arrays and should be wrapped with `CpuOperatorWrapper` for use in GPU pipelines. + ## Usage Example ```julia @@ -121,6 +128,32 @@ dct_complex = DCT(Complex{Float64}, (8, 8)) idct_op = IDCT(Float64, (10, 10)) ``` +#### GPU Support for DCT/IDCT + +To use `DCT` and `IDCT` on GPU arrays, load `AcceleratedDCTs` before constructing the operator: + +```julia +using FFTWOperators, CUDA +import AcceleratedDCTs # explicitly import to activate GPU support for DCT/IDCT + +x_gpu = CUDA.randn(Float32, 10, 10) +dct_gpu = DCT(x_gpu) +y_gpu = dct_gpu * x_gpu + +idct_gpu = IDCT(x_gpu) +x_rec_gpu = idct_gpu * y_gpu +``` + +If `AcceleratedDCTs` is not imported, `DCT` and `IDCT` will use CPU FFTW plans and will not work with GPU arrays. In that case, wrap them with `CpuOperatorWrapper` for use in GPU pipelines: + +```julia +using FFTWOperators, CUDA + +x_gpu = CUDA.randn(Float32, 10, 10) +dct_gpu = CpuOperatorWrapper(DCT(Float32, (10, 10)); array_type = CuArray{Float32}) # CPU operator wrapped for GPU use +y_gpu = dct_gpu * x_gpu # GPU in → CPU DCT → GPU out +``` + ### 5. **Shift** - Frequency Shift Operator Applies frequency shifts and zero-padding to signals. diff --git a/FFTWOperators/ext/AcceleratedDCTsExt.jl b/FFTWOperators/ext/AcceleratedDCTsExt.jl new file mode 100644 index 0000000..90e2b78 --- /dev/null +++ b/FFTWOperators/ext/AcceleratedDCTsExt.jl @@ -0,0 +1,105 @@ +module AcceleratedDCTsExt + +using AbstractFFTs +using AcceleratedDCTs +using FFTWOperators +using GPUArrays + +import Base: eltype, inv, ndims, size +import LinearAlgebra: mul! +import FFTWOperators: DCT, IDCT +import AbstractOperators: AdjointOperator, check + +struct OrthoDCTPlan{T, N, P, S, B} <: AbstractFFTs.Plan{T} + raw::P + scale::S + temp::B + pinv::Base.RefValue{Any} +end + +struct OrthoIDCTPlan{T, N, P, S, B} <: AbstractFFTs.Plan{T} + raw::P + scale::S + temp::B + pinv::Base.RefValue{Any} +end + +function FFTWOperators.DCT(x::AbstractGPUArray{T, N}) where {T <: Real, N} + A, At = _ortho_plan_pair(x) + buf = similar(x, 0) # zero-length marker for storage type inference + return DCT{N, T, typeof(A), typeof(At), typeof(buf)}(size(x), A, At, buf) +end + +function FFTWOperators.IDCT(x::AbstractGPUArray{T, N}) where {T <: Real, N} + At, A = _ortho_plan_pair(x) + buf = similar(x, 0) # zero-length marker for storage type inference + return IDCT{N, T, typeof(A), typeof(At), typeof(buf)}(size(x), A, At, buf) +end + +Base.size(p::OrthoDCTPlan) = size(p.scale) +Base.size(p::OrthoIDCTPlan) = size(p.scale) +Base.ndims(::OrthoDCTPlan{T, N}) where {T, N} = N +Base.ndims(::OrthoIDCTPlan{T, N}) where {T, N} = N +Base.eltype(::OrthoDCTPlan{T}) where {T} = T +Base.eltype(::OrthoIDCTPlan{T}) where {T} = T + +Base.inv(p::OrthoDCTPlan) = p.pinv[] +Base.inv(p::OrthoIDCTPlan) = p.pinv[] + +function mul!(y::AbstractGPUArray{T}, p::OrthoDCTPlan{T}, x::AbstractGPUArray{T}) where {T} + mul!(y, p.raw, x) + y ./= p.scale + return y +end + +function mul!(y::AbstractGPUArray{T}, p::OrthoIDCTPlan{T}, x::AbstractGPUArray{T}) where {T} + p.temp .= x .* p.scale + mul!(y, p.raw, p.temp) + return y +end + +# AcceleratedDCTs plans do not modify input, so no scratch copy is needed. +# These overrides bypass the copyto!(buf, b) in the CPU mul! methods. +function mul!(y::AbstractGPUArray, A::IDCT, b::AbstractGPUArray) + check(y, A, b) + return mul!(y, A.A, b) +end + +function mul!(y::AbstractGPUArray, A::AdjointOperator{<:DCT}, b::AbstractGPUArray) + check(y, A, b) + return mul!(y, A.A.At, b) +end + +function _backend_vector(x::AbstractGPUArray, values::Vector{T}) where {T} + ArrayT = Base.typename(typeof(x)).wrapper + return ArrayT(values) +end + +function _ortho_scale(x::AbstractGPUArray{T, N}) where {T <: Real, N} + scale = similar(x) + fill!(scale, one(T)) + for d in 1:N + axis = fill(sqrt(T(size(x, d)) / T(2)), size(x, d)) + axis[1] = sqrt(T(size(x, d))) + axis_gpu = _backend_vector(x, axis) + shape = ntuple(i -> i == d ? size(x, d) : 1, N) + scale .*= reshape(axis_gpu, shape) + end + return scale +end + +function _ortho_plan_pair(x::AbstractGPUArray{T, N}) where {T <: Real, N} + scale = _ortho_scale(x) + dct_raw = AcceleratedDCTs.plan_dct(x) + idct_raw = AcceleratedDCTs.plan_idct(x) + fwd = OrthoDCTPlan{T, N, typeof(dct_raw), typeof(scale), typeof(similar(x))}( + dct_raw, scale, similar(x), Ref{Any}(nothing), + ) + inv_plan = OrthoIDCTPlan{T, N, typeof(idct_raw), typeof(scale), typeof(similar(x))}( + idct_raw, scale, similar(x), Ref{Any}(fwd), + ) + fwd.pinv[] = inv_plan + return fwd, inv_plan +end + +end # module AcceleratedDCTsExt diff --git a/FFTWOperators/ext/GpuExt/GpuExt.jl b/FFTWOperators/ext/GpuExt/GpuExt.jl new file mode 100644 index 0000000..b16932a --- /dev/null +++ b/FFTWOperators/ext/GpuExt/GpuExt.jl @@ -0,0 +1,60 @@ +module GpuExt + +using AbstractFFTs: plan_fft, plan_bfft +using GPUArrays +using AbstractOperators: check +using FFTWOperators +import FFTWOperators: DFT, Normalization, UNNORMALIZED, _dft_scaling, SignAlternation +import LinearAlgebra: mul! + +# GPU constructor for DFT with real input — avoids FFTW-specific `flags`/`timelimit`/`num_threads` kwargs +# that are not accepted by GPU FFT backends (e.g. CUDA.CUFFT, AMDGPU.rocFFT). +function FFTWOperators.DFT( + x::AbstractGPUArray{D, N}, + dims = 1:ndims(x); + normalization::Normalization = UNNORMALIZED, + kwargs..., # silently ignore FFTW-specific kwargs (flags, timelimit, num_threads) + ) where {N, D <: Real} + xc = similar(x, Complex{D}) + A = plan_fft(xc, dims) + At = plan_bfft(xc, dims) + S = typeof(x isa SubArray ? parent(x) : x).name.wrapper + dims_t = tuple(dims...) + scaling = _dft_scaling(size(x), dims_t, normalization) + return DFT{N, Complex{D}, D, dims_t, S, typeof(A), typeof(At), D}( + size(x), A, At, normalization, scaling, + ) +end + +# GPU constructor for DFT with complex input +function FFTWOperators.DFT( + x::AbstractGPUArray{D, N}, + dims = 1:ndims(x); + normalization::Normalization = UNNORMALIZED, + kwargs..., + ) where {N, D <: Complex} + A = plan_fft(x, dims) + At = plan_bfft(x, dims) + S = typeof(x isa SubArray ? parent(x) : x).name.wrapper + dims_t = tuple(dims...) + scaling = _dft_scaling(size(x), dims_t, normalization) + return DFT{N, D, D, dims_t, S, typeof(A), typeof(At), real(D)}( + size(x), A, At, normalization, scaling, + ) +end + +# GPU mul! for SignAlternation: uses broadcast to avoid scalar indexing +function mul!(y::AbstractGPUArray, L::SignAlternation{T, N}, b::AbstractGPUArray) where {T, N} + check(y, L, b) + y .= b + for d in L.dirs + nd = size(b, d) + sv_shape = ntuple(i -> i == d ? nd : 1, N) + sv = similar(b, T, sv_shape) + copyto!(sv, T[isodd(i - 1) ? -one(T) : one(T) for i in 1:nd]) + y .*= sv + end + return y +end + +end # module GpuExt diff --git a/FFTWOperators/src/DCT.jl b/FFTWOperators/src/DCT.jl index cc6c771..ccdf2d5 100644 --- a/FFTWOperators/src/DCT.jl +++ b/FFTWOperators/src/DCT.jl @@ -29,10 +29,11 @@ julia> A*ones(3) ``` """ -struct DCT{N, C, T1 <: AbstractFFTs.Plan, T2 <: AbstractFFTs.Plan} <: CosineTransform{N, C, T1, T2} +struct DCT{N, C, T1 <: AbstractFFTs.Plan, T2 <: AbstractFFTs.Plan, B} <: CosineTransform{N, C, T1, T2} dim_in::NTuple{N, Int} A::T1 At::T2 + buf::B end """ @@ -62,10 +63,11 @@ julia> A*[1.;0.;0.] ``` """ -struct IDCT{N, C, T1 <: AbstractFFTs.Plan, T2 <: AbstractFFTs.Plan} <: CosineTransform{N, C, T1, T2} +struct IDCT{N, C, T1 <: AbstractFFTs.Plan, T2 <: AbstractFFTs.Plan, B} <: CosineTransform{N, C, T1, T2} dim_in::NTuple{N, Int} A::T1 At::T2 + buf::B end # Constructors @@ -78,7 +80,8 @@ DCT(T::Type, dim_in::Vararg{Int64}) = DCT(T, dim_in) function DCT(x::AbstractArray{C, N}) where {N, C} A, At = plan_dct(x), plan_idct(x) - return DCT{N, C, typeof(A), typeof(At)}(size(x), A, At) + buf = similar(x) + return DCT{N, C, typeof(A), typeof(At), typeof(buf)}(size(x), A, At, buf) end #standard constructor @@ -90,33 +93,38 @@ IDCT(T::Type, dim_in::Vararg{Int64}) = IDCT(T, dim_in) function IDCT(x::AbstractArray{C, N}) where {N, C} A, At = plan_idct(x), plan_dct(x) - return IDCT{N, C, typeof(A), typeof(At)}(size(x), A, At) + buf = similar(x) + return IDCT{N, C, typeof(A), typeof(At), typeof(buf)}(size(x), A, At, buf) end # Mappings -function mul!( - y::AbstractArray{C, N}, A::DCT{N, C, T1, T2}, b::AbstractArray{C, N} - ) where {N, C, T1, T2} - return mul!(y, A.A, b) +function mul!(y::AbstractArray, A::DCT, b::AbstractArray) + check(y, A, b) + mul!(y, A.A, b) # DCT plan (REDFT10): non-destructive to input + return y end -function mul!( - y::AbstractArray{C, N}, A::AdjointOperator{DCT{N, C, T1, T2}}, b::AbstractArray{C, N} - ) where {N, C, T1, T2} - return y .= A.A.At * b +function mul!(y::AbstractArray, A::AdjointOperator{<:DCT}, b::AbstractArray) + check(y, A, b) + # IDCT plan (REDFT01) modifies its input in-place; use scratch buffer + copyto!(A.A.buf, b) + mul!(y, A.A.At, A.A.buf) + return y end -function mul!( - y::AbstractArray{C, N}, A::IDCT{N, C, T1, T2}, b::AbstractArray{C, N} - ) where {N, C, T1, T2} - return y .= A.A * b +function mul!(y::AbstractArray, A::IDCT, b::AbstractArray) + check(y, A, b) + # IDCT plan (REDFT01) modifies its input in-place; use scratch buffer + copyto!(A.buf, b) + mul!(y, A.A, A.buf) + return y end -function mul!( - y::AbstractArray{C, N}, A::AdjointOperator{IDCT{N, C, T1, T2}}, b::AbstractArray{C, N} - ) where {N, C, T1, T2} - return mul!(y, A.A.At, b) +function mul!(y::AbstractArray, A::AdjointOperator{<:IDCT}, b::AbstractArray) + check(y, A, b) + mul!(y, A.A.At, b) # DCT plan (REDFT10): non-destructive to input + return y end # Properties @@ -126,9 +134,13 @@ size(L::CosineTransform) = (L.dim_in, L.dim_in) fun_name(A::DCT) = "ℱc" fun_name(A::IDCT) = "ℱc⁻¹" -domain_type(::CosineTransform{N, C, T1, T2}) where {N, C, T1, T2} = C -codomain_type(::CosineTransform{N, C, T1, T2}) where {N, C, T1, T2} = C -is_thread_safe(::CosineTransform) = true +domain_type(::CosineTransform{N, C}) where {N, C} = C +codomain_type(::CosineTransform{N, C}) where {N, C} = C +domain_array_type(::DCT{N, C, T1, T2, B}) where {N, C, T1, T2, B} = Base.typename(B).wrapper{C} +domain_array_type(::IDCT{N, C, T1, T2, B}) where {N, C, T1, T2, B} = Base.typename(B).wrapper{C} +codomain_array_type(::DCT{N, C, T1, T2, B}) where {N, C, T1, T2, B} = Base.typename(B).wrapper{C} +codomain_array_type(::IDCT{N, C, T1, T2, B}) where {N, C, T1, T2, B} = Base.typename(B).wrapper{C} +is_thread_safe(::CosineTransform) = false is_AcA_diagonal(L::CosineTransform) = true is_AAc_diagonal(L::CosineTransform) = true diff --git a/FFTWOperators/src/DFT.jl b/FFTWOperators/src/DFT.jl index ea82a5d..477ceb8 100644 --- a/FFTWOperators/src/DFT.jl +++ b/FFTWOperators/src/DFT.jl @@ -280,12 +280,12 @@ end # Properties size(L::DFT) = (L.dim_in, L.dim_in) -function domain_storage_type( +function domain_array_type( ::DFT{N, C, D, Dir, S} ) where {N, C, D, Dir, S} return S{D} end -function codomain_storage_type( +function codomain_array_type( ::DFT{N, C, D, Dir, S} ) where {N, C, D, Dir, S} return S{C} diff --git a/FFTWOperators/src/FFTWOperators.jl b/FFTWOperators/src/FFTWOperators.jl index 0b59cde..5a14078 100644 --- a/FFTWOperators/src/FFTWOperators.jl +++ b/FFTWOperators/src/FFTWOperators.jl @@ -7,14 +7,16 @@ import LinearAlgebra: mul! import Base: size, ndims import AbstractOperators: + _normalize_array_type, + _array_wrapper_type, domain_type, codomain_type, fun_name, get_normal_op, allocate_in_domain, allocate_in_codomain, - domain_storage_type, - codomain_storage_type, + domain_array_type, + codomain_array_type, can_be_combined, combine, is_thread_safe, diff --git a/FFTWOperators/src/IRDFT.jl b/FFTWOperators/src/IRDFT.jl index 97e7358..70978a3 100644 --- a/FFTWOperators/src/IRDFT.jl +++ b/FFTWOperators/src/IRDFT.jl @@ -17,7 +17,7 @@ julia> A = IRDFT((5,10,8),19,2) ``` """ -struct IRDFT{T <: Number, N, D, T1 <: AbstractFFTs.Plan, T2 <: AbstractFFTs.Plan, T3 <: NTuple{N, Any}} <: +struct IRDFT{T <: Number, N, D, T1 <: AbstractFFTs.Plan, T2 <: AbstractFFTs.Plan, T3 <: NTuple{N, Any}, S} <: LinearOperator dim_in::NTuple{N, Int} dim_out::NTuple{N, Int} @@ -39,7 +39,8 @@ function IRDFT(x::AbstractArray{Complex{T}, N}, d::Int, dims::Int = 1) where {T idx = i == dims ? (idx..., 2:ceil(Int, d / 2)) : (idx..., Colon()) end At = plan_rfft(similar(x, T, dim_out), dims) - return IRDFT{T, N, dims, typeof(A), typeof(At), typeof(idx)}(dim_in, dim_out, A, At, idx) + S = _array_wrapper_type(typeof(x isa SubArray ? parent(x) : x)) + return IRDFT{T, N, dims, typeof(A), typeof(At), typeof(idx), S}(dim_in, dim_out, A, At, idx) end function IRDFT(T::Type, dim_in::NTuple{N, Int}, d::Int, dims::Int = 1) where {N} @@ -54,12 +55,15 @@ end function mul!( y::C1, L::IRDFT{T, N, D, T1, T2, T3}, b::C2 ) where {N, T, D, T1, T2, T3, C1 <: AbstractArray{T, N}, C2 <: AbstractArray{Complex{T}, N}} - return mul!(y, L.A, b) + check(y, L, b) + mul!(y, L.A, b) + return y end function mul!( - y::C2, L::AdjointOperator{IRDFT{T, N, D, T1, T2, T3}}, b::C1 - ) where {N, T, D, T1, T2, T3, C1 <: AbstractArray{T, N}, C2 <: AbstractArray{Complex{T}, N}} + y::C2, L::AdjointOperator{<:IRDFT{T, N, D}}, b::C1 + ) where {N, T, D, C1 <: AbstractArray{T, N}, C2 <: AbstractArray{Complex{T}, N}} + check(y, L, b) A = L.A mul!(y, A.At, b) y ./= size(b, D) @@ -77,6 +81,13 @@ domain_type(::IRDFT{T}) where {T} = Complex{T} codomain_type(::IRDFT{T}) where {T} = T is_thread_safe(::IRDFT) = true +function domain_array_type(::IRDFT{T, N, D, T1, T2, T3, S}) where {T, N, D, T1, T2, T3, S} + return S{Complex{T}} +end +function codomain_array_type(::IRDFT{T, N, D, T1, T2, T3, S}) where {T, N, D, T1, T2, T3, S} + return S{T} +end + is_AAc_diagonal(L::IRDFT) = false #TODO but might be true? is_invertible(L::IRDFT) = true is_full_row_rank(L::IRDFT) = true diff --git a/FFTWOperators/src/RDFT.jl b/FFTWOperators/src/RDFT.jl index 0f0628d..e88e9d7 100644 --- a/FFTWOperators/src/RDFT.jl +++ b/FFTWOperators/src/RDFT.jl @@ -57,14 +57,18 @@ RDFT(T::Type, dim_in::Vararg{Int}) = RDFT(T, dim_in) function mul!( y::T3, L::RDFT{T, N, T1, T2, T3}, b::T4 ) where {N, T, T1, T2, T3, T4 <: AbstractArray{T, N}} - return mul!(y, L.A, b) + check(y, L, b) + mul!(y, L.A, b) + return y end function mul!( y::T4, L::AdjointOperator{RDFT{T, N, T1, T2, T3}}, b::T3 ) where {N, T, T1, T2, T3, T4 <: AbstractArray{T, N}} + check(y, L, b) A = L.A - mul!(A.b2, A.Zp, b) + fill!(A.b2, zero(eltype(A.b2))) + copyto!(view(A.b2, ntuple(i -> Base.OneTo(size(b, i)), N)...), b) mul!(A.y2, A.At, A.b2) y .= real.(A.y2) return y @@ -80,6 +84,13 @@ domain_type(::RDFT{T}) where {T} = T codomain_type(::RDFT{T}) where {T} = Complex{T} is_thread_safe(::RDFT) = false +function domain_array_type(L::RDFT{T, N, T1, T2, T3}) where {T, N, T1, T2, T3} + return T3.name.wrapper{T} +end +function codomain_array_type(L::RDFT{T, N, T1, T2, T3}) where {T, N, T1, T2, T3} + return T3.name.wrapper{Complex{T}} +end + is_AAc_diagonal(L::RDFT) = false #TODO but might be true? is_invertible(L::RDFT) = true is_full_row_rank(L::RDFT) = true diff --git a/FFTWOperators/src/Shift.jl b/FFTWOperators/src/Shift.jl index 67c720e..4463de1 100644 --- a/FFTWOperators/src/Shift.jl +++ b/FFTWOperators/src/Shift.jl @@ -25,13 +25,13 @@ julia> A * x julia> B = IFFTShift((4,)); B * (A * x) == x true -julia> A2 = FFTShift((2,2), (1,2)); A2 * reshape(1.0:4.0, 2, 2) +julia> A2 = FFTShift((2,2), (1,2)); A2 * collect(reshape(1.0:4.0, 2, 2)) 2×2 Matrix{Float64}: 4.0 2.0 3.0 1.0 ``` """ -struct FFTShift{T, N, M} <: ShiftOp +struct FFTShift{T, N, M, S <: AbstractArray{T}} <: ShiftOp dim_in::NTuple{N, Int} dirs::NTuple{M, Int} end @@ -55,7 +55,7 @@ julia> B2 = IFFTShift((2,2), (1,2)); B2 * [4.0 3.0; 2.0 1.0] == [1.0 2.0; 3.0 4. true ``` """ -struct IFFTShift{T, N, M} <: ShiftOp +struct IFFTShift{T, N, M, S <: AbstractArray{T}} <: ShiftOp dim_in::NTuple{N, Int} dirs::NTuple{M, Int} end @@ -92,7 +92,7 @@ julia> S2 = SignAlternation((2,2), (1,2)); S2 * ones(2,2) -1.0 1.0 ``` """ -struct SignAlternation{S, N, M, Th} <: LinearOperator +struct SignAlternation{T, N, M, Th, S <: AbstractArray{T}} <: LinearOperator dim_in::NTuple{N, Int} dirs::NTuple{M, Int} end @@ -104,45 +104,92 @@ function _normalize_dirs(::NTuple{N, Int}, dirs) where {N} return d end -function FFTShift(::Type{T}, dim_in::NTuple{N, Int}, dirs) where {T <: Number, N} +function _make_shift(::Type{Op}, ::Type{T}, dim_in::NTuple{N, Int}, dirs, array_type::Type) where {Op, T <: Number, N} d = _normalize_dirs(dim_in, dirs) - return FFTShift{T, N, length(d)}(dim_in, d) + S = _normalize_array_type(array_type, T) + return Op{T, N, length(d), S}(dim_in, d) end -function FFTShift(::Type{T}, dim_in::NTuple{N, Int}) where {T <: Number, N} - return FFTShift(T, dim_in, Tuple(1:N)) + +function FFTShift( + ::Type{T}, dim_in::NTuple{N, Int}, dirs; + array_type::Type{<:AbstractArray} = Array{T}, + ) where {T <: Number, N} + return _make_shift(FFTShift, T, dim_in, dirs, array_type) +end +function FFTShift( + ::Type{T}, dim_in::NTuple{N, Int}; + array_type::Type{<:AbstractArray} = Array{T}, + ) where {T <: Number, N} + return FFTShift(T, dim_in, Tuple(1:N); array_type) +end +function FFTShift(dim_in::NTuple{N, Int}, dirs; array_type::Type{<:AbstractArray} = Array{Float64}) where {N} + return FFTShift(Float64, dim_in, dirs; array_type) +end +function FFTShift(dim_in::NTuple{N, Int}; array_type::Type{<:AbstractArray} = Array{Float64}) where {N} + return FFTShift(Float64, dim_in, Tuple(1:N); array_type) end -FFTShift(dim_in::NTuple{N, Int}, dirs) where {N} = FFTShift(Float64, dim_in, dirs) -FFTShift(dim_in::NTuple{N, Int}) where {N} = FFTShift(Float64, dim_in, Tuple(1:N)) FFTShift(dim_in::Vararg{Int}) = FFTShift(dim_in) +function FFTShift(x::A) where {A <: AbstractArray} + return FFTShift(eltype(x), size(x); array_type = typeof(x isa SubArray ? parent(x) : x)) +end -function IFFTShift(::Type{T}, dim_in::NTuple{N, Int}, dirs) where {T <: Number, N} - d = _normalize_dirs(dim_in, dirs) - return IFFTShift{T, N, length(d)}(dim_in, d) +function IFFTShift( + ::Type{T}, dim_in::NTuple{N, Int}, dirs; + array_type::Type{<:AbstractArray} = Array{T}, + ) where {T <: Number, N} + return _make_shift(IFFTShift, T, dim_in, dirs, array_type) end -function IFFTShift(::Type{T}, dim_in::NTuple{N, Int}) where {T <: Number, N} - return IFFTShift(T, dim_in, Tuple(1:N)) +function IFFTShift( + ::Type{T}, dim_in::NTuple{N, Int}; + array_type::Type{<:AbstractArray} = Array{T}, + ) where {T <: Number, N} + return IFFTShift(T, dim_in, Tuple(1:N); array_type) +end +function IFFTShift(dim_in::NTuple{N, Int}, dirs; array_type::Type{<:AbstractArray} = Array{Float64}) where {N} + return IFFTShift(Float64, dim_in, dirs; array_type) +end +function IFFTShift(dim_in::NTuple{N, Int}; array_type::Type{<:AbstractArray} = Array{Float64}) where {N} + return IFFTShift(Float64, dim_in, Tuple(1:N); array_type) end -IFFTShift(dim_in::NTuple{N, Int}, dirs) where {N} = IFFTShift(Float64, dim_in, dirs) -IFFTShift(dim_in::NTuple{N, Int}) where {N} = IFFTShift(Float64, dim_in, Tuple(1:N)) IFFTShift(dim_in::Vararg{Int}) = IFFTShift(dim_in) +function IFFTShift(x::A) where {A <: AbstractArray} + return IFFTShift(eltype(x), size(x); array_type = typeof(x isa SubArray ? parent(x) : x)) +end function SignAlternation( - ::Type{S}, dim_in::NTuple{N, Int}, dirs; threaded::Bool = true - ) where {S <: Number, N} + ::Type{T}, dim_in::NTuple{N, Int}, dirs; + threaded::Bool = true, array_type::Type{<:AbstractArray} = Array{T}, + ) where {T <: Number, N} d = _normalize_dirs(dim_in, dirs) threaded = threaded && Threads.nthreads() > 1 - return SignAlternation{S, N, length(d), threaded}(dim_in, d) + S = _normalize_array_type(array_type, T) + return SignAlternation{T, N, length(d), threaded, S}(dim_in, d) end -function SignAlternation(dim_in::NTuple{N, Int}, dirs; threaded::Bool = true) where {N} - return SignAlternation(Float64, dim_in, dirs; threaded) +function SignAlternation( + dim_in::NTuple{N, Int}, dirs; + threaded::Bool = true, array_type::Type{<:AbstractArray} = Array{Float64}, + ) where {N} + return SignAlternation(Float64, dim_in, dirs; threaded, array_type) +end +function SignAlternation(x::A, dirs; threaded::Bool = true) where {A <: AbstractArray} + return SignAlternation( + eltype(x), size(x), dirs; + threaded, array_type = typeof(x isa SubArray ? parent(x) : x), + ) end domain_type(::FFTShift{T}) where {T} = T codomain_type(::FFTShift{T}) where {T} = T +domain_array_type(::FFTShift{T, N, M, S}) where {T, N, M, S} = S +codomain_array_type(::FFTShift{T, N, M, S}) where {T, N, M, S} = S domain_type(::IFFTShift{T}) where {T} = T codomain_type(::IFFTShift{T}) where {T} = T -domain_type(::SignAlternation{S}) where {S} = S -codomain_type(::SignAlternation{S}) where {S} = S +domain_array_type(::IFFTShift{T, N, M, S}) where {T, N, M, S} = S +codomain_array_type(::IFFTShift{T, N, M, S}) where {T, N, M, S} = S +domain_type(::SignAlternation{T}) where {T} = T +codomain_type(::SignAlternation{T}) where {T} = T +domain_array_type(::SignAlternation{T, N, M, Th, S}) where {T, N, M, Th, S} = S +codomain_array_type(::SignAlternation{T, N, M, Th, S}) where {T, N, M, Th, S} = S size(L::FFTShift) = (L.dim_in, L.dim_in) size(L::IFFTShift) = (L.dim_in, L.dim_in) @@ -354,15 +401,15 @@ function alternate_sign(x::AbstractArray, dirs::Int...; threaded::Bool = true) end function mul!( - y::AbstractArray, L::SignAlternation{S, N, M, Th}, b::AbstractArray - ) where {S, N, M, Th} + y::AbstractArray, L::SignAlternation{T, N, M, Th}, b::AbstractArray + ) where {T, N, M, Th} check(y, L, b) return _alternate_sign!(y, b, L.dirs; threaded = Th) end has_optimized_normalop(::Union{FFTShift, IFFTShift, SignAlternation}) = true function get_normal_op(L::Union{FFTShift, IFFTShift, SignAlternation}) - return Eye(domain_type(L), size(L, 1)) + return Eye(domain_type(L), size(L, 1); array_type = domain_array_type(L)) end LinearAlgebra.adjoint(L::SignAlternation) = L @@ -384,7 +431,7 @@ function _check_shift_dirs(::NTuple{N, Int}, dirs::NTuple{M, Int}) where {N, M} throw(ArgumentError("dirs must be sorted in ascending order")) end end - return + return nothing end function _is_dft_op(op, side) @@ -412,10 +459,10 @@ function _shift_op( _check_shift_dirs(size(op, 2), domain_shifts) shifted_domain_dims_shape = size(op, 2)[collect(domain_shifts)] if all(iseven, shifted_domain_dims_shape) && _is_dft_op(op, :domain) - domain_op = SignAlternation(codomain_type(op), size(op, 1), domain_shifts) + domain_op = SignAlternation(codomain_type(op), size(op, 1), domain_shifts; array_type = codomain_array_type(op)) op = domain_op * op else - domain_op = shift_op_type(domain_type(op), size(op, 2), domain_shifts) + domain_op = shift_op_type(domain_type(op), size(op, 2), domain_shifts; array_type = domain_array_type(op)) op = op * domain_op end end @@ -423,10 +470,10 @@ function _shift_op( _check_shift_dirs(size(op, 1), codomain_shifts) shifted_codomain_dims_shape = size(op, 1)[collect(codomain_shifts)] if all(iseven, shifted_codomain_dims_shape) && _is_dft_op(op, :codomain) - codomain_op = SignAlternation(domain_type(op), size(op, 2), codomain_shifts) + codomain_op = SignAlternation(domain_type(op), size(op, 2), codomain_shifts; array_type = domain_array_type(op)) op = op * codomain_op else - codomain_op = shift_op_type(codomain_type(op), size(op, 1), codomain_shifts) + codomain_op = shift_op_type(codomain_type(op), size(op, 1), codomain_shifts; array_type = codomain_array_type(op)) op = codomain_op * op end end diff --git a/FFTWOperators/src/combination_rules.jl b/FFTWOperators/src/combination_rules.jl index 0ae2da9..d12ce35 100644 --- a/FFTWOperators/src/combination_rules.jl +++ b/FFTWOperators/src/combination_rules.jl @@ -27,13 +27,13 @@ function combine(T1::DFT{N, C, D, Dir}, T2::AdjointOperator{<:DFT}) where {N, C, scaling = _dft_scaling(T1.dim_in, Dir, FORWARD) scaling /= _dft_scaling(T1.dim_in, Dir, T1.normalization) scaling /= _dft_scaling(T1.dim_in, Dir, T2.A.normalization) - return scaling * Eye(domain_type(T2), T1.dim_in, domain_storage_type(T2)) + return scaling * Eye(domain_type(T2), T1.dim_in; array_type = domain_array_type(T2)) end function combine(T1::AdjointOperator{<:DFT}, T2::DFT{N, C, D, Dir}) where {N, C, D, Dir} scaling = _dft_scaling(T2.dim_in, Dir, FORWARD) scaling /= _dft_scaling(T2.dim_in, Dir, T1.A.normalization) scaling /= _dft_scaling(T2.dim_in, Dir, T2.normalization) - return scaling * Eye(domain_type(T2), T2.dim_in, domain_storage_type(T2)) + return scaling * Eye(domain_type(T2), T2.dim_in; array_type = domain_array_type(T2)) end # FFTShift/IFFTShift with DFT @@ -62,28 +62,28 @@ function can_be_combined(T1::AdjointOperator{<:ShiftOp}, T2::AdjointOperator{<:D return all(iseven, size(T2, 1)[collect(T1.dirs)]) end function combine(T1::DFT, T2::ShiftOp) - return SignAlternation(codomain_type(T1), size(T1, 1), T2.dirs) * T1 + return SignAlternation(codomain_type(T1), size(T1, 1), T2.dirs; array_type = codomain_array_type(T1)) * T1 end function combine(T1::ShiftOp, T2::DFT) - return T2 * SignAlternation(domain_type(T2), size(T2, 2), T1.dirs) + return T2 * SignAlternation(domain_type(T2), size(T2, 2), T1.dirs; array_type = domain_array_type(T2)) end function combine(T1::AdjointOperator{<:DFT}, T2::ShiftOp) - return SignAlternation(codomain_type(T1), size(T1, 1), T2.dirs) * T1 + return SignAlternation(codomain_type(T1), size(T1, 1), T2.dirs; array_type = codomain_array_type(T1)) * T1 end function combine(T1::ShiftOp, T2::AdjointOperator{<:DFT}) - return T2 * SignAlternation(domain_type(T2), size(T2, 2), T1.dirs) + return T2 * SignAlternation(domain_type(T2), size(T2, 2), T1.dirs; array_type = domain_array_type(T2)) end function combine(T1::DFT, T2::AdjointOperator{<:ShiftOp}) - return SignAlternation(codomain_type(T1), size(T1, 1), T2.dirs) * T1 + return SignAlternation(codomain_type(T1), size(T1, 1), T2.dirs; array_type = codomain_array_type(T1)) * T1 end function combine(T1::AdjointOperator{<:ShiftOp}, T2::DFT) - return T2 * SignAlternation(domain_type(T2), size(T2, 2), T1.dirs) + return T2 * SignAlternation(domain_type(T2), size(T2, 2), T1.dirs; array_type = domain_array_type(T2)) end function combine(T1::AdjointOperator{<:DFT}, T2::AdjointOperator{<:ShiftOp}) - return SignAlternation(codomain_type(T1), size(T1, 1), T2.dirs) * T1 + return SignAlternation(codomain_type(T1), size(T1, 1), T2.dirs; array_type = codomain_array_type(T1)) * T1 end function combine(T1::AdjointOperator{<:ShiftOp}, T2::AdjointOperator{<:DFT}) - return T2 * SignAlternation(domain_type(T2), size(T2, 2), T1.dirs) + return T2 * SignAlternation(domain_type(T2), size(T2, 2), T1.dirs; array_type = domain_array_type(T2)) end # FFTShift/IFFTShift with DFT and SignAlternation @@ -155,36 +155,36 @@ function can_be_combined(T1::IFFTShift, T2::FFTShift) end function combine(T1::FFTShift, T2::FFTShift) new_dirs = Tuple(d for d in 1:ndims(T1, 2) if (d in T1.dirs) || (d in T2.dirs)) - return FFTShift(domain_type(T1), size(T1, 2), new_dirs) + return FFTShift(domain_type(T1), size(T1, 2), new_dirs; array_type = domain_array_type(T1)) end function combine(T1::IFFTShift, T2::IFFTShift) new_dirs = Tuple(d for d in 1:ndims(T1, 2) if (d in T1.dirs) || (d in T2.dirs)) - return IFFTShift(domain_type(T1), size(T1, 2), new_dirs) + return IFFTShift(domain_type(T1), size(T1, 2), new_dirs; array_type = domain_array_type(T1)) end function combine(T1::FFTShift, T2::IFFTShift) if does_fully_cover(T1, T2) new_dirs = Tuple(d for d in T1.dirs if !(d in T2.dirs)) if isempty(new_dirs) - return Eye(domain_type(T1), size(T1, 2)) + return Eye(domain_type(T1), size(T1, 2); array_type = domain_array_type(T1)) else - return FFTShift(domain_type(T1), size(T1, 2), new_dirs) + return FFTShift(domain_type(T1), size(T1, 2), new_dirs; array_type = domain_array_type(T1)) end else new_dirs = Tuple(d for d in T2.dirs if !(d in T1.dirs)) - return IFFTShift(domain_type(T1), size(T1, 2), new_dirs) + return IFFTShift(domain_type(T1), size(T1, 2), new_dirs; array_type = domain_array_type(T1)) end end function combine(T1::IFFTShift, T2::FFTShift) if does_fully_cover(T1, T2) new_dirs = Tuple(d for d in T1.dirs if !(d in T2.dirs)) if isempty(new_dirs) - return Eye(domain_type(T1), size(T1, 2)) + return Eye(domain_type(T1), size(T1, 2); array_type = domain_array_type(T1)) else - return IFFTShift(domain_type(T1), size(T1, 2), new_dirs) + return IFFTShift(domain_type(T1), size(T1, 2), new_dirs; array_type = domain_array_type(T1)) end else new_dirs = Tuple(d for d in T2.dirs if !(d in T1.dirs)) - return FFTShift(domain_type(T1), size(T1, 2), new_dirs) + return FFTShift(domain_type(T1), size(T1, 2), new_dirs; array_type = domain_array_type(T1)) end end @@ -193,9 +193,9 @@ can_be_combined(T1::SignAlternation, T2::SignAlternation) = true function combine(T1::SignAlternation, T2::SignAlternation) new_dirs = Tuple(d for d in 1:ndims(T1, 1) if (d in T1.dirs) != (d in T2.dirs)) if isempty(new_dirs) - return Eye(domain_type(T1), size(T1, 2)) + return Eye(domain_type(T1), size(T1, 2); array_type = domain_array_type(T1)) else - return SignAlternation(domain_type(T1), size(T1, 2), new_dirs) + return SignAlternation(domain_type(T1), size(T1, 2), new_dirs; array_type = domain_array_type(T1)) end end diff --git a/NFFTOperators/Project.toml b/NFFTOperators/Project.toml index 21324eb..6075a65 100644 --- a/NFFTOperators/Project.toml +++ b/NFFTOperators/Project.toml @@ -11,10 +11,19 @@ NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" NFFTTools = "7424e34d-94f7-41d6-98a0-85abaf1b6c91" Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588" +[weakdeps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" + +[extensions] +NFFTOperatorsGPUArraysExt = ["Adapt", "GPUArrays"] + [compat] AbstractOperators = "0.4" +Adapt = "4" FFTW = "1.8.1" FastBroadcast = "0.3.5" +GPUArrays = "11" LinearAlgebra = "1.10.0" NFFT = "0.13.7" NFFTTools = "0.2.6" diff --git a/NFFTOperators/README.md b/NFFTOperators/README.md index 0d229f4..9c9297c 100644 --- a/NFFTOperators/README.md +++ b/NFFTOperators/README.md @@ -19,6 +19,10 @@ NFFTOperators.jl is a **subpackage** of the AbstractOperators.jl ecosystem. Whil pkg> add NFFTOperators ``` +## GPU Support + +NFFTOperators.jl supports GPU execution only with CUDA.jl. The GPU implementation relies on CUDA-specific FFT planning, so the package does not currently support AMDGPU.jl, oneAPI.jl, or OpenCL.jl backends. + ## Usage Example ```julia diff --git a/NFFTOperators/ext/NFFTOperatorsGPUArraysExt.jl b/NFFTOperators/ext/NFFTOperatorsGPUArraysExt.jl new file mode 100644 index 0000000..6c2e197 --- /dev/null +++ b/NFFTOperators/ext/NFFTOperatorsGPUArraysExt.jl @@ -0,0 +1,35 @@ +module NFFTOperatorsGPUArraysExt + +using NFFTOperators +import NFFTOperators: _nfft_plan, _nfft_adapt, NFFTOp +import NFFTOperators.NFFT: NFFT +using GPUArrays +using Adapt + +""" + _nfft_plan(array_type, trajectory, image_size, threaded; kwargs...) + +GPU override: creates a GPU NFFT plan via `NFFT.plan_nfft(array_type, ...)`. +The trajectory must be a CPU `Matrix{T}`; only the computation buffers are on GPU. +Threading is ignored for GPU plans (GPU parallelism is used instead). +""" +function _nfft_plan( + array_type::Type{<:AbstractGPUArray}, + trajectory::AbstractArray{T}, + image_size, + threaded; + kwargs..., + ) where {T} + traj = Matrix{T}(reshape(trajectory, size(trajectory, 1), :)) + N = image_size + return NFFT.plan_nfft(array_type, traj, N; kwargs...) +end + +""" + _nfft_adapt(array_type, arr) + +GPU override: adapts a CPU array to the target GPU array type using Adapt.jl. +""" +_nfft_adapt(array_type::Type{<:AbstractGPUArray}, arr::AbstractArray) = adapt(array_type, arr) + +end # module NFFTOperatorsGPUArraysExt diff --git a/NFFTOperators/src/NFFTOp.jl b/NFFTOperators/src/NFFTOp.jl index d3567f2..a4e2dae 100644 --- a/NFFTOperators/src/NFFTOp.jl +++ b/NFFTOperators/src/NFFTOp.jl @@ -68,33 +68,45 @@ function NFFTOp( trajectory::AbstractArray{T}, dcf::AbstractArray; threaded::Bool = true, + array_type::Type = Array{T}, kwargs..., ) where {T, D} check_traj_and_dcf(trajectory, dcf, D) - plan = create_plan(trajectory, image_size, threaded; kwargs...) - ksp_buffer = similar( - trajectory, complex(eltype(trajectory)), size(trajectory)[2:end]... - ) - return NFFTOp{T, D, typeof(plan), typeof(ksp_buffer), typeof(dcf)}(plan, ksp_buffer, dcf, threaded) + arr_wrapper = _array_wrapper_type(array_type) + plan = _nfft_plan(arr_wrapper, trajectory, image_size, threaded; kwargs...) + ksp_shape = size(trajectory)[2:end] + ksp_buffer = _nfft_adapt(arr_wrapper, zeros(complex(T), ksp_shape...)) + adapted_dcf = _nfft_adapt(arr_wrapper, collect(dcf)) + threaded_flag = threaded && arr_wrapper === Array + return NFFTOp{T, D, typeof(plan), typeof(ksp_buffer), typeof(adapted_dcf)}(plan, ksp_buffer, adapted_dcf, threaded_flag) end function NFFTOp( image_size::NTuple{D, Int}, trajectory::AbstractArray{T}; threaded::Bool = true, + array_type::Type = Array{T}, dcf_estimation_iterations::Int = 20, dcf_correction_function::Function = identity, kwargs..., ) where {T, D} check_traj(trajectory, D) - plan = create_plan(trajectory, image_size, threaded; kwargs...) - ksp_buffer = similar( - trajectory, complex(eltype(trajectory)), size(trajectory)[2:end]... - ) + arr_wrapper = _array_wrapper_type(array_type) + plan = _nfft_plan(arr_wrapper, trajectory, image_size, threaded; kwargs...) + ksp_shape = size(trajectory)[2:end] + ksp_buffer = _nfft_adapt(arr_wrapper, zeros(complex(T), ksp_shape...)) raw_dcf = NFFTTools.sdc(plan; iters = dcf_estimation_iterations) - dcf = dcf_correction_function(reshape(raw_dcf, size(ksp_buffer))) - return NFFTOp{T, D, typeof(plan), typeof(ksp_buffer), typeof(dcf)}(plan, ksp_buffer, dcf, threaded) + dcf_cpu = dcf_correction_function(reshape(raw_dcf, ksp_shape)) + adapted_dcf = _nfft_adapt(arr_wrapper, collect(dcf_cpu)) + threaded_flag = threaded && arr_wrapper === Array + return NFFTOp{T, D, typeof(plan), typeof(ksp_buffer), typeof(adapted_dcf)}(plan, ksp_buffer, adapted_dcf, threaded_flag) +end + +# Default (CPU) implementations — overridden by NFFTOperatorsGPUArraysExt for GPU types +function _nfft_plan(::Type{Array}, trajectory, image_size, threaded; kwargs...) + return create_plan(trajectory, image_size, threaded; kwargs...) end +_nfft_adapt(::Type{Array}, arr::AbstractArray) = collect(arr) function set_nfft_threading_expr(threading_state_expr, thread_count_expr, body_expr) return quote @@ -126,18 +138,19 @@ end function mul!( img::AbstractArray, - op::AbstractOperators.AdjointOperator{<:NFFTOp}, + adjop::AbstractOperators.AdjointOperator{<:NFFTOp}, ksp::AbstractArray, ) - op = op.A - AbstractOperators.check(ksp, op, img) - return if op.threaded + AbstractOperators.check(img, adjop, ksp) + op = adjop.A + if op.threaded @.. thread = true op.ksp_buffer = ksp * op.dcf @enable_nfft_threading mul!(img, op.plan', vec(op.ksp_buffer)) else @.. op.ksp_buffer = ksp * op.dcf @disable_nfft_threading mul!(img, op.plan', vec(op.ksp_buffer)) end + return img end # Properties @@ -146,6 +159,8 @@ size(L::NFFTOp) = size(L.ksp_buffer), NFFT.size_in(L.plan) fun_name(::NFFTOp) = "𝒩" domain_type(::NFFTOp{T}) where {T} = complex(T) codomain_type(::NFFTOp{T}) where {T} = complex(T) +domain_array_type(op::NFFTOp) = typeof(op.plan.tmpVec) +codomain_array_type(op::NFFTOp{T, D, P, K}) where {T, D, P, K} = K # Utility diff --git a/NFFTOperators/src/NFFTOperators.jl b/NFFTOperators/src/NFFTOperators.jl index 420a39a..a760d97 100644 --- a/NFFTOperators/src/NFFTOperators.jl +++ b/NFFTOperators/src/NFFTOperators.jl @@ -18,10 +18,12 @@ import AbstractOperators: get_normal_op, allocate_in_domain, allocate_in_codomain, - domain_storage_type, - codomain_storage_type, + domain_array_type, + codomain_array_type, + _array_wrapper_type, AdjointOperator import Base.Threads: nthreads + import FFTW: FFTW include("NFFTOp.jl") diff --git a/NFFTOperators/src/NormalNfftOp.jl b/NFFTOperators/src/NormalNfftOp.jl index 5a2612c..609f820 100644 --- a/NFFTOperators/src/NormalNfftOp.jl +++ b/NFFTOperators/src/NormalNfftOp.jl @@ -26,7 +26,7 @@ # SOFTWARE. struct NfftNormalOp{N, T, A <: AbstractArray, F, I} <: AbstractOperators.LinearOperator - storage_type::Type{T} + array_type::Type{T} shape::NTuple{N, Int} λ::A fftplan::F @@ -36,6 +36,7 @@ struct NfftNormalOp{N, T, A <: AbstractArray, F, I} <: AbstractOperators.LinearO end function mul!(y::AbstractArray, op::NfftNormalOp, x::AbstractArray) + AbstractOperators.check(y, op, x) return if op.threaded @enable_nfft_threading _mul!(y, op, x) else @@ -67,8 +68,8 @@ function _get_normal_op(op::NFFTOp) shape_ext = 2 .* shape buf = allocate_in_domain(op, shape_ext...) - buf .= 0 - tmp = allocate_in_codomain(op, size(op.plan.k, 2)...) + fill!(buf, 0) + tmp = allocate_in_codomain(op, size(op.plan.k, 2)) tmp .= vec(op.dcf) fftplan = FFTW.plan_fft!(buf) @@ -85,7 +86,7 @@ function _get_normal_op(op::NFFTOp) mul!(buf, adjoint(p), tmp) λ = fftplan * FFTW.fftshift(buf) # create a new array by fftshift and apply in-place FFT - return NfftNormalOp(domain_storage_type(op), shape, λ, fftplan, inv(fftplan), buf, op.threaded) + return NfftNormalOp(domain_array_type(op), shape, λ, fftplan, inv(fftplan), buf, op.threaded) end # properties @@ -94,7 +95,7 @@ Base.size(op::NfftNormalOp) = op.shape, op.shape AbstractOperators.fun_name(::NfftNormalOp) = "(𝒩ᵃ𝒩)" domain_type(::NfftNormalOp{N, T}) where {N, T} = eltype(T) codomain_type(::NfftNormalOp{N, T}) where {N, T} = eltype(T) -domain_storage_type(op::NfftNormalOp{N, T}) where {N, T} = op.storage_type -codomain_storage_type(op::NfftNormalOp{N, T}) where {N, T} = op.storage_type +domain_array_type(op::NfftNormalOp{N, T}) where {N, T} = op.array_type +codomain_array_type(op::NfftNormalOp{N, T}) where {N, T} = op.array_type is_symmetric(::NfftNormalOp) = true AdjointOperator(op::NfftNormalOp) = op diff --git a/Project.toml b/Project.toml index 93b472e..1943dd0 100644 --- a/Project.toml +++ b/Project.toml @@ -15,15 +15,17 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" [weakdeps] GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e" [extensions] -GpuExt = "GPUArrays" +GpuExt = ["GPUArrays", "KernelAbstractions"] LinearMapsExt = "LinearMaps" [compat] FastBroadcast = "0.3.5" -GPUArrays = "11.2.3" +GPUArrays = "10, 11" +KernelAbstractions = "0.9" LinearAlgebra = "1" LinearMaps = "3.11.4" OperatorCore = "0.1" diff --git a/README.md b/README.md index 3d77ce8..91818e5 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ [![Build status](https://github.com/kul-optec/AbstractOperators.jl/workflows/CI/badge.svg)](https://github.com/kul-optec/AbstractOperators.jl/actions?query=workflow%3ACI) [![codecov](https://codecov.io/gh/kul-optec/AbstractOperators.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/kul-optec/AbstractOperators.jl) [![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl) -[![code style blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle) +[![Runic formatter](https://img.shields.io/badge/code_style-%E1%9A%B1%E1%9A%A2%E1%9A%BE%E1%9B%81%E1%9A%B2-black)](https://github.com/fredrikekre/Runic.jl) [![](https://img.shields.io/badge/docs-stable-blue.svg)](https://kul-optec.github.io/AbstractOperators.jl/stable) [![](https://img.shields.io/badge/docs-latest-blue.svg)](https://kul-optec.github.io/AbstractOperators.jl/latest) @@ -123,6 +123,16 @@ A list of the available `AbstractOperators` and calculus rules can be found in t `AbstractOperators.jl` is distinguished by its support for multi-dimensional array domains and codomains, efficient in-place implementations of both linear and nonlinear operators, and seamless integration with optimization algorithms in related packages: [ProximalOperators.jl](https://github.com/kul-forbes/ProximalOperators.jl), [ProximalAlgorithms.jl](https://github.com/kul-forbes/ProximalAlgorithms.jl), and [StructuredOptimization.jl](https://github.com/kul-forbes/StructuredOptimization.jl). It has built-in threading support for many operators, and partial (and extending) GPU support. +## GPU Support + +Most operators in the AbstractOperators.jl ecosystem work with GPU arrays on CUDA.jl, AMDGPU.jl, oneAPI.jl, and OpenCL.jl. + +Subpackage-specific exceptions are: + +- `NFFTOperators.jl` supports GPU execution only with CUDA.jl +- `WaveletOperators.jl` currently works on CPU only +- `DCT` and `IDCT` (in FFTWOperators.jl) use CPU FFTW plans by default, but load `AcceleratedDCTs` to activate GPU support for those operators + ## Credits AbstractOperators.jl is developed by diff --git a/WaveletOperators/README.md b/WaveletOperators/README.md index 5843ed6..1d01b11 100644 --- a/WaveletOperators/README.md +++ b/WaveletOperators/README.md @@ -19,6 +19,10 @@ WaveletOperators.jl is a **subpackage** of the AbstractOperators.jl ecosystem. W pkg> add WaveletOperators ``` +## GPU Support + +WaveletOperators.jl is currently CPU-only. Unlike most of the AbstractOperators.jl ecosystem, its operators do not yet support CUDA.jl, AMDGPU.jl, oneAPI.jl, or OpenCL.jl arrays. + ## Usage Example ```julia diff --git a/WaveletOperators/src/WaveletOperators.jl b/WaveletOperators/src/WaveletOperators.jl index f1f4715..1c75ad4 100644 --- a/WaveletOperators/src/WaveletOperators.jl +++ b/WaveletOperators/src/WaveletOperators.jl @@ -9,6 +9,8 @@ import Base: size import AbstractOperators: domain_type, codomain_type, + domain_array_type, + codomain_array_type, fun_name, is_thread_safe, has_fast_opnorm @@ -85,12 +87,14 @@ end # Mappings function mul!(y::AbstractArray{T}, L::WaveletOp{T}, x::AbstractArray{T}) where {T} + AbstractOperators.check(y, L, x) return dwt!(y, x, L.wavelet, L.levels) end function mul!( y::AbstractArray{T}, L::AdjointOperator{<:WaveletOp{T}}, x::AbstractArray{T} ) where {T} + AbstractOperators.check(y, L, x) return idwt!(y, x, L.A.wavelet, L.A.levels) end @@ -102,6 +106,8 @@ size(L::WaveletOp) = (L.dim_in, L.dim_in) domain_type(::WaveletOp{T}) where {T} = T codomain_type(::WaveletOp{T}) where {T} = T +domain_array_type(::WaveletOp{T}) where {T} = Array{T} +codomain_array_type(::WaveletOp{T}) where {T} = Array{T} is_AcA_diagonal(L::WaveletOp) = true is_AAc_diagonal(L::WaveletOp) = true diff --git a/benchmark/Project.toml b/benchmark/Project.toml index 61a2a25..fea1db0 100644 --- a/benchmark/Project.toml +++ b/benchmark/Project.toml @@ -4,6 +4,7 @@ ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" DSPOperators = "d5a72628-6e2f-430e-82f5-561df0bb8116" FFTWOperators = "c59a084b-ba08-4f3f-af9e-f4298d6caa94" +GPUEnv = "78a0b619-6146-4252-b244-0f81c54be577" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NFFTOperators = "35a93202-2abf-4a13-bc94-0aee2f861ee0" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 1c0deb8..8a40c98 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -5,6 +5,8 @@ using Pkg using Random using RecursiveArrayTools: ArrayPartition +make_rng() = MersenneTwister(1234) + function _load_local_subpackage(pkg::Symbol, relpath::String) try @eval using $pkg @@ -18,7 +20,7 @@ function _load_local_subpackage(pkg::Symbol, relpath::String) return false end try - Pkg.develop(path = local_path) + Pkg.develop(; path = local_path) @eval using $pkg return true catch @@ -72,8 +74,6 @@ const BENCH_NFFT_IMAGE = (48, 48) const BENCH_NFFT_NSAMP = 48 const BENCH_NFFT_NPROF = 24 -make_rng() = MersenneTwister(1234) - function linear_state(op) rng = make_rng() x = randn(rng, domain_type(op), size(op, 2)...) @@ -84,14 +84,22 @@ end function nonlinear_state(op; positive = false) rng = make_rng() - x = positive ? abs.(randn(rng, domain_type(op), size(op, 2)...)) : randn(rng, domain_type(op), size(op, 2)...) + x = if positive + abs.(randn(rng, domain_type(op), size(op, 2)...)) + else + randn(rng, domain_type(op), size(op, 2)...) + end y = zeros(codomain_type(op), size(op, 1)...) return (op = op, x = x, y = y) end function jacobian_state(op; positive = false) rng = make_rng() - x = positive ? abs.(randn(rng, domain_type(op), size(op, 2)...)) : randn(rng, domain_type(op), size(op, 2)...) + x = if positive + abs.(randn(rng, domain_type(op), size(op, 2)...)) + else + randn(rng, domain_type(op), size(op, 2)...) + end jac = Jacobian(op, x) b = randn(rng, codomain_type(jac), size(jac, 1)...) y = zeros(domain_type(jac), size(jac, 2)...) @@ -177,26 +185,38 @@ hadamardprod_jacobian_state() = jacobian_state(HadamardProd(Sin((BENCH_CALC_N,)) function ax_mul_bxt_state() rng = make_rng() # Ax_mul_Bxt(A,B): computes A(x) * B(x)'. Requires same domain and A,B output same codomain. - op = Ax_mul_Bxt(MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ)) + op = Ax_mul_Bxt( + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + ) return nonlinear_state(op) end function ax_mul_bxt_jacobian_state() rng = make_rng() - op = Ax_mul_Bxt(MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ)) + op = Ax_mul_Bxt( + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + ) return jacobian_state(op) end function axt_mul_bx_state() rng = make_rng() # Axt_mul_Bx(A,B): computes A(x)' * B(x). Requires A and B share domain; rows(A)==rows(B). - op = Axt_mul_Bx(MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ)) + op = Axt_mul_Bx( + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + ) return nonlinear_state(op) end function axt_mul_bx_jacobian_state() rng = make_rng() - op = Axt_mul_Bx(MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ)) + op = Axt_mul_Bx( + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + ) return jacobian_state(op) end @@ -204,13 +224,19 @@ function ax_mul_bx_state() rng = make_rng() # Ax_mul_Bx(A,B): computes A(x)*B(x). Requires size(A,1)[2] == size(B,1)[1]. # Requires square codomain shape where col(A) == row(B). - op = Ax_mul_Bx(MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ)) + op = Ax_mul_Bx( + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + ) return nonlinear_state(op) end function ax_mul_bx_jacobian_state() rng = make_rng() - op = Ax_mul_Bx(MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ)) + op = Ax_mul_Bx( + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + MatrixOp(randn(rng, BENCH_CALC_SQ, BENCH_CALC_SQ), BENCH_CALC_SQ), + ) return jacobian_state(op) end @@ -294,13 +320,19 @@ function nfft_state() end function normal_state(op) - nop = AbstractOperators.get_normal_op(op) rng = make_rng() - x = randn(rng, domain_type(nop), size(nop, 2)...) - y = zeros(codomain_type(nop), size(nop, 1)...) + nop = AbstractOperators.get_normal_op(op) + dT = domain_type(nop) + cT = codomain_type(nop) + dEl = dT <: AbstractArray ? eltype(dT) : dT + cEl = cT <: AbstractArray ? eltype(cT) : cT + x = randn(rng, dEl, size(nop, 2)...) + y = zeros(cEl, size(nop, 1)...) return (op = nop, x = x, y = y) end +nfft_normal_state() = normal_state(nfft_state().op) + linear = SUITE["linearoperators"] = BenchmarkGroup() calculus = SUITE["calculus"] = BenchmarkGroup() nonlinear = SUITE["nonlinearoperators"] = BenchmarkGroup() @@ -315,16 +347,16 @@ linear["Eye"] = BenchmarkGroup() linear["Eye"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(Eye(Float64, (BENCH_LINEAR_EYE_N,)))) linear["DiagOp"] = BenchmarkGroup() -linear["DiagOp"]["forward-single"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(DiagOp(randn(make_rng(), BENCH_LINEAR_DIAG_N); threaded = false))) -linear["DiagOp"]["adjoint-single"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(DiagOp(randn(make_rng(), BENCH_LINEAR_DIAG_N); threaded = false))) +linear["DiagOp"]["forward-single"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = linear_state(DiagOp(randn(rng, BENCH_LINEAR_DIAG_N); threaded = false))) +linear["DiagOp"]["adjoint-single"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (rng = make_rng(); state = linear_state(DiagOp(randn(rng, BENCH_LINEAR_DIAG_N); threaded = false))) if Threads.nthreads() > 1 - linear["DiagOp"]["forward-threaded"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(DiagOp(randn(make_rng(), BENCH_LINEAR_DIAG_N); threaded = true))) - linear["DiagOp"]["adjoint-threaded"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(DiagOp(randn(make_rng(), BENCH_LINEAR_DIAG_N); threaded = true))) + linear["DiagOp"]["forward-threaded"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = linear_state(DiagOp(randn(rng, BENCH_LINEAR_DIAG_N); threaded = true))) + linear["DiagOp"]["adjoint-threaded"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (rng = make_rng(); state = linear_state(DiagOp(randn(rng, BENCH_LINEAR_DIAG_N); threaded = true))) end linear["MatrixOp"] = BenchmarkGroup() -linear["MatrixOp"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(MatrixOp(randn(make_rng(), BENCH_LINEAR_MATRIX_SHAPE...), BENCH_LINEAR_MATRIX_DOMAIN))) -linear["MatrixOp"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(MatrixOp(randn(make_rng(), BENCH_LINEAR_MATRIX_SHAPE...), BENCH_LINEAR_MATRIX_DOMAIN))) +linear["MatrixOp"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = linear_state(MatrixOp(randn(rng, BENCH_LINEAR_MATRIX_SHAPE...), BENCH_LINEAR_MATRIX_DOMAIN))) +linear["MatrixOp"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (rng = make_rng(); state = linear_state(MatrixOp(randn(rng, BENCH_LINEAR_MATRIX_SHAPE...), BENCH_LINEAR_MATRIX_DOMAIN))) linear["FiniteDiff"] = BenchmarkGroup() linear["FiniteDiff"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(FiniteDiff(Float64, (BENCH_LINEAR_FD_N,), 1))) @@ -350,8 +382,8 @@ linear["Zeros"] = BenchmarkGroup() linear["Zeros"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(Zeros(Float64, (BENCH_LINEAR_ZEROS_N,), Float64, (BENCH_LINEAR_ZEROS_N,)))) linear["LMatrixOp"] = BenchmarkGroup() -linear["LMatrixOp"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(LMatrixOp(randn(make_rng(), BENCH_LINEAR_LMATRIX_N), BENCH_LINEAR_LMATRIX_N))) -linear["LMatrixOp"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(LMatrixOp(randn(make_rng(), BENCH_LINEAR_LMATRIX_N), BENCH_LINEAR_LMATRIX_N))) +linear["LMatrixOp"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = linear_state(LMatrixOp(randn(rng, BENCH_LINEAR_LMATRIX_N), BENCH_LINEAR_LMATRIX_N))) +linear["LMatrixOp"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (rng = make_rng(); state = linear_state(LMatrixOp(randn(rng, BENCH_LINEAR_LMATRIX_N), BENCH_LINEAR_LMATRIX_N))) linear["MyLinOp"] = BenchmarkGroup() linear["MyLinOp"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = mylinop_state()) @@ -362,8 +394,8 @@ linear["LBFGS"]["update"] = @benchmarkable update!(state.op, state.x, state.x_pr linear["LBFGS"]["mul"] = @benchmarkable mul!(state.d, state.op, state.grad) setup = (state = lbfgs_mul_state()) calculus["Compose"] = BenchmarkGroup() -calculus["Compose"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(Compose(DiagOp(randn(make_rng(), BENCH_CALC_N)), Eye(Float64, (BENCH_CALC_N,))))) -calculus["Compose"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(Compose(DiagOp(randn(make_rng(), BENCH_CALC_N)), Eye(Float64, (BENCH_CALC_N,))))) +calculus["Compose"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = linear_state(Compose(DiagOp(randn(rng, BENCH_CALC_N)), Eye(Float64, (BENCH_CALC_N,))))) +calculus["Compose"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (rng = make_rng(); state = linear_state(Compose(DiagOp(randn(rng, BENCH_CALC_N)), Eye(Float64, (BENCH_CALC_N,))))) calculus["Reshape"] = BenchmarkGroup() calculus["Reshape"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(Reshape(Eye(Float64, (BENCH_CALC_N,)), BENCH_CALC_2D...))) @@ -374,8 +406,8 @@ calculus["Scale"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) s calculus["Scale"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(Scale(2.0, Eye(Float64, (BENCH_CALC_N,))))) calculus["Sum"] = BenchmarkGroup() -calculus["Sum"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(Sum(Eye(Float64, (BENCH_CALC_N,)), DiagOp(randn(make_rng(), BENCH_CALC_N))))) -calculus["Sum"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(Sum(Eye(Float64, (BENCH_CALC_N,)), DiagOp(randn(make_rng(), BENCH_CALC_N))))) +calculus["Sum"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = linear_state(Sum(Eye(Float64, (BENCH_CALC_N,)), DiagOp(randn(rng, BENCH_CALC_N))))) +calculus["Sum"]["adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (rng = make_rng(); state = linear_state(Sum(Eye(Float64, (BENCH_CALC_N,)), DiagOp(randn(rng, BENCH_CALC_N))))) calculus["HCAT"] = BenchmarkGroup() calculus["HCAT"]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = hcat_state()) @@ -394,11 +426,11 @@ calculus["BroadCast"]["identity-single"] = @benchmarkable mul!(state.y, state.op if Threads.nthreads() > 2 calculus["BroadCast"]["identity-threaded"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(BroadCast(Eye(Float64, (BENCH_CALC_N,)), (BENCH_CALC_N, 8); threaded = true))) end -calculus["BroadCast"]["operator-single-forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(BroadCast(DiagOp(randn(make_rng(), 256)), (256, 8); threaded = false))) -calculus["BroadCast"]["operator-single-adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(BroadCast(DiagOp(randn(make_rng(), 256)), (256, 8); threaded = false))) +calculus["BroadCast"]["operator-single-forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = linear_state(BroadCast(DiagOp(randn(rng, 256)), (256, 8); threaded = false))) +calculus["BroadCast"]["operator-single-adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (rng = make_rng(); state = linear_state(BroadCast(DiagOp(randn(rng, 256)), (256, 8); threaded = false))) if Threads.nthreads() > 2 - calculus["BroadCast"]["operator-threaded-forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = linear_state(BroadCast(DiagOp(randn(make_rng(), 256)), (256, 8); threaded = true))) - calculus["BroadCast"]["operator-threaded-adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = linear_state(BroadCast(DiagOp(randn(make_rng(), 256)), (256, 8); threaded = true))) + calculus["BroadCast"]["operator-threaded-forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = linear_state(BroadCast(DiagOp(randn(rng, 256)), (256, 8); threaded = true))) + calculus["BroadCast"]["operator-threaded-adjoint"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (rng = make_rng(); state = linear_state(BroadCast(DiagOp(randn(rng, 256)), (256, 8); threaded = true))) end calculus["AffineAdd"] = BenchmarkGroup() @@ -437,8 +469,12 @@ for (name, builder, positive) in [ ("SoftPlus", () -> SoftPlus(Float64, (BENCH_NONLIN_N["SoftPlus"],)), false), ] nonlinear[name] = BenchmarkGroup() - nonlinear[name]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = nonlinear_state($builder(); positive = $positive)) - nonlinear[name]["jacobian-adjoint"] = @benchmarkable mul!(state.y, state.adj, state.b) setup = (state = jacobian_state($builder(); positive = $positive)) + nonlinear[name]["forward"] = @benchmarkable mul!(state.y, state.op, state.x) setup = ( + state = nonlinear_state($builder(); positive = ($positive)) + ) + nonlinear[name]["jacobian-adjoint"] = @benchmarkable mul!(state.y, state.adj, state.b) setup = ( + state = jacobian_state($builder(); positive = ($positive)) + ) end batching["SimpleBatchOp"] = BenchmarkGroup() @@ -455,8 +491,12 @@ batching["SpreadingBatchOp"]["adjoint-single"] = @benchmarkable mul!(state.z, st if Threads.nthreads() > 2 for strategy in (ThreadingStrategy.COPYING, ThreadingStrategy.LOCKING, ThreadingStrategy.FIXED_OPERATOR) strategy_name = String(Symbol(strategy)) - batching["SpreadingBatchOp"]["forward-$(strategy_name)"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = spreading_batch_state(true; strategy = $strategy)) - batching["SpreadingBatchOp"]["adjoint-$(strategy_name)"] = @benchmarkable mul!(state.z, state.adj, state.y) setup = (state = spreading_batch_state(true; strategy = $strategy)) + batching["SpreadingBatchOp"]["forward-$(strategy_name)"] = @benchmarkable mul!( + state.y, state.op, state.x + ) setup = (state = spreading_batch_state(true; strategy = ($strategy))) + batching["SpreadingBatchOp"]["adjoint-$(strategy_name)"] = @benchmarkable mul!( + state.z, state.adj, state.y + ) setup = (state = spreading_batch_state(true; strategy = ($strategy))) end end @@ -493,7 +533,7 @@ if HAS_WAVELET end normal["DiagOp"] = BenchmarkGroup() -normal["DiagOp"]["mul"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = normal_state(DiagOp(randn(make_rng(), BENCH_LINEAR_DIAG_N); threaded = false))) +normal["DiagOp"]["mul"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (rng = make_rng(); state = normal_state(DiagOp(randn(rng, BENCH_LINEAR_DIAG_N); threaded = false))) if HAS_FFTW normal["DFT"] = BenchmarkGroup() @@ -502,7 +542,7 @@ end if HAS_NFFT normal["NFFTOp"] = BenchmarkGroup() - normal["NFFTOp"]["mul"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = normal_state(nfft_state().op)) + normal["NFFTOp"]["mul"] = @benchmarkable mul!(state.y, state.op, state.x) setup = (state = nfft_normal_state()) end # Cap run time so CI / ASV comparisons complete in reasonable time diff --git a/benchmark/gpu_crossover.jl b/benchmark/gpu_crossover.jl new file mode 100644 index 0000000..8773814 --- /dev/null +++ b/benchmark/gpu_crossover.jl @@ -0,0 +1,853 @@ +# Run with: julia --project=benchmark benchmark/gpu_crossover.jl +using GPUEnv + +if VERSION < v"1.11" + println("AbstractOperators GPU crossover benchmark requires Julia 1.11 or newer. Skipping benchmark.") + exit(0) +end + +GPUEnv.activate(persist = true, include_jlarrays = false, only_first = true, supports_fftw = true) + +using AbstractOperators +using BenchmarkTools +using DSPOperators +using FFTWOperators +using LinearAlgebra +using Random +using RecursiveArrayTools: ArrayPartition + +# ── GPU backend detection ──────────────────────────────────────────────────── +# GPUEnv.activate installed available native backends; pick the first one. +const _backend = let + bs = gpu_backends() + if isempty(bs) + @error "No functional GPU backend found" + exit(0) + end + first(bs) +end + +@info "GPU benchmark" backend = _backend.name + +# ── Timing utilities ───────────────────────────────────────────────────────── +function _measure!(y, op, x, is_gpu::Bool; samples::Int) + mul!(y, op, x) + is_gpu && GPUEnv.synchronize_backend(_backend) + return @belapsed begin + mul!($y, $op, $x) + $is_gpu && GPUEnv.synchronize_backend($_backend) + end samples = samples evals = 1 +end + +function _threaded_time(threaded_builder, n; samples) + threaded_builder === nothing && return missing + op, x, y = threaded_builder(n) + return _measure!(y, op, x, false; samples) +end + +function _gpu_time(gpu_builder, n; samples) + op, x, y = gpu_builder(n) + return _measure!(y, op, x, true; samples) +end + +function _shape2(n) + m = max(8, round(Int, sqrt(n))) + return (m, max(8, cld(n, m))) +end + +function _broad_shape(n) + base = max(128, n) + return (base, 8) +end + +function _diag_cpu(n; threaded) + d = randn(Float32, n) + op = DiagOp(d; threaded) + x = randn(Float32, n) + y = zeros(Float32, n) + return op, x, y +end + +function _diag_gpu(n) + d = gpu_randn(_backend, Float32, n) + op = DiagOp(d; threaded = false) + x = gpu_randn(_backend, Float32, n) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +function _finitediff_cpu(n; threaded) + op = FiniteDiff(Float32, (n,), 1) + x = randn(Float32, n) + y = zeros(Float32, n - 1) + return op, x, y +end + +function _finitediff_gpu(n) + x = gpu_randn(_backend, Float32, n) + op = FiniteDiff(x, 1) + y = gpu_zeros(_backend, Float32, n - 1) + return op, x, y +end + +function _getindex_cpu(n; threaded) + dims = _shape2(n) + x = randn(Float32, dims...) + idx = (2:(dims[1] - 1), 2:(dims[2] - 1)) + op = GetIndex(Float32, dims, idx) + y = zeros(Float32, size(op, 1)...) + return op, x, y +end + +function _getindex_gpu(n) + dims = _shape2(n) + x = gpu_randn(_backend, Float32, dims...) + idx = (2:(dims[1] - 1), 2:(dims[2] - 1)) + op = GetIndex(x, idx) + y = gpu_zeros(_backend, Float32, size(op, 1)...) + return op, x, y +end + +# GetIndex: bool-mask variant +function _getindex_boolmask_cpu(n; threaded) + x = randn(Float32, n) + mask = [i % 3 == 0 for i in 1:n] + op = GetIndex(Float32, (n,), mask) + y = zeros(Float32, size(op, 1)...) + return op, x, y +end + +function _getindex_boolmask_gpu(n) + x = gpu_randn(_backend, Float32, n) + mask = [i % 3 == 0 for i in 1:n] + op = GetIndex(x, mask) + y = gpu_zeros(_backend, Float32, size(op, 1)...) + return op, x, y +end + +# GetIndex: integer-vector variant (select every 2nd element → output ≈ n/2) +function _getindex_intvec_cpu(n; threaded) + x = randn(Float32, n) + idx = collect(1:2:n) + op = GetIndex(Float32, (n,), idx) + y = zeros(Float32, size(op, 1)...) + return op, x, y +end + +function _getindex_intvec_gpu(n) + x = gpu_randn(_backend, Float32, n) + idx = collect(1:2:n) + op = GetIndex(x, idx) + y = gpu_zeros(_backend, Float32, size(op, 1)...) + return op, x, y +end + +# Eye +function _eye_cpu(n; threaded) + x = randn(Float32, n) + op = Eye(Float32, (n,)) + y = zeros(Float32, n) + return op, x, y +end + +function _eye_gpu(n) + x = gpu_randn(_backend, Float32, n) + op = Eye(x) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +# MatrixOp (dense matrix-vector product) +function _matrixop_cpu(n; threaded) + A = randn(Float32, n, n) + op = MatrixOp(Float32, (n,), A) + x = randn(Float32, n) + y = zeros(Float32, n) + return op, x, y +end + +function _matrixop_gpu(n) + A = gpu_randn(_backend, Float32, n, n) + op = MatrixOp(A) + x = gpu_randn(_backend, Float32, n) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +# Scale (scales an Eye operator by 2) +function _scale_cpu(n; threaded) + A = Eye(Float32, (n,)) + op = Scale(Float32(2.0), A; threaded) + x = randn(Float32, n) + y = zeros(Float32, n) + return op, x, y +end + +function _scale_gpu(n) + x = gpu_randn(_backend, Float32, n) + A = Eye(x) + op = Scale(Float32(2.0), A; threaded = false) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +# HadamardProd (Sin .* Cos) +function _hadamardprod_cpu(n; threaded) + A = Sin(Float32, (n,)) + B = Cos(Float32, (n,)) + op = HadamardProd(A, B) + x = randn(Float32, n) + y = zeros(Float32, n) + return op, x, y +end + +function _hadamardprod_gpu(n) + x = gpu_randn(_backend, Float32, n) + op = HadamardProd(Sin(x), Cos(x)) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +# Nonlinear: Sin +function _sin_cpu(n; threaded) + x = randn(Float32, n) + op = Sin(Float32, (n,)) + y = zeros(Float32, n) + return op, x, y +end + +function _sin_gpu(n) + x = gpu_randn(_backend, Float32, n) + op = Sin(x) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +# Nonlinear: Tanh +function _tanh_cpu(n; threaded) + x = randn(Float32, n) + op = Tanh(Float32, (n,)) + y = zeros(Float32, n) + return op, x, y +end + +function _tanh_gpu(n) + x = gpu_randn(_backend, Float32, n) + op = Tanh(x) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +# Nonlinear: Exp +function _exp_cpu(n; threaded) + x = randn(Float32, n) + op = Exp(Float32, (n,)) + y = zeros(Float32, n) + return op, x, y +end + +function _exp_gpu(n) + x = gpu_randn(_backend, Float32, n) + op = Exp(x) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +# DCAT +function _dcat_cpu(n; threaded) + x1 = randn(Float32, n) + x2 = randn(Float32, n) + x = ArrayPartition(x1, x2) + op = DCAT(Eye(Float32, (n,)), DiagOp(randn(Float32, n); threaded = false)) + y = ArrayPartition(zeros(Float32, n), zeros(Float32, n)) + return op, x, y +end + +function _dcat_gpu(n) + x1 = gpu_randn(_backend, Float32, n) + x2 = gpu_randn(_backend, Float32, n) + x = ArrayPartition(x1, x2) + op = DCAT(Eye(x1), DiagOp(gpu_randn(_backend, Float32, n); threaded = false)) + y = ArrayPartition(gpu_zeros(_backend, Float32, n), gpu_zeros(_backend, Float32, n)) + return op, x, y +end + +# AffineAdd +function _affineadd_cpu(n; threaded) + x = randn(Float32, n) + b = randn(Float32, n) + op = AffineAdd(DiagOp(randn(Float32, n); threaded = false), b) + y = zeros(Float32, n) + return op, x, y +end + +function _affineadd_gpu(n) + x = gpu_randn(_backend, Float32, n) + b = gpu_randn(_backend, Float32, n) + op = AffineAdd(DiagOp(gpu_randn(_backend, Float32, n); threaded = false), b) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +function _zeropad_cpu(n; threaded) + dims = _shape2(n) + x = randn(Float32, dims...) + op = ZeroPad(Float32, dims, (0, 8)) + y = zeros(Float32, size(op, 1)...) + return op, x, y +end + +function _zeropad_gpu(n) + dims = _shape2(n) + x = gpu_randn(_backend, Float32, dims...) + op = ZeroPad(x, (0, 8)) + y = gpu_zeros(_backend, Float32, size(op, 1)...) + return op, x, y +end + +function _variation_cpu(n; threaded) + dims = _shape2(n) + x = randn(Float32, dims...) + op = Variation(Float32, dims; threaded) + y = zeros(Float32, size(op, 1)...) + return op, x, y +end + +function _variation_gpu(n) + dims = _shape2(n) + x = gpu_randn(_backend, Float32, dims...) + op = Variation(x; threaded = false) + y = gpu_zeros(_backend, Float32, size(op, 1)...) + return op, x, y +end + +function _variation_adj_cpu(n; threaded) + dims = _shape2(n) + op = adjoint(Variation(Float32, dims; threaded)) + b = zeros(Float32, prod(dims), length(dims)) + y = zeros(Float32, dims...) + return op, b, y +end + +function _variation_adj_gpu(n) + dims = _shape2(n) + x = gpu_randn(_backend, Float32, dims...) + op = adjoint(Variation(x; threaded = false)) + b = gpu_zeros(_backend, Float32, prod(dims), length(dims)) + y = gpu_zeros(_backend, Float32, dims...) + return op, b, y +end + +function _compose_cpu(n; threaded) + d = randn(Float32, n - 1) + fd = FiniteDiff(Float32, (n,), 1) + op = Compose(DiagOp(d; threaded = false), fd) + x = randn(Float32, n) + y = zeros(Float32, n - 1) + return op, x, y +end + +function _compose_gpu(n) + d = gpu_randn(_backend, Float32, n - 1) + x = gpu_randn(_backend, Float32, n) + op = Compose(DiagOp(d; threaded = false), FiniteDiff(x, 1)) + y = gpu_zeros(_backend, Float32, n - 1) + return op, x, y +end + +function _sum_cpu(n; threaded) + d = randn(Float32, n) + op = Sum(Eye(Float32, (n,)), DiagOp(d; threaded = false)) + x = randn(Float32, n) + y = zeros(Float32, n) + return op, x, y +end + +function _sum_gpu(n) + x = gpu_randn(_backend, Float32, n) + d1 = gpu_randn(_backend, Float32, n) + d2 = gpu_randn(_backend, Float32, n) + op = Sum(DiagOp(d1; threaded = false), DiagOp(d2; threaded = false)) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +function _broadcast_cpu(n; threaded) + dims = _broad_shape(n) + x = randn(Float32, dims[1]) + op = BroadCast(Eye(Float32, (dims[1],)), dims; threaded) + y = zeros(Float32, dims...) + return op, x, y +end + +function _broadcast_gpu(n) + dims = _broad_shape(n) + x = gpu_randn(_backend, Float32, dims[1]) + op = BroadCast(Eye(x), dims; threaded = false) + y = gpu_zeros(_backend, Float32, dims...) + return op, x, y +end + +function _hcat_cpu(n; threaded) + x1 = randn(Float32, n) + x2 = randn(Float32, n) + x = ArrayPartition(x1, x2) + op = HCAT(Eye(Float32, (n,)), DiagOp(randn(Float32, n); threaded = false)) + y = zeros(Float32, n) + return op, x, y +end + +function _hcat_gpu(n) + x1 = gpu_randn(_backend, Float32, n) + x2 = gpu_randn(_backend, Float32, n) + x = ArrayPartition(x1, x2) + op = HCAT(DiagOp(gpu_randn(_backend, Float32, n); threaded = false), DiagOp(gpu_randn(_backend, Float32, n); threaded = false)) + y = gpu_zeros(_backend, Float32, n) + return op, x, y +end + +function _vcat_cpu(n; threaded) + op = VCAT(Eye(Float32, (n,)), DiagOp(randn(Float32, n); threaded = false)) + x = randn(Float32, n) + y = ArrayPartition(zeros(Float32, n), zeros(Float32, n)) + return op, x, y +end + +function _vcat_gpu(n) + x = gpu_randn(_backend, Float32, n) + op = VCAT(Eye(x), DiagOp(gpu_randn(_backend, Float32, n); threaded = false)) + y = ArrayPartition(gpu_zeros(_backend, Float32, n), gpu_zeros(_backend, Float32, n)) + return op, x, y +end + +function _filt_cpu(n; threaded) + taps = randn(Float32, 7) + op = Filt(Float32, (n,), taps) + x = randn(Float32, n) + y = zeros(Float32, n) + return op, x, y +end + +function _filt_gpu(n) + x_cpu = randn(Float64, n) + x = to_gpu(_fftw_backend, x_cpu) + taps = randn(Float64, 7) + op = Filt(x, taps) + y = gpu_zeros(_fftw_backend, Float64, n) + return op, x, y +end + +function _mimofilt_cpu(n; threaded) + len = max(64, n) + taps = [randn(Float32, 5), randn(Float32, 3), randn(Float32, 4), randn(Float32, 5)] + coeffs = [Float32[1] for _ in taps] + op = MIMOFilt(Float32, (len, 2), taps, coeffs) + x = randn(Float32, len, 2) + y = zeros(Float32, len, 2) + return op, x, y +end + +function _mimofilt_gpu(n) + len = max(64, n) + taps = [randn(Float32, 5), randn(Float32, 3), randn(Float32, 4), randn(Float32, 5)] + coeffs = [Float32[1] for _ in taps] + x = gpu_randn(_backend, Float32, len, 2) + op = MIMOFilt(x, taps, coeffs) + y = gpu_zeros(_backend, Float32, len, 2) + return op, x, y +end + +function _conv_cpu(n; threaded) + h = randn(Float32, 17) + op = Conv(Float32, (n,), h) + x = randn(Float32, n) + y = zeros(Float32, size(op, 1)...) + return op, x, y +end + +function _conv_gpu(n) + h = gpu_randn(_backend, Float32, 17) + op = Conv(Float32, (n,), h) + x = gpu_randn(_backend, Float32, n) + y = gpu_zeros(_backend, Float32, size(op, 1)...) + return op, x, y +end + +function _xcorr_cpu(n; threaded) + h = randn(Float32, 17) + op = Xcorr(Float32, (n,), h) + x = randn(Float32, n) + y = zeros(Float32, size(op, 1)...) + return op, x, y +end + +function _xcorr_gpu(n) + h = gpu_randn(_backend, Float32, 17) + op = Xcorr(Float32, (n,), h) + x = gpu_randn(_backend, Float32, n) + y = gpu_zeros(_backend, Float32, size(op, 1)...) + return op, x, y +end + +function _dft_cpu(n; threaded) + dims = _shape2(n) + x = randn(ComplexF32, dims...) + op = DFT(x) + y = zeros(ComplexF32, dims...) + return op, x, y +end + +function _dft_gpu(n) + dims = _shape2(n) + x = gpu_randn(_backend, ComplexF32, dims...) + op = DFT(x) + y = similar(x, ComplexF32) + return op, x, y +end + +const CONFIGS = [ + ( + name = "DiagOp", + sizes = [2^k for k in 10:20], + cpu_single = n -> _diag_cpu(n; threaded = false), + cpu_threaded = Threads.nthreads() > 1 ? (n -> _diag_cpu(n; threaded = true)) : nothing, + gpu = _diag_gpu, + ), + ( + name = "FiniteDiff", + sizes = [2^k for k in 10:20], + cpu_single = n -> _finitediff_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _finitediff_gpu, + ), + ( + name = "GetIndex", + sizes = [2^k for k in 10:18], + cpu_single = n -> _getindex_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _getindex_gpu, + ), + ( + name = "GetIndex_boolmask", + sizes = [2^k for k in 10:20], + cpu_single = n -> _getindex_boolmask_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _getindex_boolmask_gpu, + ), + ( + name = "GetIndex_intvec", + sizes = [2^k for k in 10:20], + cpu_single = n -> _getindex_intvec_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _getindex_intvec_gpu, + ), + ( + name = "Eye", + sizes = [2^k for k in 10:20], + cpu_single = n -> _eye_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _eye_gpu, + ), + ( + name = "MatrixOp", + sizes = [2^k for k in 6:12], + cpu_single = n -> _matrixop_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _matrixop_gpu, + ), + ( + name = "Scale", + sizes = [2^k for k in 10:20], + cpu_single = n -> _scale_cpu(n; threaded = false), + cpu_threaded = Threads.nthreads() > 1 ? (n -> _scale_cpu(n; threaded = true)) : nothing, + gpu = _scale_gpu, + ), + ( + name = "HadamardProd", + sizes = [2^k for k in 10:20], + cpu_single = n -> _hadamardprod_cpu(n; threaded = false), + cpu_threaded = Threads.nthreads() > 1 ? (n -> _hadamardprod_cpu(n; threaded = true)) : nothing, + gpu = _hadamardprod_gpu, + ), + ( + name = "DCAT", + sizes = [2^k for k in 10:18], + cpu_single = n -> _dcat_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _dcat_gpu, + ), + ( + name = "AffineAdd", + sizes = [2^k for k in 10:20], + cpu_single = n -> _affineadd_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _affineadd_gpu, + ), + ( + name = "Sin", + sizes = [2^k for k in 10:20], + cpu_single = n -> _sin_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _sin_gpu, + ), + ( + name = "Tanh", + sizes = [2^k for k in 10:20], + cpu_single = n -> _tanh_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _tanh_gpu, + ), + ( + name = "Exp", + sizes = [2^k for k in 10:20], + cpu_single = n -> _exp_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _exp_gpu, + ), + ( + name = "ZeroPad", + sizes = [2^k for k in 10:18], + cpu_single = n -> _zeropad_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _zeropad_gpu, + ), + ( + name = "Variation", + sizes = [2^k for k in 10:18], + cpu_single = n -> _variation_cpu(n; threaded = false), + cpu_threaded = Threads.nthreads() > 1 ? (n -> _variation_cpu(n; threaded = true)) : nothing, + gpu = _variation_gpu, + ), + ( + name = "Variation_adj", + sizes = [2^k for k in 10:18], + cpu_single = n -> _variation_adj_cpu(n; threaded = false), + cpu_threaded = Threads.nthreads() > 1 ? (n -> _variation_adj_cpu(n; threaded = true)) : nothing, + gpu = _variation_adj_gpu, + ), + ( + name = "Compose", + sizes = [2^k for k in 10:20], + cpu_single = n -> _compose_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _compose_gpu, + ), + ( + name = "Sum", + sizes = [2^k for k in 10:20], + cpu_single = n -> _sum_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _sum_gpu, + ), + ( + name = "BroadCast", + sizes = [2^k for k in 10:18], + cpu_single = n -> _broadcast_cpu(n; threaded = false), + cpu_threaded = Threads.nthreads() > 1 ? (n -> _broadcast_cpu(n; threaded = true)) : nothing, + gpu = _broadcast_gpu, + ), + ( + name = "HCAT", + sizes = [2^k for k in 10:18], + cpu_single = n -> _hcat_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _hcat_gpu, + ), + ( + name = "VCAT", + sizes = [2^k for k in 10:18], + cpu_single = n -> _vcat_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _vcat_gpu, + ), + ( + name = "Filt", + sizes = [2^k for k in 10:20], + cpu_single = n -> _filt_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _filt_gpu, + ), + ( + name = "MIMOFilt", + sizes = [2^k for k in 8:16], + cpu_single = n -> _mimofilt_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _mimofilt_gpu, + ), + ( + name = "Conv", + sizes = [2^k for k in 10:20], + cpu_single = n -> _conv_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _conv_gpu, + ), + ( + name = "Xcorr", + sizes = [2^k for k in 10:20], + cpu_single = n -> _xcorr_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _xcorr_gpu, + ), + ( + name = "DFT", + sizes = [2^k for k in 10:18], + cpu_single = n -> _dft_cpu(n; threaded = false), + cpu_threaded = nothing, + gpu = _dft_gpu, + ), +] + +function main() + Random.seed!(1234) + + selected_names = split(get(ENV, "ABSTRACTOPERATORS_GPU_BENCH_OPERATORS", ""), ',') + selected_names = filter!(!isempty, selected_names) + configs = isempty(selected_names) ? CONFIGS : filter(c -> c.name in selected_names, CONFIGS) + samples = parse(Int, get(ENV, "ABSTRACTOPERATORS_GPU_BENCH_SAMPLES", "20")) + size_override = split(get(ENV, "ABSTRACTOPERATORS_GPU_BENCH_SIZES", ""), ',') + size_override = + if isempty(size_override) || (length(size_override) == 1 && isempty(size_override[1])) + nothing + else + parse.(Int, size_override) + end + + out_dir = joinpath(@__DIR__, "..", ".temp", "gpu-crossover") + mkpath(out_dir) + out_file = joinpath(out_dir, "results.tsv") + md_file = joinpath(out_dir, "report.md") + + # Collect per-operator summary for the markdown report + summary_rows = NamedTuple[] + + open(out_file, "w") do io + println( + io, + "operator\tsize\tcpu_single_s\tcpu_threaded_s\tgpu_s\tbest_cpu_s\tgpu_faster\tnote", + ) + for config in configs + crossover = nothing + last_gpu = nothing # for "GPU faster at max size?" check + last_best_cpu = nothing + sizes = size_override === nothing ? config.sizes : size_override + println("=== ", config.name, " ===") + for n in sizes + try + op_cpu, x_cpu, y_cpu = config.cpu_single(n) + cpu_single = _measure!(y_cpu, op_cpu, x_cpu, false; samples) + cpu_threaded = _threaded_time(config.cpu_threaded, n; samples) + gpu = _gpu_time(config.gpu, n; samples) + best_cpu = cpu_threaded === missing ? cpu_single : min(cpu_single, cpu_threaded) + gpu_faster = gpu < best_cpu + if crossover === nothing && gpu_faster + crossover = n + end + last_gpu = gpu + last_best_cpu = best_cpu + println( + io, + join( + ( + config.name, + n, + cpu_single, + cpu_threaded, + gpu, + best_cpu, + gpu_faster, + "", + ), + '\t', + ), + ) + println( + config.name, + " n=", + n, + " cpu_single=", + cpu_single, + " cpu_threaded=", + cpu_threaded, + " gpu=", + gpu, + ) + catch err + note = sprint(showerror, err) + println( + io, + join( + ( + config.name, + n, + missing, + missing, + missing, + missing, + missing, + note, + ), + '\t', + ), + ) + println(config.name, " n=", n, " ERROR: ", note) + end + end + range_str = "n=$(first(sizes))..$(last(sizes))" + if crossover === nothing + gpu_at_max = if (last_gpu !== nothing && last_best_cpu !== nothing) + (last_gpu < last_best_cpu) + else + missing + end + msg = + "no crossover detected in range [$range_str]" * ( + if gpu_at_max === true + " (GPU faster at max)" + elseif gpu_at_max === false + " (CPU still faster at max)" + else + "" + end + ) + println(config.name, " → ", msg) + else + println(config.name, " → crossover at n=", crossover) + end + push!( + summary_rows, + ( + operator = config.name, + crossover = crossover, + range = range_str, + gpu_at_max = if (last_gpu !== nothing && last_best_cpu !== nothing) + (last_gpu < last_best_cpu) + else + missing + end, + ), + ) + end + end + + # Write markdown report + open(md_file, "w") do io + println(io, "# GPU Crossover Benchmark Report\n") + println( + io, + "Each row shows the smallest problem size at which the GPU becomes faster than the best CPU time.", + ) + println(io, "\"—\" means no crossover was found in the tested range.\n") + println(io, "| Operator | Crossover size | Tested range | GPU faster at max size |") + println(io, "|----------|---------------|-------------|----------------------|") + for row in summary_rows + xover = row.crossover === nothing ? "—" : string(row.crossover) + at_max = row.gpu_at_max === missing ? "n/a" : (row.gpu_at_max ? "yes" : "no") + println(io, "| $(row.operator) | $xover | $(row.range) | $at_max |") + end + end + + println("\nWrote results to ", out_file) + return println("Wrote markdown report to ", md_file) +end + +main() diff --git a/docs/make.jl b/docs/make.jl index c0c9136..119d151 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -17,6 +17,8 @@ makedocs(; "Calculus rules" => "calculus.md", "Properties" => "properties.md", "Batch operators" => "batching.md", + "GPU Support" => "gpu.md", + "Performance" => "performance.md", "Custom Operators" => "custom.md", "LinearMaps wrapper" => "linearmaps.md", ], diff --git a/docs/src/custom.md b/docs/src/custom.md index 1bbca8b..edd0329 100644 --- a/docs/src/custom.md +++ b/docs/src/custom.md @@ -39,8 +39,8 @@ function LinearAlgebra.mul!(y::AbstractArray, L::MyCustomLinOp, x::AbstractArray # - eltype(y) == codomain_type(L) # - size(x) == size(L, 2) # - size(y) == size(L, 1) - # - x isa domain_storage_type(L) - # - y isa codomain_storage_type(L) + # - x isa domain_array_type(L) + # - y isa codomain_array_type(L) AbstractOperators.check(y, L, x) # Implement your forward operation here # Example: y .= some_function(x) @@ -200,7 +200,7 @@ All custom operators must implement: Beyond the mandatory functions, you can optionally define various properties and traits to enable optimizations and provide additional information about your operator. These include: - **Thread safety**: `is_thread_safe(L::YourOperator) = true` -- **Storage types**: `domain_storage_type`, `codomain_storage_type` +- **Storage types**: `domain_array_type`, `codomain_array_type` - **Algebraic properties**: `is_diagonal`, `is_symmetric`, `is_invertible`, etc. - **Operator norm**: `opnorm(L::YourOperator)` - And many more... @@ -248,3 +248,28 @@ y_out = similar(y) mul!(y_out, L, x) @test y_out ≈ L * x ``` + +## Performance and GPU Storage + +### Threading heuristic + +`_should_thread(x)` controls whether constructors enable threaded kernels. +GPU extensions should override `_should_thread(::AbstractGPUArray)=false` to avoid mixing host threading with device arrays. + +### Storage-type correctness + +When implementing custom operators, define both `domain_array_type` and `codomain_array_type` whenever buffers are allocated internally. +This ensures that `op * x` and `allocate_in_*` keep outputs on the same backend (Array, JLArray, CuArray, ROCArray, etc.). + +### Generic FFT backend compatibility + +Prefer backend-generic APIs (`inv(plan)`, AbstractFFTs dispatch) instead of backend-specific functions (e.g. `FFTW.plan_inv`) in code paths that must support CUDA/AMDGPU. + +### Test conventions for backend coverage + +Use utility backends from `test/utils.jl`: +- `ALL_BACKENDS`: CPU + JLArray + available real GPU backends +- `GPU_BACKENDS`: JLArray + available real GPU backends +- tagged checks via `get_backend(:cuda)` / `get_backend(:amdgpu)` + +For backend-specific testitems, add tags `:cuda` / `:amdgpu` and include safety guards with `@test_skip` when unavailable. diff --git a/docs/src/gpu.md b/docs/src/gpu.md new file mode 100644 index 0000000..b86f5e0 --- /dev/null +++ b/docs/src/gpu.md @@ -0,0 +1,176 @@ +# GPU Support + +AbstractOperators.jl provides GPU compatibility through a lightweight extension model. Most operators work transparently with GPU arrays on any of the four supported backends: CUDA.jl, AMDGPU.jl, oneAPI.jl, and OpenCL.jl. + +The main exceptions are: + +- `NFFTOperators.jl`, which currently supports GPU execution only with CUDA.jl +- `WaveletOperators.jl`, which currently works on CPU only +- `DCT` and `IDCT` (in FFTWOperators.jl) needs explicit loading of `AcceleratedDCTs` to activate GPU support + +## Using GPU Arrays + +Simply pass a GPU array to the `mul!` function or the `*` operator: + +```julia +using AbstractOperators, CUDA # or AMDGPU, oneAPI, or OpenCL + +x_gpu = CuArray(randn(Float32, 10)) +y_gpu = CuArray(zeros(Float32, 9)) + +F = FiniteDiff(Float32, (10,)) # CPU operator, GPU arrays as input +mul!(y_gpu, F, x_gpu) # works transparently +``` + +For operators that carry internal arrays (like `DiagOp`), construct them from GPU arrays to get a GPU-typed operator: + +```julia +d_gpu = CuArray(rand(Float32, 10)) +D = DiagOp(d_gpu) # GPU DiagOp — threading disabled automatically +y_gpu = similar(x_gpu) +mul!(y_gpu, D, x_gpu) +``` + +## Threading and GPU + +CPU threading (via [Polyester.jl](https://github.com/JuliaSIMD/Polyester.jl)'s `@batch`) is automatically disabled when GPU arrays are detected. This happens through the `_should_thread` mechanism: + +- For operators with array fields (e.g., `DiagOp`), the constructor checks `_should_thread(d)` which returns `false` for `AbstractGPUArray`s. +- For GPU arrays passed to operators at construction time (e.g., `Variation(gpu_x)`), threading is disabled automatically. +- The `GpuExt` extension overrides `_should_thread(::AbstractGPUArray) = false`. + +## Storage Type Tracking + +Every operator tracks its storage type via `domain_array_type` and `codomain_array_type`. This enables correct intermediate buffer allocation in composed operators: + +```julia +using AbstractOperators, CUDA + +# Create GPU operators +D = DiagOp(CuArray(rand(Float32, 5))) +domain_array_type(D) # CuArray{Float32} +codomain_array_type(D) # CuArray{Float32} + +# Composition allocates GPU buffers automatically +F = FiniteDiff(CuArray{Float32}, (10,)) # array_type-aware FiniteDiff +domain_array_type(F) # CuArray{Float32} +``` + +For operators without an array at construction time, pass a GPU array to the array-based constructor: + +```julia +F = FiniteDiff(CuArray(ones(Float32, 5, 4)), 1) # deduces storage type from input +V = Variation(CuArray(ones(Float32, 4, 3))) # threaded=false automatically +``` + +Or use the `array_type` keyword argument (for operators that support it): + +```julia +Z = Zeros(Float32, (4,), Float32, (3,); array_type=CuArray) +``` + +## GpuExt Extension + +GPU-specific overrides live in `ext/GpuExt/`. The extension is loaded automatically when `GPUArrays.jl` is loaded (directly or via any GPU backend like CUDA.jl). + +The extension provides: +- `_should_thread(::AbstractGPUArray) = false` — disables CPU threading for GPU arrays +- `array_type_display_string(::Type{<:AbstractGPUArray}) = "ᵍᵖᵘ"` — shows ᵍᵖᵘ superscript in operator display +- Variation adjoint GPU override — vectorized reshape-based stencil (no scalar indexing) +- BroadCast threading GPU override — falls back to non-threaded path + +## Testing with JLArrays + +The test suite includes GPU tests using [JLArrays.jl](https://github.com/JuliaGPU/JLArrays.jl), a simple CPU-backed array type that mimics GPU behavior (no scalar indexing, GPUArrays interface). GPU tests are tagged `:gpu`: + +```bash +julia --project=test -e ' + using TestItemRunner + @run_package_tests filter=ti -> :gpu in ti.tags +' +``` + +## Backend-Specific Notes + +Most operators in AbstractOperators.jl, DSPOperators.jl, and the GPU-compatible parts of FFTWOperators.jl follow the same generic GPU-array execution model and therefore work with CUDA.jl, AMDGPU.jl, oneAPI.jl, and OpenCL.jl. + +### NFFTOperators GPU Support (CUDA only) + +NFFTOperators.jl supports GPU via the `array_type` keyword argument. This requires loading CUDA.jl (or any package that loads GPUArrays + Adapt): + +```julia +using NFFTOperators, CUDA + +image_size = (128, 128) +trajectory = rand(Float32, 2, 128, 50) .- 0.5f0 +dcf = rand(Float32, 128, 50) + +# Create GPU NFFT operator — ksp_buffer and dcf are CuArrays +op_gpu = NFFTOp(image_size, trajectory, dcf; array_type=CuArray, threaded=false) + +x_gpu = CUDA.randn(ComplexF32, image_size) +y_gpu = op_gpu * x_gpu # forward: image → k-space (on GPU) +img_rec = op_gpu' * y_gpu # adjoint: k-space → image (on GPU) +``` + +**Note:** GPU NFFT requires CUDA.jl because it relies on cuFFT plans. JLArrays (used in tests) does not support FFT plans and cannot be used with NFFTOperators GPU. + +The trajectory `k` is kept on CPU (as required by NFFT.jl's GPU plan); only the computation buffers (`ksp_buffer`, `dcf`) and input/output arrays are on GPU. + +### FFTWOperators DCT/IDCT (AcceleratedDCTs) + +`DCT` and `IDCT` are CPU-only unless `AcceleratedDCTs` is loaded. When `AcceleratedDCTs` is available and imported, the package activates GPU support for those operators through its extension: + +```julia +using AbstractOperators, FFTWOperators, CUDA +import AcceleratedDCTs # explicitly import to activate GPU support for DCT/IDCT + +x_gpu = CUDA.randn(Float32, 64) +dct_op = DCT(x_gpu) +y_gpu = dct_op * x_gpu + +idct_op = IDCT(x_gpu) +x_rec_gpu = idct_op * y_gpu +``` + +If `AcceleratedDCTs` is not imported, `DCT` and `IDCT` continue to use the CPU FFTW implementation, +and will not work with GPU arrays. In that case, wrap them with `CpuOperatorWrapper` to use in GPU pipelines: + +```julia +using AbstractOperators, FFTWOperators, CUDA + +x_gpu = CUDA.randn(Float32, 64) +dct_op = CpuOperatorWrapper(DCT(Float32, (64,)); array_type = CuArray{Float32}) # CPU operator wrapped for GPU use +y_gpu = dct_op * x_gpu # GPU in → CPU DCT → GPU out +``` + +### WaveletOperators CPU-only status + +WaveletOperators.jl currently relies on CPU execution. Its operators do not yet support GPU arrays, so wavelet transforms should remain on CPU or be wrapped explicitly as CPU operators when building mixed CPU/GPU pipelines. + +## CpuOperatorWrapper + +For operators that do not natively support GPU arrays (e.g., FFTWOperators DCT, custom CPU-only operators), use `CpuOperatorWrapper`. This wrapper preallocates CPU buffers for the operator's domain and codomain, allowing GPU arrays to be passed in and out while the computation happens on CPU. + +```@docs +OperatorWrapper +OperatorWrapper(::AbstractOperator) +``` + +### Example: + +```julia +using AbstractOperators, CUDA + +# Any CPU operator +op = FiniteDiff(Float32, (64,)) # or any FFTWOperators, etc. + +# Wrap it — preallocates CPU buffers for domain and codomain +wrapper = CpuOperatorWrapper(op; array_type = CuArray{Float32}) # specify GPU array type for buffers + +x_gpu = CUDA.randn(Float32, 64) +y_gpu = similar(x_gpu, 63) + +mul!(y_gpu, wrapper, x_gpu) # GPU in → CPU compute → GPU out +mul!(x_gpu, wrapper', y_gpu) # GPU in → CPU adjoint → GPU out +``` diff --git a/docs/src/index.md b/docs/src/index.md index a7a504f..a2d3892 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -115,6 +115,31 @@ A list of the available `AbstractOperators` and calculus rules can be found in t `AbstractOperators.jl` is distinguished by its support for multi-dimensional array domains and codomains, efficient in-place implementations of both linear and nonlinear operators, and seamless integration with optimization algorithms in related packages: [ProximalOperators.jl](https://github.com/kul-forbes/ProximalOperators.jl), [ProximalAlgorithms.jl](https://github.com/kul-forbes/ProximalAlgorithms.jl), and [StructuredOptimization.jl](https://github.com/kul-forbes/StructuredOptimization.jl). It has built-in threading support for many operators, and partial (and extending) GPU support. +## GPU Support + +AbstractOperators.jl supports GPU arrays through a lightweight extension. Most operators work transparently with GPU arrays (e.g., `CuArray` from CUDA.jl): + +```julia +using AbstractOperators, CUDA +x_gpu = CuArray(randn(Float32, 100)) +F = FiniteDiff(Float32, (100,)) +y_gpu = similar(x_gpu, 99) +mul!(y_gpu, F, x_gpu) # works transparently +``` + +CPU threading is disabled automatically for GPU arrays. Storage types propagate through composed operators. See the [GPU documentation](gpu.md) for details. + +## Smart Operator Copying + +For parallel or multi-threaded use, copy operators efficiently with [`copy_operator`](@ref): + +```julia +op = DiagOp(rand(1000)) + DiagOp(rand(1000)) # Sum with mutable buffers +op2 = copy_operator(op) # shares immutable data, copies buffers only +``` + +See [Performance](performance.md) for user-facing tips and [Custom Operators](custom.md) for developer internals (threading/storage/backend guidance). + ## Credits AbstractOperators.jl is developed by diff --git a/docs/src/performance.md b/docs/src/performance.md new file mode 100644 index 0000000..832c41f --- /dev/null +++ b/docs/src/performance.md @@ -0,0 +1,9 @@ +# Performance + +For end users, the main performance recommendations are: + +1. Prefer `mul!` over `*` to avoid allocations. +2. Preallocate with `allocate_in_domain(op)` / `allocate_in_codomain(op)`. +3. Use operator constructors with array inputs to preserve storage type (CPU vs GPU). + +Developer-oriented performance internals (threading heuristics, storage traits, backend caveats) were moved to [Custom Operators](@ref) to keep this page concise. diff --git a/docs/src/properties.md b/docs/src/properties.md index 6ef2548..a3b2dc9 100644 --- a/docs/src/properties.md +++ b/docs/src/properties.md @@ -8,8 +8,8 @@ ndims ndoms domain_type codomain_type -domain_storage_type -codomain_storage_type +domain_array_type +codomain_array_type ``` ## Traits @@ -57,3 +57,9 @@ remove_slicing is_thread_safe estimate_opnorm ``` + +## Operator Copying + +```@docs +copy_operator +``` diff --git a/ext/GpuExt.jl b/ext/GpuExt.jl deleted file mode 100644 index a772ab8..0000000 --- a/ext/GpuExt.jl +++ /dev/null @@ -1,8 +0,0 @@ -module GpuExt - -using GPUArrays -import AbstractOperators: storage_type_display_string - -storage_type_display_string(::Type{<:AbstractGPUArray}) = "ᵍᵖᵘ" - -end # module GpuExt diff --git a/ext/GpuExt/GpuExt.jl b/ext/GpuExt/GpuExt.jl new file mode 100644 index 0000000..337d18e --- /dev/null +++ b/ext/GpuExt/GpuExt.jl @@ -0,0 +1,21 @@ +module GpuExt + +using GPUArrays +using KernelAbstractions +using AbstractOperators +import LinearAlgebra: mul! +import AbstractOperators: + _should_thread, array_type_display_string, check, + AdjointOperator, Variation, NoOperatorBroadCast, + domain_type, allocate_in_domain, allocate_in_codomain, + GetIndex, ZeroPad, OperatorWrapper +using RecursiveArrayTools: ArrayPartition + +include("properties.jl") +include("linearoperators/getindex.jl") +include("linearoperators/zeropad.jl") +include("linearoperators/variation.jl") +include("cpuwrapper.jl") +include("guards.jl") + +end # module GpuExt diff --git a/ext/GpuExt/cpuwrapper.jl b/ext/GpuExt/cpuwrapper.jl new file mode 100644 index 0000000..3502727 --- /dev/null +++ b/ext/GpuExt/cpuwrapper.jl @@ -0,0 +1,26 @@ +# GPU mul! overrides for OperatorWrapper. +# check(y, A, x) passes because domain_array_type/codomain_array_type reflect +# the constructor-supplied array_type (e.g. CuArray), not the inner CPU op's type. + +# Forward: GPU → CPU → op → CPU → GPU +function mul!(y::AbstractGPUArray, A::OperatorWrapper, x::AbstractGPUArray) + check(y, A, x) + copyto!(A.dom_buf, x) + mul!(A.cod_buf, A.op, A.dom_buf) + copyto!(y, A.cod_buf) + return y +end + +# Adjoint: GPU → CPU → op' → CPU → GPU +function mul!( + y::AbstractGPUArray, + Ac::AdjointOperator{<:OperatorWrapper}, + x::AbstractGPUArray, + ) + check(y, Ac, x) + A = Ac.A + copyto!(A.cod_buf, x) + mul!(A.dom_buf, A.op', A.cod_buf) + copyto!(y, A.dom_buf) + return y +end diff --git a/ext/GpuExt/guards.jl b/ext/GpuExt/guards.jl new file mode 100644 index 0000000..0cb5a02 --- /dev/null +++ b/ext/GpuExt/guards.jl @@ -0,0 +1,45 @@ +# CPU/GPU mutual rejection via check overrides. +# These ensure that mixing CPU and GPU arrays in mul! gives a clear error message. +# +# GPU-GPU and GPU-ArrayPartition overloads use `invoke` to call the base +# AbstractOperators.check directly, bypassing these specialised methods. +# Without `invoke`, the GPU-specific methods would re-dispatch to themselves +# (infinite recursion) because AbstractGPUArray <: AbstractArray. + + +# GPU-GPU: delegate to the base check (handles ArrayPartition too) +function check(y::AbstractGPUArray, A, b::AbstractGPUArray) + invoke(AbstractOperators.check, Tuple{Any, Any, Any}, y, A, b) + return nothing +end + +# ArrayPartition containers may hold GPU arrays — delegate to base check +function check(y::AbstractGPUArray, A, b::ArrayPartition) + invoke(AbstractOperators.check, Tuple{Any, Any, Any}, y, A, b) + return nothing +end + +function check(y::ArrayPartition, A, b::AbstractGPUArray) + invoke(AbstractOperators.check, Tuple{Any, Any, Any}, y, A, b) + return nothing +end + +# GPU output + CPU input: error (plain arrays only, not ArrayPartition) +function check(y::AbstractGPUArray, A, b::AbstractArray) + throw( + ArgumentError( + "Cannot use CPU input $(typeof(b)) with GPU output $(typeof(y)). " * + "Ensure both arrays have the same storage type.", + ), + ) +end + +# CPU output + GPU input: error (plain arrays only, not ArrayPartition) +function check(y::AbstractArray, A, b::AbstractGPUArray) + throw( + ArgumentError( + "Cannot use GPU input $(typeof(b)) with CPU output $(typeof(y)). " * + "Ensure both arrays have the same storage type.", + ), + ) +end diff --git a/ext/GpuExt/linearoperators/getindex.jl b/ext/GpuExt/linearoperators/getindex.jl new file mode 100644 index 0000000..805688d --- /dev/null +++ b/ext/GpuExt/linearoperators/getindex.jl @@ -0,0 +1,72 @@ +function _to_gpu_indices(ref_array::AbstractGPUArray, cpu_idx::AbstractVector{<:Integer}) + ArrayT = Base.typename(typeof(ref_array)).wrapper + return ArrayT(Vector{Int}(cpu_idx)) +end + +function _to_gpu_indices(ref_array::AbstractGPUArray, gpu_idx::AbstractGPUArray{<:Integer}) + return gpu_idx +end + +function _to_gpu_indices(array_type::Type, cpu_idx::AbstractVector{<:Integer}) + ArrayT = Base.typename(array_type).wrapper + return ArrayT(Vector{Int}(cpu_idx)) +end + +function _mask_to_linear_indices(mask::AbstractArray{Bool}) + return findall(vec(mask)) +end + +function _mask_to_linear_indices(mask::AbstractGPUArray{Bool}) + return findall(vec(Array(mask))) +end + +function AbstractOperators._prepare_getindex_intvec( + idx::Vector{Int}, array_type::Type{<:AbstractGPUArray} + ) + return _to_gpu_indices(array_type, idx) +end + +function AbstractOperators._prepare_getindex_boolmask( + mask::AbstractArray{Bool}, array_type::Type{<:AbstractGPUArray} + ) + return _to_gpu_indices(array_type, _mask_to_linear_indices(mask)) +end + +function GetIndex(x::AbstractGPUArray, idx::AbstractVector{Int}) + dim_in = size(x) + dim_out = AbstractOperators.get_dim_out(dim_in, idx) + if dim_out == dim_in + return AbstractOperators.Eye(eltype(x), dim_in; array_type = typeof(x)) + end + S = AbstractOperators._array_wrapper(x){eltype(x)} + gpu_idx = _to_gpu_indices(x, idx) + return AbstractOperators.GetIndex(eltype(x), S, dim_out, dim_in, gpu_idx) +end + +function GetIndex(x::AbstractGPUArray, mask::AbstractArray{Bool}) + dim_in = size(x) + dim_out = AbstractOperators.get_dim_out(dim_in, mask) + if dim_out[1] == prod(dim_in) + return reshape(AbstractOperators.Eye(eltype(x), dim_in; array_type = typeof(x)), dim_out) + end + S = AbstractOperators._array_wrapper(x){eltype(x)} + gpu_idx = _to_gpu_indices(x, _mask_to_linear_indices(mask)) + return AbstractOperators.GetIndex(eltype(x), S, dim_out, dim_in, gpu_idx) +end + +function mul!( + y::AbstractGPUArray, L::GetIndex{I}, b::AbstractGPUArray + ) where {K, I <: NTuple{K, Any}} + check(y, L, b) + y .= view(b, L.idx...) + return y +end + +function mul!( + y::AbstractGPUArray, Lc::AdjointOperator{<:GetIndex{I}}, b::AbstractGPUArray + ) where {K, I <: NTuple{K, Any}} + check(y, Lc, b) + fill!(y, zero(eltype(y))) + view(y, Lc.A.idx...) .= b + return y +end diff --git a/ext/GpuExt/linearoperators/variation.jl b/ext/GpuExt/linearoperators/variation.jl new file mode 100644 index 0000000..6ad9b37 --- /dev/null +++ b/ext/GpuExt/linearoperators/variation.jl @@ -0,0 +1,109 @@ +# Variation: KernelAbstractions @kernel implementations. + +# JLBackend (JLArrays) is synchronous — KA provides no synchronize method for it. +# Use this guard to avoid a MethodError on JLBackend while still synchronizing +# asynchronous GPU backends (CUDA, AMDGPU). +_ka_synchronize(backend) = + hasmethod(KernelAbstractions.synchronize, (typeof(backend),)) && + KernelAbstractions.synchronize(backend) + +@kernel function _var_fwd_dim1_kernel!(y_col, @Const(b_flat), dim_stride) + cnt = @index(Global, Linear) + n = length(b_flat) + i1 = (cnt - 1) % dim_stride + 1 + if i1 == 1 + next = cnt + 1 <= n ? cnt + 1 : cnt + y_col[cnt] = b_flat[next] - b_flat[cnt] + else + y_col[cnt] = b_flat[cnt] - b_flat[cnt - 1] + end +end + +@kernel function _var_fwd_dimd_kernel!(y_col, @Const(b_flat), dim_stride, nd) + cnt = @index(Global, Linear) + k = (cnt - 1) ÷ dim_stride + pos = k % nd + if pos == 0 + y_col[cnt] = b_flat[cnt + dim_stride] - b_flat[cnt] + else + y_col[cnt] = b_flat[cnt] - b_flat[cnt - dim_stride] + end +end + +@kernel function _var_adj_kernel!(y_flat, @Const(b), @Const(dim_in_ntuple)) + cnt = @index(Global, Linear) + dim_in = dim_in_ntuple + N = length(dim_in) + T = eltype(y_flat) + acc = zero(T) + dim_stride = 1 + for d in 1:N + nd = dim_in[d] + i_d = ((cnt - 1) ÷ dim_stride) % nd + 1 + acc += if i_d == 1 + -(b[cnt, d] + b[cnt + dim_stride, d]) + elseif i_d == nd + b[cnt, d] + elseif i_d == 2 + b[cnt, d] + b[cnt - dim_stride, d] - b[cnt + dim_stride, d] + else + b[cnt, d] - b[cnt + dim_stride, d] + end + dim_stride *= nd + end + y_flat[cnt] = acc +end + +function _var_fwd_gpu!(y::AbstractGPUArray, A::Variation, b::AbstractGPUArray) + backend = KernelAbstractions.get_backend(b) + n = length(b) + b_flat = reshape(b, n) + + dim1_extent = size(b, 1) + ker1 = _var_fwd_dim1_kernel!(backend, 256) + ker1(view(y, :, 1), b_flat, dim1_extent; ndrange = n) + + N = ndims(b) + dim_stride = dim1_extent + for d in 2:N + nd = size(b, d) + kerd = _var_fwd_dimd_kernel!(backend, 256) + kerd(view(y, :, d), b_flat, dim_stride, nd; ndrange = n) + dim_stride *= nd + end + _ka_synchronize(backend) + return y +end + +function _var_adj_gpu!(y::AbstractGPUArray, A::AdjointOperator{<:Variation}, b::AbstractGPUArray) + backend = KernelAbstractions.get_backend(y) + n = length(y) + ker = _var_adj_kernel!(backend, 256) + ker(reshape(y, n), b, A.A.dim_in; ndrange = n) + _ka_synchronize(backend) + return y +end + +function mul!(y::AbstractGPUArray, A::Variation{T, N, false}, b::AbstractGPUArray) where {T, N} + check(y, A, b) + return _var_fwd_gpu!(y, A, b) +end + +function mul!(y::AbstractGPUArray, A::Variation{T, N, true}, b::AbstractGPUArray) where {T, N} + check(y, A, b) + return _var_fwd_gpu!(y, A, b) +end + +function mul!( + y::AbstractGPUArray, A::AdjointOperator{<:Variation{T, N, false}}, b::AbstractGPUArray + ) where {T, N} + check(y, A, b) + return _var_adj_gpu!(y, A, b) +end + +function mul!( + y::AbstractGPUArray, A::AdjointOperator{<:Variation{T, N, true}}, b::AbstractGPUArray + ) where {T, N} + check(y, A, b) + return _var_adj_gpu!(y, A, b) +end diff --git a/ext/GpuExt/linearoperators/zeropad.jl b/ext/GpuExt/linearoperators/zeropad.jl new file mode 100644 index 0000000..800ad5d --- /dev/null +++ b/ext/GpuExt/linearoperators/zeropad.jl @@ -0,0 +1,18 @@ +# ZeroPad: vectorized GPU mul! to avoid scalar indexing in @generated loop + +function mul!(y::AbstractGPUArray, L::ZeroPad, b::AbstractGPUArray) + check(y, L, b) + fill!(y, zero(eltype(y))) + N = length(L.dim_in) + dst = view(y, ntuple(i -> 1:L.dim_in[i], N)...) + copyto!(dst, b) + return y +end + +function mul!(y::AbstractGPUArray, Lc::AdjointOperator{<:ZeroPad}, b::AbstractGPUArray) + check(y, Lc, b) + N = length(Lc.A.dim_in) + src = view(b, ntuple(i -> 1:Lc.A.dim_in[i], N)...) + copyto!(y, src) + return y +end diff --git a/ext/GpuExt/properties.jl b/ext/GpuExt/properties.jl new file mode 100644 index 0000000..ed72027 --- /dev/null +++ b/ext/GpuExt/properties.jl @@ -0,0 +1,4 @@ +array_type_display_string(::Type{<:AbstractGPUArray}) = "ᵍᵖᵘ" + +_should_thread(::AbstractGPUArray) = false +_should_thread(::Type{<:AbstractGPUArray}) = false diff --git a/src/AbstractOperators.jl b/src/AbstractOperators.jl index 8e0b119..000608d 100644 --- a/src/AbstractOperators.jl +++ b/src/AbstractOperators.jl @@ -58,6 +58,8 @@ include("linearoperators/FiniteDiff.jl") include("linearoperators/Variation.jl") include("linearoperators/LBFGS.jl") +# Calculus rules + # Batch operators include("batching/BatchOp.jl") include("batching/SimpleBatchOp.jl") @@ -78,6 +80,7 @@ include("calculus/Axt_mul_Bx.jl") include("calculus/Ax_mul_Bxt.jl") include("calculus/Ax_mul_Bx.jl") include("calculus/HadamardProd.jl") +include("calculus/OperatorWrapper.jl") # Non-Linear operators include("nonlinearoperators/Pow.jl") diff --git a/src/batching/BatchOp.jl b/src/batching/BatchOp.jl index dc0bb6a..099057e 100644 --- a/src/batching/BatchOp.jl +++ b/src/batching/BatchOp.jl @@ -25,15 +25,6 @@ end get_domain_batch_dim_mask(::Type{<:BatchOp{T1, T2, dM, cM}}) where {T1, T2, dM, cM} = dM get_codomain_batch_dim_mask(::Type{<:BatchOp{T1, T2, dM, cM}}) where {T1, T2, dM, cM} = cM -copy_op(x) = deepcopy(x) -function copy_op(x::T) where {T <: AbstractOperator} - return if AbstractOperators.is_thread_safe(x) - x - else - T.name.wrapper([copy_op(getfield(x, k)) for k in fieldnames(T)]...) - end -end - function prepare_batch_op( operator::AbstractOperators.AbstractOperator, domain_size::NTuple{M1, Int}, diff --git a/src/batching/SimpleBatchOp.jl b/src/batching/SimpleBatchOp.jl index 05d1b08..a05c54f 100644 --- a/src/batching/SimpleBatchOp.jl +++ b/src/batching/SimpleBatchOp.jl @@ -132,51 +132,6 @@ function create_BatchOp( ) end -function _build_threaded_operators(operator::AbstractOperators.AbstractOperator, batch_size::Tuple) - batch_length = prod(batch_size) - copies = min(nthreads(), batch_length) - return tuple([i == 1 ? operator : copy_op(operator) for i in 1:copies]...) -end - -function _build_single_threaded_batchop( - operator, - domain_size, - codomain_size, - batch_size, - domain_batch_dim_mask, - codomain_batch_dim_mask, - dType, - cdType, - ) - B = length(batch_size) - opType = typeof(operator) - N = length(domain_size) - M = length(codomain_size) - return SimpleBatchOpSingleThreaded{ - dType, cdType, domain_batch_dim_mask, codomain_batch_dim_mask, opType, N, M, B, - }(operator, domain_size, codomain_size, batch_size) -end - -function _build_multi_threaded_batchop( - operator, - domain_size, - codomain_size, - batch_size, - domain_batch_dim_mask, - codomain_batch_dim_mask, - dType, - cdType, - ) - operators = _build_threaded_operators(operator, batch_size) - C = length(operators) - opType = typeof(operator) - N = length(domain_size) - M = length(codomain_size) - return SimpleBatchOpMultiThreaded{ - dType, cdType, domain_batch_dim_mask, codomain_batch_dim_mask, opType, N, M, C, - }(operators, domain_size, codomain_size, CartesianIndices(batch_size)) -end - function create_BatchOp( operator::AbstractOperators.AbstractOperator, domain_size::NTuple{N, Int}, @@ -190,30 +145,30 @@ function create_BatchOp( batch_size, dType, cdType = prepare_batch_op( operator, domain_size, domain_batch_dim_mask, codomain_size, codomain_batch_dim_mask ) - - if threaded && nthreads() > 1 - return _build_multi_threaded_batchop( - operator, - domain_size, - codomain_size, - batch_size, - domain_batch_dim_mask, - codomain_batch_dim_mask, - dType, - cdType, + opType = typeof(operator) + threaded = threaded && _should_thread(operator) + return if threaded && nthreads() > 1 + batch_length = prod(batch_size) + operators = tuple( + [ + i == 1 ? operator : AbstractOperators.copy_operator(operator) for + i in 1:min(nthreads(), batch_length) + ]..., + ) + C = length(operators) + SimpleBatchOpMultiThreaded{ + dType, cdType, domain_batch_dim_mask, codomain_batch_dim_mask, opType, N, M, C, + }( + operators, domain_size, codomain_size, CartesianIndices(batch_size) + ) + else + B = length(batch_size) + SimpleBatchOpSingleThreaded{ + dType, cdType, domain_batch_dim_mask, codomain_batch_dim_mask, opType, N, M, B, + }( + operator, domain_size, codomain_size, batch_size ) end - - return _build_single_threaded_batchop( - operator, - domain_size, - codomain_size, - batch_size, - domain_batch_dim_mask, - codomain_batch_dim_mask, - dType, - cdType, - ) end # mul! implementations @@ -283,10 +238,10 @@ fun_name(L::SimpleBatchOpMultiThreaded) = "⟳" * fun_name(L.operator[1]) size(L::SimpleBatchOp) = L.codomain_size, L.domain_size -domain_storage_type(L::SimpleBatchOpSingleThreaded) = domain_storage_type(L.operator) -domain_storage_type(L::SimpleBatchOpMultiThreaded) = domain_storage_type(L.operator[1]) -codomain_storage_type(L::SimpleBatchOpSingleThreaded) = codomain_storage_type(L.operator) -codomain_storage_type(L::SimpleBatchOpMultiThreaded) = codomain_storage_type(L.operator[1]) +domain_array_type(L::SimpleBatchOpSingleThreaded) = domain_array_type(L.operator) +domain_array_type(L::SimpleBatchOpMultiThreaded) = domain_array_type(L.operator[1]) +codomain_array_type(L::SimpleBatchOpSingleThreaded) = codomain_array_type(L.operator) +codomain_array_type(L::SimpleBatchOpMultiThreaded) = codomain_array_type(L.operator[1]) is_linear(L::SimpleBatchOpSingleThreaded) = is_linear(L.operator) is_linear(L::SimpleBatchOpMultiThreaded) = is_linear(L.operator[1]) @@ -392,7 +347,7 @@ function get_normal_op( L::SimpleBatchOpMultiThreaded{dT, cT, dM, cM, opT, N, M, C} ) where {dT, cT, dM, cM, opT, N, M, C} new_op = get_normal_op(L.operator[1]) - new_ops = tuple([i == 1 ? new_op : copy_op(new_op) for i in 1:length(L.operator)]...) + new_ops = tuple([i == 1 ? new_op : AbstractOperators.copy_operator(new_op) for i in 1:length(L.operator)]...) return SimpleBatchOpMultiThreaded{dT, cT, dM, dM, typeof(new_op), N, N, C}( new_ops, L.domain_size, L.domain_size, L.batch_indices ) diff --git a/src/batching/SpreadingBatchOp.jl b/src/batching/SpreadingBatchOp.jl index 8da457e..3a7ee32 100644 --- a/src/batching/SpreadingBatchOp.jl +++ b/src/batching/SpreadingBatchOp.jl @@ -242,6 +242,7 @@ function create_BatchOp( codomain_batch_dim_mask, spreading_dims, ) + threaded = threaded && _should_thread(operators[1]) if threaded && nthreads() > 1 type_args = ( dType, @@ -291,10 +292,10 @@ function prepare_SpreadingBatchOp( @assert all(op -> domain_type(op) == domain_type(operators[1]), operators) "All operators must have the same domain type" @assert all(op -> codomain_type(op) == codomain_type(operators[1]), operators) "All operators must have the same codomain type" @assert all( - op -> domain_storage_type(op) == domain_storage_type(operators[1]), operators + op -> domain_array_type(op) == domain_array_type(operators[1]), operators ) "All operators must have the same storage type" @assert all( - op -> codomain_storage_type(op) == codomain_storage_type(operators[1]), operators + op -> codomain_array_type(op) == codomain_array_type(operators[1]), operators ) "All operators must have the same storage type" @assert all(op -> size(op, 2) == size(operators[1], 2), operators) "All operators must have the same input size" @assert all(op -> size(op, 1) == size(operators[1], 1), operators) "All operators must have the same output size" @@ -318,8 +319,23 @@ function guess_optimal_threading_strategy(operators, batch_size) copied_ops_size = Base.summarysize(operators) * min(nthreads(), prod(batch_size)) domain_array_size = sizeof(domain_type(operators[1])) * prod(size(operators[1], 2)) codomain_array_size = sizeof(codomain_type(operators[1])) * prod(size(operators[1], 1)) - if copied_ops_size < 10.0e6 || - copied_ops_size < max(domain_array_size, codomain_array_size) * 2 #= 10MB =# + if copied_ops_size < 10.0e6 || copied_ops_size < max(domain_array_size, codomain_array_size) * 2 + # Public constructors + + # Public constructors + + # Public constructors + + #= 10MB =# + + # Public constructors + + # Public constructors + + # Public constructors + + # Public constructors + threading_strategy = ThreadingStrategy.COPYING else threading_strategy = ThreadingStrategy.LOCKING @@ -340,7 +356,7 @@ function create_threaded_SpreadingBatchOp( end if threading_strategy == ThreadingStrategy.COPYING operators = [ - i == 1 ? operators : [copy_op(op) for op in operators] for + i == 1 ? operators : [AbstractOperators.copy_operator(op) for op in operators] for i in 1:min(nthreads(), prod(batch_size)) ] return SpreadingBatchOpCopying{type_args...}( @@ -574,10 +590,24 @@ fun_name(L::SpreadingBatchOpCopying) = "⟳" * fun_name(L.operators[1][1]) size(L::SpreadingBatchOp) = L.codomain_size, L.domain_size -domain_storage_type(L::SpreadingBatchOp) = domain_storage_type(L.operators[1]) -domain_storage_type(L::SpreadingBatchOpCopying) = domain_storage_type(L.operators[1][1]) -codomain_storage_type(L::SpreadingBatchOp) = codomain_storage_type(L.operators[1]) -codomain_storage_type(L::SpreadingBatchOpCopying) = codomain_storage_type(L.operators[1][1]) +domain_array_type(L::SpreadingBatchOpSingleThreaded) = domain_array_type(L.operators[1]) +domain_array_type(L::SpreadingBatchOpThreadSafe) = domain_array_type(L.operators[1]) +domain_array_type(L::SpreadingBatchOpLocking) = domain_array_type(L.operators[1]) +domain_array_type(L::SpreadingBatchOpFixedOperator) = domain_array_type(L.operators[1]) +domain_array_type(L::SpreadingBatchOpCopying) = domain_array_type(L.operators[1][1]) +codomain_array_type(L::SpreadingBatchOpSingleThreaded) = codomain_array_type(L.operators[1]) +codomain_array_type(L::SpreadingBatchOpThreadSafe) = codomain_array_type(L.operators[1]) +codomain_array_type(L::SpreadingBatchOpLocking) = codomain_array_type(L.operators[1]) +codomain_array_type(L::SpreadingBatchOpFixedOperator) = codomain_array_type(L.operators[1]) +codomain_array_type(L::SpreadingBatchOpCopying) = codomain_array_type(L.operators[1][1]) +AbstractOperators._check_domain_storage(domain_array, ::SpreadingBatchOp) = nothing +AbstractOperators._check_codomain_storage(codomain_array, ::SpreadingBatchOp) = nothing +AbstractOperators._check_domain_storage(domain_array::ArrayPartition, ::SpreadingBatchOp) = nothing +function AbstractOperators._check_codomain_storage( + codomain_array::ArrayPartition, ::SpreadingBatchOp + ) + return nothing +end is_linear(L::SpreadingBatchOp) = is_linear(L.operators[1]) is_linear(L::SpreadingBatchOpCopying) = is_linear(L.operators[1][1]) diff --git a/src/calculus/AdjointOperator.jl b/src/calculus/AdjointOperator.jl index 1ef729a..e0c563c 100644 --- a/src/calculus/AdjointOperator.jl +++ b/src/calculus/AdjointOperator.jl @@ -44,8 +44,8 @@ size(L::AdjointOperator) = size(L.A, 2), size(L.A, 1) domain_type(L::AdjointOperator) = codomain_type(L.A) codomain_type(L::AdjointOperator) = domain_type(L.A) -domain_storage_type(L::AdjointOperator) = codomain_storage_type(L.A) -codomain_storage_type(L::AdjointOperator) = domain_storage_type(L.A) +domain_array_type(L::AdjointOperator) = codomain_array_type(L.A) +codomain_array_type(L::AdjointOperator) = domain_array_type(L.A) is_thread_safe(L::AdjointOperator) = is_thread_safe(L.A) fun_name(L::AdjointOperator) = fun_name(L.A) * "ᵃ" diff --git a/src/calculus/AffineAdd.jl b/src/calculus/AffineAdd.jl index 06f850f..f522a43 100644 --- a/src/calculus/AffineAdd.jl +++ b/src/calculus/AffineAdd.jl @@ -85,8 +85,8 @@ size(L::AffineAdd) = size(L.A) domain_type(L::AffineAdd) = domain_type(L.A) codomain_type(L::AffineAdd) = codomain_type(L.A) -domain_storage_type(L::AffineAdd) = domain_storage_type(L.A) -codomain_storage_type(L::AffineAdd) = codomain_storage_type(L.A) +domain_array_type(L::AffineAdd) = domain_array_type(L.A) +codomain_array_type(L::AffineAdd) = codomain_array_type(L.A) is_thread_safe(L::AffineAdd) = is_thread_safe(L.A) is_linear(L::AffineAdd) = is_linear(L.A) diff --git a/src/calculus/Ax_mul_Bx.jl b/src/calculus/Ax_mul_Bx.jl index 1996c83..b53c9fd 100644 --- a/src/calculus/Ax_mul_Bx.jl +++ b/src/calculus/Ax_mul_Bx.jl @@ -95,8 +95,8 @@ fun_name(L::Union{Ax_mul_Bx, Ax_mul_BxJac}) = fun_name(L.A) * "*" * fun_name(L.B domain_type(L::Union{Ax_mul_Bx, Ax_mul_BxJac}) = domain_type(L.A) codomain_type(L::Union{Ax_mul_Bx, Ax_mul_BxJac}) = codomain_type(L.A) -domain_storage_type(L::Union{Ax_mul_Bx, Ax_mul_BxJac}) = domain_storage_type(L.A) -codomain_storage_type(L::Union{Ax_mul_Bx, Ax_mul_BxJac}) = codomain_storage_type(L.B) +domain_array_type(L::Union{Ax_mul_Bx, Ax_mul_BxJac}) = domain_array_type(L.A) +codomain_array_type(L::Union{Ax_mul_Bx, Ax_mul_BxJac}) = codomain_array_type(L.B) # utils function permute( diff --git a/src/calculus/Ax_mul_Bxt.jl b/src/calculus/Ax_mul_Bxt.jl index 035755b..432a958 100644 --- a/src/calculus/Ax_mul_Bxt.jl +++ b/src/calculus/Ax_mul_Bxt.jl @@ -95,8 +95,12 @@ function mul!(y::AbstractArray, J::AdjointOperator{<:Ax_mul_BxtJac}, b::Abstract end # Properties -Base.:(==)(P1::Ax_mul_Bxt{L1, L2, C, D}, P2::Ax_mul_Bxt{L1, L2, C, D}) where {L1, L2, C, D} = P1.A == P2.A && P1.B == P2.B -Base.:(==)(P1::Ax_mul_BxtJac{L1, L2, C, D}, P2::Ax_mul_BxtJac{L1, L2, C, D}) where {L1, L2, C, D} = P1.A == P2.A && P1.B == P2.B +function Base.:(==)(P1::Ax_mul_Bxt{L1, L2, C, D}, P2::Ax_mul_Bxt{L1, L2, C, D}) where {L1, L2, C, D} + return P1.A == P2.A && P1.B == P2.B +end +function Base.:(==)(P1::Ax_mul_BxtJac{L1, L2, C, D}, P2::Ax_mul_BxtJac{L1, L2, C, D}) where {L1, L2, C, D} + return P1.A == P2.A && P1.B == P2.B +end function size(P::Union{Ax_mul_Bxt, Ax_mul_BxtJac}) return ((size(P.A, 1)[1], size(P.B, 1)[1]), size(P.A, 2)) @@ -106,8 +110,8 @@ fun_name(L::Union{Ax_mul_Bxt, Ax_mul_BxtJac}) = fun_name(L.A) * "*" * fun_name(L domain_type(L::Union{Ax_mul_Bxt, Ax_mul_BxtJac}) = domain_type(L.A) codomain_type(L::Union{Ax_mul_Bxt, Ax_mul_BxtJac}) = codomain_type(L.A) -domain_storage_type(L::Union{Ax_mul_Bxt, Ax_mul_BxtJac}) = domain_storage_type(L.A) -codomain_storage_type(L::Union{Ax_mul_Bxt, Ax_mul_BxtJac}) = codomain_storage_type(L.B) +domain_array_type(L::Union{Ax_mul_Bxt, Ax_mul_BxtJac}) = domain_array_type(L.A) +codomain_array_type(L::Union{Ax_mul_Bxt, Ax_mul_BxtJac}) = codomain_array_type(L.B) # utils function permute( diff --git a/src/calculus/Axt_mul_Bx.jl b/src/calculus/Axt_mul_Bx.jl index d2c036d..369b48b 100644 --- a/src/calculus/Axt_mul_Bx.jl +++ b/src/calculus/Axt_mul_Bx.jl @@ -84,7 +84,8 @@ function mul!(y::AbstractArray, P::Axt_mul_Bx{1}, b::AbstractArray) check(y, P, b) mul!(P.bufA, P.A, b) mul!(P.bufB, P.B, b) - return y[1] = dot(P.bufA, P.bufB) + fill!(y, dot(P.bufA, P.bufB)) + return y end function mul!(y::AbstractArray, J::AdjointOperator{<:Axt_mul_BxJac{1}}, b::AbstractArray) @@ -118,8 +119,14 @@ end # Properties -Base.:(==)(P1::Axt_mul_Bx{1, L1, L2, C, D}, P2::Axt_mul_Bx{1, L1, L2, C, D}) where {L1, L2, C, D} = P1.A == P2.A && P1.B == P2.B -Base.:(==)(P1::Axt_mul_BxJac{1, L1, L2, C, D}, P2::Axt_mul_BxJac{1, L1, L2, C, D}) where {L1, L2, C, D} = P1.A == P2.A && P1.B == P2.B +function Base.:(==)(P1::Axt_mul_Bx{1, L1, L2, C, D}, P2::Axt_mul_Bx{1, L1, L2, C, D}) where {L1, L2, C, D} + return P1.A == P2.A && P1.B == P2.B +end +function Base.:(==)( + P1::Axt_mul_BxJac{1, L1, L2, C, D}, P2::Axt_mul_BxJac{1, L1, L2, C, D} + ) where {L1, L2, C, D} + return P1.A == P2.A && P1.B == P2.B +end size(P::Union{Axt_mul_Bx{1}, Axt_mul_BxJac{1}}) = ((1,), size(P.A, 2)) function size(P::Union{Axt_mul_Bx{2}, Axt_mul_BxJac{2}}) @@ -130,8 +137,8 @@ fun_name(L::Union{Axt_mul_Bx, Axt_mul_BxJac}) = fun_name(L.A) * "*" * fun_name(L domain_type(L::Union{Axt_mul_Bx, Axt_mul_BxJac}) = domain_type(L.A) codomain_type(L::Union{Axt_mul_Bx, Axt_mul_BxJac}) = codomain_type(L.A) -domain_storage_type(L::Union{Axt_mul_Bx, Axt_mul_BxJac}) = domain_storage_type(L.A) -codomain_storage_type(L::Union{Axt_mul_Bx, Axt_mul_BxJac}) = codomain_storage_type(L.B) +domain_array_type(L::Union{Axt_mul_Bx, Axt_mul_BxJac}) = domain_array_type(L.A) +codomain_array_type(L::Union{Axt_mul_Bx, Axt_mul_BxJac}) = codomain_array_type(L.B) is_thread_safe(::Axt_mul_Bx) = false # utils diff --git a/src/calculus/BroadCast.jl b/src/calculus/BroadCast.jl index 39a6912..17c75f9 100644 --- a/src/calculus/BroadCast.jl +++ b/src/calculus/BroadCast.jl @@ -11,7 +11,7 @@ struct NoOperatorBroadCast{T, N, M, Threaded, S} <: AbstractBroadCast{T, N, M, T ) where {N, M} Base.Broadcast.check_broadcast_shape(dim_out, reshaped_dim_in) compact = all(reshaped_dim_in[d] == dim_out[d] for d in 1:N) - threaded = threaded && Threads.nthreads() > 1 && compact && prod(dim_in) * sizeof(T) > 2^16 + threaded = threaded && _should_thread(S) && Threads.nthreads() > 1 && compact && prod(dim_in) * sizeof(T) > 2^16 return new{T, N, M, threaded, S}(dim_in, reshaped_dim_in, dim_out) end end @@ -26,13 +26,21 @@ struct OperatorBroadCast{T, N, M, Threaded, Compact, Imask, L, C, D, K} <: Abstr A, dim_out::NTuple{M, Int}; threaded::Bool = true ) where {M} Base.Broadcast.check_broadcast_shape(dim_out, size(A, 1)) - threaded = threaded && Threads.nthreads() > 1 + threaded = threaded && _should_thread(A) && Threads.nthreads() > 1 N = ndims(A, 1) T = codomain_type(A) dim_in = size(A, 1) - Imask, broadcast_dims, idxs, compact = _setup_broadcast_indices(dim_out, dim_in, N) + Imask = Tuple(d ≤ N && (dim_out[d] == dim_in[d]) for d in 1:M) + broadcast_dims = Tuple(Imask[d] ? 1 : dim_out[d] for d in eachindex(dim_out)) + idxs = CartesianIndices(broadcast_dims) + compact = all(Imask[1:N]) bufC = allocate_in_codomain(A) - A, bufD = _setup_broadcast_buffers(A, threaded) + if threaded + bufD = [allocate_in_domain(A) for _ in 1:Threads.nthreads()] + A = [i == 1 ? A : AbstractOperators.copy_operator(A) for i in 1:Threads.nthreads()] + else + bufD = allocate_in_domain(A) + end L = typeof(A) C = typeof(bufC) D = typeof(bufD) @@ -41,23 +49,6 @@ struct OperatorBroadCast{T, N, M, Threaded, Compact, Imask, L, C, D, K} <: Abstr end end -function _setup_broadcast_indices(dim_out::NTuple{M, Int}, dim_in::NTuple{N, Int}, nd::Int) where {M, N} - imask = Tuple(d <= nd && (dim_out[d] == dim_in[d]) for d in 1:M) - broadcast_dims = Tuple(imask[d] ? 1 : dim_out[d] for d in eachindex(dim_out)) - idxs = CartesianIndices(broadcast_dims) - compact = all(imask[1:nd]) - return imask, broadcast_dims, idxs, compact -end - -function _setup_broadcast_buffers(A::AbstractOperator, threaded::Bool) - if threaded - ops = [i == 1 ? A : copy_op(A) for i in 1:Threads.nthreads()] - bufs = [allocate_in_domain(A) for _ in 1:Threads.nthreads()] - return ops, bufs - end - return A, allocate_in_domain(A) -end - # Constructors """ @@ -90,7 +81,7 @@ function BroadCast( elseif is_eye(A) dim_in = size(A, 2) reshaped_dim_in = ntuple(d -> d <= ndims(A, 1) ? size(A, 1)[d] : 1, length(dim_out)) - return NoOperatorBroadCast(domain_type(A), domain_storage_type(A), dim_in, reshaped_dim_in, dim_out; threaded) + return NoOperatorBroadCast(domain_type(A), domain_array_type(A), dim_in, reshaped_dim_in, dim_out; threaded) else return OperatorBroadCast(A, dim_out; threaded) end @@ -109,18 +100,18 @@ function tbroadcast!(y, x) end # NoOperatorBroadCast -function mul!(y::AbstractArray, A::NoOperatorBroadCast{T, N, M, false}, b::AbstractArray) where {T, N, M} +function mul!(y, A::NoOperatorBroadCast{T, N, M, false}, b) where {T, N, M} check(y, A, b) - b_reshaped = reshape(b, A.reshaped_dim_in) - return y .= b_reshaped + b = reshape(b, A.reshaped_dim_in) + return y .= b # non-threaded broadcasting end -function mul!(y::AbstractArray, A::NoOperatorBroadCast{T, N, M, true}, b::AbstractArray) where {T, N, M} +function mul!(y, A::NoOperatorBroadCast{T, N, M, true}, b) where {T, N, M} check(y, A, b) return tbroadcast!(y, b) # threaded broadcasting, handles reshaping (threading is only enabled when compact) end -function mul!(y::AbstractArray, A::AdjointOperator{<:NoOperatorBroadCast}, b::AbstractArray) # there is no threaded option here +function mul!(y, A::AdjointOperator{<:NoOperatorBroadCast}, b) # there is no threaded option here check(y, A, b) y = reshape(y, A.A.reshaped_dim_in) return sum!(y, b) @@ -128,23 +119,23 @@ end # OperatorBroadCast -function mul!(y::AbstractArray, R::OperatorBroadCast, b::AbstractArray) # Non-threaded +function mul!(y, R::OperatorBroadCast, b) # Non-threaded check(y, R, b) mul!(R.bufC, R.A, b) - return y .= R.bufC + return y .= R.bufC # non-threaded broadcasting end -function mul!(y::AbstractArray, R::OperatorBroadCast{T, N, M, true, Compact}, b::AbstractArray) where {T, N, M, Compact} # Threaded +function mul!(y, R::OperatorBroadCast{T, N, M, true, Compact}, b) where {T, N, M, Compact} # Threaded check(y, R, b) mul!(R.bufC, R.A[1], b) if Compact return tbroadcast!(y, R.bufC) # threaded broadcasting else - return y .= R.bufC + return y .= R.bufC # non-threaded broadcasting end end -function mul!(y::AbstractArray, A::AdjointOperator{<:OperatorBroadCast{T, N, M, false}}, b::AbstractArray) where {T, N, M} # Non-threaded +function mul!(y, A::AdjointOperator{<:OperatorBroadCast{T, N, M, false}}, b) where {T, N, M} # Non-threaded check(y, A, b) R = A.A for idx in R.idxs @@ -162,21 +153,16 @@ function mul!(y::AbstractArray, A::AdjointOperator{<:OperatorBroadCast{T, N, M, return y end -function mul!(y::AbstractArray, A::AdjointOperator{<:OperatorBroadCast{T, N, M, true}}, b::AbstractArray) where {T, N, M} # Threaded +function mul!(y, A::AdjointOperator{<:OperatorBroadCast{T, N, M, true}}, b) where {T, N, M} # Threaded check(y, A, b) - _threaded_broadcast_adjoint!(y, A.A, b) - return y -end - -function _threaded_broadcast_adjoint!(y, R::OperatorBroadCast, b) + R = A.A fill!(y, 0) lock = ReentrantLock() thread_count = min(Threads.nthreads(), length(R.idxs)) - chunk = cld(length(R.idxs), thread_count) - + batch_size = length(R.idxs) / thread_count @threads for t in 1:thread_count - idx_start = (t - 1) * chunk + 1 - idx_end = min(t * chunk, length(R.idxs)) + idx_start = max(1, floor(Int, (t - 1) * batch_size + 1)) + idx_end = min(length(R.idxs), floor(Int, t * batch_size)) for i in idx_start:idx_end b_slice = get_input_slice(R, R.idxs[i], b) if size(b_slice) != size(R.A[t], 1) @@ -186,7 +172,6 @@ function _threaded_broadcast_adjoint!(y, R::OperatorBroadCast, b) @lock lock y .+= R.bufD[t] end end - return y end @@ -210,12 +195,12 @@ domain_type(R::OperatorBroadCast{T, N, M, true}) where {T, N, M} = domain_type(R codomain_type(::NoOperatorBroadCast{T}) where {T} = T codomain_type(R::OperatorBroadCast{T, N, M, false}) where {T, N, M} = codomain_type(R.A) codomain_type(R::OperatorBroadCast{T, N, M, true}) where {T, N, M} = codomain_type(R.A[1]) -domain_storage_type(::NoOperatorBroadCast{T, N, M, Threaded, S}) where {T, N, M, Threaded, S} = S -domain_storage_type(R::OperatorBroadCast{T, N, M, false}) where {T, N, M} = domain_storage_type(R.A) -domain_storage_type(R::OperatorBroadCast{T, N, M, true}) where {T, N, M} = domain_storage_type(R.A[1]) -codomain_storage_type(::NoOperatorBroadCast{T, N, M, Threaded, S}) where {T, N, M, Threaded, S} = S -codomain_storage_type(R::OperatorBroadCast{T, N, M, false}) where {T, N, M} = codomain_storage_type(R.A) -codomain_storage_type(R::OperatorBroadCast{T, N, M, true}) where {T, N, M} = codomain_storage_type(R.A[1]) +domain_array_type(::NoOperatorBroadCast{T, N, M, Threaded, S}) where {T, N, M, Threaded, S} = S +domain_array_type(R::OperatorBroadCast{T, N, M, false}) where {T, N, M} = domain_array_type(R.A) +domain_array_type(R::OperatorBroadCast{T, N, M, true}) where {T, N, M} = domain_array_type(R.A[1]) +codomain_array_type(::NoOperatorBroadCast{T, N, M, Threaded, S}) where {T, N, M, Threaded, S} = S +codomain_array_type(R::OperatorBroadCast{T, N, M, false}) where {T, N, M} = codomain_array_type(R.A) +codomain_array_type(R::OperatorBroadCast{T, N, M, true}) where {T, N, M} = codomain_array_type(R.A[1]) is_thread_safe(::NoOperatorBroadCast) = true is_thread_safe(::OperatorBroadCast) = false @@ -233,11 +218,11 @@ fun_name(R::OperatorBroadCast{T, N, M, true}) where {T, N, M} = "." * fun_name(R remove_displacement(R::NoOperatorBroadCast) = R function remove_displacement(R::OperatorBroadCast{T, N, M, false, Imask}) where {T, N, M, Imask} new_A = remove_displacement(R.A) - return OperatorBroadCast(new_A, R.dim_out, threaded = false) + return OperatorBroadCast(new_A, R.dim_out; threaded = false) end function remove_displacement(R::OperatorBroadCast{T, N, M, true, Imask}) where {T, N, M, Imask} new_A = remove_displacement(R.A[1]) - return OperatorBroadCast(new_A, R.dim_out, threaded = true) + return OperatorBroadCast(new_A, R.dim_out; threaded = true) end has_fast_opnorm(::NoOperatorBroadCast) = true @@ -258,6 +243,23 @@ function permute(R::OperatorBroadCast{T, N, M, true}, p::AbstractVector{Int}) wh return BroadCast(permute(R.A[1], p), R.dim_out; threaded = true) end +function _copy_operator_impl( + op::NoOperatorBroadCast{T, N, M, Th, S}; array_type = nothing, threaded = nothing + ) where {T, N, M, Th, S} + new_threaded = threaded === nothing ? Th : threaded + new_S = array_type === nothing ? S : array_type + return NoOperatorBroadCast(T, new_S, op.dim_in, op.reshaped_dim_in, op.dim_out; threaded = new_threaded) +end + +function _copy_operator_impl( + op::OperatorBroadCast{T, N, M, Th}; array_type = nothing, threaded = nothing + ) where {T, N, M, Th} + new_threaded = threaded === nothing ? Th : threaded + inner_op = Th ? op.A[1] : op.A + new_op = copy_operator(inner_op; array_type, threaded) + return BroadCast(new_op, op.dim_out; threaded = new_threaded) +end + @generated function get_input_slice( ::OperatorBroadCast{T, N, M, Threaded, Compact, Imask}, idx::CartesianIndex, b ) where {T, N, M, Threaded, Compact, Imask} diff --git a/src/calculus/Compose.jl b/src/calculus/Compose.jl index 47353b2..ebe7d8e 100644 --- a/src/calculus/Compose.jl +++ b/src/calculus/Compose.jl @@ -35,12 +35,10 @@ struct Compose{N, M, L <: Tuple, T <: Tuple} <: AbstractOperator while i < length(A) should_be_combined = false triple_combination = false - if ( - A[i + 1] isa AdjointOperator && - A[i + 1].A == A[i] && - has_optimized_normalop(A[i]) + if (A[i + 1] isa AdjointOperator && A[i + 1].A == A[i] && has_optimized_normalop(A[i])) + DEBUG_COMPOSE[] && print( + "Replacing $(typeof((A[i])).name.wrapper) and $(typeof((A[i + 1])).name.wrapper) with normal operator", ) - DEBUG_COMPOSE[] && print("Replacing $(typeof((A[i])).name.wrapper) and $(typeof((A[i + 1])).name.wrapper) with normal operator") new_op = get_normal_op(A[i]) should_be_combined = true elseif can_be_combined(A[i + 1], A[i]) @@ -113,9 +111,9 @@ function Compose(L1::AbstractOperator, L2::AbstractOperator) msg = "cannot compose operators with different domain and codomain types" throw(DomainError((domain_type(L1), codomain_type(L2)), msg)) end - if domain_storage_type(L1) != codomain_storage_type(L2) + if domain_array_type(L1) != codomain_array_type(L2) msg = "cannot compose operators with different input and output storage types" - throw(DomainError((domain_storage_type(L1), codomain_storage_type(L2)), msg)) + throw(DomainError((domain_array_type(L1), codomain_array_type(L2)), msg)) end if L1 isa AdjointOperator && L1.A == L2 && has_optimized_normalop(L2) return get_normal_op(L2) @@ -132,7 +130,7 @@ function Compose(L1::AbstractOperator, L2::AbstractOperator) end compatible_bufs = x -> - x isa codomain_storage_type(L1) && + x isa codomain_array_type(L1) && size(x) == size(L2, 1) && eltype(x) == codomain_type(L2) new_buf_pos = findfirst(compatible_bufs, available_bufs) @@ -198,6 +196,7 @@ _ndoms_from_type(::Type{<:Compose{N, M, L, T}}, dim::Int) where {N, M, L <: Tupl end end return ex = quote + check(y, L, b) $ex mul!(y, L.A[N], L.buf[M]) return y @@ -215,6 +214,7 @@ end end end return ex = quote + check(y, L, b) $ex mul!(y, L.A.A[1]', L.A.buf[1]) return y @@ -243,8 +243,8 @@ fun_name(L::Compose) = length(L.A) == 2 ? fun_name(L.A[2]) * "*" * fun_name(L.A[ domain_type(L::Compose) = domain_type(L.A[1]) codomain_type(L::Compose) = codomain_type(L.A[end]) -domain_storage_type(L::Compose) = domain_storage_type(L.A[1]) -codomain_storage_type(L::Compose) = codomain_storage_type(L.A[end]) +domain_array_type(L::Compose) = domain_array_type(L.A[1]) +codomain_array_type(L::Compose) = codomain_array_type(L.A[end]) is_linear(L::Compose) = all(is_linear.(L.A)) function is_diagonal(L::Compose) @@ -302,3 +302,9 @@ end remove_displacement(C::Compose) = Compose(remove_displacement.(C.A), C.buf) get_operators(C::Compose) = C.A + +function _copy_operator_impl(op::Compose; array_type = nothing, threaded = nothing) + new_bufs = tuple([_convert_buffer(b, array_type) for b in op.buf]...) + new_ops = tuple([copy_operator(a; array_type, threaded) for a in op.A]...) + return Compose(new_ops, new_bufs) +end diff --git a/src/calculus/DCAT.jl b/src/calculus/DCAT.jl index 8e40099..061b451 100644 --- a/src/calculus/DCAT.jl +++ b/src/calculus/DCAT.jl @@ -26,11 +26,58 @@ struct DCAT{ L <: NTuple{N, AbstractOperator}, P1 <: NTuple{N, Union{Int, Tuple}}, P2 <: NTuple{N, Union{Int, Tuple}}, + DS <: AbstractArray, # domain storage type (fixed at construction) + CS <: AbstractArray, # codomain storage type (fixed at construction) } <: AbstractOperator A::L idxD::P1 idxC::P2 - DCAT(A::L, idxD::P1, idxC::P2) where {L, P1, P2} = new{length(A), L, P1, P2}(A, idxD, idxC) + function DCAT(A::L, idxD::P1, idxC::P2) where {L, P1, P2} + DS = _compute_dcat_ds(A, idxD) + CS = _compute_dcat_cs(A, idxC) + return new{length(A), L, P1, P2, DS, CS}(A, idxD, idxC) + end +end + +function _compute_dcat_ds(A, idxD) + ds_list = [d <: ArrayPartition ? [d.parameters[2].types...] : d for d in domain_array_type.(A)] + domain = vcat(ds_list...) + p = vcat([[idx...] for idx in idxD]...) + invpermute!(domain, p) + T = promote_type(map(_storage_eltype, domain)...) + return ArrayPartition{T, Tuple{domain...}} +end + +function _compute_dcat_cs(A, idxC) + cs_list = [d <: ArrayPartition ? [d.parameters[2].types...] : d for d in codomain_array_type.(A)] + codomain = vcat(cs_list...) + p = vcat([[idx...] for idx in idxC]...) + invpermute!(codomain, p) + T = promote_type(map(_storage_eltype, codomain)...) + return ArrayPartition{T, Tuple{codomain...}} +end + +# Flatten DCAT index structure (tuple of ints/tuples) into a flat tuple of ints. +# Used in _dcat_apply_invperm to avoid allocating intermediate Vectors. +@inline _dcat_flatten_idxs(::Tuple{}) = () +@inline function _dcat_flatten_idxs(idxs::Tuple) + head = first(idxs) + tail = _dcat_flatten_idxs(Base.tail(idxs)) + return head isa Integer ? (Int(head), tail...) : (head..., tail...) +end + +# Linear scan for value v in tuple t, returning its 1-based position. +@inline _dcat_find_in_tuple(::Tuple{}, ::Int, ::Int) = 0 +@inline function _dcat_find_in_tuple(t::Tuple, v::Int, i::Int = 1) + return first(t) == v ? i : _dcat_find_in_tuple(Base.tail(t), v, i + 1) +end + +# Apply the inverse permutation encoded in idxs to natural, using only tuple +# operations so that no Vector is allocated on the hot path. +function _dcat_apply_invperm(natural::Tuple, idxs::Tuple) + N = length(natural) + p = _dcat_flatten_idxs(idxs) + return ntuple(j -> natural[_dcat_find_in_tuple(p, j)], Val(N)) end # Constructors @@ -83,70 +130,70 @@ end ) where {N, L, P1, P2} # extract stuff - ex = :(y = yy.x; b = bb.x) + ex = :(check(yy, H, bb); y = yy.x; b = bb.x) for i in 1:N if fieldtype(P2, i) <: Int # flatten operator # build mul!(y[H.idxC[i]], H.A[i], b) - yy = :(y[H.idxC[$i]]) + yyi = :(y[H.idxC[$i]]) else # stacked operator # build mul!(( y[H.idxC[i][1]], y[H.idxC[i][2]] ... ), H.A[i], b) - yy = [:(y[H.idxC[$i][$ii]]) for ii in eachindex(fieldnames(fieldtype(P2, i)))] - yy = :(ArrayPartition($(yy...))) + yyi = [:(y[H.idxC[$i][$ii]]) for ii in eachindex(fieldnames(fieldtype(P2, i)))] + yyi = :(ArrayPartition($(yyi...))) end if fieldtype(P1, i) <: Int # flatten operator # build mul!(H.buf, H.A[i], b[H.idxD[i]]) - bb = :(b[H.idxD[$i]]) + bbi = :(b[H.idxD[$i]]) else # stacked operator # build mul!(H.buf, H.A[i],( b[H.idxD[i][1]], b[H.idxD[i][2]] ... )) - bb = [:(b[H.idxD[$i][$ii]]) for ii in eachindex(fieldnames(fieldtype(P1, i)))] - bb = :(ArrayPartition($(bb...))) + bbi = [:(b[H.idxD[$i][$ii]]) for ii in eachindex(fieldnames(fieldtype(P1, i)))] + bbi = :(ArrayPartition($(bbi...))) end - ex = :($ex; mul!($yy, H.A[$i], $bb)) + ex = :($ex; mul!($yyi, H.A[$i], $bbi)) end - ex = :($ex; return y) + ex = :($ex; return yy) return ex end @generated function mul!( - yy::ArrayPartition, A::AdjointOperator{DCAT{N, L, P1, P2}}, bb::ArrayPartition + yy::ArrayPartition, A::AdjointOperator{<:DCAT{N, L, P1, P2}}, bb::ArrayPartition ) where {N, L, P1, P2} # extract stuff - ex = :(H = A.A; y = yy.x; b = bb.x) + ex = :(check(yy, A, bb); H = A.A; y = yy.x; b = bb.x) for i in 1:N if fieldtype(P1, i) <: Int # flatten operator # build mul!(y[H.idxD[i]], H.A[i]', b) - yy = :(y[H.idxD[$i]]) + yyi = :(y[H.idxD[$i]]) else # stacked operator # build mul!(( y[H.idxD[i][1]], y[H.idxD[i][2]] ... ), H.A[i]', b) - yy = [:(y[H.idxD[$i][$ii]]) for ii in eachindex(fieldnames(fieldtype(P1, i)))] - yy = :(ArrayPartition($(yy...))) + yyi = [:(y[H.idxD[$i][$ii]]) for ii in eachindex(fieldnames(fieldtype(P1, i)))] + yyi = :(ArrayPartition($(yyi...))) end if fieldtype(P2, i) <: Int # flatten operator # build mul!(H.buf, H.A[i]', b[H.idxC[i]]) - bb = :(b[H.idxC[$i]]) + bbi = :(b[H.idxC[$i]]) else # stacked operator # build mul!(H.buf, H.A[i]',( b[H.idxC[i][1]], b[H.idxC[i][2]] ... )) - bb = [:(b[H.idxC[$i][$ii]]) for ii in eachindex(fieldnames(fieldtype(P2, i)))] - bb = :(ArrayPartition($(bb...))) + bbi = [:(b[H.idxC[$i][$ii]]) for ii in eachindex(fieldnames(fieldtype(P2, i)))] + bbi = :(ArrayPartition($(bbi...))) end - ex = :($ex; mul!($yy, H.A[$i]', $bb)) + ex = :($ex; mul!($yyi, H.A[$i]', $bbi)) end - ex = :($ex; return y) + ex = :($ex; return yy) return ex end @@ -158,17 +205,36 @@ end # Properties Base.:(==)(H1::DCAT{N, L1, P1, P2}, H2::DCAT{N, L2, P1, P2}) where {N, L1, L2, P1, P2} = H1.A == H2.A && H1.idxD == H2.idxD && H1.idxC == H2.idxC -size(H::DCAT) = size(H, 1), size(H, 2) -function size(H::DCAT, i::Int) - sz = [] - for s in size.(H.A, i) - eltype(s) <: Int ? push!(sz, s) : push!(sz, s...) +@generated function size(H::DCAT{N, L, P1, P2}) where {N, L, P1, P2} + cod_exprs = [] + for i in 1:N + Pi = fieldtype(P2, i) + if Pi <: Integer + push!(cod_exprs, :(size(H.A[$i], 1))) + else + for ii in eachindex(fieldnames(Pi)) + push!(cod_exprs, :(size(H.A[$i], 1)[$ii])) + end + end end - p = vcat([[idx...] for idx in (i == 1 ? H.idxC : H.idxD)]...) - invpermute!(sz, p) - - return (sz...,) + dom_exprs = [] + for i in 1:N + Pi = fieldtype(P1, i) + if Pi <: Integer + push!(dom_exprs, :(size(H.A[$i], 2))) + else + for ii in eachindex(fieldnames(Pi)) + push!(dom_exprs, :(size(H.A[$i], 2)[$ii])) + end + end + end + natural_cod = Expr(:tuple, cod_exprs...) + natural_dom = Expr(:tuple, dom_exprs...) + return :( + _dcat_apply_invperm($natural_cod, H.idxC), + _dcat_apply_invperm($natural_dom, H.idxD), + ) end function fun_name(L::DCAT) @@ -179,31 +245,56 @@ function fun_name(L::DCAT) end end -function domain_type(H::DCAT) - domain = vcat([typeof(d) <: Tuple ? [d...] : d for d in domain_type.(H.A)]...) - p = vcat([[idx...] for idx in H.idxD]...) - invpermute!(domain, p) - return (domain...,) +@generated function ndoms(H::DCAT{N, L}) where {N, L} + nc = sum(_ndoms_from_type(fieldtype(L, i), 1) for i in 1:N) + nd = sum(_ndoms_from_type(fieldtype(L, i), 2) for i in 1:N) + return :(($(nc), $(nd))) end -function domain_storage_type(H::DCAT) - domain = vcat([d <: ArrayPartition ? [d.parameters[2].types...] : d for d in domain_storage_type.(H.A)]...) - p = vcat([[idx...] for idx in H.idxD]...) - invpermute!(domain, p) - T = promote_type(domain_type(H)...) - return ArrayPartition{T, Tuple{domain...}} + +@generated function ndoms(H::DCAT{N, L}, dim::Int) where {N, L} + nc = sum(_ndoms_from_type(fieldtype(L, i), 1) for i in 1:N) + nd = sum(_ndoms_from_type(fieldtype(L, i), 2) for i in 1:N) + return :(dim == 1 ? $(nc) : $(nd)) end -function codomain_type(H::DCAT) - codomain = vcat([typeof(d) <: Tuple ? [d...] : d for d in codomain_type.(H.A)]...) - p = vcat([[idx...] for idx in H.idxC]...) - invpermute!(codomain, p) - return (codomain...,) + +@generated function domain_type(H::DCAT{N, L, P1, P2}) where {N, L, P1, P2} + exprs = [] + for i in 1:N + Pi = fieldtype(P1, i) + if Pi <: Integer + push!(exprs, :(domain_type(H.A[$i]))) + else + for ii in eachindex(fieldnames(Pi)) + push!(exprs, :(domain_type(H.A[$i])[$ii])) + end + end + end + natural_expr = Expr(:tuple, exprs...) + return :(_dcat_apply_invperm($natural_expr, H.idxD)) end -function codomain_storage_type(H::DCAT) - codomain = vcat([d <: ArrayPartition ? [d.parameters[2].types...] : d for d in codomain_storage_type.(H.A)]...) - p = vcat([[idx...] for idx in H.idxC]...) - invpermute!(codomain, p) - T = promote_type(codomain_type(H)...) - return ArrayPartition{T, Tuple{codomain...}} + +function domain_array_type(::DCAT{N, L, P1, P2, DS, CS}) where {N, L, P1, P2, DS, CS} + return DS +end + +@generated function codomain_type(H::DCAT{N, L, P1, P2}) where {N, L, P1, P2} + exprs = [] + for i in 1:N + Pi = fieldtype(P2, i) + if Pi <: Integer + push!(exprs, :(codomain_type(H.A[$i]))) + else + for ii in eachindex(fieldnames(Pi)) + push!(exprs, :(codomain_type(H.A[$i])[$ii])) + end + end + end + natural_expr = Expr(:tuple, exprs...) + return :(_dcat_apply_invperm($natural_expr, H.idxC)) +end + +function codomain_array_type(::DCAT{N, L, P1, P2, DS, CS}) where {N, L, P1, P2, DS, CS} + return CS end is_thread_safe(H::DCAT) = all(is_thread_safe.(H.A)) diff --git a/src/calculus/HCAT.jl b/src/calculus/HCAT.jl index 7f0be7c..ba1ce10 100644 --- a/src/calculus/HCAT.jl +++ b/src/calculus/HCAT.jl @@ -44,6 +44,7 @@ struct HCAT{ L <: NTuple{N, AbstractOperator}, P <: Tuple, C <: AbstractArray, + DS <: AbstractArray, # domain storage type (fixed at construction) } <: AbstractOperator A::L # tuple of AbstractOperators idxs::P # indices @@ -60,10 +61,20 @@ struct HCAT{ if any([codomain_type(A[1]) != codomain_type(a) for a in A]) throw(error("operators must all share the same codomain_type!")) end - return new{N, L, P, C}(A, idxs, buf) + DS = _compute_hcat_ds(A, idxs) + return new{N, L, P, C, DS}(A, idxs, buf) end end +function _compute_hcat_ds(A, idxs) + ds_list = [d <: ArrayPartition ? [d.parameters[2].types...] : d for d in domain_array_type.(A)] + domain = vcat(ds_list...) + p = vcat([[idx...] for idx in idxs]...) + invpermute!(domain, p) + T = promote_type(map(_storage_eltype, domain)...) + return ArrayPartition{T, Tuple{domain...}} +end + function HCAT(A::Vararg{AbstractOperator}) if any((<:).(typeof.(A), HCAT)) #there are HCATs in A AA = () @@ -305,14 +316,8 @@ end return :(_hcat_apply_invperm($natural_expr, H.idxs)) end codomain_type(L::HCAT) = codomain_type.(Ref(L.A[1])) -function domain_storage_type(H::HCAT) - domain = vcat([d <: ArrayPartition ? [d.parameters[2].types...] : d for d in domain_storage_type.(H.A)]...) - p = vcat([[idx...] for idx in H.idxs]...) - invpermute!(domain, p) - T = promote_type(domain_type(H)...) - return ArrayPartition{T, Tuple{domain...}} -end -codomain_storage_type(L::HCAT) = codomain_storage_type.(Ref(L.A[1])) +domain_array_type(::HCAT{N, L, P, C, DS}) where {N, L, P, C, DS} = DS +codomain_array_type(L::HCAT) = codomain_array_type.(Ref(L.A[1])) is_linear(L::HCAT) = all(is_linear.(L.A)) is_AAc_diagonal(L::HCAT) = all(is_AAc_diagonal.(L.A)) @@ -358,3 +363,9 @@ function permute(H::HCAT, p::AbstractVector{Int}) end remove_displacement(H::HCAT) = HCAT(remove_displacement.(H.A), H.idxs, H.buf) + +function _copy_operator_impl(op::HCAT; array_type = nothing, threaded = nothing) + new_buf = _convert_buffer(op.buf, array_type) + new_ops = tuple([copy_operator(a; array_type, threaded) for a in op.A]...) + return HCAT(new_ops, op.idxs, new_buf) +end diff --git a/src/calculus/HadamardProd.jl b/src/calculus/HadamardProd.jl index 4229c5e..712b929 100644 --- a/src/calculus/HadamardProd.jl +++ b/src/calculus/HadamardProd.jl @@ -85,15 +85,17 @@ function mul!(y::AbstractArray, J::AdjointOperator{<:HadamardProdJac}, b::Abstra end # Properties -Base.:(==)(P1::HadamardProd{L1, L2, C, D}, P2::HadamardProd{L1, L2, C, D}) where {L1, L2, C, D} = P1.A == P2.A && P1.B == P2.B +function Base.:(==)(P1::HadamardProd{L1, L2, C, D}, P2::HadamardProd{L1, L2, C, D}) where {L1, L2, C, D} + return P1.A == P2.A && P1.B == P2.B +end size(P::Union{HadamardProd, HadamardProdJac}) = (size(P.A, 1), size(P.A, 2)) fun_name(L::Union{HadamardProd, HadamardProdJac}) = fun_name(L.A) * ".*" * fun_name(L.B) domain_type(L::Union{HadamardProd, HadamardProdJac}) = domain_type(L.A) codomain_type(L::Union{HadamardProd, HadamardProdJac}) = codomain_type(L.A) -domain_storage_type(L::Union{HadamardProd, HadamardProdJac}) = domain_storage_type(L.A) -codomain_storage_type(L::Union{HadamardProd, HadamardProdJac}) = codomain_storage_type(L.A) +domain_array_type(L::Union{HadamardProd, HadamardProdJac}) = domain_array_type(L.A) +codomain_array_type(L::Union{HadamardProd, HadamardProdJac}) = codomain_array_type(L.A) # utils function permute( @@ -109,3 +111,23 @@ function remove_displacement(P::HadamardProd) remove_displacement(P.A), remove_displacement(P.B), P.bufA, P.bufB, P.bufD ) end + +function _copy_operator_impl(op::HadamardProd; array_type = nothing, threaded = nothing) + new_bufA = _convert_buffer(op.bufA, array_type) + new_bufB = _convert_buffer(op.bufB, array_type) + new_bufD = _convert_buffer(op.bufD, array_type) + new_A = copy_operator(op.A; array_type, threaded) + new_B = copy_operator(op.B; array_type, threaded) + return HadamardProd(new_A, new_B, new_bufA, new_bufB, new_bufD) +end + +function _copy_operator_impl(op::HadamardProdJac; array_type = nothing, threaded = nothing) + new_bufA = _convert_buffer(op.bufA, array_type) + new_bufB = _convert_buffer(op.bufB, array_type) + new_bufD = _convert_buffer(op.bufD, array_type) + new_A = copy_operator(op.A; array_type, threaded) + new_B = copy_operator(op.B; array_type, threaded) + return HadamardProdJac{typeof(new_A), typeof(new_B), typeof(new_bufA), typeof(new_bufD)}( + new_A, new_B, new_bufA, new_bufB, new_bufD + ) +end diff --git a/src/calculus/Jacobian.jl b/src/calculus/Jacobian.jl index 96ff7f3..22b86fa 100644 --- a/src/calculus/Jacobian.jl +++ b/src/calculus/Jacobian.jl @@ -94,10 +94,10 @@ Jacobian(T::Transpose{<:AbstractOperator}, ::AbstractArray) = T #Jacobian of BroadCast Jacobian(L::NoOperatorBroadCast, ::AbstractArray) = L function Jacobian(B::OperatorBroadCast{T, N, M, false}, x::AbstractArray) where {T, N, M} - return OperatorBroadCast(Jacobian(B.A, x), B.dim_out, threaded = false) + return OperatorBroadCast(Jacobian(B.A, x), B.dim_out; threaded = false) end function Jacobian(B::OperatorBroadCast{T, N, M, true}, x::AbstractArray) where {T, N, M} - return OperatorBroadCast(Jacobian(B.A[1], x), B.dim_out, threaded = true) + return OperatorBroadCast(Jacobian(B.A[1], x), B.dim_out; threaded = true) end #Jacobian of AffineAdd Jacobian(B::AffineAdd, x) = Jacobian(B.A, x) @@ -109,5 +109,5 @@ size(L::Jacobian) = size(L.A, 1), size(L.A, 2) domain_type(L::Jacobian) = domain_type(L.A) codomain_type(L::Jacobian) = codomain_type(L.A) -domain_storage_type(L::Jacobian) = domain_storage_type(L.A) -codomain_storage_type(L::Jacobian) = codomain_storage_type(L.A) +domain_array_type(L::Jacobian) = domain_array_type(L.A) +codomain_array_type(L::Jacobian) = codomain_array_type(L.A) diff --git a/src/calculus/OperatorWrapper.jl b/src/calculus/OperatorWrapper.jl new file mode 100644 index 0000000..e7a1e26 --- /dev/null +++ b/src/calculus/OperatorWrapper.jl @@ -0,0 +1,136 @@ +export OperatorWrapper + +""" + OperatorWrapper{Op, DB, CB, DS, CS} + +A wrapper that adapts any operator's storage type independently of the wrapped operator. +Useful for making CPU operators transparently usable with GPU arrays. + +When `mul!(y, wrapper, x)` is called the wrapper: +1. Copies `x` to the preallocated domain buffer (same storage as the inner op) +2. Executes the inner operator on its native buffers +3. Copies the result to `y` + +GPU `mul!` overrides are provided by the GpuExt extension (requires loading GPUArrays). + +# Construction + OperatorWrapper(op::AbstractOperator; array_type = Array) + +`array_type` sets the *outer* storage type reported by the wrapper (used by `check`). +This is independent of the wrapped operator's storage type, which governs the internal +CPU buffers. + +# Notes +- Thread safety: each wrapper has its own buffers; use `copy_operator` for parallel use. + +# Example +```jldoctest +julia> using AbstractOperators + +julia> op = FiniteDiff(Float64, (8,), 1) +δx ℝ^8 -> ℝ^7 + +julia> w = OperatorWrapper(op) +CPU[δx] ℝ^8 -> ℝ^7 + +julia> size(w) == size(op) +true +``` +""" +struct OperatorWrapper{ + Op <: AbstractOperator, + DB <: AbstractArray, # inner domain buffer (always CPU-backed) + CB <: AbstractArray, # inner codomain buffer (always CPU-backed) + DS <: AbstractArray, # outer domain storage type (reported to check) + CS <: AbstractArray, # outer codomain storage type (reported to check) + } <: LinearOperator + op::Op + dom_buf::DB + cod_buf::CB +end + +""" + OperatorWrapper(op::AbstractOperator; array_type = Array) + +Wrap `op`, preallocating internal CPU buffers from its domain/codomain. +`array_type` sets the outer storage type used in `check` (default: `Array`). +""" +function OperatorWrapper(op::AbstractOperator; array_type::Type = Array) + dom_buf = allocate_in_domain(op) + cod_buf = allocate_in_codomain(op) + S = _array_wrapper_type(array_type) + T_dom = domain_type(op) + T_cod = codomain_type(op) + N_dom = ndims(dom_buf) + N_cod = ndims(cod_buf) + DS = S{T_dom} + CS = S{T_cod} + return OperatorWrapper{typeof(op), typeof(dom_buf), typeof(cod_buf), DS, CS}(op, dom_buf, cod_buf) +end + +# CPU mul! — copies through internal buffers +function mul!(y::AbstractArray, A::OperatorWrapper, x::AbstractArray) + check(y, A, x) + copyto!(A.dom_buf, x) + mul!(A.cod_buf, A.op, A.dom_buf) + copyto!(y, A.cod_buf) + return y +end + +function mul!(y::AbstractArray, Ac::AdjointOperator{<:OperatorWrapper}, x::AbstractArray) + check(y, Ac, x) + A = Ac.A + copyto!(A.cod_buf, x) + mul!(A.dom_buf, A.op', A.cod_buf) + copyto!(y, A.dom_buf) + return y +end + +# Properties — size/types delegate to inner op + +Base.size(A::OperatorWrapper) = size(A.op) +fun_name(A::OperatorWrapper) = "CPU[$(fun_name(A.op))]" + +domain_type(A::OperatorWrapper) = domain_type(A.op) +codomain_type(A::OperatorWrapper) = codomain_type(A.op) + +# Outer storage types are fixed at construction via DS/CS type parameters. +domain_array_type(::OperatorWrapper{Op, DB, CB, DS, CS}) where {Op, DB, CB, DS, CS} = DS +codomain_array_type(::OperatorWrapper{Op, DB, CB, DS, CS}) where {Op, DB, CB, DS, CS} = CS + +# Forward all predicates to the wrapped operator. +import OperatorCore: + is_linear, is_eye, is_null, is_diagonal, + is_AcA_diagonal, is_AAc_diagonal, diag_AcA, diag_AAc, + is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank, + is_symmetric, is_positive_definite, is_positive_semidefinite + +is_linear(A::OperatorWrapper) = is_linear(A.op) +is_eye(A::OperatorWrapper) = is_eye(A.op) +is_null(A::OperatorWrapper) = is_null(A.op) +is_diagonal(A::OperatorWrapper) = is_diagonal(A.op) +is_AcA_diagonal(A::OperatorWrapper) = is_AcA_diagonal(A.op) +is_AAc_diagonal(A::OperatorWrapper) = is_AAc_diagonal(A.op) +diag_AcA(A::OperatorWrapper) = diag_AcA(A.op) +diag_AAc(A::OperatorWrapper) = diag_AAc(A.op) +is_orthogonal(A::OperatorWrapper) = is_orthogonal(A.op) +is_full_row_rank(A::OperatorWrapper) = is_full_row_rank(A.op) +is_full_column_rank(A::OperatorWrapper) = is_full_column_rank(A.op) +is_symmetric(A::OperatorWrapper) = is_symmetric(A.op) +is_positive_definite(A::OperatorWrapper) = is_positive_definite(A.op) +is_positive_semidefinite(A::OperatorWrapper) = is_positive_semidefinite(A.op) + +# OperatorWrapper has mutable buffers — never thread-safe regardless of inner op. +is_thread_safe(::OperatorWrapper) = false + +displacement(A::OperatorWrapper) = displacement(A.op) +remove_displacement(A::OperatorWrapper) = OperatorWrapper(remove_displacement(A.op)) + +function _copy_operator_impl( + A::OperatorWrapper{Op, DB, CB, DS, CS}; array_type = nothing, threaded = nothing + ) where {Op, DB, CB, DS, CS} + new_op = copy_operator(A.op; array_type = nothing, threaded) + return OperatorWrapper{typeof(new_op), DB, CB, DS, CS}( + new_op, similar(A.dom_buf), similar(A.cod_buf) + ) +end diff --git a/src/calculus/Reshape.jl b/src/calculus/Reshape.jl index 43df9ea..b1077f7 100644 --- a/src/calculus/Reshape.jl +++ b/src/calculus/Reshape.jl @@ -57,13 +57,15 @@ has_optimized_normalop(R::Reshape) = true get_normal_op(R::Reshape) = get_normal_op(R.A) # Properties -Base.:(==)(R1::Reshape{N, L}, R2::Reshape{N, L}) where {N, L} = R1.A == R2.A && R1.dim_out == R2.dim_out +function Base.:(==)(R1::Reshape{N, L}, R2::Reshape{N, L}) where {N, L} + return R1.A == R2.A && R1.dim_out == R2.dim_out +end size(R::Reshape) = (R.dim_out, size(R.A, 2)) domain_type(R::Reshape) = domain_type(R.A) codomain_type(R::Reshape) = codomain_type(R.A) -domain_storage_type(R::Reshape) = domain_storage_type(R.A) -codomain_storage_type(R::Reshape) = codomain_storage_type(R.A) +domain_array_type(R::Reshape) = domain_array_type(R.A) +codomain_array_type(R::Reshape) = codomain_array_type(R.A) is_thread_safe(R::Reshape) = is_thread_safe(R.A) is_linear(R::Reshape) = is_linear(R.A) diff --git a/src/calculus/Scale.jl b/src/calculus/Scale.jl index 9826d29..9b7d1af 100644 --- a/src/calculus/Scale.jl +++ b/src/calculus/Scale.jl @@ -73,7 +73,7 @@ function mul!(y::Tuple, L::Scale{Th}, x::AbstractArray) where {Th} for k in eachindex(y) @.. thread = Th y[k] *= L.coeff end - return + return y end function mul!( @@ -92,7 +92,7 @@ function mul!(y::Tuple, S::AdjointOperator{<:Scale{Th}}, x::AbstractArray) where for k in eachindex(y) @.. thread = Th y[k] .*= L.coeff_conj end - return + return y end has_optimized_normalop(L::Scale) = is_linear(L.A) && has_optimized_normalop(L.A) @@ -113,8 +113,8 @@ size(L::Scale) = size(L.A) domain_type(L::Scale) = domain_type(L.A) codomain_type(L::Scale) = codomain_type(L.A) -domain_storage_type(L::Scale) = domain_storage_type(L.A) -codomain_storage_type(L::Scale) = codomain_storage_type(L.A) +domain_array_type(L::Scale) = domain_array_type(L.A) +codomain_array_type(L::Scale) = codomain_array_type(L.A) is_thread_safe(L::Scale) = is_thread_safe(L.A) is_linear(L::Scale) = is_linear(L.A) diff --git a/src/calculus/Sum.jl b/src/calculus/Sum.jl index 08470e3..e51f373 100644 --- a/src/calculus/Sum.jl +++ b/src/calculus/Sum.jl @@ -154,10 +154,10 @@ domain_type(S::Sum{K, C, D, L}) where {K, C, D <: AbstractArray, L} = domain_typ domain_type(S::Sum{K, C, D, L}) where {K, C, D <: Tuple, L} = domain_type.(Ref(S.A[1])) codomain_type(S::Sum{K, C, D, L}) where {K, C <: AbstractArray, D, L} = codomain_type(S.A[1]) codomain_type(S::Sum{K, C, D, L}) where {K, C <: Tuple, D, L} = codomain_type.(Ref(S.A[1])) -domain_storage_type(S::Sum{K, C, D, L}) where {K, C <: AbstractArray, D, L} = domain_storage_type(S.A[1]) -domain_storage_type(S::Sum{K, C, D, L}) where {K, C <: Tuple, D, L} = domain_storage_type.(Ref(S.A[1])) -codomain_storage_type(S::Sum{K, C, D, L}) where {K, C <: AbstractArray, D, L} = codomain_storage_type(S.A[1]) -codomain_storage_type(S::Sum{K, C, D, L}) where {K, C <: Tuple, D, L} = codomain_storage_type.(Ref(S.A[1])) +domain_array_type(S::Sum{K, C, D, L}) where {K, C <: AbstractArray, D, L} = domain_array_type(S.A[1]) +domain_array_type(S::Sum{K, C, D, L}) where {K, C <: Tuple, D, L} = domain_array_type.(Ref(S.A[1])) +codomain_array_type(S::Sum{K, C, D, L}) where {K, C <: AbstractArray, D, L} = codomain_array_type(S.A[1]) +codomain_array_type(S::Sum{K, C, D, L}) where {K, C <: Tuple, D, L} = codomain_array_type.(Ref(S.A[1])) fun_domain(S::Sum) = fun_domain(S.A[1]) fun_codomain(S::Sum) = fun_codomain(S.A[1]) @@ -185,3 +185,12 @@ function permute(S::Sum, p::AbstractVector{Int}) end remove_displacement(S::Sum) = Sum(remove_displacement.(S.A), S.bufC, S.bufD) + +function _copy_operator_impl(op::Sum; array_type = nothing, threaded = nothing) + new_bufC = _convert_buffer(op.bufC, array_type) + new_bufD = _convert_buffer(op.bufD, array_type) + new_ops = tuple([copy_operator(a; array_type, threaded) for a in op.A]...) + K = length(new_ops) + L = typeof(new_ops) + return Sum{K, typeof(new_bufC), typeof(new_bufD), L}(new_ops, new_bufC, new_bufD) +end diff --git a/src/calculus/VCAT.jl b/src/calculus/VCAT.jl index 9b092f6..6579b7f 100644 --- a/src/calculus/VCAT.jl +++ b/src/calculus/VCAT.jl @@ -32,6 +32,7 @@ struct VCAT{ L <: NTuple{N, AbstractOperator}, P <: Tuple, C <: AbstractArray, + CS <: AbstractArray, # codomain storage type (fixed at construction) } <: AbstractOperator A::L # tuple of AbstractOperators idxs::P # indices @@ -50,10 +51,20 @@ struct VCAT{ if any([domain_type(A[1]) != domain_type(a) for a in A]) throw(error("operators must all share the same domain_type!")) end - return new{N, L, P, C}(A, idxs, buf) + CS = _compute_vcat_cs(A, idxs) + return new{N, L, P, C, CS}(A, idxs, buf) end end +function _compute_vcat_cs(A, idxs) + cs_list = [d <: ArrayPartition ? [d.parameters[2].types...] : d for d in codomain_array_type.(A)] + codomain = vcat(cs_list...) + p = vcat([[idx...] for idx in idxs]...) + invpermute!(codomain, p) + T = promote_type(map(_storage_eltype, codomain)...) + return ArrayPartition{T, Tuple{codomain...}} +end + function VCAT(A::Vararg{AbstractOperator}) if any((<:).(typeof.(A), VCAT)) #there are VCATs in A AA = () @@ -159,7 +170,9 @@ end # Properties -Base.:(==)(H1::VCAT{N, L1, P1, C}, H2::VCAT{N, L2, P2, C}) where {N, L1, L2, P1, P2, C} = H1.A == H2.A && H1.idxs == H2.idxs +function Base.:(==)(H1::VCAT{N, L1, P1, C}, H2::VCAT{N, L2, P2, C}) where {N, L1, L2, P1, P2, C} + return H1.A == H2.A && H1.idxs == H2.idxs +end @generated function size(H::VCAT{N, L, P}) where {N, L, P} exprs = [] @@ -204,14 +217,8 @@ domain_type(L::VCAT) = domain_type.(Ref(L.A[1])) natural_expr = Expr(:tuple, exprs...) return :(_vcat_apply_invperm($natural_expr, H.idxs)) end -domain_storage_type(L::VCAT) = domain_storage_type.(Ref(L.A[1])) -function codomain_storage_type(H::VCAT) - codomain = vcat([d <: ArrayPartition ? [d.parameters[2].types...] : d for d in codomain_storage_type.(H.A)]...) - p = vcat([[idx...] for idx in H.idxs]...) - invpermute!(codomain, p) - T = promote_type(codomain_type(H)...) - return ArrayPartition{T, Tuple{codomain...}} -end +domain_array_type(L::VCAT) = domain_array_type.(Ref(L.A[1])) +codomain_array_type(::VCAT{N, L, P, C, CS}) where {N, L, P, C, CS} = CS is_linear(L::VCAT) = all(is_linear.(L.A)) is_AcA_diagonal(L::VCAT) = all(is_AcA_diagonal.(L.A)) @@ -276,3 +283,9 @@ function permute(H::VCAT{N, L, P, C}, p::AbstractVector{Int}) where {N, L, P, C} end remove_displacement(V::VCAT) = VCAT(remove_displacement.(V.A), V.idxs, V.buf) + +function _copy_operator_impl(op::VCAT; array_type = nothing, threaded = nothing) + new_buf = _convert_buffer(op.buf, array_type) + new_ops = tuple([copy_operator(a; array_type, threaded) for a in op.A]...) + return VCAT(new_ops, op.idxs, new_buf) +end diff --git a/src/linearoperators/DiagOp.jl b/src/linearoperators/DiagOp.jl index c7a4aad..f24fc80 100644 --- a/src/linearoperators/DiagOp.jl +++ b/src/linearoperators/DiagOp.jl @@ -25,32 +25,48 @@ end # Constructors ###standard constructor Operator{N}(D::Type, domain_dim::NTuple{N,Int}) -function DiagOp(D::Type, domain_dim::NTuple{N, Int}, d::T; threaded::Bool = true) where {N, T <: AbstractArray} +function DiagOp( + D::Type, domain_dim::NTuple{N, Int}, d::T; + threaded::Bool = true, array_type::Type = _array_wrapper(d), + ) where {N, T <: AbstractArray} size(d) != domain_dim && error("dimension of d must coincide with domain_dim") C = promote_type(eltype(d), D) - threaded = threaded && Threads.nthreads() > 1 && length(d) * sizeof(D) > 2^16 + threaded = threaded && _should_thread(d) B = threaded ? FastBroadcast.True() : FastBroadcast.False() - dS = typeof(d isa SubArray ? parent(d) : d).name.wrapper{D} - cS = typeof(d isa SubArray ? parent(d) : d).name.wrapper{C} + dS = _normalize_array_type(array_type, D) + cS = _normalize_array_type(array_type, C) return DiagOp{B, D, C, N, dS, cS, T}(domain_dim, d) end ###standard constructor with Scalar -function DiagOp(D::Type, domain_dim::NTuple{N, Int}, d::T; threaded::Bool = true) where {N, T <: Number} +function DiagOp( + D::Type, domain_dim::NTuple{N, Int}, d::T; + threaded::Bool = true, array_type::Type = Array{D}, + ) where {N, T <: Number} C = promote_type(eltype(d), D) - threaded = threaded && Threads.nthreads() > 1 && length(d) > 2^16 + threaded = threaded && _should_thread(d) B = threaded ? FastBroadcast.True() : FastBroadcast.False() - return DiagOp{B, D, C, N, Array{D}, Array{C}, T}(domain_dim, d) + dS = _normalize_array_type(array_type, D) + cS = _normalize_array_type(array_type, C) + return DiagOp{B, D, C, N, dS, cS, T}(domain_dim, d) end # other constructors -function DiagOp(d::AbstractArray{T, N}; threaded::Bool = true) where {N, T <: Number} +function DiagOp( + d::AbstractArray{T, N}; + threaded::Bool = true, array_type::Type = _array_wrapper(d), + ) where {N, T <: Number} C = eltype(d) - B = (threaded && length(d) > 2^16) ? FastBroadcast.True() : FastBroadcast.False() - S = typeof(d isa SubArray ? parent(d) : d).name.wrapper{T} + threaded = threaded && _should_thread(d) + B = threaded ? FastBroadcast.True() : FastBroadcast.False() + S = _normalize_array_type(array_type, T) return DiagOp{B, eltype(d), C, N, S, S, typeof(d)}(size(d), d) end -DiagOp(domain_dim::NTuple{N, Int}, d::A; threaded::Bool = true) where {N, A <: Number} = DiagOp(Float64, domain_dim, d; threaded) +function DiagOp( + domain_dim::NTuple{N, Int}, d::A; threaded::Bool = true, array_type::Type = Array{Float64} + ) where {N, A <: Number} + return DiagOp(Float64, domain_dim, d; threaded, array_type) +end # scale of DiagOp function Scale(coeff::T, L::DiagOp{B, D, C, N, dS, cS}) where {T <: Number, B, D, C, N, dS, cS} @@ -84,8 +100,8 @@ end # Properties -domain_storage_type(::DiagOp{<:Any, <:Any, <:Any, <:Any, dS}) where {dS} = dS -codomain_storage_type(::DiagOp{<:Any, <:Any, <:Any, <:Any, <:Any, cS}) where {cS} = cS +domain_array_type(::DiagOp{<:Any, <:Any, <:Any, <:Any, dS}) where {dS} = dS +codomain_array_type(::DiagOp{<:Any, <:Any, <:Any, <:Any, <:Any, cS}) where {cS} = cS diag(L::DiagOp) = L.d diag_AAc(L::DiagOp{B}) where {B} = @.. thread = B L.d * conj(L.d) @@ -95,6 +111,13 @@ domain_type(::DiagOp{<:Any, D}) where {D} = D codomain_type(::DiagOp{<:Any, <:Any, C}) where {C} = C is_thread_safe(::DiagOp) = true +function _copy_operator_impl(op::DiagOp{B}; array_type = nothing, threaded = nothing) where {B} + new_threaded = threaded === nothing ? (B == FastBroadcast.True()) : threaded + new_d = array_type === nothing ? op.d : _convert_buffer(op.d, array_type) + new_at = array_type === nothing ? _array_wrapper(op.d) : array_type + return DiagOp(domain_type(op), op.dim_in, new_d; threaded = new_threaded, array_type = new_at) +end + size(L::DiagOp) = (L.dim_in, L.dim_in) fun_name(L::DiagOp) = "╲" diff --git a/src/linearoperators/Eye.jl b/src/linearoperators/Eye.jl index 43b6793..dbf956a 100644 --- a/src/linearoperators/Eye.jl +++ b/src/linearoperators/Eye.jl @@ -1,7 +1,5 @@ export Eye -abstract type AbstractEye{T, N, S <: AbstractArray} <: LinearOperator end - """ Eye([domain_type=Float64::Type,] dim_in::Tuple) Eye([domain_type=Float64::Type,] dims...) @@ -20,60 +18,71 @@ true ``` """ -struct Eye{T, N, S <: AbstractArray{T}} <: AbstractEye{T, N, S} +struct Eye{T, N, S <: AbstractArray{T}} <: LinearOperator dim::NTuple{N, Int} end # Constructors -###standard constructor Operator{N}(domain_type::Type, DomainDim::NTuple{N,Int}) +@inline _make_eye(::Type{T}, dims::NTuple{N, Int}, ::Type{S}) where {T, N, S <: AbstractArray} = + Eye{T, N, _normalize_array_type(S, T)}(dims) + function Eye( - domain_type::Type{T}, domainDim::NTuple{N, <:Integer}, storageType::Type{S} = Array{T} - ) where {N, T, S <: AbstractArray{T}} - return Eye{domain_type, N, storageType}(map(Int, domainDim)) + domain_type::Type{T}, domainDim::NTuple{N, <:Integer}; array_type::Type{<:AbstractArray} = Array{T} + ) where {N, T} + return _make_eye(domain_type, map(Int, domainDim), array_type) end -### -Eye(t::Type, dims::Vararg{Integer}) = Eye(t, dims) -Eye(dims::NTuple{N, Integer}) where {N} = Eye(Float64, dims) -Eye(dims::Vararg{Integer}) = Eye(Float64, dims) -Eye(x::A) where {A <: AbstractArray} = Eye(eltype(x), size(x), Array{eltype(x)}) +function Eye( + t::Type{T}, dims::Vararg{Integer}; array_type::Type{<:AbstractArray} = Array{T} + ) where {T} + return _make_eye(t, map(Int, dims), array_type) +end +function Eye(dims::NTuple{N, Integer}; array_type::Type{<:AbstractArray} = Array{Float64}) where {N} + return _make_eye(Float64, map(Int, dims), array_type) +end +function Eye(dims::Vararg{Integer}; array_type::Type{<:AbstractArray} = Array{Float64}) + return _make_eye(Float64, map(Int, dims), array_type) +end +function Eye(x::A) where {A <: AbstractArray} + return _make_eye(eltype(x), size(x), typeof(x isa SubArray ? parent(x) : x)) +end # Mappings -function mul!(y::AbstractArray, L::AbstractEye, b::AbstractArray) +function mul!(y::AbstractArray, L::Eye, b::AbstractArray) check(y, L, b) y .= b return y end # Properties -diag(::AbstractEye) = 1.0 -diag_AcA(::AbstractEye) = 1.0 -diag_AAc(::AbstractEye) = 1.0 - -domain_type(::AbstractEye{T, N}) where {T, N} = T -codomain_type(::AbstractEye{T, N}) where {T, N} = T -domain_storage_type(::AbstractEye{T, N, S}) where {T, N, S} = S -codomain_storage_type(::AbstractEye{T, N, S}) where {T, N, S} = S +diag(::Eye) = 1.0 +diag_AcA(::Eye) = 1.0 +diag_AAc(::Eye) = 1.0 + +domain_type(::Eye{T, N}) where {T, N} = T +codomain_type(::Eye{T, N}) where {T, N} = T +domain_array_type(::Eye{T, N, S}) where {T, N, S} = S +codomain_array_type(::Eye{T, N, S}) where {T, N, S} = S is_thread_safe(::Eye) = true -size(L::AbstractEye) = (L.dim, L.dim) +size(L::Eye) = (L.dim, L.dim) -fun_name(::AbstractEye) = "I" +fun_name(::Eye) = "I" -is_eye(::AbstractEye) = true -is_diagonal(::AbstractEye) = true -is_orthogonal(::AbstractEye) = true -is_invertible(::AbstractEye) = true -is_full_row_rank(::AbstractEye) = true -is_full_column_rank(::AbstractEye) = true -is_symmetric(::AbstractEye) = true -is_positive_definite(::AbstractEye) = true -is_positive_semidefinite(::AbstractEye) = true +is_eye(::Eye) = true +is_diagonal(::Eye) = true +is_orthogonal(::Eye) = true +is_invertible(::Eye) = true +is_full_row_rank(::Eye) = true +is_full_column_rank(::Eye) = true +is_symmetric(::Eye) = true +is_positive_definite(::Eye) = true +is_positive_semidefinite(::Eye) = true -has_optimized_normalop(::AbstractEye) = true -get_normal_op(L::AbstractEye) = L +has_optimized_normalop(::Eye) = true +get_normal_op(L::Eye) = L -has_fast_opnorm(::AbstractEye) = true -LinearAlgebra.opnorm(L::AbstractEye) = one(real(domain_type(L))) -AdjointOperator(L::AbstractEye) = L +has_fast_opnorm(::Eye) = true +LinearAlgebra.opnorm(L::Eye) = one(real(domain_type(L))) +AdjointOperator(L::Eye) = L diff --git a/src/linearoperators/FiniteDiff.jl b/src/linearoperators/FiniteDiff.jl index ad7ece1..83613fc 100644 --- a/src/linearoperators/FiniteDiff.jl +++ b/src/linearoperators/FiniteDiff.jl @@ -20,32 +20,44 @@ true ``` """ -struct FiniteDiff{N, D, T} <: LinearOperator +struct FiniteDiff{N, D, T, S <: AbstractArray{T}} <: LinearOperator dim_in::NTuple{N, Int} - function FiniteDiff{N, D, T}(dim_in) where {N, D, T} + function FiniteDiff{N, D, T, S}(dim_in) where {N, D, T, S <: AbstractArray{T}} D > N && error("direction is bigger the number of dimension $N") - return new{N, D, T}(dim_in) + return new{N, D, T, S}(dim_in) end end # Constructors # Val-dispatch constructor — fully type-stable (D is known at compile time) -function FiniteDiff(::Type{T}, dim_in::NTuple{N, Int}, ::Val{D}) where {T, N, D} - return FiniteDiff{N, D, T}(dim_in) +function FiniteDiff( + ::Type{T}, dim_in::NTuple{N, Int}, ::Val{D}; array_type::Type = Array{T} + ) where {T, N, D} + S = _normalize_array_type(array_type, T) + return FiniteDiff{N, D, T, S}(dim_in) end # Specialized no-direction constructor: D=1 is a compile-time literal — fully type-stable -FiniteDiff(dim_in::NTuple{N, Int}) where {N} = FiniteDiff{N, 1, Float64}(dim_in) +function FiniteDiff(dim_in::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + S = _normalize_array_type(array_type, Float64) + return FiniteDiff{N, 1, Float64, S}(dim_in) +end #default constructor (direction as runtime Int — delegates to Val) -function FiniteDiff(domain_type::Type, dim_in::NTuple{N, Int}, dir::Int = 1) where {N} - return FiniteDiff(domain_type, dim_in, Val(dir)) +function FiniteDiff( + domain_type::Type{T}, dim_in::NTuple{N, Int}, dir::Int = 1; + array_type::Type = Array{T} + ) where {T, N} + return FiniteDiff(domain_type, dim_in, Val(dir); array_type) end -FiniteDiff(dim_in::NTuple{N, Int}, dir::Int) where {N} = FiniteDiff(Float64, dim_in, Val(dir)) +function FiniteDiff(dim_in::NTuple{N, Int}, dir::Int; array_type::Type = Array{Float64}) where {N} + return FiniteDiff(Float64, dim_in, Val(dir); array_type) +end function FiniteDiff(x::AbstractArray{T, N}, dir::Int = 1) where {T, N} - return FiniteDiff(T, size(x), Val(dir)) + S = _normalize_array_type(_array_wrapper(x), T) + return FiniteDiff{N, dir, T, S}(size(x)) end # Mappings @@ -76,6 +88,8 @@ end domain_type(::FiniteDiff{<:Any, <:Any, T}) where {T} = T codomain_type(::FiniteDiff{<:Any, <:Any, T}) where {T} = T +domain_array_type(::FiniteDiff{N, D, T, S}) where {N, D, T, S} = S +codomain_array_type(::FiniteDiff{N, D, T, S}) where {N, D, T, S} = S is_thread_safe(::FiniteDiff) = true function size(L::FiniteDiff{N, D}) where {N, D} diff --git a/src/linearoperators/GetIndex.jl b/src/linearoperators/GetIndex.jl index 2786801..e1c7259 100644 --- a/src/linearoperators/GetIndex.jl +++ b/src/linearoperators/GetIndex.jl @@ -13,7 +13,7 @@ Supported indices are: - `Tuple`: a tuple of indices, where each element can be one of the above types. ```jldoctest -julia> x = (1:10) .* 1.0; +julia> x = collect(1.0:10.0); julia> G = GetIndex(Float64,(10,), 1:3) ↓ ℝ^10 -> ℝ^3 @@ -44,36 +44,51 @@ struct GetIndex{I, N, M, T, S} <: LinearOperator end end +_prepare_getindex_intvec(idx, ::Type{<:AbstractArray}) = idx +_prepare_getindex_boolmask(mask, ::Type{<:AbstractArray}) = mask + # Constructors # default -function GetIndex(domain_type::Type, dim_in::Dims, idx::Tuple) +function GetIndex( + domain_type::Type{T}, dim_in::Dims, idx::Tuple; array_type::Type = Array{T} + ) where {T} dim_out = get_dim_out(dim_in, idx...) if dim_out == dim_in return Eye(domain_type, dim_in) else - return GetIndex(domain_type, Array{domain_type}, dim_out, dim_in, idx) + S = _normalize_array_type(array_type, domain_type) + return GetIndex(domain_type, S, dim_out, dim_in, idx) end end function GetIndex( - domain_type::Type, + domain_type::Type{T}, dim_in::Dims, idx::Union{AbstractVector{Int}, AbstractVector{<:CartesianIndex}}, - ) + ; array_type::Type = Array{T}, + ) where {T} + idx = _prepare_getindex_intvec(idx, array_type) dim_out = (length(idx),) - return GetIndex(domain_type, Array{domain_type}, dim_out, dim_in, idx) + S = _normalize_array_type(array_type, domain_type) + return GetIndex(domain_type, S, dim_out, dim_in, idx) end -function GetIndex(domain_type::Type, dim_in::Dims, mask::AbstractArray{Bool}) - dim_out = (sum(mask),) +function GetIndex( + domain_type::Type{T}, dim_in::Dims, mask::AbstractArray{Bool}; array_type::Type = Array{T} + ) where {T} + mask = _prepare_getindex_boolmask(mask, array_type) + dim_out = mask isa AbstractArray{Bool} ? (sum(mask),) : (length(mask),) if dim_out[1] == prod(dim_in) return reshape(Eye(domain_type, dim_in), dim_out) else - return GetIndex(domain_type, Array{domain_type}, dim_out, dim_in, mask) + S = _normalize_array_type(array_type, domain_type) + return GetIndex(domain_type, S, dim_out, dim_in, mask) end end -GetIndex(domain_type::Type, dim_in::Dims, idx...) = GetIndex(domain_type, dim_in, idx) +function GetIndex(domain_type::Type{T}, dim_in::Dims, idx...; array_type::Type = Array{T}) where {T} + return GetIndex(domain_type, dim_in, idx; array_type) +end GetIndex(dim_in::Dims, idx...) = GetIndex(Float64, dim_in, idx) GetIndex(dim_in::Dims, idx::Tuple) = GetIndex(Float64, dim_in, idx) GetIndex(dim_in::Dims, idx) = GetIndex(Float64, dim_in, idx) @@ -83,7 +98,7 @@ function GetIndex(x::AbstractArray, idx::Tuple) if dim_out == dim_in return Eye(eltype(x), dim_in) else - S = typeof(x isa SubArray ? parent(x) : x).name.wrapper{eltype(x)} + S = _array_wrapper(x){eltype(x)} return GetIndex(eltype(x), S, dim_out, dim_in, idx) end end @@ -96,23 +111,21 @@ function GetIndex( if dim_out == dim_in return Eye(eltype(x), dim_in) else - S = typeof(x isa SubArray ? parent(x) : x).name.wrapper{eltype(x)} + S = _array_wrapper(x){eltype(x)} return GetIndex(eltype(x), S, dim_out, dim_in, idx) end end # Mappings -@generated function mul!( - y::AbstractArray, L::GetIndex{I}, b::AbstractArray - ) where {K, I <: Tuple{Vararg{Any, K}}} - return if K == 1 +@generated function mul!(y, L::GetIndex{I}, b) where {K, I <: Tuple{Vararg{Any, K}}} + if K == 1 if I.parameters[1] <: AbstractArray{Bool} indices = :(to_indices(b, L.idx)[1]) else indices = :(L.idx[1]) end - quote + return quote check(y, L, b) for (i, j) in enumerate($indices) @inbounds y[i] = b[j] @@ -120,7 +133,7 @@ end return y end else - quote + return quote check(y, L, b) indices = to_indices(b, L.idx) j = 1 @@ -170,12 +183,12 @@ end # Properties domain_type(::GetIndex{I, N, M, T}) where {I, N, M, T} = T domain_type(::NormalGetIndex{I, N, T}) where {I, N, T} = T -domain_storage_type(::GetIndex{I, N, M, T, S}) where {I, N, M, T, S} = S -domain_storage_type(::NormalGetIndex{I, N, T, S}) where {I, N, T, S} = S +domain_array_type(::GetIndex{I, N, M, T, S}) where {I, N, M, T, S} = S +domain_array_type(::NormalGetIndex{I, N, T, S}) where {I, N, T, S} = S codomain_type(::GetIndex{I, N, M, T}) where {I, N, M, T} = T codomain_type(::NormalGetIndex{I, N, T}) where {I, N, T} = T -codomain_storage_type(::GetIndex{I, N, M, T, S}) where {I, N, M, T, S} = S -codomain_storage_type(::NormalGetIndex{I, N, T, S}) where {I, N, T, S} = S +codomain_array_type(::GetIndex{I, N, M, T, S}) where {I, N, M, T, S} = S +codomain_array_type(::NormalGetIndex{I, N, T, S}) where {I, N, T, S} = S is_thread_safe(L::GetIndex) = true is_thread_safe(L::NormalGetIndex) = true diff --git a/src/linearoperators/LBFGS.jl b/src/linearoperators/LBFGS.jl index d3d81eb..efc65fc 100644 --- a/src/linearoperators/LBFGS.jl +++ b/src/linearoperators/LBFGS.jl @@ -34,16 +34,14 @@ end #default constructor -function LBFGS(x::T, M::I) where {T <: AbstractArray, I <: Integer} +function LBFGS(x::T, M::I; array_type::Type = _array_wrapper(x)) where {T <: AbstractArray, I <: Integer} s_M = [zero(x) for i in 1:M] y_M = [zero(x) for i in 1:M] s = zero(x) y = zero(x) R = real(eltype(x)) - ys_M = similar(x, R, M) - fill!(ys_M, one(R)) - alphas = similar(x, R, M) - fill!(alphas, one(R)) + ys_M = fill(one(R), M) # always CPU: only used with scalar indexing + alphas = fill(one(R), M) # always CPU: only used with scalar indexing return LBFGS{R, T, M, I}(0, 0, s, y, s_M, y_M, ys_M, alphas, one(R)) end @@ -93,6 +91,7 @@ mul!(x::AbstractArray, L::AdjointOperator{<:LBFGS}, y::AbstractArray) = mul!(x, # Two-loop recursion function mul!(d::AbstractArray, L::LBFGS, gradx::AbstractArray) + check(d, L, gradx) d .= gradx idx = loop1!(d, L) d .*= L.H @@ -131,6 +130,8 @@ domain_type(L::LBFGS) = eltype(L.y_M[1]) domain_type(L::LBFGS{R, T}) where {R, T <: ArrayPartition} = eltype.(L.y_M[1].x) codomain_type(L::LBFGS) = eltype(L.y_M[1]) codomain_type(L::LBFGS{R, T}) where {R, T <: ArrayPartition} = eltype.(L.y_M[1].x) +domain_array_type(L::LBFGS) = typeof(L.s) +codomain_array_type(L::LBFGS) = typeof(L.s) is_thread_safe(L::LBFGS) = false size(A::LBFGS{R, T}) where {R, T <: ArrayPartition} = (size.(A.s.x), size.(A.s.x)) diff --git a/src/linearoperators/LMatrixOp.jl b/src/linearoperators/LMatrixOp.jl index 9ec2cf2..0e4bd14 100644 --- a/src/linearoperators/LMatrixOp.jl +++ b/src/linearoperators/LMatrixOp.jl @@ -21,7 +21,7 @@ julia> op*ones(3,4) ``` """ -struct LMatrixOp{T, A <: Union{AbstractVector, AbstractMatrix}, B <: AbstractMatrix} <: +struct LMatrixOp{T, A <: Union{AbstractVector, AbstractMatrix}, B <: AbstractMatrix, dS, cS} <: LinearOperator b::A bt::B @@ -32,15 +32,18 @@ end # Constructors function LMatrixOp( domain_type::Type, DomainDim::Tuple{Int, Int}, b::A + ; array_type::Type = _array_wrapper_type(A), ) where {A <: Union{AbstractVector, AbstractMatrix}} bt = b' - return LMatrixOp{domain_type, A, typeof(bt)}(b, bt, DomainDim[1]) + dS = _normalize_array_type(array_type, domain_type) + cS = _normalize_array_type(array_type, domain_type) + return LMatrixOp{domain_type, A, typeof(bt), dS, cS}(b, bt, DomainDim[1]) end function LMatrixOp( b::A, n_row_in::Int ) where {T, A <: Union{AbstractVector{T}, AbstractMatrix{T}}} - return LMatrixOp(T, (n_row_in, size(b, 1)), b) + return LMatrixOp(T, (n_row_in, size(b, 1)), b; array_type = _array_wrapper_type(A)) end # Mappings @@ -62,14 +65,16 @@ end # Properties domain_type(::LMatrixOp{T}) where {T} = T codomain_type(::LMatrixOp{T}) where {T} = T +domain_array_type(::LMatrixOp{T, A, B, dS}) where {T, A, B, dS} = dS +codomain_array_type(::LMatrixOp{T, A, B, dS, cS}) where {T, A, B, dS, cS} = cS is_thread_safe(::LMatrixOp) = true fun_name(L::LMatrixOp) = "(⋅)b" -function size(L::LMatrixOp{T, A, B}) where {T, A <: AbstractVector, B <: Adjoint} +function size(L::LMatrixOp{T, A, B, dS, cS}) where {T, A <: AbstractVector, B <: Adjoint, dS, cS} return (L.n_row_in,), (L.n_row_in, length(L.b)) end -function size(L::LMatrixOp{T, A, B}) where {T, A <: AbstractMatrix, B <: AbstractMatrix} +function size(L::LMatrixOp{T, A, B, dS, cS}) where {T, A <: AbstractMatrix, B <: AbstractMatrix, dS, cS} return (L.n_row_in, size(L.b, 2)), (L.n_row_in, size(L.b, 1)) end diff --git a/src/linearoperators/MatrixOp.jl b/src/linearoperators/MatrixOp.jl index 61738f6..05fd83b 100644 --- a/src/linearoperators/MatrixOp.jl +++ b/src/linearoperators/MatrixOp.jl @@ -24,7 +24,7 @@ julia> MatrixOp(randn(20,10),4) ``` """ -struct MatrixOp{D, T, M <: AbstractMatrix{T}, NC} <: LinearOperator +struct MatrixOp{D, T, M <: AbstractMatrix{T}, NC, dS, cS} <: LinearOperator A::M end @@ -33,27 +33,35 @@ end ###standard constructor Operator{N}(domain_type::Type, DomainDim::NTuple{N,Int}) function MatrixOp( domain_type::Type, DomainDim::NTuple{N, Int}, A::M + ; array_type::Type = _array_wrapper_type(M), ) where {N, T, M <: AbstractMatrix{T}} N > 2 && error("cannot multiply a Matrix by a n-dimensional Variable with n > 2") size(A, 2) != DomainDim[1] && error("wrong input dimensions") + codomainT = domain_type <: Real && T <: Complex ? T : domain_type + dS = _normalize_array_type(array_type, domain_type) + cS = _normalize_array_type(array_type, codomainT) return if N == 1 - MatrixOp{domain_type, T, M, 1}(A) + MatrixOp{domain_type, T, M, 1, dS, cS}(A) else - MatrixOp{domain_type, T, M, DomainDim[2]}(A) + MatrixOp{domain_type, T, M, DomainDim[2], dS, cS}(A) end end ### -MatrixOp(A::M) where {M <: AbstractMatrix} = MatrixOp(eltype(A), (size(A, 2),), A) -MatrixOp(D::Type, A::M) where {M <: AbstractMatrix} = MatrixOp(D, (size(A, 2),), A) +function MatrixOp(A::M; array_type::Type = _array_wrapper_type(M)) where {M <: AbstractMatrix} + return MatrixOp(eltype(A), (size(A, 2),), A; array_type) +end +function MatrixOp(D::Type, A::M; array_type::Type = _array_wrapper_type(M)) where {M <: AbstractMatrix} + return MatrixOp(D, (size(A, 2),), A; array_type) +end function MatrixOp(A::M, n::Integer) where {M <: AbstractMatrix} - return MatrixOp(eltype(A), (size(A, 2), n), A) + return MatrixOp(eltype(A), (size(A, 2), n), A; array_type = _array_wrapper_type(M)) end function MatrixOp(D::Type, A::M, n::Integer) where {M <: AbstractMatrix} - return MatrixOp(D, (size(A, 2), n), A) + return MatrixOp(D, (size(A, 2), n), A; array_type = _array_wrapper_type(M)) end -function Scale(coeff::Number, A::MatrixOp{D, T, M, NC}) where {D, T, M, NC} +function Scale(coeff::Number, A::MatrixOp{D, T, M, NC, dS, cS}) where {D, T, M, NC, dS, cS} if coeff == 1 return A end @@ -64,7 +72,7 @@ function Scale(coeff::Number, A::MatrixOp{D, T, M, NC}) where {D, T, M, NC} "Cannot Scale AbstractOperator with real codomain with complex scalar. Use `DiagOp` instead.", ) end - return MatrixOp(coeff * A.A, NC) + return MatrixOp(domain_type(A), size(A, 2), coeff * A.A; array_type = _array_wrapper_type(dS)) end function Scale(coeff::Number, A::AdjointOperator{<:MatrixOp}) if coeff == 1 @@ -81,7 +89,7 @@ function Scale(coeff::Number, A::AdjointOperator{<:MatrixOp}) end import Base: convert -convert(::Type{LinearOperator}, L::M) where {T, M <: AbstractMatrix{T}} = MatrixOp{T, T, M, 1}(L) +convert(::Type{LinearOperator}, L::M) where {T, M <: AbstractMatrix{T}} = MatrixOp(T, (size(L, 2),), L) function convert(::Type{LinearOperator}, L::M, n::Integer) where {T, M <: AbstractMatrix{T}} return MatrixOp(L, n) end @@ -104,18 +112,21 @@ function mul!(y::AbstractArray, L::AdjointOperator{<:MatrixOp}, b::AbstractArray return mul!(y, L.A.A', b) end # NC=1 adjoint implicit batching: matrix input accepted -function mul!(y::AbstractArray, L::AdjointOperator{<:MatrixOp{<:Any, <:Any, <:Any, 1}}, b::AbstractArray) +function mul!( + y::AbstractArray, L::AdjointOperator{<:MatrixOp{<:Any, <:Any, <:Any, 1}}, b::AbstractArray + ) return mul!(y, L.A.A', b) end # Special Case, real b, complex matrix: accepts real b (type mismatch is intentional) function mul!(y::AbstractArray, L::AdjointOperator{<:MatrixOp{D, T}}, b::AbstractArray) where {D <: Real, T <: Complex} + check(y, L, b) yc = similar(y, T, size(y)) mul!(yc, L.A.A', b) return y .= real.(yc) end # Resolves ambiguity: NC=1 real-domain complex-matrix adjoint with matrix batching -function mul!(y::AbstractArray, L::AdjointOperator{<:MatrixOp{D, T, M, 1}}, b::AbstractArray) where {D <: Real, T <: Complex, M} +function mul!(y::AbstractArray, L::AdjointOperator{<:MatrixOp{D, T, M, 1, <:Any, <:Any}}, b::AbstractArray) where {D <: Real, T <: Complex, M} yc = similar(y, T, size(y)) mul!(yc, L.A.A', b) return y .= real.(yc) @@ -125,11 +136,13 @@ end domain_type(::MatrixOp{D}) where {D} = D codomain_type(::MatrixOp{D, T}) where {D, T} = D <: Real && T <: Complex ? T : D +domain_array_type(::MatrixOp{D, T, M, NC, dS}) where {D, T, M, NC, dS} = dS +codomain_array_type(::MatrixOp{D, T, M, NC, dS, cS}) where {D, T, M, NC, dS, cS} = cS is_thread_safe(::MatrixOp) = true # Type-stable size dispatch: NC=1 → 1D, NC>1 → 2D -size(L::MatrixOp{D, T, M, 1}) where {D, T, M} = ((size(L.A, 1),), (size(L.A, 2),)) -size(L::MatrixOp{D, T, M, NC}) where {D, T, M, NC} = ((size(L.A, 1), NC), (size(L.A, 2), NC)) +size(L::MatrixOp{D, T, M, 1, dS, cS}) where {D, T, M, dS, cS} = ((size(L.A, 1),), (size(L.A, 2),)) +size(L::MatrixOp{D, T, M, NC, dS, cS}) where {D, T, M, NC, dS, cS} = ((size(L.A, 1), NC), (size(L.A, 2), NC)) fun_name(L::MatrixOp) = "▒" @@ -139,10 +152,13 @@ is_AAc_diagonal(L::MatrixOp) = isdiag(L.A * L.A') is_AcA_diagonal(L::MatrixOp) = isdiag(L.A' * L.A) is_null(L::MatrixOp) = L.A == 0 * I is_eye(L::MatrixOp) = L.A == I -is_invertible(L::MatrixOp) = - size(L.A, 1) == size(L.A, 2) && - !isapprox(det(BigFloat.(L.A)), 0, atol = eps(eltype(L.A)) * 10) -is_orthogonal(L::MatrixOp) = size(L.A, 1) == size(L.A, 2) && all(<(eps(eltype(L.A)) * 10), L.A' * L.A - I) +function is_invertible(L::MatrixOp) + return size(L.A, 1) == size(L.A, 2) && + !isapprox(det(BigFloat.(L.A)), 0; atol = eps(eltype(L.A)) * 10) +end +function is_orthogonal(L::MatrixOp) + return size(L.A, 1) == size(L.A, 2) && all(<(eps(eltype(L.A)) * 10), L.A' * L.A - I) +end is_full_row_rank(L::MatrixOp) = rank(L.A) == size(L.A, 1) is_full_column_rank(L::MatrixOp) = rank(L.A) == size(L.A, 2) is_positive_definite(L::MatrixOp) = isposdef(L.A) diff --git a/src/linearoperators/MyLinOp.jl b/src/linearoperators/MyLinOp.jl index faf9615..78fcfdd 100644 --- a/src/linearoperators/MyLinOp.jl +++ b/src/linearoperators/MyLinOp.jl @@ -18,7 +18,7 @@ A ℝ^4 -> ℝ^5 ``` """ -struct MyLinOp{N, M, C, D, F <: Function, G <: Function} <: LinearOperator +struct MyLinOp{N, M, C, D, F <: Function, G <: Function, dS, cS} <: LinearOperator dim_out::NTuple{N, Int} dim_in::NTuple{M, Int} Fwd!::F @@ -28,35 +28,42 @@ end # Constructors function MyLinOp( - domain_type::Type, + domain_type::Type{T}, dim_in::NTuple{N, Int}, dim_out::NTuple{M, Int}, Fwd!::F, Adj!::G, - ) where {N, M, F <: Function, G <: Function} - return MyLinOp{N, M, domain_type, domain_type, F, G}(dim_out, dim_in, Fwd!, Adj!) + ; array_type::Type = Array{T}, + ) where {T, N, M, F <: Function, G <: Function} + dS = _normalize_array_type(array_type, domain_type) + return MyLinOp{N, M, domain_type, domain_type, F, G, dS, dS}(dim_out, dim_in, Fwd!, Adj!) end function MyLinOp( - domain_type::Type, + domain_type::Type{T}, dim_in::NTuple{N, Int}, - codomain_type::Type, + codomain_type::Type{C}, dim_out::NTuple{M, Int}, Fwd!::F, Adj!::G, - ) where {N, M, F <: Function, G <: Function} - return MyLinOp{N, M, domain_type, codomain_type, F, G}(dim_out, dim_in, Fwd!, Adj!) + ; array_type::Type = Array{T}, + ) where {T, C, N, M, F <: Function, G <: Function} + dS = _normalize_array_type(array_type, domain_type) + cS = _normalize_array_type(array_type, codomain_type) + return MyLinOp{N, M, domain_type, codomain_type, F, G, dS, cS}(dim_out, dim_in, Fwd!, Adj!) end # Mappings function mul!(y::AbstractArray, L::MyLinOp, b::AbstractArray) check(y, L, b) - return L.Fwd!(y, b) + L.Fwd!(y, b) + return y end function mul!(y::AbstractArray, L::AdjointOperator{<:MyLinOp}, b::AbstractArray) check(y, L, b) - return L.A.Adj!(y, b) + L.A.Adj!(y, b) + return y end # Properties @@ -65,5 +72,7 @@ size(L::MyLinOp) = (L.dim_out, L.dim_in) codomain_type(::MyLinOp{N, M, C}) where {N, M, C} = C domain_type(::MyLinOp{N, M, C, D}) where {N, M, C, D} = D +domain_array_type(::MyLinOp{N, M, C, D, F, G, dS}) where {N, M, C, D, F, G, dS} = dS +codomain_array_type(::MyLinOp{N, M, C, D, F, G, dS, cS}) where {N, M, C, D, F, G, dS, cS} = cS fun_name(L::MyLinOp) = "A" diff --git a/src/linearoperators/Variation.jl b/src/linearoperators/Variation.jl index ed515ea..a067e8a 100644 --- a/src/linearoperators/Variation.jl +++ b/src/linearoperators/Variation.jl @@ -23,31 +23,48 @@ julia> Variation(ones(2,2))*[1. 2.; 1. 2.] ``` """ -struct Variation{T, N, Th} <: LinearOperator +struct Variation{T, N, Th, S <: AbstractArray{T}} <: LinearOperator dim_in::NTuple{N, Int} end # Constructors #default constructor -function Variation(domain_type::Type, dim_in::NTuple{N, Int}; threaded = true) where {N} +function Variation( + domain_type::Type{T}, dim_in::NTuple{N, Int}; + threaded = true, array_type::Type = Array{T} + ) where {T, N} N == 1 && error("use FiniteDiff instead!") threaded = threaded && Threads.nthreads() > 1 && prod(dim_in) * sizeof(domain_type) > 2^16 - return Variation{domain_type, N, threaded}(dim_in) + S = _normalize_array_type(array_type, domain_type) + return Variation{domain_type, N, threaded, S}(dim_in) end -function Variation(domain_type::Type, dim_in::Vararg{Int}; threaded = true) - return Variation(domain_type, dim_in; threaded) +function Variation( + domain_type::Type{T}, dim_in::Vararg{Int}; threaded = true, array_type::Type = Array{T} + ) where {T} + return Variation(domain_type, dim_in; threaded, array_type) +end +function Variation( + dim_in::NTuple{N, Int}; threaded = true, array_type::Type = Array{Float64} + ) where {N} + return Variation(Float64, dim_in; threaded, array_type) +end +function Variation(dim_in::Vararg{Int}; threaded = true, array_type::Type = Array{Float64}) + return Variation(dim_in; threaded, array_type) end -function Variation(dim_in::NTuple{N, Int}; threaded = true) where {N} - return Variation(Float64, dim_in; threaded) +function Variation(x::AbstractArray; threaded = true) + S = _array_wrapper(x){eltype(x)} + T = eltype(x) + N = ndims(x) + threaded = threaded && _should_thread(x) + N == 1 && error("use FiniteDiff instead!") + return Variation{T, N, threaded, S}(size(x)) end -Variation(dim_in::Vararg{Int}; threaded = true) = Variation(dim_in; threaded) -Variation(x::AbstractArray; threaded = true) = Variation(eltype(x), size(x); threaded) # Mappings # Non-threaded forward @inbounds function LinearAlgebra.mul!( - y::AbstractArray, A::Variation{T, N, false}, b::AbstractArray + y::AbstractArray, A::Variation{T, N, false, <:Any}, b::AbstractArray ) where {T, N} check(y, A, b) @assert firstindex(b) == 1 "Only support 1-based arrays" @@ -69,12 +86,12 @@ Variation(x::AbstractArray; threaded = true) = Variation(eltype(x), size(x); thr next_slice_start = (k + 1) * batch_length + 1 next_slice_end = (k + 2) * batch_length next_slicing = next_slice_start:next_slice_end - y[slicing, d] = b[next_slicing] - b[slicing] + @views y[slicing, d] .= b[next_slicing] .- b[slicing] else prev_slice_start = (k - 1) * batch_length + 1 prev_slice_end = k * batch_length prev_slicing = prev_slice_start:prev_slice_end - y[slicing, d] = b[slicing] - b[prev_slicing] + @views y[slicing, d] .= b[slicing] .- b[prev_slicing] end end batch_count ÷= size(b, d) @@ -84,7 +101,7 @@ Variation(x::AbstractArray; threaded = true) = Variation(eltype(x), size(x); thr end @inbounds function LinearAlgebra.mul!( - y::AbstractArray, A::Variation{T, N, true}, b::AbstractArray + y::AbstractArray, A::Variation{T, N, true, <:Any}, b::AbstractArray ) where {T, N} check(y, A, b) @assert firstindex(b) == 1 "Only support 1-based arrays" @@ -106,12 +123,12 @@ end next_slice_start = (k + 1) * batch_length + 1 next_slice_end = (k + 2) * batch_length next_slicing = next_slice_start:next_slice_end - y[slicing, d] = b[next_slicing] - b[slicing] + @views y[slicing, d] .= b[next_slicing] .- b[slicing] else prev_slice_start = (k - 1) * batch_length + 1 prev_slice_end = k * batch_length prev_slicing = prev_slice_start:prev_slice_end - y[slicing, d] = b[slicing] - b[prev_slicing] + @views y[slicing, d] .= b[slicing] .- b[prev_slicing] end end batch_count ÷= size(b, d) @@ -122,7 +139,7 @@ end # Non-threaded adjoint @inbounds function LinearAlgebra.mul!( - y::AbstractArray, A::AdjointOperator{Variation{T, N, false}}, b::AbstractArray + y::AbstractArray, A::AdjointOperator{<:Variation{T, N, false}}, b::AbstractArray ) where {T, N} check(y, A, b) for cnt in LinearIndices(size(y)) @@ -156,7 +173,7 @@ end # Threaded adjoint @inbounds function LinearAlgebra.mul!( - y::AbstractArray, A::AdjointOperator{Variation{T, N, true}}, b::AbstractArray + y::AbstractArray, A::AdjointOperator{<:Variation{T, N, true}}, b::AbstractArray ) where {T, N} check(y, A, b) @batch for cnt in LinearIndices(size(y)) @@ -192,8 +209,18 @@ end domain_type(::Variation{T}) where {T} = T codomain_type(::Variation{T}) where {T} = T +domain_array_type(::Variation{T, N, Th, S}) where {T, N, Th, S} = S +codomain_array_type(::Variation{T, N, Th, S}) where {T, N, Th, S} = S is_thread_safe(::Variation) = true +function _copy_operator_impl( + op::Variation{T, N, Th, S}; array_type = nothing, threaded = nothing + ) where {T, N, Th, S} + new_threaded = threaded === nothing ? Th : threaded + new_at = array_type === nothing ? _array_wrapper_type(S) : array_type + return Variation(T, op.dim_in; threaded = new_threaded, array_type = new_at) +end + size(L::Variation{T, N}) where {T, N} = ((prod(L.dim_in), N), L.dim_in) fun_name(L::Variation) = "Ʋ" diff --git a/src/linearoperators/ZeroPad.jl b/src/linearoperators/ZeroPad.jl index f7fc2c9..988ef73 100644 --- a/src/linearoperators/ZeroPad.jl +++ b/src/linearoperators/ZeroPad.jl @@ -17,118 +17,87 @@ julia> Z*ones(2,2) ``` """ -struct ZeroPad{N, T} <: LinearOperator +struct ZeroPad{N, T, S <: AbstractArray{T}} <: LinearOperator dim_in::NTuple{N, Int} zp::NTuple{N, Int} end # Constructors #standard constructor -function ZeroPad(domain_type::Type, dim_in::NTuple{N, Int}, zp::NTuple{M, Int}) where {N, M} +function ZeroPad( + domain_type::Type{T}, dim_in::NTuple{N, Int}, zp::NTuple{M, Int}; + array_type::Type = Array{T} + ) where {T, N, M} M != N && error("dim_in and zp must have the same length") any([zp...] .< 0) && error("zero padding cannot be negative") - return ZeroPad{N, domain_type}(dim_in, zp) + S = _normalize_array_type(array_type, T) + return ZeroPad{N, T, S}(dim_in, zp) end -ZeroPad(dim_in::Tuple, zp::NTuple{N, Int}) where {N} = ZeroPad(Float64, dim_in, zp) -function ZeroPad(domain_type::Type, dim_in::Tuple, zp::Vararg{Int, N}) where {N} - return ZeroPad(domain_type, dim_in, zp) +function ZeroPad(dim_in::Tuple, zp::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return ZeroPad(Float64, dim_in, zp; array_type) +end +function ZeroPad( + domain_type::Type{T}, dim_in::Tuple, zp::Vararg{Int, N}; array_type::Type = Array{T} + ) where {T, N} + return ZeroPad(domain_type, dim_in, zp; array_type) +end +function ZeroPad(dim_in::Tuple, zp::Vararg{Int, N}; array_type::Type = Array{Float64}) where {N} + return ZeroPad(Float64, dim_in, zp; array_type) +end +function ZeroPad(x::AbstractArray{T}, zp::NTuple{N, Int}) where {T, N} + S = _normalize_array_type(_array_wrapper(x), T) + return ZeroPad{N, T, S}(size(x), zp) +end +function ZeroPad(x::AbstractArray{T}, zp::Vararg{Int, N}) where {T, N} + S = _normalize_array_type(_array_wrapper(x), T) + return ZeroPad{N, T, S}(size(x), zp) end -ZeroPad(dim_in::Tuple, zp::Vararg{Int, N}) where {N} = ZeroPad(Float64, dim_in, zp) -ZeroPad(x::AbstractArray, zp::NTuple{N, Int}) where {N} = ZeroPad(eltype(x), size(x), zp) -ZeroPad(x::AbstractArray, zp::Vararg{Int, N}) where {N} = ZeroPad(eltype(x), size(x), zp) # Mappings -@generated function mul!( - y::AbstractArray, L::ZeroPad{N}, b::AbstractArray - ) where {N} - - # builds - #for i1 =1:size(y,1), i2 =1:size(y,2) - # y[i1,i2] = i1 <= size(b,1) && i2 <= size(b,2) ? b[i1,i2] : 0. - #end - - ex = "for " - for i in 1:N - ex *= "i$i =1:size(y,$i)," +@generated function mul!(y::AbstractArray, L::ZeroPad{N, T}, b::AbstractArray) where {N, T} + z = zero(T) # compile-time literal zero for type T + vars = [Symbol("i$i") for i in 1:N] + # Condition: i1 <= size(b,1) && i2 <= size(b,2) && ... + cond = if N == 1 + :($(vars[1]) <= size(b, 1)) + else + Expr(:&&, [:($(vars[i]) <= size(b, $i)) for i in 1:N]...) end - - ex = ex[1:(end - 1)] #remove , - - ex *= " y[" - for i in 1:N - ex *= "i$i," - end - - ex = ex[1:(end - 1)] #remove , - - ex *= "] = " - for i in 1:N - ex *= " i$i <= size(b,$i) &&" - end - - ex = ex[1:(end - 2)] #remove && - - ex *= " ? b[" + # Inner assignment: @inbounds y[i1,i2,...] = cond ? b[i1,i2,...] : z + assign = :(@inbounds y[$(vars...)] = $cond ? b[$(vars...)] : $z) + # Wrap in nested for loops: outermost=dim N, innermost=dim 1 (column-major) + body = assign for i in 1:N - ex *= "i$i," + body = :( + for $(vars[i]) in 1:size(y, $i) + ;$body + end + ) end - - ex = ex[1:(end - 1)] #remove , - ex *= "] : 0. end" - - loop_expr = Meta.parse(ex) return quote check(y, L, b) - $loop_expr + $body return y end end -@generated function mul!( - y::AbstractArray, L::AdjointOperator{<:ZeroPad{N}}, b::AbstractArray - ) where {N} - - #builds - #for l = 1:size(y,1), m = 1:size(y,2) - # y[l,m] = b[l,m] - #end - - ex = "for " - for i in 1:N - ex *= "i$i =1:size(y,$i)," - end - - ex *= " y[" - idx = "" - for i in 1:N - idx *= "i$i," - end - - idx = idx[1:(end - 1)] #remove , - - ex *= idx - - ex *= "] = b[" - ex *= idx - ex *= "] end" - - loop_expr = Meta.parse(ex) - return quote - check(y, L, b) - $loop_expr - return y - end +function mul!(y::AbstractArray, L::AdjointOperator{<:ZeroPad}, b::AbstractArray) + check(y, L, b) + copyto!(y, view(b, ntuple(i -> 1:L.A.dim_in[i], length(L.A.dim_in))...)) + return y end -function get_normal_op(L::ZeroPad{N, T}) where {N, T} - return Eye(domain_type(L), size(L, 2), domain_storage_type(L)) +function get_normal_op(L::ZeroPad{N, T, S}) where {N, T, S} + return Eye(domain_type(L), size(L, 2); array_type = S) end # Properties domain_type(::ZeroPad{<:Any, T}) where {T} = T codomain_type(::ZeroPad{<:Any, T}) where {T} = T +domain_array_type(::ZeroPad{N, T, S}) where {N, T, S} = S +codomain_array_type(::ZeroPad{N, T, S}) where {N, T, S} = S is_thread_safe(::ZeroPad) = true size(L::ZeroPad) = L.dim_in .+ L.zp, L.dim_in diff --git a/src/linearoperators/Zeros.jl b/src/linearoperators/Zeros.jl index d228127..8628191 100644 --- a/src/linearoperators/Zeros.jl +++ b/src/linearoperators/Zeros.jl @@ -16,7 +16,7 @@ julia> Zeros([Eye(10,20) Eye(10,20)]) ``` """ -struct Zeros{C, N, D, M} <: LinearOperator +struct Zeros{C, N, D, M, dS, cS} <: LinearOperator dim_out::NTuple{N, Int} dim_in::NTuple{M, Int} end @@ -24,13 +24,25 @@ end # Constructors #default function Zeros( - domain_type::Type, dim_in::NTuple{M, Int}, codomain_type::Type, dim_out::NTuple{N, Int} - ) where {N, M} - return Zeros{codomain_type, N, domain_type, M}(dim_out, dim_in) + domain_type::Type{D}, + dim_in::NTuple{M, Int}, + codomain_type::Type{C}, + dim_out::NTuple{N, Int}; + array_type::Type = Array{D}, + ) where {N, M, D, C} + dS = _normalize_array_type(array_type, domain_type) + cS = _normalize_array_type(array_type, codomain_type) + return Zeros{codomain_type, N, domain_type, M, dS, cS}(dim_out, dim_in) end -function Zeros(domain_type::Type, dim_in::NTuple{M, Int}, dim_out::NTuple{N, Int}) where {N, M} - return Zeros{domain_type, N, domain_type, M}(dim_out, dim_in) +function Zeros( + domain_type::Type{T}, + dim_in::NTuple{M, Int}, + dim_out::NTuple{N, Int}; + array_type::Type = Array{T}, + ) where {N, M, T} + dS = _normalize_array_type(array_type, domain_type) + return Zeros{domain_type, N, domain_type, M, dS, dS}(dim_out, dim_in) end function Zeros( @@ -69,6 +81,8 @@ end domain_type(::Zeros{C, N, D, M}) where {C, N, D, M} = D codomain_type(::Zeros{C, N, D, M}) where {C, N, D, M} = C +domain_array_type(::Zeros{C, N, D, M, dS}) where {C, N, D, M, dS} = dS +codomain_array_type(::Zeros{C, N, D, M, dS, cS}) where {C, N, D, M, dS, cS} = cS is_thread_safe(::Zeros) = true size(L::Zeros) = (L.dim_out, L.dim_in) diff --git a/src/nonlinearoperators/Atan.jl b/src/nonlinearoperators/Atan.jl index 04b2cf8..2ec7b5f 100644 --- a/src/nonlinearoperators/Atan.jl +++ b/src/nonlinearoperators/Atan.jl @@ -9,16 +9,26 @@ Creates an inverse tangent non-linear operator with input dimensions `dim_in`: ``` """ -struct Atan{T, N} <: NonLinearOperator +struct Atan{T, N, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} end -function Atan(domain_type::Type, DomainDim::NTuple{N, Int}) where {N} - return Atan{domain_type, N}(DomainDim) +function Atan( + domain_type::Type{T}, DomainDim::NTuple{N, Int}; array_type::Type = Array{T} + ) where {T, N} + S = _normalize_array_type(array_type, T) + return Atan{T, N, S}(DomainDim) end -Atan(DomainDim::NTuple{N, Int}) where {N} = Atan{Float64, N}(DomainDim) -Atan(DomainDim::Vararg{Int}) = Atan{Float64, length(DomainDim)}(DomainDim) +function Atan(DomainDim::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return Atan(Float64, DomainDim; array_type) +end +Atan(DomainDim::Vararg{Int}; array_type::Type = Array{Float64}) = Atan(Float64, DomainDim; array_type) + +function Atan(x::AbstractArray{T}; array_type::Type = _array_wrapper(x)) where {T} + S = _normalize_array_type(array_type, T) + return Atan{T, ndims(x), S}(size(x)) +end function mul!(y::AbstractArray, L::Atan, x::AbstractArray) check(y, L, x) @@ -37,4 +47,6 @@ size(L::Atan) = (L.dim, L.dim) domain_type(::Atan{T, N}) where {T, N} = T codomain_type(::Atan{T, N}) where {T, N} = T +domain_array_type(::Atan{T, N, S}) where {T, N, S} = S +codomain_array_type(::Atan{T, N, S}) where {T, N, S} = S is_thread_safe(::Atan) = true diff --git a/src/nonlinearoperators/Cos.jl b/src/nonlinearoperators/Cos.jl index 0b10f56..9041ee8 100644 --- a/src/nonlinearoperators/Cos.jl +++ b/src/nonlinearoperators/Cos.jl @@ -9,16 +9,26 @@ Creates a cosine non-linear operator with input dimensions `dim_in`: ``` """ -struct Cos{T, N} <: NonLinearOperator +struct Cos{T, N, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} end -function Cos(domain_type::Type, DomainDim::NTuple{N, Int}) where {N} - return Cos{domain_type, N}(DomainDim) +function Cos( + domain_type::Type{T}, DomainDim::NTuple{N, Int}; array_type::Type = Array{T} + ) where {T, N} + S = _normalize_array_type(array_type, T) + return Cos{T, N, S}(DomainDim) end -Cos(DomainDim::NTuple{N, Int}) where {N} = Cos{Float64, N}(DomainDim) -Cos(DomainDim::Vararg{Int}) = Cos{Float64, length(DomainDim)}(DomainDim) +function Cos(DomainDim::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return Cos(Float64, DomainDim; array_type) +end +Cos(DomainDim::Vararg{Int}; array_type::Type = Array{Float64}) = Cos(Float64, DomainDim; array_type) + +function Cos(x::AbstractArray{T}; array_type::Type = _array_wrapper(x)) where {T} + S = _normalize_array_type(array_type, T) + return Cos{T, ndims(x), S}(size(x)) +end function mul!(y::AbstractArray, L::Cos, x::AbstractArray) check(y, L, x) @@ -37,4 +47,6 @@ size(L::Cos) = (L.dim, L.dim) domain_type(::Cos{T, N}) where {T, N} = T codomain_type(::Cos{T, N}) where {T, N} = T +domain_array_type(::Cos{T, N, S}) where {T, N, S} = S +codomain_array_type(::Cos{T, N, S}) where {T, N, S} = S is_thread_safe(::Cos) = true diff --git a/src/nonlinearoperators/Exp.jl b/src/nonlinearoperators/Exp.jl index 34cf846..0777fee 100644 --- a/src/nonlinearoperators/Exp.jl +++ b/src/nonlinearoperators/Exp.jl @@ -9,16 +9,26 @@ e^{ \\mathbf{x} }. ``` """ -struct Exp{T, N} <: NonLinearOperator +struct Exp{T, N, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} end -function Exp(domain_type::Type, DomainDim::NTuple{N, Int}) where {N} - return Exp{domain_type, N}(DomainDim) +function Exp( + domain_type::Type{T}, DomainDim::NTuple{N, Int}; array_type::Type = Array{T} + ) where {T, N} + S = _normalize_array_type(array_type, T) + return Exp{T, N, S}(DomainDim) end -Exp(DomainDim::NTuple{N, Int}) where {N} = Exp{Float64, N}(DomainDim) -Exp(DomainDim::Vararg{Int}) = Exp{Float64, length(DomainDim)}(DomainDim) +function Exp(DomainDim::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return Exp(Float64, DomainDim; array_type) +end +Exp(DomainDim::Vararg{Int}; array_type::Type = Array{Float64}) = Exp(Float64, DomainDim; array_type) + +function Exp(x::AbstractArray{T}; array_type::Type = _array_wrapper(x)) where {T} + S = _normalize_array_type(array_type, T) + return Exp{T, ndims(x), S}(size(x)) +end function mul!(y::AbstractArray, L::Exp, x::AbstractArray) check(y, L, x) @@ -37,4 +47,6 @@ size(L::Exp) = (L.dim, L.dim) domain_type(::Exp{T, N}) where {T, N} = T codomain_type(::Exp{T, N}) where {T, N} = T +domain_array_type(::Exp{T, N, S}) where {T, N, S} = S +codomain_array_type(::Exp{T, N, S}) where {T, N, S} = S is_thread_safe(::Exp) = true diff --git a/src/nonlinearoperators/Pow.jl b/src/nonlinearoperators/Pow.jl index 5bda1c0..a351855 100644 --- a/src/nonlinearoperators/Pow.jl +++ b/src/nonlinearoperators/Pow.jl @@ -6,16 +6,26 @@ export Pow Elementwise power `p` non-linear operator with input dimensions `dim_in`. """ -struct Pow{T, N, I <: Real} <: NonLinearOperator +struct Pow{T, N, I <: Real, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} p::I end -function Pow(domain_type::Type, DomainDim::NTuple{N, Int}, p::I) where {N, I <: Real} - return Pow{domain_type, N, I}(DomainDim, p) +function Pow( + domain_type::Type{T}, DomainDim::NTuple{N, Int}, p::I; array_type::Type = Array{T} + ) where {T, N, I <: Real} + S = _normalize_array_type(array_type, T) + return Pow{T, N, I, S}(DomainDim, p) end -Pow(DomainDim::NTuple{N, Int}, p::I) where {N, I <: Real} = Pow{Float64, N, I}(DomainDim, p) +function Pow(DomainDim::NTuple{N, Int}, p::I; array_type::Type = Array{Float64}) where {N, I <: Real} + return Pow(Float64, DomainDim, p; array_type) +end + +function Pow(x::AbstractArray{T}, p::I; array_type::Type = _array_wrapper(x)) where {T, I <: Real} + S = _normalize_array_type(array_type, T) + return Pow{T, ndims(x), I, S}(size(x), p) +end function mul!(y::AbstractArray, L::Pow, x::AbstractArray) check(y, L, x) @@ -34,4 +44,6 @@ size(L::Pow) = (L.dim, L.dim) domain_type(::Pow{T, N}) where {T, N} = T codomain_type(::Pow{T, N}) where {T, N} = T +domain_array_type(::Pow{T, N, I, S}) where {T, N, I, S} = S +codomain_array_type(::Pow{T, N, I, S}) where {T, N, I, S} = S is_thread_safe(::Pow) = true diff --git a/src/nonlinearoperators/Sech.jl b/src/nonlinearoperators/Sech.jl index cdc15b3..c61bf5b 100644 --- a/src/nonlinearoperators/Sech.jl +++ b/src/nonlinearoperators/Sech.jl @@ -9,16 +9,26 @@ Creates an hyperbolic secant non-linear operator with input dimensions `dim_in`: ``` """ -struct Sech{T, N} <: NonLinearOperator +struct Sech{T, N, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} end -function Sech(domain_type::Type, DomainDim::NTuple{N, Int}) where {N} - return Sech{domain_type, N}(DomainDim) +function Sech( + domain_type::Type{T}, DomainDim::NTuple{N, Int}; array_type::Type = Array{T} + ) where {T, N} + S = _normalize_array_type(array_type, T) + return Sech{T, N, S}(DomainDim) end -Sech(DomainDim::NTuple{N, Int}) where {N} = Sech{Float64, N}(DomainDim) -Sech(DomainDim::Vararg{Int}) = Sech{Float64, length(DomainDim)}(DomainDim) +function Sech(DomainDim::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return Sech(Float64, DomainDim; array_type) +end +Sech(DomainDim::Vararg{Int}; array_type::Type = Array{Float64}) = Sech(Float64, DomainDim; array_type) + +function Sech(x::AbstractArray{T}; array_type::Type = _array_wrapper(x)) where {T} + S = _normalize_array_type(array_type, T) + return Sech{T, ndims(x), S}(size(x)) +end function mul!(y::AbstractArray, L::Sech, x::AbstractArray) check(y, L, x) @@ -37,4 +47,6 @@ size(L::Sech) = (L.dim, L.dim) domain_type(::Sech{T, N}) where {T, N} = T codomain_type(::Sech{T, N}) where {T, N} = T +domain_array_type(::Sech{T, N, S}) where {T, N, S} = S +codomain_array_type(::Sech{T, N, S}) where {T, N, S} = S is_thread_safe(::Sech) = true diff --git a/src/nonlinearoperators/Sigmoid.jl b/src/nonlinearoperators/Sigmoid.jl index 20c7701..a310c3a 100644 --- a/src/nonlinearoperators/Sigmoid.jl +++ b/src/nonlinearoperators/Sigmoid.jl @@ -9,17 +9,32 @@ Creates the sigmoid non-linear operator with input dimensions `dim_in`. ``` """ -struct Sigmoid{T, N, G <: Real} <: NonLinearOperator +struct Sigmoid{T, N, G <: Real, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} gamma::G end -function Sigmoid(domain_type::Type, DomainDim::NTuple{N, Int}, gamma::G = 1.0) where {N, G <: Real} - return Sigmoid{domain_type, N, G}(DomainDim, gamma) +function Sigmoid( + domain_type::Type{T}, + DomainDim::NTuple{N, Int}, + gamma::G = 1.0; + array_type::Type = Array{T}, + ) where {T, N, G <: Real} + S = _normalize_array_type(array_type, T) + return Sigmoid{T, N, G, S}(DomainDim, gamma) end -function Sigmoid(DomainDim::NTuple{N, Int}, gamma::G = 1.0) where {N, G} - return Sigmoid{Float64, N, G}(DomainDim, gamma) +function Sigmoid( + DomainDim::NTuple{N, Int}, gamma::G = 1.0; array_type::Type = Array{Float64} + ) where {N, G} + return Sigmoid(Float64, DomainDim, gamma; array_type) +end + +function Sigmoid( + x::AbstractArray{T}; gamma::G = 1.0, array_type::Type = _array_wrapper(x) + ) where {T, G <: Real} + S = _normalize_array_type(array_type, T) + return Sigmoid{T, ndims(x), G, S}(size(x), gamma) end function mul!(y::AbstractArray, L::Sigmoid, x::AbstractArray) @@ -42,4 +57,6 @@ size(L::Sigmoid) = (L.dim, L.dim) domain_type(::Sigmoid{T, N, D}) where {T, N, D} = T codomain_type(::Sigmoid{T, N, D}) where {T, N, D} = T +domain_array_type(::Sigmoid{T, N, D, S}) where {T, N, D, S} = S +codomain_array_type(::Sigmoid{T, N, D, S}) where {T, N, D, S} = S is_thread_safe(::Sigmoid) = true diff --git a/src/nonlinearoperators/Sin.jl b/src/nonlinearoperators/Sin.jl index 98188ac..e9f6bda 100644 --- a/src/nonlinearoperators/Sin.jl +++ b/src/nonlinearoperators/Sin.jl @@ -9,16 +9,26 @@ Creates a sinusoid non-linear operator with input dimensions `dim_in`: ``` """ -struct Sin{T, N} <: NonLinearOperator +struct Sin{T, N, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} end -function Sin(domain_type::Type, DomainDim::NTuple{N, Int}) where {N} - return Sin{domain_type, N}(DomainDim) +function Sin( + domain_type::Type{T}, DomainDim::NTuple{N, Int}; array_type::Type = Array{T} + ) where {T, N} + S = _normalize_array_type(array_type, T) + return Sin{T, N, S}(DomainDim) end -Sin(DomainDim::NTuple{N, Int}) where {N} = Sin{Float64, N}(DomainDim) -Sin(DomainDim::Vararg{Int}) = Sin{Float64, length(DomainDim)}(DomainDim) +function Sin(DomainDim::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return Sin(Float64, DomainDim; array_type) +end +Sin(DomainDim::Vararg{Int}; array_type::Type = Array{Float64}) = Sin(Float64, DomainDim; array_type) + +function Sin(x::AbstractArray{T}; array_type::Type = _array_wrapper(x)) where {T} + S = _normalize_array_type(array_type, T) + return Sin{T, ndims(x), S}(size(x)) +end function mul!(y::AbstractArray, L::Sin, x::AbstractArray) check(y, L, x) @@ -37,4 +47,6 @@ size(L::Sin) = (L.dim, L.dim) domain_type(::Sin{T, N}) where {T, N} = T codomain_type(::Sin{T, N}) where {T, N} = T +domain_array_type(::Sin{T, N, S}) where {T, N, S} = S +codomain_array_type(::Sin{T, N, S}) where {T, N, S} = S is_thread_safe(::Sin) = true diff --git a/src/nonlinearoperators/SoftMax.jl b/src/nonlinearoperators/SoftMax.jl index 7c8f9a0..eef73da 100644 --- a/src/nonlinearoperators/SoftMax.jl +++ b/src/nonlinearoperators/SoftMax.jl @@ -14,17 +14,24 @@ struct SoftMax{T, N, B <: AbstractArray{T, N}} <: NonLinearOperator buf::B end -function SoftMax(x::AbstractArray{T, N}) where {T, N} - buf = similar(x) +function SoftMax(x::AbstractArray{T, N}; array_type::Type = _array_wrapper(x)) where {T, N} + S = _normalize_array_type(array_type, T) + buf = similar(S, size(x)) return SoftMax{T, N, typeof(buf)}(size(x), buf) end -function SoftMax(domain_type::Type, DomainDim::NTuple{N, Int}) where {N} - buf = zeros(domain_type, DomainDim) - return SoftMax{domain_type, N, typeof(buf)}(DomainDim, buf) +function SoftMax( + domain_type::Type{T}, DomainDim::NTuple{N, Int}; array_type::Type = Array{T} + ) where {T, N} + S = _normalize_array_type(array_type, T) + buf = similar(S, DomainDim) + fill!(buf, zero(T)) + return SoftMax{T, N, typeof(buf)}(DomainDim, buf) end -SoftMax(DomainDim::NTuple{N, Int}) where {N} = SoftMax(Float64, DomainDim) +function SoftMax(DomainDim::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return SoftMax(Float64, DomainDim; array_type) +end function mul!(y::AbstractArray, L::SoftMax, x::AbstractArray) check(y, L, x) @@ -35,13 +42,10 @@ end function mul!(y::AbstractArray, J::AdjointOperator{<:Jacobian{<:SoftMax}}, b::AbstractArray) check(y, J, b) L = J.A - fill!(y, zero(eltype(y))) L.A.buf .= exp.(L.x .- maximum(L.x)) L.A.buf ./= sum(L.A.buf) - for i in eachindex(y) - y[i] = -L.A.buf[i] * dot(L.A.buf, b) - y[i] += L.A.buf[i] * b[i] - end + d = dot(L.A.buf, b) + y .= L.A.buf .* (b .- d) return y end @@ -51,4 +55,6 @@ size(L::SoftMax) = (L.dim, L.dim) domain_type(::SoftMax{T}) where {T} = T codomain_type(::SoftMax{T}) where {T} = T -is_thread_safe(::SoftMax) = true +domain_array_type(::SoftMax{T, N, B}) where {T, N, B} = _array_wrapper_type(B){T} +codomain_array_type(::SoftMax{T, N, B}) where {T, N, B} = _array_wrapper_type(B){T} +is_thread_safe(::SoftMax) = false diff --git a/src/nonlinearoperators/SoftPlus.jl b/src/nonlinearoperators/SoftPlus.jl index 3ac6553..2614d18 100644 --- a/src/nonlinearoperators/SoftPlus.jl +++ b/src/nonlinearoperators/SoftPlus.jl @@ -9,15 +9,25 @@ Creates the softplus non-linear operator with input dimensions `dim_in`. ``` """ -struct SoftPlus{T, N} <: NonLinearOperator +struct SoftPlus{T, N, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} end -function SoftPlus(domain_type::Type, DomainDim::NTuple{N, Int}) where {N} - return SoftPlus{domain_type, N}(DomainDim) +function SoftPlus( + domain_type::Type{T}, DomainDim::NTuple{N, Int}; array_type::Type = Array{T} + ) where {T, N} + S = _normalize_array_type(array_type, T) + return SoftPlus{T, N, S}(DomainDim) end -SoftPlus(DomainDim::NTuple{N, Int}) where {N} = SoftPlus{Float64, N}(DomainDim) +function SoftPlus(DomainDim::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return SoftPlus(Float64, DomainDim; array_type) +end + +function SoftPlus(x::AbstractArray{T}; array_type::Type = _array_wrapper(x)) where {T} + S = _normalize_array_type(array_type, T) + return SoftPlus{T, ndims(x), S}(size(x)) +end function mul!(y::AbstractArray, L::SoftPlus, x::AbstractArray) check(y, L, x) @@ -36,4 +46,6 @@ size(L::SoftPlus) = (L.dim, L.dim) domain_type(::SoftPlus{T, N}) where {T, N} = T codomain_type(::SoftPlus{T, N}) where {T, N} = T +domain_array_type(::SoftPlus{T, N, S}) where {T, N, S} = S +codomain_array_type(::SoftPlus{T, N, S}) where {T, N, S} = S is_thread_safe(::SoftPlus) = true diff --git a/src/nonlinearoperators/Tanh.jl b/src/nonlinearoperators/Tanh.jl index a7789a7..3e6b2b8 100644 --- a/src/nonlinearoperators/Tanh.jl +++ b/src/nonlinearoperators/Tanh.jl @@ -9,16 +9,26 @@ Creates an hyperbolic tangent non-linear operator with input dimensions `dim_in` ``` """ -struct Tanh{T, N} <: NonLinearOperator +struct Tanh{T, N, S <: AbstractArray{T}} <: NonLinearOperator dim::NTuple{N, Int} end -function Tanh(domain_type::Type, DomainDim::NTuple{N, Int}) where {N} - return Tanh{domain_type, N}(DomainDim) +function Tanh( + domain_type::Type{T}, DomainDim::NTuple{N, Int}; array_type::Type = Array{T} + ) where {T, N} + S = _normalize_array_type(array_type, T) + return Tanh{T, N, S}(DomainDim) end -Tanh(DomainDim::NTuple{N, Int}) where {N} = Tanh{Float64, N}(DomainDim) -Tanh(DomainDim::Vararg{Int}) = Tanh{Float64, length(DomainDim)}(DomainDim) +function Tanh(DomainDim::NTuple{N, Int}; array_type::Type = Array{Float64}) where {N} + return Tanh(Float64, DomainDim; array_type) +end +Tanh(DomainDim::Vararg{Int}; array_type::Type = Array{Float64}) = Tanh(Float64, DomainDim; array_type) + +function Tanh(x::AbstractArray{T}; array_type::Type = _array_wrapper(x)) where {T} + S = _normalize_array_type(array_type, T) + return Tanh{T, ndims(x), S}(size(x)) +end function mul!(y::AbstractArray, L::Tanh, x::AbstractArray) check(y, L, x) @@ -37,4 +47,6 @@ size(L::Tanh) = (L.dim, L.dim) domain_type(::Tanh{T, N}) where {T, N} = T codomain_type(::Tanh{T, N}) where {T, N} = T +domain_array_type(::Tanh{T, N, S}) where {T, N, S} = S +codomain_array_type(::Tanh{T, N, S}) where {T, N, S} = S is_thread_safe(::Tanh) = true diff --git a/src/properties.jl b/src/properties.jl index e91ed9c..0d878a0 100644 --- a/src/properties.jl +++ b/src/properties.jl @@ -1,11 +1,13 @@ import Base: size, ndims, similar, copy import LinearAlgebra: diag, opnorm +export copy_operator +# _should_thread is internal, not exported export ndoms, domain_type, codomain_type, - domain_storage_type, - codomain_storage_type, + domain_array_type, + codomain_array_type, is_linear, is_eye, is_null, @@ -59,74 +61,82 @@ julia> codomain_type(vcat(Eye(Complex{Float64},(10,)),DiagOp(rand(ComplexF64, 10 codomain_type """ - domain_storage_type(L::AbstractOperator) + domain_array_type(L::AbstractOperator) Returns the type of the storage for the domain of the operator. ```jldoctest -julia> domain_storage_type(DiagOp(rand(10))) +julia> domain_array_type(DiagOp(rand(10))) Array{Float64} -julia> domain_storage_type(hcat(Eye(Complex{Float64},(10,)),DiagOp(rand(ComplexF64, 10)))) +julia> domain_array_type(hcat(Eye(Complex{Float64},(10,)),DiagOp(rand(ComplexF64, 10)))) RecursiveArrayTools.ArrayPartition{ComplexF64, Tuple{Array{ComplexF64}, Array{ComplexF64}}} ``` """ -function domain_storage_type(L::AbstractOperator) - dt = domain_type(L) - return if dt isa Tuple - arrayTypes = Tuple{[Array{t} for t in dt]...} - ArrayPartition{promote_type(dt...), arrayTypes} - else - Array{dt} - end +function domain_array_type(L::AbstractOperator) + return _array_type_for_elem(domain_type(L)) end """ - codomain_storage_type(L::AbstractOperator) + codomain_array_type(L::AbstractOperator) Returns the type of the storage of for the codomain of the operator. ```jldoctest -julia> codomain_storage_type(DiagOp(rand(ComplexF64,10))) +julia> codomain_array_type(DiagOp(rand(ComplexF64,10))) Array{ComplexF64} -julia> codomain_storage_type(vcat(Eye(Complex{Float64},(10,)),DiagOp(rand(ComplexF64,10)))) +julia> codomain_array_type(vcat(Eye(Complex{Float64},(10,)),DiagOp(rand(ComplexF64,10)))) RecursiveArrayTools.ArrayPartition{ComplexF64, Tuple{Array{ComplexF64}, Array{ComplexF64}}} ``` """ -function codomain_storage_type(L::AbstractOperator) - dt = codomain_type(L) - return if dt isa Tuple - arrayTypes = Tuple{[Array{t} for t in dt]...} - ArrayPartition{promote_type(dt...), arrayTypes} - else - Array{dt} - end +function codomain_array_type(L::AbstractOperator) + return _array_type_for_elem(codomain_type(L)) +end + +_array_type_for_elem(T::Type) = Array{T} +function _array_type_for_elem(dt::Tuple) + arrayTypes = Tuple{[Array{t} for t in dt]...} + return ArrayPartition{promote_type(dt...), arrayTypes} +end + +# Get the array container type (Array, CuArray, etc.) from an array instance +_array_wrapper_type(::Type{A}) where {A <: AbstractArray} = Base.typename(A).wrapper +function _array_wrapper(x::A) where {T, A <: AbstractArray{T}} + return _array_wrapper_type(typeof(x isa SubArray ? parent(x) : x)) +end +function _normalize_array_type(array_type::Type{A}, elem_type::Type{T}) where {A <: AbstractArray, T} + return _array_wrapper_type(A){T} end +_storage_eltype(::Type{<:AbstractArray{T}}) where {T} = T function allocate_in_domain(L::AbstractOperator, dims... = size(L, 2)...) - dS = domain_storage_type(L) + dS = domain_array_type(L) if dS <: ArrayPartition S = dS.parameters[2] return ArrayPartition([similar(s, d...) for (s, d) in zip(S.parameters, dims)]...) + elseif dS <: Array + return Array{domain_type(L), length(dims)}(undef, dims...) else return similar(dS, dims...) end end function allocate_in_codomain(L::AbstractOperator, dims... = size(L, 1)...) - cS = codomain_storage_type(L) + cS = codomain_array_type(L) if cS <: ArrayPartition S = cS.parameters[2] return ArrayPartition([similar(s, d...) for (s, d) in zip(S.parameters, dims)]...) + elseif cS <: Array + return Array{codomain_type(L), length(dims)}(undef, dims...) else return similar(cS, dims...) end end -storage_type_display_string(::Type{T}) where {T <: AbstractArray} = "" +array_type_display_string(::Type{T}) where {T <: AbstractArray} = "" function storage_display_string(L::AbstractOperator) - return storage_type_display_string(codomain_storage_type(L)) + return array_type_display_string(codomain_array_type(L)) end """ @@ -162,7 +172,7 @@ size(L::AbstractOperator, i::Int) = size(L)[i] Returns a `Tuple` with the number of dimensions of the codomain and domain of an `AbstractOperator`. Type `ndims(A,1)` for the number of dimensions of the codomain and `ndims(A,2)` for the number of dimensions of the codomain. -```jldoctest +```jldoctest; setup = :(using AbstractOperators) julia> V = Variation((2,3,4)) Ʋ ℝ^(2, 3, 4) -> ℝ^(24, 3) @@ -289,12 +299,10 @@ remove_displacement(A::AbstractOperator) = A import Base: convert function convert(::Type{T}, dom::Type, dim_in::Tuple, L::T) where {T <: AbstractOperator} - domain_type(L) != dom && error( - "cannot convert operator with domain $(domain_type(L)) to operator with domain $dom ", - ) - size(L, 1) != dim_in && error( - "cannot convert operator with size $(size(L, 1)) to operator with domain $dim_in ", - ) + domain_type(L) != dom && + error("cannot convert operator with domain $(domain_type(L)) to operator with domain $dom ") + size(L, 1) != dim_in && + error("cannot convert operator with size $(size(L, 1)) to operator with domain $dim_in ") return L end @@ -305,7 +313,7 @@ end Returns whether the operators `L` and `R` can be merged when they are multiplied. Examples: -```jldoctest +```jldoctest; setup = :(using AbstractOperators) julia> AbstractOperators.can_be_combined(DiagOp(rand(10)), FiniteDiff((10,))) false @@ -313,7 +321,12 @@ julia> AbstractOperators.can_be_combined(Eye(10), FiniteDiff((11,))) true ``` """ -can_be_combined(L, R) = is_eye(L) || is_eye(R) || is_null(L) || (is_null(R) && is_linear(L) && all(displacement(L) .== 0)) +function can_be_combined(L, R) + return is_eye(L) || + is_eye(R) || + is_null(L) || + (is_null(R) && is_linear(L) && all(displacement(L) .== 0)) +end can_be_combined(L, M, R) = false """ @@ -321,7 +334,7 @@ can_be_combined(L, M, R) = false Returns the combined operator of `L` and `R`. The combined operator is defined as `L * R` where `L` and `R` are the operators to be combined. Examples: -```jldoctest +```jldoctest; setup = :(using AbstractOperators) julia> AbstractOperators.combine(Eye(10), DiagOp(rand(10))) ╲ ℝ^10 -> ℝ^10 ``` @@ -471,3 +484,43 @@ function string_dom(dm::Tuple, sz::Tuple) s = string_dom.(dm, sz) return length(s) > 3 ? s[1] * "..." * s[end] : *(s...) end + +""" + copy_operator(op::AbstractOperator; array_type=nothing, threaded=nothing) + +Create a copy of `op` suitable for parallel use. + +- Immutable fields (operator arrays, type params) are **shared** (no copy). +- Mutable buffer fields are **deep-copied**. +- `array_type`: if provided (e.g., `CuArray`), convert buffer arrays to that storage. +- `threaded`: if provided (`true`/`false`), toggle threading for operators that support it. + +When `array_type` is `nothing` and `threaded` is `nothing`, equivalent to the old `copy_op` +but more efficient (shares immutable data). +""" +function copy_operator(op::AbstractOperator; array_type = nothing, threaded = nothing) + if is_thread_safe(op) && threaded === nothing && array_type === nothing + return op # safe to share + end + return _copy_operator_impl(op; array_type, threaded) +end + +# Default implementation: just deepcopy (fallback) +function _copy_operator_impl( + op::T; array_type = nothing, threaded = nothing + ) where {T <: AbstractOperator} + return deepcopy(op) +end + +# Helper: convert a buffer array to the target storage type +function _convert_buffer(buf::AbstractArray, ::Nothing) + return similar(buf) # same type, new allocation +end +function _convert_buffer(buf::AbstractArray{T}, array_type::Type) where {T} + return similar(array_type{T}, size(buf)) +end + +_should_thread(::Number) = false +_should_thread(d::AbstractArray) = length(d) > 2^16 && Threads.nthreads() > 1 +_should_thread(::Type{<:AbstractArray}) = Threads.nthreads() > 1 +_should_thread(op::AbstractOperator) = _should_thread(domain_array_type(op)) diff --git a/src/utils.jl b/src/utils.jl index 3dafa0d..b980fbc 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -19,12 +19,14 @@ end for pair in thread_count_functions[] pair.second(n) end + return end @noinline function _restore_thread_counts(prev::Vector) for (i, pair) in enumerate(thread_count_functions[]) pair.second(prev[i]) end + return end function set_thread_counts_expr(thread_count_expr, body_expr) @@ -57,13 +59,51 @@ macro restrict_threading(expr) return set_thread_counts_expr(1, expr) end -function check(codomain_array, op, domain_array) - if domain_array isa AbstractArray === false - throw(ArgumentError("Input must be an AbstractArray")) +_storage_parent(a) = a +_storage_parent(a::SubArray) = _storage_parent(parent(a)) +_storage_parent(a::Base.ReshapedArray) = _storage_parent(parent(a)) + +_is_storage_compatible(a, ::Type{<:Array}) = _storage_parent(a) isa Array +_is_storage_compatible(a, ::Type{T}) where {T} = _storage_parent(a) isa T + +@generated function _is_storage_compatible(a::ArrayPartition, ::Type{T}) where {T <: ArrayPartition} + Tp = T.parameters[2] + if !(Tp isa DataType && Tp <: Tuple) + return :(a isa T) end - if codomain_array isa AbstractArray === false - throw(ArgumentError("Output must be an AbstractArray")) + types = Tp.parameters + checks = [:(a.x[$i] isa $(types[i])) for i in 1:length(types)] + cond = isempty(checks) ? :(true) : reduce((x, y) -> :($x && $y), checks) + return :(length(a.x) == $(length(types)) && $cond) +end + +function _check_domain_storage(domain_array, op) + if !_is_storage_compatible(domain_array, domain_array_type(op)) + throw( + ArgumentError( + "Input storage type $(typeof(domain_array)) is not compatible with " * + "operator's expected domain storage $(domain_array_type(op))", + ), + ) end + return +end + +function _check_codomain_storage(codomain_array, op) + if !_is_storage_compatible(codomain_array, codomain_array_type(op)) + throw( + ArgumentError( + "Output storage type $(typeof(codomain_array)) is not compatible with " * + "operator's expected codomain storage $(codomain_array_type(op))", + ), + ) + end + return +end + +function check(codomain_array, op, domain_array) + _check_domain_storage(domain_array, op) + _check_codomain_storage(codomain_array, op) if (ndoms(op, 2) > 1) != (domain_array isa ArrayPartition) throw(ArgumentError("Input must be an ArrayPartition if and only if operator has multiple input domains")) end @@ -89,7 +129,11 @@ function check(codomain_array, op, domain_array) ) end if (ndoms(op, 1) > 1) != (codomain_array isa ArrayPartition) - throw(ArgumentError("Output must be an ArrayPartition if and only if operator has multiple output domains")) + throw( + ArgumentError( + "Output must be an ArrayPartition if and only if operator has multiple output domains", + ), + ) end if codomain_array isa ArrayPartition dtype = eltype.(codomain_array.x) @@ -104,11 +148,12 @@ function check(codomain_array, op, domain_array) ) end dim_out = codomain_array isa ArrayPartition ? size.(codomain_array.x) : size(codomain_array) - return if !isequal(dim_out, size(op, 1)) + if !isequal(dim_out, size(op, 1)) throw( DimensionMismatch( "Output size $(dim_out) does not match operator output size $(size(op, 1))", ), ) end + return nothing end diff --git a/test/Project.toml b/test/Project.toml index c15a0d6..be10c12 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -7,6 +7,7 @@ DSPOperators = "d5a72628-6e2f-430e-82f5-561df0bb8116" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" FFTWOperators = "c59a084b-ba08-4f3f-af9e-f4298d6caa94" +GPUEnv = "78a0b619-6146-4252-b244-0f81c54be577" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -30,7 +31,6 @@ Aqua = "0.8" BenchmarkTools = "1" Documenter = "1.8" JET = "0.8, 0.9, 0.10, 0.11" -JLArrays = "0.2.0" LinearAlgebra = "1" LinearMaps = "3.11.4" Pkg = "1.10.0" diff --git a/test/batching/test_SimpleBatchOp.jl b/test/batching/test_SimpleBatchOp.jl index 3fd8314..d50d69d 100644 --- a/test/batching/test_SimpleBatchOp.jl +++ b/test/batching/test_SimpleBatchOp.jl @@ -1,6 +1,5 @@ -@testitem "SimpleBatchOp" tags = [:batching, :SimpleBatchOp] setup = [TestUtils] begin - using Random, BenchmarkTools, AbstractOperators - Random.seed!(0) +@testmodule SimpleBatchOpHelpers begin + using Random, BenchmarkTools, LinearAlgebra, AbstractOperators, Test function test_simple_batchop(op, batch_op, x, y, z, threaded) if threaded && Threads.nthreads() > 1 @@ -70,17 +69,13 @@ function other_tests(threaded) op = DiagOp([1.0, 2.0]) batch_op = BatchOp(op, (2,); threaded) - # show (fun_name) io = IOBuffer(); show(io, batch_op); s = String(take!(io)) @test occursin("⟳", s) - # == and isequal - batch_op_copy = AbstractOperators.copy_op(batch_op) + batch_op_copy = AbstractOperators.copy_operator(batch_op) @test batch_op == batch_op_copy @test isequal(batch_op, batch_op_copy) - # storage types - @test domain_storage_type(batch_op) == domain_storage_type(op) - @test codomain_storage_type(batch_op) == codomain_storage_type(op) - # property queries + @test domain_array_type(batch_op) == domain_array_type(op) + @test codomain_array_type(batch_op) == codomain_array_type(op) @test is_linear(batch_op) == is_linear(op) @test is_eye(batch_op) == is_eye(op) @test is_null(batch_op) == is_null(op) @@ -92,89 +87,114 @@ @test is_full_column_rank(batch_op) == is_full_column_rank(op) @test is_sliced(batch_op) == is_sliced(op) @test is_thread_safe(batch_op) == is_thread_safe(op) - # has_optimized_normalop, get_normal_op @test AbstractOperators.has_optimized_normalop(batch_op) == AbstractOperators.has_optimized_normalop(op) n_op = AbstractOperators.get_normal_op(batch_op) @test typeof(n_op) <: typeof(batch_op) - # opnorm, estimate_opnorm -- exact solution is expected for both @test AbstractOperators.has_fast_opnorm(batch_op) == AbstractOperators.has_fast_opnorm(op) @test opnorm(batch_op) == opnorm(op) @test estimate_opnorm(batch_op) == estimate_opnorm(op) @test estimate_opnorm(batch_op) == opnorm(batch_op) - # diag, diag_AcA, diag_AAc @test diag(batch_op) == [diag(op)'; diag(op)']' @test diag_AcA(batch_op) == [diag_AcA(op)'; diag_AcA(op)']' @test diag_AAc(batch_op) == [diag_AAc(op)'; diag_AAc(op)']' - # error path: mismatched input type x_bad = rand(Int, 2, 2) y_bad = zeros(2, 2) @test_throws ArgumentError mul!(y_bad, batch_op, x_bad) - # error path: mismatched input size x_bad2 = rand(2, 3) @test_throws DimensionMismatch mul!(y_bad, batch_op, x_bad2) - # error path: mismatched output type y_bad2 = rand(Int, 2, 2) x_good = rand(2, 2) @test_throws ArgumentError mul!(y_bad2, batch_op, x_good) - # error path: mismatched output size y_bad3 = zeros(3, 2) @test_throws DimensionMismatch mul!(y_bad3, batch_op, x_good) - - # scalar-diag path (Eye reports scalar diagonal) eye_batch = BatchOp(Eye(Float64, (2,)), (2,); threaded) @test diag(eye_batch) == 1.0 @test diag_AcA(eye_batch) == 1.0 @test diag_AAc(eye_batch) == 1.0 return end +end - @testset "SimpleBatchOp" begin - @testset "Shape-keeping op (DiagOp)" begin - @testset "non-threaded" begin - test_shape_keeping_simple_batch_op(false) - end - if Threads.nthreads() > 1 - @testset "threaded (thread-safe)" begin - test_shape_keeping_simple_batch_op(true) - end - end - end - @testset "Dimension count changing op (Variation)" begin - @testset "non-threaded" begin - test_variation_simple_batch_op(false) - end - if Threads.nthreads() > 1 - @testset "threaded (thread-safe)" begin - test_variation_simple_batch_op(true) - end - end - end - @testset "Shape-changing op (DiagOp∘FiniteDiff)" begin - @testset "non-threaded" begin - test_shape_changing_simple_batch_op(false) - end - if Threads.nthreads() > 1 - @testset "threaded (thread-safe)" begin - test_shape_changing_simple_batch_op(true) - end - end - end - @testset "Other tests" begin - @testset "non-threaded" begin - other_tests(false) - end - if Threads.nthreads() > 1 - @testset "threaded (thread-safe)" begin - other_tests(true) - end - end - end - if Threads.nthreads() > 1 && get(ENV, "CI", "false") == "false" - @testset "Benchmark" begin - t_single_threaded = benchmark_threading(false) - t_multi_threaded = benchmark_threading(true) - @test t_multi_threaded < t_single_threaded - end - end +@testitem "SimpleBatchOp shape-keeping non-threaded" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + SimpleBatchOpHelpers.test_shape_keeping_simple_batch_op(false) +end + +@testitem "SimpleBatchOp shape-keeping threaded" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + if Threads.nthreads() > 1 + SimpleBatchOpHelpers.test_shape_keeping_simple_batch_op(true) + end +end + +@testitem "SimpleBatchOp variation non-threaded" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + SimpleBatchOpHelpers.test_variation_simple_batch_op(false) +end + +@testitem "SimpleBatchOp variation threaded" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + if Threads.nthreads() > 1 + SimpleBatchOpHelpers.test_variation_simple_batch_op(true) + end +end + +@testitem "SimpleBatchOp shape-changing non-threaded" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + SimpleBatchOpHelpers.test_shape_changing_simple_batch_op(false) +end + +@testitem "SimpleBatchOp shape-changing threaded" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + if Threads.nthreads() > 1 + SimpleBatchOpHelpers.test_shape_changing_simple_batch_op(true) + end +end + +@testitem "SimpleBatchOp other tests non-threaded" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + SimpleBatchOpHelpers.other_tests(false) +end + +@testitem "SimpleBatchOp other tests threaded" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + if Threads.nthreads() > 1 + SimpleBatchOpHelpers.other_tests(true) + end +end + +@testitem "SimpleBatchOp benchmark" tags = [:batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random + Random.seed!(0) + if Threads.nthreads() > 1 && get(ENV, "CI", "false") == "false" + t_single_threaded = SimpleBatchOpHelpers.benchmark_threading(false) + t_multi_threaded = SimpleBatchOpHelpers.benchmark_threading(true) + @test t_multi_threaded < t_single_threaded + end +end + +@testitem "SimpleBatchOp (GPU)" tags = [:gpu, :batching, :SimpleBatchOp] setup = [TestUtils, SimpleBatchOpHelpers] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + op = DiagOp(to_gpu(backend, [1.0, 2.0])) + batch_op = BatchOp(op, (3, 4), (:_, :b, :b)) + x = gpu_ones(backend, Float64, 2, 3, 4) + y_gpu = batch_op * x + @test size(Array(y_gpu)) == (2, 3, 4) + @test all(Array(y_gpu)[1, :, :] .≈ 1.0) + @test all(Array(y_gpu)[2, :, :] .≈ 2.0) + y_gpu2 = similar(y_gpu) + mul!(y_gpu2, batch_op, x) + @test Array(y_gpu2) ≈ Array(y_gpu) end -end # @testitem "SimpleBatchOp" +end diff --git a/test/batching/test_SpreadingBatchOp.jl b/test/batching/test_SpreadingBatchOp.jl index d00d85f..02bfc7e 100644 --- a/test/batching/test_SpreadingBatchOp.jl +++ b/test/batching/test_SpreadingBatchOp.jl @@ -1,6 +1,5 @@ -@testitem "SpreadingBatchOp" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils] begin - using Random, BenchmarkTools, AbstractOperators - Random.seed!(0) +@testmodule SpreadingBatchOpHelpers begin + using Random, BenchmarkTools, LinearAlgebra, AbstractOperators, Test function test_spreading_batchop(operators, batch_op, x, y, z, threaded) if !threaded @@ -66,7 +65,7 @@ num_ops = Threads.nthreads() + 5 op = GetIndex(Float64, (m - 1,), 1:6) * FiniteDiff((m,)) ops = [reshape(i * op, 2, 3) for i in 1:num_ops] - return @test_throws ArgumentError BatchOp(ops, n, (:b, :s, :_) => (:b, :s, :_, :_); threaded = true, threading_strategy = ThreadingStrategy.FIXED_OPERATOR) + return @test_throws ArgumentError BatchOp(ops, n, (:b, :s, :_) => (:b, :s, :_, :_); threaded = true, threading_strategy = AbstractOperators.ThreadingStrategy.FIXED_OPERATOR) end function benchmark_threading_strategy(threaded, threading_strategy) @@ -81,16 +80,12 @@ function other_spreadingbatchop_tests(threaded) ops = [DiagOp([1.0, 2.0]) for _ in 1:3] bop = BatchOp(ops, 4; threaded = threaded) - # show output (fun_name indirectly) io = IOBuffer(); show(io, bop); s = String(take!(io)) @test occursin("⟳", s) - # size consistency cod, dom = size(bop) @test cod == size(bop, 1) && dom == size(bop, 2) - # storage types - @test domain_storage_type(bop) == domain_storage_type(ops[1]) - @test codomain_storage_type(bop) == codomain_storage_type(ops[1]) - # property queries mirror first operator + @test domain_array_type(bop) == domain_array_type(ops[1]) + @test codomain_array_type(bop) == codomain_array_type(ops[1]) @test is_linear(bop) == is_linear(ops[1]) @test is_eye(bop) == is_eye(ops[1]) @test is_null(bop) == is_null(ops[1]) @@ -102,19 +97,15 @@ @test is_full_column_rank(bop) == is_full_column_rank(ops[1]) @test is_sliced(bop) == is_sliced(ops[1]) @test is_thread_safe(bop) == is_thread_safe(ops[1]) - # normal op @test AbstractOperators.has_optimized_normalop(bop) == AbstractOperators.has_optimized_normalop(ops[1]) nbop = AbstractOperators.get_normal_op(bop) @test size(nbop, 1) == size(bop, 1) && size(nbop, 2) == size(bop, 2) - # opnorm / estimate_opnorm aggregate as maximum -- exact solution is expected for both @test opnorm(bop) == maximum(opnorm.(ops)) @test estimate_opnorm(bop) == maximum(estimate_opnorm.(ops)) @test estimate_opnorm(bop) == opnorm(bop) - # diag family collapses to first operator entries (identical) @test diag(bop) == repeat(diag(ops[1]), outer = (1, 3, 4)) @test diag_AcA(bop) == repeat(diag_AcA(ops[1]), outer = (1, 3, 4)) @test diag_AAc(bop) == repeat(diag_AAc(ops[1]), outer = (1, 3, 4)) - # Functional sanity: apply and compare with manual broadcast x = rand(2, 3, 4) y1 = bop * x y2 = similar(x) @@ -122,205 +113,209 @@ mul!(@view(y2[:, i, j]), ops[i], @view(x[:, i, j])) end @test y1 == y2 - # Error paths (type / size) x_bad_type = rand(Int, 2, 3, 4) y = zeros(2, 3, 4) @test_throws ArgumentError mul!(y, bop, x_bad_type) - x_bad_size = rand(2, 3, 5) # wrong last dim + x_bad_size = rand(2, 3, 5) @test_throws DimensionMismatch mul!(y, bop, x_bad_size) y_bad_type = rand(Int, 2, 3, 4) @test_throws ArgumentError mul!(y_bad_type, bop, x) y_bad_size = zeros(3, 3, 4) @test_throws DimensionMismatch mul!(y_bad_size, bop, x) - # Constructor assertion error: mismatched operator shapes (different size) bad_ops = [DiagOp([1.0, 2.0]), DiagOp([1.0, 2.0, 3.0]), DiagOp([1.0, 2.0])] return @test_throws AssertionError BatchOp(bad_ops, 4) end +end - @testset "SpreadingBatchOp" begin - @testset "Shape-keeping op (DiagOp)" begin - @testset "non-threaded" begin - test_shape_keeping_threadsafe_spreading_batch_op(false) - end - if Threads.nthreads() > 1 - @testset "threaded (thread-safe)" begin - test_shape_keeping_threadsafe_spreading_batch_op(true) - end - end - end - @testset "Shape-changing op (Variation)" begin - @testset "non-threaded" begin - test_shape_changing_threadsafe_spreading_batch_op(false) - end - if Threads.nthreads() > 1 - @testset "threaded (thread-safe)" begin - test_shape_changing_threadsafe_spreading_batch_op(true) - end - end - end - @testset "Non-threadsafe op (Compose)" begin - @testset "non-threaded" begin - test_nonthreadsafe_spreading_batch_op(false, ThreadingStrategy.AUTO) - end - if Threads.nthreads() > 1 - @testset "threaded (copying)" begin - test_nonthreadsafe_spreading_batch_op(true, ThreadingStrategy.COPYING) - end - @testset "threaded (locking)" begin - test_nonthreadsafe_spreading_batch_op(true, ThreadingStrategy.LOCKING) - end - @testset "threaded (fixed operator)" begin - test_nonthreadsafe_spreading_batch_op(true, ThreadingStrategy.FIXED_OPERATOR) - test_failing_nonthreadsafe_spreading_batch_op() - end - end - end - @testset "Other tests" begin - @testset "non-threaded" begin - other_spreadingbatchop_tests(false) - end - if Threads.nthreads() > 1 - @testset "threaded (thread-safe)" begin - other_spreadingbatchop_tests(true) - end - end - end - if Threads.nthreads() >= 4 && get(ENV, "CI", "false") == "false" - @testset "Benchmark" begin - t_single_threaded = benchmark_threading_strategy(false, ThreadingStrategy.AUTO) - t_copying = benchmark_threading_strategy(true, ThreadingStrategy.COPYING) - t_fixed_operator = benchmark_threading_strategy(true, ThreadingStrategy.FIXED_OPERATOR) - @test t_copying < t_single_threaded - @test t_fixed_operator < t_single_threaded - end - end - - @testset "SpreadingBatchOpCopying property delegations" begin - if Threads.nthreads() > 1 - # Create non-threadsafe operators to trigger COPYING strategy - ops = [DiagOp(rand(5)) * FiniteDiff((6,)) for i in 1:3] - bop = BatchOp(ops, 4, (:_, :s, :b); threaded = true, threading_strategy = ThreadingStrategy.COPYING) - - # fun_name for Copying variant - io = IOBuffer(); show(io, bop); s = String(take!(io)) - @test occursin("⟳", s) +@testitem "SpreadingBatchOp shape-keeping non-threaded" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + SpreadingBatchOpHelpers.test_shape_keeping_threadsafe_spreading_batch_op(false) +end - # Storage types for Copying - @test domain_storage_type(bop) == domain_storage_type(ops[1]) - @test codomain_storage_type(bop) == codomain_storage_type(ops[1]) +@testitem "SpreadingBatchOp shape-keeping threaded" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + SpreadingBatchOpHelpers.test_shape_keeping_threadsafe_spreading_batch_op(true) + end +end - # Properties for Copying - @test is_linear(bop) == is_linear(ops[1]) - @test is_eye(bop) == is_eye(ops[1]) - @test is_AAc_diagonal(bop) == is_AAc_diagonal(ops[1]) - @test is_AcA_diagonal(bop) == is_AcA_diagonal(ops[1]) - @test is_full_row_rank(bop) == is_full_row_rank(ops[1]) - @test is_full_column_rank(bop) == is_full_column_rank(ops[1]) - @test is_sliced(bop) == is_sliced(ops[1]) - @test is_null(bop) == is_null(ops[1]) - @test is_diagonal(bop) == is_diagonal(ops[1]) - @test is_invertible(bop) == is_invertible(ops[1]) - @test is_orthogonal(bop) == is_orthogonal(ops[1]) - @test is_thread_safe(bop) == is_thread_safe(ops[1]) +@testitem "SpreadingBatchOp variation non-threaded" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + SpreadingBatchOpHelpers.test_shape_changing_threadsafe_spreading_batch_op(false) +end - # Normal op for Copying - requires optimized normal op - @test AbstractOperators.has_optimized_normalop(bop) == AbstractOperators.has_optimized_normalop(ops[1]) +@testitem "SpreadingBatchOp variation threaded" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + SpreadingBatchOpHelpers.test_shape_changing_threadsafe_spreading_batch_op(true) + end +end - # opnorm methods for Copying (use approximate equality for floating point) - @test AbstractOperators.has_fast_opnorm(bop) == AbstractOperators.has_fast_opnorm(ops[1]) - operator_norm = opnorm(bop) - @test operator_norm ≈ maximum(opnorm.(ops)) rtol = 5.0e-6 - @test estimate_opnorm(bop) ≈ operator_norm rtol = 0.05 +@testitem "SpreadingBatchOp non-threadsafe auto" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + SpreadingBatchOpHelpers.test_nonthreadsafe_spreading_batch_op(false, AbstractOperators.ThreadingStrategy.AUTO) +end - # diag methods for Copying - use operators with diag - ops2 = [DiagOp(rand(5)) for i in 1:3] - bop2 = BatchOp(ops2, 4, (:_, :s, :b); threaded = true, threading_strategy = ThreadingStrategy.COPYING) - d = diag(bop2) - @test size(d) == (5, 3, 4) - daca = diag_AcA(bop2) - @test size(daca) == (5, 3, 4) - daac = diag_AAc(bop2) - @test size(daac) == (5, 3, 4) - end - end +@testitem "SpreadingBatchOp threaded copying" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + SpreadingBatchOpHelpers.test_nonthreadsafe_spreading_batch_op(true, AbstractOperators.ThreadingStrategy.COPYING) + end +end - @testset "Locking get_normal_op and reused operators" begin - if Threads.nthreads() > 1 - # Create scenario with reused operator to trigger haskey branch - op = DiagOp(rand(6)) * FiniteDiff((7,)) - ops = [op, op, DiagOp(rand(6)) * FiniteDiff((7,))] # First two are same instance - bop = BatchOp(ops, 4, (:_, :s, :b); threaded = true, threading_strategy = ThreadingStrategy.LOCKING) +@testitem "SpreadingBatchOp threaded locking" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + SpreadingBatchOpHelpers.test_nonthreadsafe_spreading_batch_op(true, AbstractOperators.ThreadingStrategy.LOCKING) + end +end - # Verify it works - x = rand(7, 3, 4) - y = bop * x - @test size(y) == (6, 3, 4) - end - end +@testitem "SpreadingBatchOp threaded fixed operator" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + SpreadingBatchOpHelpers.test_nonthreadsafe_spreading_batch_op(true, AbstractOperators.ThreadingStrategy.FIXED_OPERATOR) + SpreadingBatchOpHelpers.test_failing_nonthreadsafe_spreading_batch_op() + end +end - @testset "FixedOperator get_normal_op and get_spreading_dims" begin - if Threads.nthreads() > 1 - ops = [DiagOp(rand(6)) * FiniteDiff((7,)) for i in 1:3] - bop = BatchOp(ops, 4, (:_, :s, :b); threaded = true, threading_strategy = ThreadingStrategy.FIXED_OPERATOR) +@testitem "SpreadingBatchOp other tests non-threaded" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + SpreadingBatchOpHelpers.other_spreadingbatchop_tests(false) +end - # Verify get_spreading_dims is called (indirectly via operations) - x = rand(7, 3, 4) - y = bop * x - @test size(y) == (6, 3, 4) - end - end +@testitem "SpreadingBatchOp other tests threaded" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + SpreadingBatchOpHelpers.other_spreadingbatchop_tests(true) + end +end - @testset "Orthogonal property for SpreadingBatchOp" begin - # Use Identity operators which are orthogonal - ops = [Eye(Float64, 5) for i in 1:3] - bop = BatchOp(ops, 4; threaded = false) - @test is_orthogonal(bop) == true +@testitem "SpreadingBatchOp benchmark" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() >= 4 && get(ENV, "CI", "false") == "false" + t_single_threaded = SpreadingBatchOpHelpers.benchmark_threading_strategy(false, AbstractOperators.ThreadingStrategy.AUTO) + t_copying = SpreadingBatchOpHelpers.benchmark_threading_strategy(true, AbstractOperators.ThreadingStrategy.COPYING) + t_fixed_operator = SpreadingBatchOpHelpers.benchmark_threading_strategy(true, AbstractOperators.ThreadingStrategy.FIXED_OPERATOR) + @test t_copying < t_single_threaded + @test t_fixed_operator < t_single_threaded + end +end - if Threads.nthreads() > 1 - # Test with threaded version - bop_threaded = BatchOp(ops, 4; threaded = true) - @test is_orthogonal(bop_threaded) == true +@testitem "SpreadingBatchOpCopying property delegations" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + ops = [DiagOp(rand(5)) * FiniteDiff((6,)) for i in 1:3] + bop = BatchOp(ops, 4, (:_, :s, :b); threaded = true, threading_strategy = AbstractOperators.ThreadingStrategy.COPYING) + io = IOBuffer(); show(io, bop); s = String(take!(io)); @test occursin("⟳", s) + @test domain_array_type(bop) == domain_array_type(ops[1]) + @test codomain_array_type(bop) == codomain_array_type(ops[1]) + @test is_linear(bop) == is_linear(ops[1]) + @test is_eye(bop) == is_eye(ops[1]) + @test is_AAc_diagonal(bop) == is_AAc_diagonal(ops[1]) + @test is_AcA_diagonal(bop) == is_AcA_diagonal(ops[1]) + @test is_full_row_rank(bop) == is_full_row_rank(ops[1]) + @test is_full_column_rank(bop) == is_full_column_rank(ops[1]) + @test is_sliced(bop) == is_sliced(ops[1]) + @test is_null(bop) == is_null(ops[1]) + @test is_diagonal(bop) == is_diagonal(ops[1]) + @test is_invertible(bop) == is_invertible(ops[1]) + @test is_orthogonal(bop) == is_orthogonal(ops[1]) + @test is_thread_safe(bop) == is_thread_safe(ops[1]) + @test AbstractOperators.has_optimized_normalop(bop) == AbstractOperators.has_optimized_normalop(ops[1]) + @test AbstractOperators.has_fast_opnorm(bop) == AbstractOperators.has_fast_opnorm(ops[1]) + operator_norm = opnorm(bop) + @test operator_norm ≈ maximum(opnorm.(ops)) rtol = 5.0e-6 + @test estimate_opnorm(bop) ≈ operator_norm rtol = 0.05 + ops2 = [DiagOp(rand(5)) for i in 1:3] + bop2 = BatchOp(ops2, 4, (:_, :s, :b); threaded = true, threading_strategy = AbstractOperators.ThreadingStrategy.COPYING) + @test size(diag(bop2)) == (5, 3, 4) + @test size(diag_AcA(bop2)) == (5, 3, 4) + @test size(diag_AAc(bop2)) == (5, 3, 4) + end +end - # Test with Copying variant - need non-threadsafe orthogonal operator - # Use Compose with Eye operators (still orthogonal but might not be thread-safe depending on implementation) - ops2 = [Eye(Float64, 5) for i in 1:3] - bop2 = BatchOp(ops2, 4; threaded = true) - @test is_orthogonal(bop2) == is_orthogonal(ops2[1]) - end - end +@testitem "Locking get_normal_op and reused operators" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + op = DiagOp(rand(6)) * FiniteDiff((7,)) + ops = [op, op, DiagOp(rand(6)) * FiniteDiff((7,))] + bop = BatchOp(ops, 4, (:_, :s, :b); threaded = true, threading_strategy = AbstractOperators.ThreadingStrategy.LOCKING) + y = bop * rand(7, 3, 4) + @test size(y) == (6, 3, 4) + end +end - @testset "AUTO threading strategy triggering" begin - if Threads.nthreads() > 1 - # Create scenario where AUTO should pick a strategy - # Small operators with FiniteDiff (non-threadsafe) - ops = [FiniteDiff((11,)) for i in 1:3] - bop = BatchOp(ops, 4, (:_, :s, :b); threaded = true, threading_strategy = ThreadingStrategy.AUTO) +@testitem "FixedOperator get_normal_op and get_spreading_dims" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + ops = [DiagOp(rand(6)) * FiniteDiff((7,)) for i in 1:3] + bop = BatchOp(ops, 4, (:_, :s, :b); threaded = true, threading_strategy = AbstractOperators.ThreadingStrategy.FIXED_OPERATOR) + y = bop * rand(7, 3, 4) + @test size(y) == (6, 3, 4) + end +end - # Should work regardless of chosen strategy - x = rand(11, 3, 4) - y = bop * x - @test size(y) == (10, 3, 4) - end - end +@testitem "Orthogonal property for SpreadingBatchOp" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + ops = [Eye(Float64, 5) for i in 1:3] + @test is_orthogonal(BatchOp(ops, 4; threaded = false)) == true + if Threads.nthreads() > 1 + bop_threaded = BatchOp(ops, 4; threaded = true) + @test is_orthogonal(bop_threaded) == true + bop2 = BatchOp([Eye(Float64, 5) for i in 1:3], 4; threaded = true) + @test is_orthogonal(bop2) == is_orthogonal(ops[1]) + end +end - @testset "Scalar diagonal return paths" begin - # Create operators where all have identical scalar diagonals - scale_val = 2.0 - ops = [scale_val * Eye(Float64, 5) for i in 1:3] - bop = BatchOp(ops, 4; threaded = false) +@testitem "AUTO threading strategy triggering" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + if Threads.nthreads() > 1 + ops = [FiniteDiff((11,)) for i in 1:3] + bop = BatchOp(ops, 4, (:_, :s, :b); threaded = true, threading_strategy = AbstractOperators.ThreadingStrategy.AUTO) + @test size(bop * rand(11, 3, 4)) == (10, 3, 4) + end +end - # These should return scalars when all operators have identical scalar diagonals - d = diag(bop) - @test d isa Number # Expect scalar - @test d == scale_val +@testitem "Scalar diagonal return paths" tags = [:batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators + Random.seed!(0) + scale_val = 2.0 + ops = [scale_val * Eye(Float64, 5) for i in 1:3] + bop = BatchOp(ops, 4; threaded = false) + @test diag(bop) isa Number + @test diag(bop) == scale_val + @test diag_AcA(bop) isa Number + @test diag_AcA(bop) == scale_val^2 + @test diag_AAc(bop) isa Number + @test diag_AAc(bop) == scale_val^2 +end - daca = diag_AcA(bop) - @test daca isa Number # Expect scalar - @test daca == scale_val^2 +@testitem "SpreadingBatchOp (GPU)" tags = [:gpu, :batching, :SpreadingBatchOp] setup = [TestUtils, SpreadingBatchOpHelpers] begin + using Random, AbstractOperators, GPUEnv - daac = diag_AAc(bop) - @test daac isa Number # Expect scalar - @test daac == scale_val^2 - end + for backend in gpu_backends() + Random.seed!(0) + ops = [DiagOp(to_gpu(backend, [1.0, 2.0])) for _ in 1:4] + bop = BatchOp(ops, 3, (:_, :s, :b)) + y_gpu = bop * gpu_ones(backend, Float64, 2, 4, 3) + @test size(Array(y_gpu)) == (2, 4, 3) + @test all(Array(y_gpu)[1, :, :] .≈ 1.0) + @test all(Array(y_gpu)[2, :, :] .≈ 2.0) end -end # @testitem "SpreadingBatchOp" +end diff --git a/test/calculus/test_Ax_mul_Bx.jl b/test/calculus/test_Ax_mul_Bx.jl index acafdc1..1b0e987 100644 --- a/test/calculus/test_Ax_mul_Bx.jl +++ b/test/calculus/test_Ax_mul_Bx.jl @@ -1,7 +1,6 @@ @testitem "Ax_mul_Bx" tags = [:calculus, :Ax_mul_Bx] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Ax_mul_Bx --- ") n = 3 A, B = Eye(n, n), Eye(n, n) P = Ax_mul_Bx(A, B) @@ -101,3 +100,26 @@ @test Ax_mul_Bx(A, B) == Ax_mul_Bx(A, B) @test Jacobian(Ax_mul_Bx(A, B), x) == Jacobian(Ax_mul_Bx(A, B), x) end + +@testitem "Ax_mul_Bx (GPU)" tags = [:gpu, :calculus, :Ax_mul_Bx] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n = 3 + P = Ax_mul_Bx( + Eye(Float64, (n, n); array_type = gpu_wrapper(backend, Float64, n, n)), + Eye(Float64, (n, n); array_type = gpu_wrapper(backend, Float64, n, n)), + ) + x = gpu_randn(backend, n, n) + r = gpu_randn(backend, n, n) + test_NLop_gpu(P, x, r, false) + + n2 = 3 + P2 = Ax_mul_Bx(Sin(gpu_zeros(backend, Float64, n2, n2)), Cos(gpu_zeros(backend, Float64, n2, n2))) + x2 = gpu_randn(backend, n2, n2) + r2 = gpu_randn(backend, n2, n2) + test_NLop_gpu(P2, x2, r2, false) + end +end diff --git a/test/calculus/test_Ax_mul_Bxt.jl b/test/calculus/test_Ax_mul_Bxt.jl index 09e8480..45fdfc9 100644 --- a/test/calculus/test_Ax_mul_Bxt.jl +++ b/test/calculus/test_Ax_mul_Bxt.jl @@ -1,7 +1,7 @@ -@testitem "Ax_mul_Bxt" tags = [:calculus, :Ax_mul_Bxt] setup = [TestUtils] begin +@testitem "Ax_mul_Bxt: basic mul" tags = [:calculus, :Ax_mul_Bxt] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Ax_mul_Bxt --- ") + n = 10 A, B = Eye(n), Sin(n) P = Ax_mul_Bxt(A, B) @@ -33,6 +33,11 @@ r = randn(n, n) y, grad = test_NLop(P, x, r, verb) @test norm((A * x) * (B * x)' - y) < 1.0e-8 +end + +@testitem "Ax_mul_Bxt: HCAT and permute" tags = [:calculus, :Ax_mul_Bxt] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) # testing with HCAT m, n = 3, 5 @@ -62,6 +67,11 @@ @test_throws Exception Ax_mul_Bxt(Eye(2, 2), Eye(2, 1)) @test_throws Exception Ax_mul_Bxt(Eye(2, 2, 2), Eye(2, 2, 2)) +end + +@testitem "Ax_mul_Bxt: error paths and equality" tags = [:calculus, :Ax_mul_Bxt] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) # ndims==2 branch with mismatched second codomain dimension struct AxDummy2D <: AbstractOperator @@ -71,8 +81,8 @@ Base.size(op::AxDummy2D) = (op.dim_out, op.dim_in) AbstractOperators.domain_type(::AxDummy2D) = Float64 AbstractOperators.codomain_type(::AxDummy2D) = Float64 - AbstractOperators.domain_storage_type(::AxDummy2D) = Array{Float64} - AbstractOperators.codomain_storage_type(::AxDummy2D) = Array{Float64} + AbstractOperators.domain_array_type(::AxDummy2D) = Array{Float64} + AbstractOperators.codomain_array_type(::AxDummy2D) = Array{Float64} struct AxDummyMixed <: AbstractOperator dim_out::Tuple{Int} @@ -81,8 +91,8 @@ Base.size(op::AxDummyMixed) = (op.dim_out, op.dim_in) AbstractOperators.domain_type(::AxDummyMixed) = Float64 AbstractOperators.codomain_type(::AxDummyMixed) = Float64 - AbstractOperators.domain_storage_type(::AxDummyMixed) = Array{Float64} - AbstractOperators.codomain_storage_type(::AxDummyMixed) = Array{Float64} + AbstractOperators.domain_array_type(::AxDummyMixed) = Array{Float64} + AbstractOperators.codomain_array_type(::AxDummyMixed) = Array{Float64} A2d = AxDummy2D((2, 3), (4, 3)) B2d_bad = AxDummy2D((2, 4), (4, 3)) @@ -93,6 +103,25 @@ # test equality n, m = 3, 4 A, B = MatrixOp(randn(n, m)), MatrixOp(randn(n, m)) + # x matches the HCAT ArrayPartition context used in original test + x = ArrayPartition(randn(3), randn(5)) @test Ax_mul_Bxt(A, B) == Ax_mul_Bxt(A, B) @test Jacobian(Ax_mul_Bxt(A, B), x) == Jacobian(Ax_mul_Bxt(A, B), x) end + +@testitem "Ax_mul_Bxt (GPU)" tags = [:gpu, :calculus, :Ax_mul_Bxt] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n = 10 + P = Ax_mul_Bxt( + Eye(Float64, (n,); array_type = gpu_wrapper(backend, Float64, n)), + Sin(gpu_zeros(backend, Float64, n)), + ) + x = gpu_randn(backend, n) + r = gpu_randn(backend, n, n) + test_NLop_gpu(P, x, r, false) + end +end diff --git a/test/calculus/test_Axt_mul_Bx.jl b/test/calculus/test_Axt_mul_Bx.jl index 7e1731a..1546321 100644 --- a/test/calculus/test_Axt_mul_Bx.jl +++ b/test/calculus/test_Axt_mul_Bx.jl @@ -1,7 +1,7 @@ -@testitem "Axt_mul_Bx" tags = [:calculus, :Axt_mul_Bx] setup = [TestUtils] begin +@testitem "Axt_mul_Bx: basic mul" tags = [:calculus, :Axt_mul_Bx] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Axt_mul_Bx --- ") + n = 10 A, B = Eye(n), Sin(n) P = Axt_mul_Bx(A, B) @@ -35,6 +35,11 @@ r = randn(m, m) y, grad = test_NLop(P, x, r, verb) @test norm((A * x)' * (B * x) - y) < 1.0e-8 +end + +@testitem "Axt_mul_Bx: HCAT and permute" tags = [:calculus, :Axt_mul_Bx] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) # testing with HCAT m, n = 3, 5 @@ -64,6 +69,11 @@ @test_throws Exception Axt_mul_Bx(Eye(2, 2), Eye(2, 1)) @test_throws Exception Axt_mul_Bx(Eye(2, 2, 2), Eye(2, 2, 2)) +end + +@testitem "Axt_mul_Bx: error paths and equality" tags = [:calculus, :Axt_mul_Bx] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) # ndims==2 branch with mismatched first codomain dimension struct AxtDummy2D <: AbstractOperator @@ -73,8 +83,8 @@ Base.size(op::AxtDummy2D) = (op.dim_out, op.dim_in) AbstractOperators.domain_type(::AxtDummy2D) = Float64 AbstractOperators.codomain_type(::AxtDummy2D) = Float64 - AbstractOperators.domain_storage_type(::AxtDummy2D) = Array{Float64} - AbstractOperators.codomain_storage_type(::AxtDummy2D) = Array{Float64} + AbstractOperators.domain_array_type(::AxtDummy2D) = Array{Float64} + AbstractOperators.codomain_array_type(::AxtDummy2D) = Array{Float64} struct AxtDummyMixed <: AbstractOperator dim_out::Tuple{Int} @@ -83,8 +93,8 @@ Base.size(op::AxtDummyMixed) = (op.dim_out, op.dim_in) AbstractOperators.domain_type(::AxtDummyMixed) = Float64 AbstractOperators.codomain_type(::AxtDummyMixed) = Float64 - AbstractOperators.domain_storage_type(::AxtDummyMixed) = Array{Float64} - AbstractOperators.codomain_storage_type(::AxtDummyMixed) = Array{Float64} + AbstractOperators.domain_array_type(::AxtDummyMixed) = Array{Float64} + AbstractOperators.codomain_array_type(::AxtDummyMixed) = Array{Float64} A2d = AxtDummy2D((2, 3), (4, 3)) B2d_bad = AxtDummy2D((3, 3), (4, 3)) @@ -95,6 +105,37 @@ # test equality n, m = 3, 4 A, B = MatrixOp(randn(n, m)), MatrixOp(randn(n, m)) + # x matches the HCAT ArrayPartition context used in original test + x = ArrayPartition(randn(3), randn(5)) @test Axt_mul_Bx(A, B) == Axt_mul_Bx(A, B) @test Jacobian(Axt_mul_Bx(A, B), x) == Jacobian(Axt_mul_Bx(A, B), x) end + +@testitem "Axt_mul_Bx (GPU)" tags = [:gpu, :calculus, :Axt_mul_Bx] setup = [TestUtils] begin + using Random, AbstractOperators, LinearAlgebra, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n, m = 3, 4 + AL = gpu_randn(backend, n, m) + BL = gpu_randn(backend, n, m) + A, B = MatrixOp(AL), MatrixOp(BL) + P = Axt_mul_Bx(A, B) + x = gpu_randn(backend, m) + y = gpu_zeros(backend, Float64, 1) + mul!(y, P, x) + Acpu, Bcpu, xcpu = Array(AL), Array(BL), Array(x) + @test Array(y)[1] ≈ dot(Acpu * xcpu, Bcpu * xcpu) + + n, m, l = 3, 5, 4 + A2, B2 = MatrixOp(gpu_randn(backend, n, m), l), MatrixOp(gpu_randn(backend, n, m), l) + P2 = Axt_mul_Bx(A2, B2) + x2 = gpu_randn(backend, m, l) + y2 = gpu_zeros(backend, Float64, l, l) + mul!(y2, P2, x2) + Ax = Array(A2.A) * Array(x2) + Bx = Array(B2.A) * Array(x2) + @test Array(y2) ≈ Ax' * Bx + end +end diff --git a/test/calculus/test_Jacobian.jl b/test/calculus/test_Jacobian.jl index e77b50d..a357302 100644 --- a/test/calculus/test_Jacobian.jl +++ b/test/calculus/test_Jacobian.jl @@ -1,6 +1,5 @@ -@testitem "Jacobian" tags = [:calculus, :Jacobian] setup = [TestUtils] begin +@testitem "Jacobian: basic HCAT" tags = [:calculus, :Jacobian] setup = [TestUtils] begin using AbstractOperators - verb && println(" --- Testing Jacobian --- ") m, n = 3, 5 x = ArrayPartition(randn(m), randn(n)) @@ -16,8 +15,10 @@ J = Jacobian(opP, xp)' verb && println(size(J, 1)) y, grad = test_NLop(opP, xp, r, verb) +end - # Additional coverage-oriented tests appended (no new testset created) +@testitem "Jacobian: LinearOperator and Scale paths" tags = [:calculus, :Jacobian] setup = [TestUtils] begin + using AbstractOperators # 1. LinearOperator path (Jacobian of a LinearOperator returns itself) n_lin = 6 @@ -35,6 +36,10 @@ @test JS isa Scale @test JS.coeff == 2.5 @test JS.A == Jacobian(base_op, x_sc) +end + +@testitem "Jacobian: AffineAdd and Transpose paths" tags = [:calculus, :Jacobian] setup = [TestUtils] begin + using AbstractOperators # 3. AffineAdd path (Jacobian should drop displacement) n_aff = 4 @@ -49,6 +54,10 @@ n_tr = 5 op_tr = Sigmoid(Float64, (n_tr,), 2) @test_throws ErrorException op_tr' +end + +@testitem "Jacobian: Compose Sum VCAT paths" tags = [:calculus, :Jacobian] setup = [TestUtils] begin + using AbstractOperators # 5. Compose path (single op) with tuple input to trigger second Compose method n_cp = 3 @@ -80,6 +89,10 @@ JV = Jacobian(Vop, x_v) @test JV.A[1] == Jacobian(op_v1, x_v) @test JV.A[2] == Jacobian(op_v2, x_v) +end + +@testitem "Jacobian: HCAT and DCAT paths" tags = [:calculus, :Jacobian] setup = [TestUtils] begin + using AbstractOperators # 8. HCAT path with multi-index block (length(idx) > 1) to cover else branch m_h1, m_h2 = 3, 2 @@ -104,6 +117,10 @@ JD = Jacobian(D, xh) @test JD.A[1] == Jacobian(op_joint, (xh.x[1], xh.x[2])) @test JD.A[2] == Jacobian(Ah1, xh.x[3]) +end + +@testitem "Jacobian: Reshape and BroadCast paths" tags = [:calculus, :Jacobian] setup = [TestUtils] begin + using AbstractOperators # 10. Reshape path n_rs = 6 @@ -121,6 +138,10 @@ JBc = Jacobian(Bc, x_bc) # ensure repeating structure consistent @test size(JBc) == ((n_bc, l_bc), (n_bc,)) +end + +@testitem "Jacobian: equality and properties" tags = [:calculus, :Jacobian] setup = [TestUtils] begin + using AbstractOperators # 12. Equality and show output n_eq = 4 @@ -140,8 +161,8 @@ # 13. Type and thread safety properties @test domain_type(J1) == domain_type(op_eq) @test codomain_type(J1) == codomain_type(op_eq) - @test domain_storage_type(J1) == domain_storage_type(op_eq) - @test codomain_storage_type(J1) == codomain_storage_type(op_eq) + @test domain_array_type(J1) == domain_array_type(op_eq) + @test codomain_array_type(J1) == codomain_array_type(op_eq) @test !is_thread_safe(J1) # Direct DCAT + ArrayPartition Jacobian path @@ -157,3 +178,30 @@ @test JJ isa HCAT @test length(JJ.A) == 2 end + +@testitem "Jacobian (GPU)" tags = [:gpu, :calculus, :Jacobian] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n = 5 + A = Sin(gpu_zeros(backend, Float64, n)) + x = gpu_randn(backend, n) + J = Jacobian(A, x) + y = gpu_randn(backend, n) + grad = J' * y + grad2 = similar(grad) + mul!(grad2, J', y) + @test collect(grad) ≈ collect(grad2) + + A2 = Exp(gpu_zeros(backend, Float64, n)) + x2 = gpu_randn(backend, n) + J2 = Jacobian(A2, x2) + y2 = gpu_randn(backend, n) + grad3 = J2' * y2 + grad4 = similar(grad3) + mul!(grad4, J2', y2) + @test collect(grad3) ≈ collect(grad4) + end +end diff --git a/test/calculus/test_adjointoperator.jl b/test/calculus/test_adjointoperator.jl index dfb1077..1e6f64c 100644 --- a/test/calculus/test_adjointoperator.jl +++ b/test/calculus/test_adjointoperator.jl @@ -1,7 +1,6 @@ @testitem "AdjointOperator" tags = [:calculus, :AdjointOperator] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing AdjointOperator --- ") m, n = 5, 7 A1 = @@ -58,8 +57,8 @@ @test op2 == opA1 # storage & thread-safety propagation - _dst = domain_storage_type(opT) - _cst = codomain_storage_type(opT) + _dst = domain_array_type(opT) + _cst = codomain_array_type(opT) @test _dst !== nothing && _cst !== nothing @test is_thread_safe(opT) == is_thread_safe(opA1) @@ -89,3 +88,21 @@ @test opA1' == opA1' end + +@testitem "AdjointOperator (GPU)" tags = [:gpu, :calculus, :AdjointOperator] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n = 5 + op = FiniteDiff(gpu_zeros(backend, Float64, n)) + opT = AdjointOperator(op) + test_op(opT, gpu_randn(backend, n - 1), gpu_randn(backend, n), false) + + n = 4 + op = DiagOp(gpu_randn(backend, n)) + opT = AdjointOperator(op) + test_op(opT, gpu_randn(backend, n), gpu_randn(backend, n), false) + end +end diff --git a/test/calculus/test_affineadd.jl b/test/calculus/test_affineadd.jl index b2b113c..8ff4276 100644 --- a/test/calculus/test_affineadd.jl +++ b/test/calculus/test_affineadd.jl @@ -1,230 +1,149 @@ -@testitem "AffineAdd" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin +@testitem "AffineAdd: basic operations" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing AffineAdd --- ") n, m = 5, 6 A = randn(n, m) opA = MatrixOp(A) d = randn(n) T = AffineAdd(opA, d) - - verb && println(T) x1 = randn(m) - y1 = T * x1 - @test norm(y1 - (A * x1 + d)) < 1.0e-9 + @test norm(T * x1 - (A * x1 + d)) < 1.0e-9 r = randn(n) @test norm(T' * r - (A' * r)) < 1.0e-9 @test displacement(T) == d @test norm(remove_displacement(T) * x1 - A * x1) < 1.0e-9 - # with sign - T = AffineAdd(opA, d, false) - @test sign(T) == -1 - - verb && println(T) - x1 = randn(m) - y1 = T * x1 - @test norm(y1 - (A * x1 - d)) < 1.0e-9 - r = randn(n) - @test norm(T' * r - (A' * r)) < 1.0e-9 - @test displacement(T) == -d - @test norm(remove_displacement(T) * x1 - A * x1) < 1.0e-9 + T_neg = AffineAdd(opA, d, false) + @test sign(T_neg) == -1 + @test norm(T_neg * x1 - (A * x1 - d)) < 1.0e-9 - # with scalar - T = AffineAdd(opA, pi) - @test sign(T) == 1 - - verb && println(T) - x1 = randn(m) - y1 = T * x1 - @test norm(y1 - (A * x1 .+ pi)) < 1.0e-9 - r = randn(n) - @test norm(T' * r - (A' * r)) < 1.0e-9 - @test displacement(T) .- pi < 1.0e-9 - @test norm(remove_displacement(T) * x1 - A * x1) < 1.0e-9 + T_scalar = AffineAdd(opA, pi) + @test norm(T_scalar * x1 - (A * x1 .+ pi)) < 1.0e-9 @test_throws DimensionMismatch AffineAdd(MatrixOp(randn(2, 5)), randn(5)) - opD = DiagOp(Float64, (4,), randn(ComplexF64, 4)) - @test_throws ErrorException AffineAdd(opD, randn(4)) - AffineAdd(opD, pi) @test_throws ErrorException AffineAdd(Eye(4), im * pi) +end - # with scalar and vector - d = randn(n) - T = AffineAdd(AffineAdd(opA, pi), d, false) - - verb && println(T) - x1 = randn(m) - y1 = T * x1 - @test norm(y1 - (A * x1 .+ pi .- d)) < 1.0e-9 - r = randn(n) - @test norm(T' * r - (A' * r)) < 1.0e-9 - @test norm(displacement(T) .- (pi .- d)) < 1.0e-9 - - T2 = remove_displacement(T) - @test norm(T2 * x1 - (A * x1)) < 1.0e-9 +@testitem "AffineAdd: nonlinear and permute" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # permute AddAffine n, m = 5, 6 A = randn(n, m) d = randn(n) opH = HCAT(Eye(n), MatrixOp(A)) x = ArrayPartition(randn(n), randn(m)) opHT = AffineAdd(opH, d) - @test norm(opHT * x - (x.x[1] + A * x.x[2] .+ d)) < 1.0e-12 p = [2; 1] - @test norm( - AbstractOperators.permute(opHT, p) * ArrayPartition(x.x[p]...) - - (x.x[1] + A * x.x[2] .+ d), - ) < 1.0e-12 + @test norm(AbstractOperators.permute(opHT, p) * ArrayPartition(x.x[p]...) - (x.x[1] + A * x.x[2] .+ d)) < 1.0e-12 - n, m = 5, 6 - A = randn(n, m) - opA = MatrixOp(A) - d = randn(n) - Tplus = AffineAdd(opA, d) - Tminus = AffineAdd(opA, d, false) - Tscalar = AffineAdd(opA, pi) - io = IOBuffer(); show(io, Tplus); splus = String(take!(io)); @test occursin("+d", splus) - io = IOBuffer(); show(io, Tminus); sminus = String(take!(io)); @test occursin("-d", sminus) - io = IOBuffer(); show(io, Tscalar); sscal = String(take!(io)); @test occursin("+d", sscal) || occursin("-d", sscal) - - # Scaling of AffineAdd (array displacement) - α = 2.0 - Tscaled = Scale(α, Tplus) - x = randn(m) - @test Tscaled * x ≈ α * (A * x + d) - # scaling of negative sign variant - Tscaled_neg = Scale(α, Tminus) - @test Tscaled_neg * x ≈ α * (A * x - d) - - # Diagonal passthrough and diag/diag_AcA/diag_AAc - dvec = randn(n) - D = DiagOp(dvec) - Tdiag = AffineAdd(D, zeros(n)) - @test is_diagonal(Tdiag) == true - @test diag(Tdiag) == diag(D) - - # is_null / is_eye cases - Zop = Zeros(Float64, (n,), Float64, (n,)) - Tzero = AffineAdd(Zop, zeros(n)) - @test is_null(Tzero) == true - Eop = Eye(n) - Teye = AffineAdd(Eop, zeros(n)) - @test is_eye(Teye) == true - - # remove_displacement idempotence (already covered for simple, add explicit check) - rd1 = remove_displacement(Tplus) - rd2 = remove_displacement(rd1) - @test rd1 * x == rd2 * x - - # AffineAdd with NonLinearOperator n = 10 d = randn(n) T = AffineAdd(Exp(n), d, false) - r = randn(n) x = randn(size(T, 2)) y, grad = test_NLop(T, x, r, verb) @test norm(y - (exp.(x) - d)) < 1.0e-8 +end - # Additional coverage-focused tests for uncovered AffineAdd branches +@testitem "AffineAdd equality operator" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n, m = 5, 6 + A = MatrixOp(randn(n, m)) + d1 = randn(n) + d2 = randn(n) + @test AffineAdd(A, d1) == AffineAdd(A, d1) + @test !(AffineAdd(A, d1) == AffineAdd(A, d2)) +end - @testset "AffineAdd equality operator" begin - n, m = 5, 6 - A = randn(n, m) - opA = MatrixOp(A) - d1 = randn(n) - d2 = randn(n) - T1 = AffineAdd(opA, d1) - T2 = AffineAdd(opA, d1) - T3 = AffineAdd(opA, d2) - @test T1 == T2 - @test !(T1 == T3) - end +@testitem "AffineAdd property delegations (invertible, rank, diagonal)" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n = 5 + E = Eye(n) + TE = AffineAdd(E, randn(n)) + @test is_invertible(TE) == is_invertible(E) + D = DiagOp(randn(n)) + TD = AffineAdd(D, randn(n)) + @test is_AcA_diagonal(TD) == is_AcA_diagonal(D) + @test is_AAc_diagonal(TD) == is_AAc_diagonal(D) + m = 6 + A = MatrixOp(randn(n, m)) + TA = AffineAdd(A, randn(n)) + @test is_full_row_rank(TA) == is_full_row_rank(A) + @test is_full_column_rank(TA) == is_full_column_rank(A) + D2 = DiagOp(randn(n)) + TD2 = AffineAdd(D2, zeros(n)) + @test diag_AcA(TD2) == diag_AcA(D2) + @test diag_AAc(TD2) == diag_AAc(D2) +end - @testset "AffineAdd property delegations (invertible, rank, diagonal)" begin - n = 5 - # is_invertible delegation - E = Eye(n) - TE = AffineAdd(E, randn(n)) - @test is_invertible(TE) == is_invertible(E) - - # is_AcA_diagonal and is_AAc_diagonal delegation - D = DiagOp(randn(n)) - TD = AffineAdd(D, randn(n)) - @test is_AcA_diagonal(TD) == is_AcA_diagonal(D) - @test is_AAc_diagonal(TD) == is_AAc_diagonal(D) - - # is_full_row_rank and is_full_column_rank delegation - m = 6 - A = randn(n, m) - opA = MatrixOp(A) - TA = AffineAdd(opA, randn(n)) - @test is_full_row_rank(TA) == is_full_row_rank(opA) - @test is_full_column_rank(TA) == is_full_column_rank(opA) - - # diag_AcA and diag_AAc delegation - dvec = randn(n) - D2 = DiagOp(dvec) - TD2 = AffineAdd(D2, zeros(n)) - @test diag_AcA(TD2) == diag_AcA(D2) - @test diag_AAc(TD2) == diag_AAc(D2) - end +@testitem "AffineAdd slicing property delegation" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + G = GetIndex(Float64, (10,), 2:5) + TG = AffineAdd(G, randn(4)) + @test is_sliced(TG) == is_sliced(G) + @test AbstractOperators.get_slicing_expr(TG) == AbstractOperators.get_slicing_expr(G) + @test AbstractOperators.get_slicing_mask(TG) == AbstractOperators.get_slicing_mask(G) + @test AbstractOperators.remove_slicing(TG) isa Eye +end - @testset "AffineAdd slicing property delegation" begin - n = 10 - G = GetIndex(Float64, (n,), 2:5) - TG = AffineAdd(G, randn(4)) - @test is_sliced(TG) == is_sliced(G) - @test AbstractOperators.get_slicing_expr(TG) == AbstractOperators.get_slicing_expr(G) - mask = AbstractOperators.get_slicing_mask(TG) - @test mask == AbstractOperators.get_slicing_mask(G) - Gremoved = AbstractOperators.remove_slicing(TG) - @test Gremoved isa Eye - end +@testitem "AffineAdd normal operator" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n, m = 5, 6 + G = GetIndex(Float64, (n, m), (1:3, :)) + d = randn(3, m) + TG = AffineAdd(G, d) + @test AbstractOperators.has_optimized_normalop(TG) == + AbstractOperators.has_optimized_normalop(G) + N = AbstractOperators.get_normal_op(TG) + x = randn(n, m) + @test N * x ≈ TG.A' * (TG.A * x + TG.d) +end - @testset "AffineAdd normal operator" begin - n, m = 5, 6 - G = GetIndex(Float64, (n, m), (1:3, :)) - d = randn(3, m) - TG = AffineAdd(G, d) - @test AbstractOperators.has_optimized_normalop(TG) == AbstractOperators.has_optimized_normalop(G) - N = AbstractOperators.get_normal_op(TG) - x = randn(n, m) - # Normal operator should compute A'*(A*x + d) for AffineAdd - expected = TG.A' * (TG.A * x + TG.d) - @test N * x ≈ expected - end +@testitem "AffineAdd is_thread_safe delegation" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + @test is_thread_safe(AffineAdd(DiagOp(randn(5)), randn(5))) == true + C = Compose(FiniteDiff((6,)), DiagOp(randn(6))) + @test is_thread_safe(AffineAdd(C, randn(5))) == false +end - @testset "AffineAdd is_thread_safe delegation" begin - # DiagOp is thread-safe, so AffineAdd(DiagOp, d) should be thread-safe - D = DiagOp(randn(5)) - d = randn(5) - TD = AffineAdd(D, d) - @test is_thread_safe(TD) == true - - # Compose with FiniteDiff is not thread-safe (has buffers) - C = Compose(FiniteDiff((6,)), DiagOp(randn(6))) - d2 = randn(5) - TC = AffineAdd(C, d2) - @test is_thread_safe(TC) == false - end +@testitem "AffineAdd sign function" tags = [:calculus, :AffineAdd] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n, m = 5, 6 + opA = MatrixOp(randn(n, m)) + d = randn(n) + @test sign(AffineAdd(opA, d, true)) == 1 + @test sign(AffineAdd(opA, d, false)) == -1 +end - @testset "AffineAdd sign function" begin - n, m = 5, 6 - A = randn(n, m) - opA = MatrixOp(A) - d = randn(n) +@testitem "AffineAdd (GPU)" tags = [:gpu, :calculus, :AffineAdd] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv - # Positive sign (addition) - Tplus = AffineAdd(opA, d, true) - @test sign(Tplus) == 1 + for backend in gpu_backends() + Random.seed!(0) - # Negative sign (subtraction) - Tminus = AffineAdd(opA, d, false) - @test sign(Tminus) == -1 + n, m = 5, 6 + A = gpu_randn(backend, n, m) + d = gpu_randn(backend, n) + T = AffineAdd(MatrixOp(A), d) + x1 = gpu_randn(backend, m) + y1 = T * x1 + y1_buf = similar(y1) + mul!(y1_buf, T, x1) + @test collect(y1) ≈ collect(y1_buf) + + r = gpu_randn(backend, n) + r_adj = T' * r + r_adj2 = similar(r_adj) + mul!(r_adj2, T', r) + @test collect(r_adj) ≈ collect(r_adj2) end end diff --git a/test/calculus/test_broadcast.jl b/test/calculus/test_broadcast.jl index 6ccd8e1..70c0ed1 100644 --- a/test/calculus/test_broadcast.jl +++ b/test/calculus/test_broadcast.jl @@ -1,346 +1,182 @@ -@testitem "BroadCast" tags = [:calculus, :BroadCast] setup = [TestUtils] begin +@testitem "BroadCast: basic mul" tags = [:calculus, :BroadCast] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing BroadCast --- ") - - function test_broadcast_op(verb, threaded) - m, n = 8, 4 - dim_out = (m, 10) - A1 = randn(m, n) - opA1 = MatrixOp(A1) - opR = BroadCast(opA1, dim_out; threaded) - x1 = randn(n) - y1 = test_op(opR, x1, randn(dim_out), verb) - y2 = zeros(dim_out) - y2 .= A1 * x1 - @test norm(y1 - y2) <= 1.0e-12 - @test AbstractOperators.has_fast_opnorm(opR) == AbstractOperators.has_fast_opnorm(opA1) - - m, n, l, k = 8, 4, 5, 7 - dim_out = (m, n, l, k) - opA1 = Eye(m, n) - opR = BroadCast(opA1, dim_out; threaded) - x1 = randn(m, n) - y1 = test_op(opR, x1, randn(dim_out), verb) - y2 = zeros(dim_out) - y2 .= x1 - @test norm(y1 - y2) <= 1.0e-12 - - @test_throws Exception BroadCast(opA1, (m, m)) - - m, l = 1, 5 - dim_out = (m, l) - opA1 = Eye(m) - opR = BroadCast(opA1, dim_out; threaded) - x1 = randn(m) - y1 = test_op(opR, x1, randn(dim_out), verb) - y2 = zeros(dim_out) - y2 .= x1 - @test norm(y1 - y2) <= 1.0e-12 - - #colum in - matrix out - m, l = 4, 5 - dim_out = (m, l) - opA1 = Eye(1, l) - opR = BroadCast(opA1, dim_out; threaded) - x1 = randn(1, l) - y1 = test_op(opR, x1, randn(dim_out), verb) - y2 = zeros(dim_out) - y2 .= x1 - @test norm(y1 - y2) <= 1.0e-12 - - op = HCAT(Eye(m, l), opR) - x1 = ArrayPartition(randn(m, l), randn(1, l)) - y1 = test_op(op, x1, randn(dim_out), verb) - y2 = x1.x[1] .+ x1.x[2] - @test norm(y1 - y2) <= 1.0e-12 - - m, n, l = 2, 5, 8 - dim_out = (m, n, l) - opA1 = Eye(m) - opR = BroadCast(opA1, dim_out; threaded) - x1 = randn(m) - y1 = test_op(opR, x1, randn(dim_out), verb) - y2 = zeros(dim_out) - y2 .= x1 - @test norm(y1 - y2) <= 1.0e-12 - - m, n, l = 1, 5, 8 - dim_out = (m, n, l) - opA1 = Eye(m) - opR = BroadCast(opA1, dim_out; threaded) - x1 = randn(m) - y1 = test_op(opR, x1, randn(dim_out), verb) - y2 = zeros(dim_out) - y2 .= x1 - @test norm(y1 - y2) <= 1.0e-12 - - m, n, l = 1, 5, 8 - dim_out = (m, n, l) - opA1 = Scale(2.4, Eye(m)) - opR = BroadCast(opA1, dim_out; threaded) - x1 = randn(m) - y1 = test_op(opR, x1, randn(dim_out), verb) - y2 = zeros(dim_out) - y2 .= 2.4 * x1 - @test norm(y1 - y2) <= 1.0e-12 - - @test is_null(opR) == is_null(opA1) - @test is_eye(opR) == false - @test is_diagonal(opR) == false - @test is_AcA_diagonal(opR) == false - @test is_AAc_diagonal(opR) == false - @test is_orthogonal(opR) == false - @test is_invertible(opR) == false - @test is_full_row_rank(opR) == false - @test is_full_column_rank(opR) == false - end - - @testset "non-threaded" test_broadcast_op(verb, false) - @testset "threaded" test_broadcast_op(verb, true) - - # NoOperatorBroadCast path (Eye case returning NoOperatorBroadCast rather than OperatorBroadCast) - # Use dim_out that differs to avoid early return of original operator - m = 3 - E = Eye(m) - SB = BroadCast(E, (m, 2)) # should be NoOperatorBroadCast - @test SB isa AbstractOperators.NoOperatorBroadCast - x = randn(m) - y = SB * x - @test y[:, 1] == x && y[:, 2] == x - # fun_name for NoOperatorBroadCast (".I") appears in show - io = IOBuffer() - show(io, SB) - s = String(take!(io)) - @test occursin(".I", s) - # opnorm / estimate consistency - @test opnorm(SB) == √2 - @test estimate_opnorm(SB) == √2 - # remove_displacement on NoOperatorBroadCast is identity - @test remove_displacement(SB) === SB - - # remove_displacement on non-threaded OperatorBroadCast returns structurally equal broadcast - A = MatrixOp(randn(m, m)) - OB = BroadCast(A, (m, 2); threaded = false) - @test remove_displacement(OB) == OB - # For a displaced inner operator - d = randn(m) - AD = AffineAdd(A, d) - OBD = BroadCast(AD, (m, 2); threaded = false) - rd = remove_displacement(OBD) - @test rd * x == BroadCast(remove_displacement(AD), (m, 2); threaded = false) * x - @test remove_displacement(rd) == rd # idempotent - - # permute test (domain permutation) using HCAT to create partition domain - for threaded in [false, true] - A1 = HCAT(Eye(m), Eye(m)) # domain is ArrayPartition - B1 = BroadCast(A1, (m, 2); threaded) - xpart = ArrayPartition(randn(m), randn(m)) - y_base = B1 * xpart - p = [2, 1] - B1p = AbstractOperators.permute(B1, p) - xpart_p = ArrayPartition(xpart.x[p]...) - y_perm = B1p * xpart_p - @test y_perm == y_base # same broadcasted sum after permutation inversion - - # Adjoint path exercise for OperatorBroadCast (non-threaded) to hit get_input_slice - r = randn(size(B1, 1)) - g = B1' * r - @test length(g) == length(xpart.x[1]) + length(xpart.x[2]) - end - - # If multi-threading available, test threaded variant basics (skip if only 1 thread) - if Threads.nthreads() > 1 - AT = BroadCast(A, (m, 3); threaded = true) - xt = randn(m) - yt = AT * xt - @test yt[:, 2] == A * xt # replicated columns - # remove_displacement idempotence for threaded case - @test remove_displacement(AT) == AT - # displaced inner operator threaded - ADT = BroadCast(AffineAdd(A, d), (m, 3); threaded = true) - rdt = remove_displacement(ADT) - @test rdt * xt == BroadCast(A, (m, 3); threaded = true) * xt - end - - # test displacement m, n = 8, 4 dim_out = (m, 10) A1 = randn(m, n) - d1 = randn(m) - opA1 = AffineAdd(MatrixOp(A1), d1) + opA1 = MatrixOp(A1) opR = BroadCast(opA1, dim_out) x1 = randn(n) - y1 = opR * x1 - y2 = zeros(dim_out) - y2 .= A1 * x1 + d1 - @test norm(y1 - y2) <= 1.0e-12 - x1 = randn(n) - y1 = remove_displacement(opR) * x1 + y1 = test_op(opR, x1, randn(dim_out), verb) y2 = zeros(dim_out) y2 .= A1 * x1 @test norm(y1 - y2) <= 1.0e-12 - # Storage types / thread safety (non-thread-safe expected) - _ds = domain_storage_type(opR) - _cs = codomain_storage_type(opR) - @test _ds !== nothing - @test _cs !== nothing - @test is_thread_safe(opR) == false + m, n, l, k = 8, 4, 5, 7 + dim_out = (m, n, l, k) + opA1 = Eye(m, n) + opR = BroadCast(opA1, dim_out) + x1 = randn(m, n) + y1 = test_op(opR, x1, randn(dim_out), verb) + y2 = zeros(dim_out) + y2 .= x1 + @test norm(y1 - y2) <= 1.0e-12 + @test_throws Exception BroadCast(opA1, (m, m)) - # In-place multiplication - m, n = 6, 3 - dim_out = (m, 5) - A1 = randn(m, n) - opA1 = MatrixOp(A1) + m, n = 8, 4 + dim_out = (m, 10) + d1 = randn(m) + opA1 = AffineAdd(MatrixOp(randn(m, n)), d1) opR = BroadCast(opA1, dim_out) x1 = randn(n) - y = zeros(dim_out) - mul!(y, opR, x1) - yref = zeros(dim_out) - yref .= A1 * x1 - @test norm(y - yref) <= 1.0e-12 - - # Adjoint of BroadCast - opR_adj = adjoint(opR) - # If defined, test shape - @test size(opR_adj, 1) == (n,) - - # Show/summary output - io = IOBuffer(); show(io, opR); str = String(take!(io)); @test occursin(".", str) + y1 = opR * x1 + y2 = zeros(dim_out) + y2 .= opA1.A.A * x1 + d1 + @test norm(y1 - y2) <= 1.0e-12 + y3 = remove_displacement(opR) * x1 + y4 = zeros(dim_out) + y4 .= opA1.A.A * x1 + @test norm(y3 - y4) <= 1.0e-12 +end - # Edge: BroadCast of Zeros - opZ = Zeros(Float64, (3,), Float64, (2,)) - opRz = BroadCast(opZ, (2, 4)) - xz = randn(3) - y = opRz * xz - @test all(y .== 0) +@testitem "BroadCast: properties and storage" tags = [:calculus, :BroadCast] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # Edge: BroadCast of Eye - opE = Eye(200) - opRe = BroadCast(opE, (200, 350); threaded = false) - xe = randn(200) - y = opRe * xe - @test all(y[:, 1] .== xe) - opRe = BroadCast(opE, (200, 350); threaded = true) - xe = randn(200) - y = opRe * xe - @test all(y[:, 1] .== xe) - @test is_linear(opRe) == true - @test is_null(opRe) == false - @test AbstractOperators.has_fast_opnorm(opRe) == true + m, n = 8, 4 + dim_out = (m, 10) + opA1 = MatrixOp(randn(m, n)) + opR = BroadCast(opA1, dim_out) + @test is_null(opR) == is_null(opA1) + @test is_eye(opR) == false + @test is_diagonal(opR) == false + @test is_AcA_diagonal(opR) == false + @test is_AAc_diagonal(opR) == false + @test is_orthogonal(opR) == false + @test is_invertible(opR) == false + @test is_full_row_rank(opR) == false + @test is_full_column_rank(opR) == false + @test is_thread_safe(opR) == false + @test domain_array_type(opR) !== nothing + @test codomain_array_type(opR) !== nothing + @test AbstractOperators.has_fast_opnorm(opR) == AbstractOperators.has_fast_opnorm(opA1) - # Edge: BroadCast of DiagOp - d = randn(2) - opD = DiagOp(d) - opRd = BroadCast(opD, (2, 3)) - xd = randn(2) - y = opRd * xd - @test all(y[:, 1] .== d .* xd) + m = 3 + E = Eye(m) + SB = BroadCast(E, (m, 2)) + @test SB isa AbstractOperators.NoOperatorBroadCast + x = randn(m) + y = SB * x + @test y[:, 1] == x && y[:, 2] == x + @test opnorm(SB) == √2 + @test remove_displacement(SB) === SB +end - # Singleton and empty array broadcast - opS = MatrixOp(randn(1, 1)) - opRs = BroadCast(opS, (1,)) - xs = [1.0] - ys = opRs * xs - @test length(ys) == 1 +@testitem "BroadCast: nonlinear" tags = [:calculus, :BroadCast] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # Testing nonlinear BroadCast n, l = 4, 7 x = randn(n) r = randn(n, l) opS = Sigmoid(Float64, (n,), 2) op = BroadCast(opS, (n, l)) - y, grad = test_NLop(op, x, r, verb) + @test norm((opS * x) .* ones(n, l) - y) < 1.0e-8 +end - Y = (opS * x) .* ones(n, l) - @test norm(Y - y) < 1.0e-8 - - n, l = 1, 7 - x = randn(n) - r = randn(n, l) - opS = Sigmoid(Float64, (n,), 2) - op = BroadCast(opS, (n, l)) +@testitem "BroadCast error cases and edge paths" tags = [:calculus, :BroadCast] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - y, grad = test_NLop(op, x, r, verb) + m, n = 4, 3 + A = MatrixOp(randn(m, n)) + @test_throws DimensionMismatch BroadCast(A, (2,)) - Y = (opS * x) .* ones(n, l) - @test norm(Y - y) < 1.0e-8 + E1 = Eye(3) + E2 = Eye(3) + B1 = BroadCast(E1, (3, 2); threaded = false) + B2 = BroadCast(E2, (3, 2); threaded = false) + B3 = BroadCast(E1, (3, 3); threaded = false) + @test B1 == B2 + @test B1 != B3 - # Additional coverage tests - @testset "BroadCast error cases and edge paths" begin - m, n = 4, 3 - A = MatrixOp(randn(m, n)) + @test AbstractOperators.has_fast_opnorm(B1) == AbstractOperators.has_fast_opnorm(E1) + A_op = MatrixOp(randn(3, 2)) + B_op = BroadCast(A_op, (3, 4); threaded = false) + @test opnorm(B_op) ≈ opnorm(A_op) - # Error: dim_out that doesn't match broadcast rules - @test_throws DimensionMismatch BroadCast(A, (2,)) + if Threads.nthreads() > 1 + B_op_t = BroadCast(A_op, (3, 4); threaded = true) + @test opnorm(B_op_t) ≈ opnorm(A_op) + end - # NoOperatorBroadCast equality - E1 = Eye(3) - E2 = Eye(3) - B1 = BroadCast(E1, (3, 2); threaded = false) - B2 = BroadCast(E2, (3, 2); threaded = false) - B3 = BroadCast(E1, (3, 3); threaded = false) - @test B1 == B2 - @test B1 != B3 + A = DiagOp(rand(4, 3, 2)) + @test_throws ErrorException BroadCast(A, (4, 2)) +end - # opnorm for OperatorBroadCast - @test AbstractOperators.has_fast_opnorm(B1) == AbstractOperators.has_fast_opnorm(E1) - A_op = MatrixOp(randn(3, 2)) - B_op = BroadCast(A_op, (3, 4); threaded = false) - @test opnorm(B_op) ≈ opnorm(A_op) +@testitem "Threaded NoOperatorBroadCast" tags = [:calculus, :BroadCast] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - if Threads.nthreads() > 1 - B_op_t = BroadCast(A_op, (3, 4); threaded = true) - @test opnorm(B_op_t) ≈ opnorm(A_op) + if Threads.nthreads() > 1 + m = 1000 + E = Eye(m) + B_threaded = BroadCast(E, (m, 100); threaded = true) + @test B_threaded isa AbstractOperators.NoOperatorBroadCast + x = randn(m) + y = B_threaded * x + for i in 1:100 + @test y[:, i] ≈ x end - - # wrong output size - A = DiagOp(rand(4, 3, 2)) - @test_throws ErrorException BroadCast(A, (4, 2)) + y_adj = randn(m, 100) + x_back = B_threaded' * y_adj + @test x_back ≈ dropdims(sum(y_adj, dims = 2), dims = 2) end +end - @testset "Threaded NoOperatorBroadCast" begin - if Threads.nthreads() > 1 - m = 1000 - E = Eye(m) - B_threaded = BroadCast(E, (m, 100); threaded = true) - @test B_threaded isa AbstractOperators.NoOperatorBroadCast - - x = randn(m) - y = B_threaded * x - - for i in 1:100 - @test y[:, i] ≈ x - end +@testitem "Non-compact threaded OperatorBroadCast" tags = [:calculus, :BroadCast] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - y_adj = randn(m, 100) - x_back = B_threaded' * y_adj - @test x_back ≈ dropdims(sum(y_adj, dims = 2), dims = 2) + if Threads.nthreads() > 1 + m, n = 3, 2 + A = reshape(MatrixOp(randn(m, n)), 1, m) + dim_out = (4, m, 5) + B_noncompact = BroadCast(A, dim_out; threaded = true) + x = randn(n) + y = B_noncompact * x + ref = A * x + for i in 1:4, j in 1:5 + @test y[i, :, j] ≈ vec(ref) end + y_test = randn(dim_out) + x_back = B_noncompact' * y_test + @test size(x_back) == (n,) + @test x_back ≈ A' * dropdims(sum(y_test, dims = (1, 3)), dims = 3) end +end - @testset "Non-compact threaded OperatorBroadCast" begin - if Threads.nthreads() > 1 - m, n = 3, 2 - A = reshape(MatrixOp(randn(m, n)), 1, m) - dim_out = (4, m, 5) - B_noncompact = BroadCast(A, dim_out; threaded = true) - - x = randn(n) - y = B_noncompact * x +@testitem "BroadCast (GPU)" tags = [:gpu, :calculus, :BroadCast] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv - ref = A * x - for i in 1:4, j in 1:5 - @test y[i, :, j] ≈ vec(ref) - end + for backend in gpu_backends() + Random.seed!(0) - y_test = randn(dim_out) - x_back = B_noncompact' * y_test - @test size(x_back) == (n,) - @test x_back ≈ A' * dropdims(sum(y_test, dims = (1, 3)), dims = 3) - end + m, n = 8, 4 + dim_out = (m, 10) + A1 = gpu_randn(backend, m, n) + opR = BroadCast(MatrixOp(A1), dim_out; threaded = false) + test_op(opR, gpu_randn(backend, n), gpu_randn(backend, dim_out...), false) + + m2, n2 = 3, 3 + dim_out2 = (m2, n2, 5) + opR2 = BroadCast( + Eye(Float64, (m2, n2); array_type = gpu_wrapper(backend, Float64, m2, n2)), + dim_out2; + threaded = false, + ) + test_op(opR2, gpu_randn(backend, m2, n2), gpu_randn(backend, dim_out2...), false) end end diff --git a/test/calculus/test_combinations.jl b/test/calculus/test_combinations.jl index 4b5e332..c40cd7e 100644 --- a/test/calculus/test_combinations.jl +++ b/test/calculus/test_combinations.jl @@ -1,8 +1,7 @@ -@testitem "Combinations" tags = [:calculus, :Combinations] setup = [TestUtils] begin +@testitem "Combinations: HCAT and Compose" tags = [:calculus, :Combinations] setup = [TestUtils] begin using Random, AbstractOperators - verb && println(" --- Testing Combinations --- ") + Random.seed!(42) - ## test Compose of HCAT m1, m2, m3, m4 = 4, 7, 3, 2 A1 = randn(m3, m1) A2 = randn(m3, m2) @@ -21,99 +20,86 @@ opCp = AbstractOperators.permute(opC, [2, 1]) y1 = test_op(opCp, ArrayPartition(x2, x1), randn(m4), verb) - @test norm(y1 - y2) < 1.0e-9 - ## test HCAT of Compose of HCAT m5 = 10 A4 = randn(m4, m5) x3 = randn(m5) opHC = HCAT(opC, MatrixOp(A4)) x = ArrayPartition(x1, x2, x3) y1 = test_op(opHC, x, randn(m4), verb) - @test norm(y1 - (y2 + A4 * x3)) < 1.0e-9 p = randperm(ndoms(opHC, 2)) opHP = AbstractOperators.permute(opHC, p) - xp = ArrayPartition(x.x[p]...) - y1 = test_op(opHP, xp, randn(m4), verb) pp = randperm(ndoms(opHC, 2)) opHPP = AbstractOperators.permute(opHC, pp) xpp = ArrayPartition(x.x[pp]...) y1 = test_op(opHPP, xpp, randn(m4), verb) +end - # test VCAT of HCAT's +@testitem "Combinations: VCAT and HCAT mixtures" tags = [:calculus, :Combinations] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(43) + + # VCAT of HCATs m1, m2, n1 = 4, 7, 3 A1 = randn(n1, m1) A2 = randn(n1, m2) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opH1 = HCAT(opA1, opA2) - + opH1 = HCAT(MatrixOp(A1), MatrixOp(A2)) m1, m2, n2 = 4, 7, 5 A3 = randn(n2, m1) A4 = randn(n2, m2) - opA3 = MatrixOp(A3) - opA4 = MatrixOp(A4) - opH2 = HCAT(opA3, opA4) - + opH2 = HCAT(MatrixOp(A3), MatrixOp(A4)) opV = VCAT(opH1, opH2) x1, x2 = randn(m1), randn(m2) y1 = test_op(opV, ArrayPartition(x1, x2), ArrayPartition(randn(n1), randn(n2)), verb) y2 = ArrayPartition(A1 * x1 + A2 * x2, A3 * x1 + A4 * x2) @test norm(y1 - y2) <= 1.0e-12 - # test VCAT of HCAT's with complex num + # VCAT of HCATs with complex m1, m2, n1 = 4, 7, 5 - A1 = randn(n1, m1) + im * randn(n1, m1) - opA1 = MatrixOp(A1) + A1c = randn(n1, m1) + im * randn(n1, m1) d1 = rand(ComplexF64, n1) - opA2 = DiagOp(Float64, (n1,), d1) - opH1 = HCAT(opA1, opA2) - + opH1c = HCAT(MatrixOp(A1c), DiagOp(Float64, (n1,), d1)) m1, m2, n2 = 4, 7, 5 - A3 = randn(n2, m1) + im * randn(n2, m1) - opA3 = MatrixOp(A3) + A3c = randn(n2, m1) + im * randn(n2, m1) d2 = rand(ComplexF64, n2) - opA4 = DiagOp(Float64, (n2,), d2) - opH2 = HCAT(opA3, opA4) - - opV = VCAT(opH1, opH2) - x1, x2 = randn(m1) + im * randn(m1), randn(n2) - y1 = test_op( - opV, - ArrayPartition(x1, x2), + opH2c = HCAT(MatrixOp(A3c), DiagOp(Float64, (n2,), d2)) + opVc = VCAT(opH1c, opH2c) + x1c = randn(m1) + im * randn(m1) + x2c = randn(n2) + y1c = test_op( + opVc, + ArrayPartition(x1c, x2c), ArrayPartition(randn(n1) + im * randn(n1), randn(n2) + im * randn(n2)), verb, ) - y2 = ArrayPartition(A1 * x1 + x2 .* d1, A3 * x1 + x2 .* d2) - @test norm(y1 - y2) <= 1.0e-12 - - # test HCAT of VCAT's + y2c = ArrayPartition(A1c * x1c + x2c .* d1, A3c * x1c + x2c .* d2) + @test norm(y1c - y2c) <= 1.0e-12 + # HCAT of VCATs n1, n2, m1, m2 = 3, 5, 4, 7 A = randn(m1, n1) - opA = MatrixOp(A) B = randn(m1, n2) - opB = MatrixOp(B) C = randn(m2, n1) - opC = MatrixOp(C) D = randn(m2, n2) - opD = MatrixOp(D) - opV = HCAT(VCAT(opA, opC), VCAT(opB, opD)) + opV2 = HCAT(VCAT(MatrixOp(A), MatrixOp(C)), VCAT(MatrixOp(B), MatrixOp(D))) x1 = randn(n1) x2 = randn(n2) - y1 = test_op(opV, ArrayPartition(x1, x2), ArrayPartition(randn(m1), randn(m2)), verb) + y1 = test_op(opV2, ArrayPartition(x1, x2), ArrayPartition(randn(m1), randn(m2)), verb) y2 = ArrayPartition(A * x1 + B * x2, C * x1 + D * x2) - @test norm(y1 - y2) <= 1.0e-12 +end - # test Sum of HCAT's +@testitem "Combinations: Sum structures" tags = [:calculus, :Combinations] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(44) + # Sum of HCATs m, n1, n2, n3 = 4, 7, 5, 3 A1 = randn(m, n1) A2 = randn(m, n2) @@ -121,29 +107,21 @@ B1 = randn(m, n1) B2 = randn(m, n2) B3 = randn(m, n3) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opA3 = MatrixOp(A3) - opB1 = MatrixOp(B1) - opB2 = MatrixOp(B2) - opB3 = MatrixOp(B3) - opHA = HCAT(opA1, opA2, opA3) - opHB = HCAT(opB1, opB2, opB3) + opHA = HCAT(MatrixOp(A1), MatrixOp(A2), MatrixOp(A3)) + opHB = HCAT(MatrixOp(B1), MatrixOp(B2), MatrixOp(B3)) opS = Sum(opHA, opHB) x1 = randn(n1) x2 = randn(n2) x3 = randn(n3) y1 = test_op(opS, ArrayPartition(x1, x2, x3), randn(m), verb) y2 = A1 * x1 + B1 * x1 + A2 * x2 + B2 * x2 + A3 * x3 + B3 * x3 - @test norm(y1 - y2) <= 1.0e-12 p = [3; 2; 1] opSp = AbstractOperators.permute(opS, p) y1 = test_op(opSp, ArrayPartition(((x1, x2, x3)[p])...), randn(m), verb) - # test Sum of VCAT's - + # Sum of VCATs m1, m2, n = 4, 7, 5 A1 = randn(m1, n) A2 = randn(m2, n) @@ -151,137 +129,108 @@ B2 = randn(m2, n) C1 = randn(m1, n) C2 = randn(m2, n) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opB1 = MatrixOp(B1) - opB2 = MatrixOp(B2) - opC1 = MatrixOp(C1) - opC2 = MatrixOp(C2) - opVA = VCAT(opA1, opA2) - opVB = VCAT(opB1, opB2) - opVC = VCAT(opC1, opC2) + opVA = VCAT(MatrixOp(A1), MatrixOp(A2)) + opVB = VCAT(MatrixOp(B1), MatrixOp(B2)) + opVC = VCAT(MatrixOp(C1), MatrixOp(C2)) opS = Sum(opVA, opVB, opVC) x = randn(n) y1 = test_op(opS, x, ArrayPartition(randn(m1), randn(m2)), verb) y2 = ArrayPartition(A1 * x + B1 * x + C1 * x, A2 * x + B2 * x + C2 * x) - @test norm(y1 - y2) .<= 1.0e-12 +end - # test Scale of DCAT +@testitem "Combinations: Scale structures" tags = [:calculus, :Combinations] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(45) + # Scale of DCAT m1, n1 = 4, 7 m2, n2 = 3, 5 A1 = randn(m1, n1) A2 = randn(m2, n2) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opD = DCAT(opA1, opA2) + opD = DCAT(MatrixOp(A1), MatrixOp(A2)) coeff = randn() opS = Scale(coeff, opD) x1 = randn(n1) x2 = randn(n2) y = test_op(opS, ArrayPartition(x1, x2), ArrayPartition(randn(m1), randn(m2)), verb) z = ArrayPartition(coeff * A1 * x1, coeff * A2 * x2) - @test norm(y - z) <= 1.0e-12 - # test Scale of VCAT - + # Scale of VCAT m1, m2, n = 4, 3, 7 A1 = randn(m1, n) A2 = randn(m2, n) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opV = VCAT(opA1, opA2) + opV = VCAT(MatrixOp(A1), MatrixOp(A2)) coeff = randn() opS = Scale(coeff, opV) x = randn(n) y = test_op(opS, x, ArrayPartition(randn(m1), randn(m2)), verb) z = ArrayPartition(coeff * A1 * x, coeff * A2 * x) - @test norm(y - z) <= 1.0e-12 - # test Scale of HCAT - + # Scale of HCAT m, n1, n2 = 4, 3, 7 A1 = randn(m, n1) A2 = randn(m, n2) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opH = HCAT(opA1, opA2) + opH = HCAT(MatrixOp(A1), MatrixOp(A2)) coeff = randn() opS = Scale(coeff, opH) x1 = randn(n1) x2 = randn(n2) y = test_op(opS, ArrayPartition(x1, x2), randn(m), verb) z = coeff * (A1 * x1 + A2 * x2) - @test norm(y - z) <= 1.0e-12 - # test DCAT of HCATs - - m1, m2, n1, n2, l1, l2, l3 = 2, 3, 4, 5, 6, 7, 8 + # DCAT of HCATs + m1, m2, n1, n2 = 2, 3, 4, 5 A1 = randn(m1, n1) A2 = randn(m1, n2) B1 = randn(m2, n1) B2 = randn(m2, n2) B3 = randn(m2, n2) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opH1 = HCAT(opA1, opA2) - opB1 = MatrixOp(B1) - opB2 = MatrixOp(B2) - opB3 = MatrixOp(B3) - opH2 = HCAT(opB1, opB2, opB3) - - op = DCAT(opA1, opH2) - x = ArrayPartition(randn.(size(op, 2))...) - y0 = ArrayPartition(randn.(size(op, 1))...) - y = test_op(op, x, y0, verb) - - op = DCAT(opH1, opH2) + opH1 = HCAT(MatrixOp(A1), MatrixOp(A2)) + opH2 = HCAT(MatrixOp(B1), MatrixOp(B2), MatrixOp(B3)) + op = DCAT(MatrixOp(A1), opH2) x = ArrayPartition(randn.(size(op, 2))...) y0 = ArrayPartition(randn.(size(op, 1))...) y = test_op(op, x, y0, verb) - - p = randperm(ndoms(op, 2)) - y2 = op[p] * ArrayPartition(x.x[p]...) - + op2 = DCAT(opH1, opH2) + x = ArrayPartition(randn.(size(op2, 2))...) + y0 = ArrayPartition(randn.(size(op2, 1))...) + y = test_op(op2, x, y0, verb) + p = randperm(ndoms(op2, 2)) + y2 = op2[p] * ArrayPartition(x.x[p]...) @test norm(y - y2) <= 1.0e-8 - # test Scale of Sum - + # Scale of Sum and Compose m, n = 5, 7 A1 = randn(m, n) A2 = randn(m, n) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opS = Sum(opA1, opA2) + opSum = Sum(MatrixOp(A1), MatrixOp(A2)) coeff = pi - opSS = Scale(coeff, opS) + opSS = Scale(coeff, opSum) x1 = randn(n) y1 = test_op(opSS, x1, randn(m), verb) y2 = coeff * (A1 * x1 + A2 * x1) @test norm(y1 - y2) <= 1.0e-12 - # test Scale of Compose - m1, m2, m3 = 4, 7, 3 - A1 = randn(m2, m1) - A2 = randn(m3, m2) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - - coeff = pi - opC = Compose(opA2, opA1) - opS = Scale(coeff, opC) + Ac1 = randn(m2, m1) + Ac2 = randn(m3, m2) + opC = Compose(MatrixOp(Ac2), MatrixOp(Ac1)) + opSC = Scale(coeff, opC) x = randn(m1) - y1 = test_op(opS, x, randn(m3), verb) - y2 = coeff * (A2 * A1 * x) + y1 = test_op(opSC, x, randn(m3), verb) + y2 = coeff * (Ac2 * Ac1 * x) @test all(norm.(y1 .- y2) .<= 1.0e-12) +end +@testitem "Combinations: Nonlinear" tags = [:calculus, :Combinations] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(46) - # Testing nonlinear HCAT of VCAT + # Nonlinear HCAT of VCAT n, m1, m2, m3 = 4, 3, 2, 7 x1 = randn(m1) x2 = randn(m2) @@ -298,13 +247,11 @@ op2 = VCAT(MatrixOp(A2), MatrixOp(B2)) op3 = VCAT(MatrixOp(A3), MatrixOp(B3)) op = HCAT(op1, op2, op3) - y, grad = test_NLop(op, x, r, verb) - Y = ArrayPartition(A1 * x1 + A2 * x2 + A3 * x3, B1 * x1 + B2 * x2 + B3 * x3) @test norm(Y - y) < 1.0e-8 - # Testing nonlinear VCAT of HCAT + # Nonlinear VCAT of HCAT m1, m2, m3, n1, n2 = 3, 4, 5, 6, 7 x1 = randn(m1) x2 = randn(n1) @@ -317,37 +264,44 @@ A2 = randn(n2, m1) B2 = randn(n2, n1) C2 = randn(n2, m3) - x = ArrayPartition(x1, x2, x3) - op1 = HCAT(MatrixOp(A1), B1, MatrixOp(C1)) - op2 = HCAT(MatrixOp(A2), MatrixOp(B2), MatrixOp(C2)) - op = VCAT(op1, op2) - + op = VCAT(HCAT(MatrixOp(A1), B1, MatrixOp(C1)), HCAT(MatrixOp(A2), MatrixOp(B2), MatrixOp(C2))) y, grad = test_NLop(op, x, r, verb) - Y = ArrayPartition(A1 * x1 + B1 * x2 + C1 * x3, A2 * x1 + B2 * x2 + C2 * x3) @test norm(Y - y) < 1.0e-8 - # Testing nonlinear AffineAdd and Compose + # Nonlinear AffineAdd and Compose n = 10 d1 = randn(n) d2 = randn(n) T = Compose(AffineAdd(Sin(n), d2), AffineAdd(Eye(n), d1)) - r = randn(n) x = randn(size(T, 2)) y, grad = test_NLop(T, x, r, verb) @test norm(y - (sin.(x + d1) + d2)) < 1.0e-8 - n = 10 - d1 = randn(n) - d2 = randn(n) d3 = pi - T = Compose( + T2 = Compose( AffineAdd(Sin(n), d3), Compose(AffineAdd(Exp(n), d2, false), AffineAdd(Eye(n), d1)) ) - r = randn(n) - x = randn(size(T, 2)) - y, grad = test_NLop(T, x, r, verb) + x = randn(size(T2, 2)) + y, grad = test_NLop(T2, x, r, verb) @test norm(y - (sin.(exp.(x + d1) - d2) .+ d3)) < 1.0e-8 end + +@testitem "Combinations (GPU)" tags = [:gpu, :calculus, :Combinations] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + m1, m2, m3, m4 = 4, 7, 3, 2 + A1 = gpu_randn(backend, m3, m1) + A2 = gpu_randn(backend, m3, m2) + A3 = gpu_randn(backend, m4, m3) + opH = HCAT(MatrixOp(A1), MatrixOp(A2)) + opC = Compose(MatrixOp(A3), opH) + x1, x2 = gpu_randn(backend, m1), gpu_randn(backend, m2) + test_op(opC, ArrayPartition(x1, x2), gpu_randn(backend, m4), false) + end +end diff --git a/test/calculus/test_compose.jl b/test/calculus/test_compose.jl index e16d20c..92b3398 100644 --- a/test/calculus/test_compose.jl +++ b/test/calculus/test_compose.jl @@ -1,143 +1,70 @@ -@testitem "Compose" tags = [:calculus, :Compose] setup = [TestUtils] begin +@testitem "Compose: basic mul" tags = [:calculus, :Compose] setup = [TestUtils] begin using Random, LinearAlgebra, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Compose --- ") m1, m2 = 4, 7 A1 = randn(m2, m1) opA1 = MatrixOp(A1) opF = FiniteDiff((m2,)) - opC = Compose(opF, opA1) x = randn(m1) y1 = test_op(opC, x, randn(m2 - 1), verb) - y2 = diff(A1 * x) - @test y1 == y2 + @test y1 == diff(A1 * x) - # test Compose longer m1, m2, m3 = 4, 7, 3 A1 = randn(m2, m1) A2 = randn(m3, m2 - 1) opA1 = MatrixOp(A1) - opF = FiniteDiff((m2,)) opA2 = MatrixOp(A2) - - opC1 = Compose(opA2, Compose(opF, opA1)) - opC2 = Compose(Compose(opA2, opF), opA1) + opC1 = Compose(opA2, Compose(FiniteDiff((m2,)), opA1)) + opC2 = Compose(Compose(opA2, FiniteDiff((m2,))), opA1) x = randn(m1) y1 = test_op(opC1, x, randn(m3), verb) y2 = test_op(opC2, x, randn(m3), verb) - y3 = A2 * diff(A1 * x) + @test all(norm.(y1 .- A2 * diff(A1 * x)) .<= 1.0e-12) @test all(norm.(y1 .- y2) .<= 1.0e-12) - @test all(norm.(y3 .- y2) .<= 1.0e-12) - #test Compose special cases @test typeof(opA1 * Eye(m1)) == typeof(opA1) @test typeof(Eye(m2) * opA1) == typeof(opA1) @test typeof(Eye(m2) * Eye(m2)) == typeof(Eye(m2)) - - opS1 = Compose(opF, opA1) - opS1c = Scale(pi, opS1) - @test opS1c isa Compose # Scaling is fused with opA1 - - # In-place multiplication coverage - opS = MyLinOp(Float64, (m2,), Float64, (m2,), (y, x) -> y .= x .* 2, (y, x) -> y .= x ./ 2) - C = Compose(opS, opA1) - x = randn(m1) - y = zeros(m2) - mul!(y, C, x) - @test y ≈ 2 .* (A1 * x) - - # Adjoint reversal property ( (A*B* C)' == C'*B'*A') - A2 = randn(m3, m2) - opA2 = MatrixOp(A2) - chain = Compose(opA2, Compose(opS, opA1)) - chain_adj = chain' - x_in = randn(m3) - y_chain = chain_adj * x_in - y_ref = opA1' * (1 // 2 .* (opA2' * x_in)) - @test y_chain ≈ y_ref - - # Dimension mismatch error @test_throws Exception Compose(MatrixOp(randn(5, 4)), MatrixOp(randn(3, 2))) +end - # Identity elimination & fusion checks - E = Eye(m2) - comp1 = opA2 * E * opA1 - @test comp1 * x ≈ A2 * A1 * x - - # Composition with Zeros yields Zeros (front) - Z = Zeros(Float64, (m2,), Float64, (m2,)) - ZC = Compose(opA2, Z) - @test is_null(ZC) - @test ZC * zeros(m2) == zeros(m3) - - # Composition with diagonal preserves diagonality when both diagonal - struct MyDiagOp <: LinearOperator end - LinearAlgebra.size(::MyDiagOp) = ((m1,), (m1,)) - AbstractOperators.domain_type(::MyDiagOp) = Float64 - AbstractOperators.codomain_type(::MyDiagOp) = Float64 - AbstractOperators.is_diagonal(::MyDiagOp) = true - AbstractOperators.diag(::MyDiagOp) = 3.0 - d = randn(m1) - D1 = DiagOp(d) - D2 = MyDiagOp() - DD = Compose(D2, D1) - @test is_diagonal(DD) - @test diag(DD) == diag(D2) .* diag(D1) - - # Scale inside composition - S = Scale(2.5, opF) - SC = Compose(S, opA1) - @test SC * x ≈ 2.5 * diff(A1 * x) - - # Show output coverage - io = IOBuffer(); show(io, chain); str = String(take!(io)); @test occursin("Π", str) - - # Displacement: nested remove_displacement idempotence - dvec = randn(m2) - Aff = AffineAdd(opA1, dvec) - comp_disp = Compose(opA2, Aff) - x = randn(m1) - y_full = comp_disp * x - y_split = A2 * (A1 * x + dvec) - @test y_full ≈ y_split - rd = remove_displacement(comp_disp) - rd2 = remove_displacement(rd) - @test rd * x ≈ A2 * (A1 * x) - @test rd2 * x ≈ rd * x - - # Sliced + diagonal detection after composition GetIndex * DiagOp - sel = 1:minimum((length(d), 3)) - Sliced1 = Compose(GetIndex((length(d),), sel), DiagOp(d)) - @test !is_sliced(Sliced1) - @test !is_diagonal(Sliced1) - @test is_AAc_diagonal(Sliced1) - Sliced2 = Compose(DiagOp(d[sel]), GetIndex((length(d),), sel)) - @test is_sliced(Sliced2) - @test is_diagonal(Sliced2) +@testitem "Compose: properties" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) - #properties + m1, m2, m3 = 4, 7, 3 + A1 = randn(m2, m1) + A2 = randn(m3, m2 - 1) + opA1 = MatrixOp(A1) + opA2 = MatrixOp(A2) + opF = FiniteDiff((m2,)) + opC = Compose(opF, opA1) + opC1 = Compose(opA2, Compose(opF, opA1)) @test is_sliced(opC) == false @test is_linear(opC1) == true @test is_null(opC1) == false @test is_eye(opC1) == false @test is_diagonal(opC1) == false - @test is_AcA_diagonal(opC1) == false - @test is_AAc_diagonal(opC1) == false @test is_orthogonal(opC1) == false @test is_invertible(opC1) == false - @test is_full_row_rank(opC1) == (is_full_row_rank(opC1.A[1]) && is_full_row_rank(opC1.A[2])) - @test is_full_column_rank(opC1) == (is_full_column_rank(opC1.A[1]) && is_full_column_rank(opC1.A[2])) - # properties special case d = randn(5) - opC = DiagOp(d) * GetIndex((10,), 1:5) - @test is_sliced(opC) == true - @test is_diagonal(opC) == true - @test diag(opC) == d + opC2 = DiagOp(d) * GetIndex((10,), 1:5) + @test is_sliced(opC2) == true + @test is_diagonal(opC2) == true + @test diag(opC2) == d + + Z = Zeros(Float64, (m2,), Float64, (m2 - 1,)) + ZC = Compose(opA2, Z) + @test is_null(ZC) +end + +@testitem "Compose: displacement" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # displacement test m1, m2, m3, m4, m5 = 4, 7, 3, 2, 11 A1 = randn(m2, m1) A2 = randn(m3, m2) @@ -151,85 +78,17 @@ opA2 = AffineAdd(MatrixOp(A2), d2) opA3 = MatrixOp(A3) opA4 = AffineAdd(MatrixOp(A4), d4, false) - opC = Compose(Compose(Compose(opA4, opA3), opA2), opA1) - x = randn(m1) - @test norm(opC * x - (A4 * (A3 * (A2 * (A1 * x + d1) .+ d2) .+ d3) - d4)) < 1.0e-9 @test norm(displacement(opC) - (A4 * (A3 * (A2 * d1 .+ d2) .+ d3) - d4)) < 1.0e-9 - - opA4 = MatrixOp(A4) - opC = AffineAdd(Compose(Compose(Compose(opA4, opA3), opA2), opA1), d4) - @test norm(opC * x - (A4 * (A3 * (A2 * (A1 * x + d1) .+ d2) .+ d3) + d4)) < 1.0e-9 - @test norm(displacement(opC) - (A4 * (A3 * (A2 * d1 .+ d2) .+ d3) + d4)) < 1.0e-9 - @test norm(remove_displacement(opC) * x - (A4 * (A3 * (A2 * (A1 * x))))) < 1.0e-9 +end - # Error paths: domain/codomain type/storage mismatch - struct ComposeDummyOp <: LinearOperator end - LinearAlgebra.size(::ComposeDummyOp) = ((2,), (2,)) - AbstractOperators.domain_type(::ComposeDummyOp) = Int - AbstractOperators.codomain_type(::ComposeDummyOp) = Int - AbstractOperators.diag(::ComposeDummyOp) = 1 - AbstractOperators.fun_name(::ComposeDummyOp) = "D2" - opint = ComposeDummyOp() - @test_throws DomainError Compose(DiagOp(rand(2)), opint) - - # Storage type mismatch path in Compose constructor - struct StorageMismatchLeft <: LinearOperator end - struct StorageMismatchRight <: LinearOperator end - LinearAlgebra.size(::StorageMismatchLeft) = ((2,), (2,)) - LinearAlgebra.size(::StorageMismatchRight) = ((2,), (2,)) - AbstractOperators.domain_type(::StorageMismatchLeft) = Float64 - AbstractOperators.codomain_type(::StorageMismatchLeft) = Float64 - AbstractOperators.domain_storage_type(::StorageMismatchLeft) = Vector{Float64} - AbstractOperators.codomain_storage_type(::StorageMismatchLeft) = Vector{Float64} - AbstractOperators.domain_type(::StorageMismatchRight) = Float64 - AbstractOperators.codomain_type(::StorageMismatchRight) = Float64 - AbstractOperators.domain_storage_type(::StorageMismatchRight) = Vector{Float64} - AbstractOperators.codomain_storage_type(::StorageMismatchRight) = Matrix{Float64} - @test_throws DomainError Compose(StorageMismatchLeft(), StorageMismatchRight()) - - # Show output patterns for Compose (2-term vs multi-term) instead of direct fun_name (non-exported) - C2 = Compose(DiagOp(rand(2)), FiniteDiff((3,))) - io_fn = IOBuffer(); show(io_fn, C2); str_fn = String(take!(io_fn)) - @test occursin("╲*δx", str_fn) - C4 = DiagOp(rand(2)) * FiniteDiff((3,)) * DiagOp(rand(3)) - io_fn = IOBuffer(); show(io_fn, C4); str_fn = String(take!(io_fn)) - @test occursin("Π", str_fn) - - # opnorm/estimate_opnorm consistency (using DiagOp for simplicity) - d = randn(2) - D1 = DiagOp(d) - D2 = FiniteDiff((3,)) - CC = Compose(D1, D2) - opnorm_CC = opnorm(CC) - @test opnorm_CC ≈ estimate_opnorm(CC) rtol = 0.05 - - # permute utility - A1 = MatrixOp(randn(2, 2)) - A2 = MatrixOp(randn(2, 2)) - A3 = MatrixOp(randn(2, 2)) - Cperm = Compose(A3, HCAT(A1, A2)) - p = [2, 1] - Cperm2 = AbstractOperators.permute(Cperm, p) - @test size(Cperm2) == size(Cperm) - x = ArrayPartition(randn(2), randn(2)) - y1 = Cperm * x - y2 = Cperm2 * ArrayPartition(x.x[p]...) - @test y1 == y2 - - # remove_slicing utility - S = Compose(DiagOp(randn(2)), GetIndex((5,), 1:2)) - S2 = AbstractOperators.remove_slicing(S) - @test S2 isa DiagOp - - # get_operators utility - ops = AbstractOperators.get_operators(CC) - @test length(ops) == 2 +@testitem "Compose: nonlinear" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # Testing nonlinear Compose l, n, m = 5, 4, 3 x = randn(m) r = randn(l) @@ -239,306 +98,261 @@ opB = Sigmoid(Float64, (n,), 2) opC = MatrixOp(C) op = Compose(opA, Compose(opB, opC)) - - y, grad = test_NLop(op, x, r, verb) - - Y = A * (opB * (opC * x)) - @test norm(Y - y) < 1.0e-8 - - ## NN - m, n, l = 4, 7, 5 - b = randn(l) - opS1 = Sigmoid(Float64, (n,), 2) - x = ArrayPartition(randn(n, l), randn(n)) - r = randn(n) - - A1 = HCAT(LMatrixOp(b, n), Eye(n)) - op = Compose(opS1, A1) y, grad = test_NLop(op, x, r, verb) + @test norm(A * (opB * (opC * x)) - y) < 1.0e-8 +end - # --- Additional coverage for Compose internals and edge branches --- - - @testset "Compose internal buffer mismatch" begin - A = MatrixOp(randn(3, 3)) - B = MatrixOp(randn(3, 3)) - # Using the low-level constructor with wrong buffer length should throw - @test_throws DimensionMismatch AbstractOperators.Compose((B, A), ()) - end +@testitem "Compose internal buffer mismatch" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = MatrixOp(randn(3, 3)) + B = MatrixOp(randn(3, 3)) + @test_throws DimensionMismatch AbstractOperators.Compose((B, A), ()) +end - @testset "Adjacent adjoint optimized normal (GetIndex*GetIndex')" begin - G = GetIndex((5,), 2:4) - # Provide a dummy buffer (won't be used because it reduces to a single op) - L = AbstractOperators.Compose((G, G'), (randn(size(G, 1)),)) - x = randn(5) - y_ref = G' * (G * x) - y = L * x - @test y == y_ref - @test is_diagonal(L) - end +@testitem "Adjacent adjoint optimized normal (GetIndex*GetIndex')" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + G = GetIndex((5,), 2:4) + L = AbstractOperators.Compose((G, G'), (randn(size(G, 1)),)) + x = randn(5) + @test L * x == G' * (G * x) + @test is_diagonal(L) +end - @testset "Combine branch producing nested Compose is inlined" begin - M = MatrixOp(randn(3, 3)) - S = Scale(2.0, Eye(3)) - # Order (S, M) triggers combine(MatrixOp, Scale) which returns a Compose - L = AbstractOperators.Compose((S, M), (zeros(3),)) - x = randn(3) - @test L * x ≈ 2.0 * (M * x) - # Ensure we got a valid operator (may fully reduce to MatrixOp) - @test L isa AbstractOperator - end +@testitem "Combine branch producing nested Compose is inlined" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + M = MatrixOp(randn(3, 3)) + S = Scale(2.0, Eye(3)) + L = AbstractOperators.Compose((S, M), (zeros(3),)) + @test L * randn(3) isa AbstractVector + x = randn(3) + @test L * x ≈ 2.0 * (M * x) + @test L isa AbstractOperator +end - @testset "remove_slicing: first op is GetIndex, compose of three operators" begin - # First operator is GetIndex, removing slicing should keep equivalent mapping on reduced domain - G = GetIndex((5,), 2:4) # 5 -> 3 - A2 = FiniteDiff((3,)) # 3 -> 2 - A3 = MatrixOp(randn(2, 2)) # 2 -> 2 - L = Compose(A3, Compose(A2, G)) # overall 5 -> 2, internal A = (G, A2, A3) - L2 = AbstractOperators.remove_slicing(L) - @test !is_sliced(L2) - @test size(L2, 2) == (3,) - # Functional equivalence when lifting a reduced-domain vector back to full domain - v = randn(3) - xfull = zeros(5); xfull[2:4] .= v - y1 = L * xfull - y2 = L2 * v - @test y1 ≈ y2 - end +@testitem "remove_slicing first op GetIndex" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + G = GetIndex((5,), 2:4) + A2 = FiniteDiff((3,)) + A3 = MatrixOp(randn(2, 2)) + L = Compose(A3, Compose(A2, G)) + L2 = AbstractOperators.remove_slicing(L) + @test !is_sliced(L2) + @test size(L2, 2) == (3,) + v = randn(3) + xfull = zeros(5) + xfull[2:4] .= v + @test L * xfull ≈ L2 * v +end - @testset "remove_slicing: sliced first op not GetIndex (Scale(GetIndex))" begin - # First operator is Scale(GetIndex), removing slicing should keep equivalent mapping on reduced domain - G = GetIndex((5,), 2:4) # 5 -> 3 - S = 3.0 * G # still 5 -> 3, sliced true - A2 = FiniteDiff((3,)) # 3 -> 2 - L = Compose(A2, S) # overall 5 -> 2, internal A = (S, A2) - L2 = AbstractOperators.remove_slicing(L) - @test !is_sliced(L2) - @test size(L2, 2) == (3,) - # Functional equivalence when lifting a reduced-domain vector back to full domain - v = randn(3) - xfull = zeros(5); xfull[2:4] .= v - y1 = L * xfull - y2 = L2 * v - @test y1 ≈ y2 +@testitem "remove_slicing first op Scale(GetIndex)" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + G = GetIndex((5,), 2:4) + S = 3.0 * G + A2 = FiniteDiff((3,)) + L = Compose(A2, S) + L2 = AbstractOperators.remove_slicing(L) + @test !is_sliced(L2) + @test size(L2, 2) == (3,) + v = randn(3) + xfull = zeros(5) + xfull[2:4] .= v + @test L * xfull ≈ L2 * v + A3 = MatrixOp(randn(2, 2)) + L3 = Compose(A3, L) + L4 = AbstractOperators.remove_slicing(L3) + @test !is_sliced(L4) + @test size(L4, 2) == (3,) + @test L3 * xfull ≈ L4 * v +end - # Compose of three operators with Scale(GetIndex) first - A3 = MatrixOp(randn(2, 2)) # 2 -> 2 - L = Compose(A3, L) # overall 5 -> 2, internal A = (S, A2, A3) - L2 = AbstractOperators.remove_slicing(L) - @test !is_sliced(L2) - @test size(L2, 2) == (3,) - # Functional equivalence when lifting a reduced-domain vector back to full domain - v = randn(3) - xfull = zeros(5); xfull[2:4] .= v - y1 = L * xfull - y2 = L2 * v - @test y1 ≈ y2 - end +@testitem "remove_slicing error path first not sliced" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + op1 = MyLinOp(Float64, (3,), Float64, (3,), (y, x) -> (y .= x), (y, x) -> (y .= x)) + op2 = MatrixOp(randn(2, 3)) + L = AbstractOperators.Compose((op1, op2), (zeros(3),)) + @test_throws ArgumentError AbstractOperators.remove_slicing(L) +end - @testset "remove_slicing error path (first not GetIndex nor sliced)" begin - # Build a Compose manually where the first operator is not sliced and not a GetIndex - op1 = MyLinOp(Float64, (3,), Float64, (3,), (y, x) -> (y .= x), (y, x) -> (y .= x)) - op2 = MatrixOp(randn(2, 3)) - L = AbstractOperators.Compose((op1, op2), (zeros(3),)) - @test_throws ArgumentError AbstractOperators.remove_slicing(L) - end +@testitem "diag_AAc on Compose ok and error" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + d = randn(5) + sel = 1:3 + L_ok = Compose(DiagOp(d[sel]), GetIndex((length(d),), sel)) + @test is_AAc_diagonal(L_ok) + @test AbstractOperators.diag_AAc(L_ok) == diag_AAc(DiagOp(d[sel])) + L_bad = Compose(MatrixOp(randn(3, 3)), MatrixOp(randn(3, 3))) + @test !is_AAc_diagonal(L_bad) + @test_throws ErrorException AbstractOperators.diag_AAc(L_bad) +end - @testset "diag_AAc on Compose (ok and error)" begin - d = randn(5) - sel = 1:3 - L_ok = Compose(DiagOp(d[sel]), GetIndex((length(d),), sel)) - @test is_AAc_diagonal(L_ok) - @test AbstractOperators.diag_AAc(L_ok) == diag_AAc(DiagOp(d[sel])) +@testitem "get_normal_op on Compose fallback path" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + L = Compose(FiniteDiff((3,)), MatrixOp(randn(3, 3))) + N = AbstractOperators.get_normal_op(L) + x = randn(3) + @test N isa AbstractOperator + @test N * x ≈ L' * (L * x) +end - L_bad = Compose(MatrixOp(randn(3, 3)), MatrixOp(randn(3, 3))) - @test !is_AAc_diagonal(L_bad) - @test_throws ErrorException AbstractOperators.diag_AAc(L_bad) - end +@testitem "Scale(coeff, L::Compose) specialized paths" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + Llin = FiniteDiff((3,)) * MatrixOp(randn(3, 3)) * FiniteDiff((4,)) + Slin = Scale(1.7, Llin) + x = randn(4) + @test Slin isa Compose + @test Slin * x ≈ 1.7 * (Llin * x) + Lnl = Compose(FiniteDiff((3,)), Sigmoid(Float64, (3,), 2)) + Snl = Scale(2.0, Lnl) + x2 = randn(3) + @test Snl isa Scale + @test Snl * x2 ≈ 2.0 * (Lnl * x2) +end - @testset "get_normal_op on Compose (fallback path)" begin - # Last operator without optimized normal => else branch in get_normal_op(Compose) - L = Compose(FiniteDiff((3,)), MatrixOp(randn(3, 3))) - N = AbstractOperators.get_normal_op(L) - @test N isa AbstractOperator - x = randn(3) - @test N * x ≈ L' * (L * x) - end +@testitem "get_normal_op(Compose) else branch" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + F1 = FiniteDiff((5,)) + F2 = FiniteDiff((6,)) + L = Compose(F1, F2) + N = AbstractOperators.get_normal_op(L) + x = randn(6) + @test N isa AbstractOperator + @test N * x ≈ L' * (L * x) +end - @testset "Scale(coeff, L::Compose) specialized paths" begin - # Linear case with combinable ops: returns a Compose (not a top-level Scale) - Llin = FiniteDiff((3,)) * MatrixOp(randn(3, 3)) * FiniteDiff((4,)) - Slin = Scale(1.7, Llin) - @test Slin isa Compose - x = randn(4) - @test Slin * x ≈ 1.7 * (Llin * x) +@testitem "Buffer reuse in 4-operator chain" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + F1 = FiniteDiff((10,)) + F2 = MatrixOp(rand(10, 9)) + F3 = FiniteDiff((10,)) + F4 = MatrixOp(rand(10, 10)) + L = F1 * F2 * F3 * F4 + @test L isa Compose + @test length(L * randn(10)) == 9 + @test length(L' * randn(9)) == 10 + @test L.buf[1] === L.buf[3] +end - # Nonlinear inside => wraps as a Scale around the Compose - Lnl = Compose(FiniteDiff((3,)), Sigmoid(Float64, (3,), 2)) - Snl = Scale(2.0, Lnl) - @test Snl isa Scale - x = randn(3) - @test Snl * x ≈ 2.0 * (Lnl * x) +@testitem "DEBUG_COMPOSE logging branches" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + old_debug = AbstractOperators.DEBUG_COMPOSE[] + try + AbstractOperators.DEBUG_COMPOSE[] = true + original_stdout = stdout + (read_pipe, write_pipe) = redirect_stdout() + G = GetIndex((5,), 2:4) + _ = AbstractOperators.Compose((G, G'), (randn(3),)) + D1 = DiagOp(randn(3)) + D2 = DiagOp(randn(3)) + _ = AbstractOperators.Compose((D1, D2), (randn(3),)) + redirect_stdout(original_stdout) + close(write_pipe) + log_str = read(read_pipe, String) + @test occursin("Replacing", log_str) || occursin("Combining", log_str) + finally + AbstractOperators.DEBUG_COMPOSE[] = old_debug end +end - @testset "get_normal_op(Compose) else branch (no optimized normal)" begin - # FiniteDiff and Variation lack optimized normal operators => hit the else path - F1 = FiniteDiff((5,)) # domain 5, codomain 4 - F2 = FiniteDiff((6,)) # domain 6, codomain 5 - L = Compose(F1, F2) # overall: domain 6, codomain 4 - N = AbstractOperators.get_normal_op(L) - @test N isa AbstractOperator +@testitem "Triple combination path" tags = [:calculus, :Compose] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) + struct TripleCombTestOp <: LinearOperator end + LinearAlgebra.size(::TripleCombTestOp) = ((5,), (5,)) + AbstractOperators.domain_type(::TripleCombTestOp) = Float64 + AbstractOperators.codomain_type(::TripleCombTestOp) = Float64 + AbstractOperators.fun_name(::TripleCombTestOp) = "TCT" + AbstractOperators.can_be_combined(::TripleCombTestOp, ::TripleCombTestOp, ::TripleCombTestOp) = true + AbstractOperators.combine(::TripleCombTestOp, ::TripleCombTestOp, ::TripleCombTestOp) = DiagOp(3 .* ones(5)) + AbstractOperators.can_be_combined(::FiniteDiff, ::DiagOp, ::TripleCombTestOp) = true + AbstractOperators.combine(L1::FiniteDiff, L2::DiagOp, ::TripleCombTestOp) = L1 * L2 + AbstractOperators.can_be_combined(::TripleCombTestOp, ::DiagOp, ::FiniteDiff) = true + AbstractOperators.combine(::TripleCombTestOp, L1::DiagOp, L2::FiniteDiff) = L1 * L2 + + T1, T2, T3 = TripleCombTestOp(), TripleCombTestOp(), TripleCombTestOp() + D1, D2 = DiagOp(randn(5)), DiagOp(randn(4)) + F1, F2 = FiniteDiff((5,)), FiniteDiff((6,)) + old_debug = AbstractOperators.DEBUG_COMPOSE[] + try + AbstractOperators.DEBUG_COMPOSE[] = true + original_stdout = stdout + (read_pipe, write_pipe) = redirect_stdout() + + L = T1 * T2 * T3 + @test L isa DiagOp + @test diag(L) == 3 .* ones(5) + L = F1 * T1 * T2 * T3 + x = randn(5) + @test L isa Compose + @test L * x ≈ F1 * (3 .* x) + L = T1 * T2 * T3 * F2 x = randn(6) - # Verify N behaves like L' * L - @test N * x ≈ L' * (L * x) - end - - @testset "Buffer reuse in 4-operator chain (FiniteDiff)" begin - # Chain 4 FiniteDiff operators to exercise buffer reuse adjacency detection - F1 = FiniteDiff((10,)) # domain 10, codomain 9 - F2 = MatrixOp(rand(10, 9)) # domain 9, codomain 10 - F3 = FiniteDiff((10,)) # domain 10, codomain 9 - F4 = MatrixOp(rand(10, 10)) # domain 10, codomain 10 - L = F1 * F2 * F3 * F4 # chains right-to-left: F4->F3->F2->F1, domain 10, codomain 9 @test L isa Compose - x = randn(10) - y = L * x - @test length(y) == 9 - # Verify adjoint chain - r = randn(9) - g = L' * r - @test length(g) == 10 - @test L.buf[1] === L.buf[3] # buffer reuse check - end - - # Enable DEBUG_COMPOSE for combination logging branches - @testset "DEBUG_COMPOSE logging branches" begin - old_debug = AbstractOperators.DEBUG_COMPOSE[] - try - AbstractOperators.DEBUG_COMPOSE[] = true - - # Capture stdout to check logging output - original_stdout = stdout - (read_pipe, write_pipe) = redirect_stdout() - - # Trigger adjacent adjoint optimization with logging - G = GetIndex((5,), 2:4) - L1 = AbstractOperators.Compose((G, G'), (randn(3),)) - - # Trigger 2-arg combination with logging - D1 = DiagOp(randn(3)) - D2 = DiagOp(randn(3)) - L2 = AbstractOperators.Compose((D1, D2), (randn(3),)) - - # Restore stdout and read captured output - redirect_stdout(original_stdout) - close(write_pipe) - log_str = read(read_pipe, String) - - # Verify logging occurred - @test occursin("Replacing", log_str) || occursin("Combining", log_str) - finally - AbstractOperators.DEBUG_COMPOSE[] = old_debug - end + @test L * x ≈ 3 .* (F2 * x) + L = F1 * D1 * T1 + x = randn(5) + @test L isa Compose + @test length(L.A) == 2 + @test L * x ≈ F1 * (D1 * x) + L = F1 * D1 * T1 * F2 + x = randn(6) + @test L isa Compose + @test length(L.A) == 3 + @test L * x ≈ F1 * (D1 * (F2 * x)) + L = D2 * F1 * D1 * T1 + x = randn(5) + @test L isa Compose + @test length(L.A) == 3 + @test L * x ≈ D2 * (F1 * (D1 * x)) + L = D2 * F1 * D1 * T1 * F2 + x = randn(6) + @test L isa Compose + @test length(L.A) == 4 + @test L * x ≈ D2 * (F1 * (D1 * (F2 * x))) + L = F1 * D1 * T1 * D1 + x = randn(5) + @test L isa Compose + @test length(L.A) == 2 + @test L * x ≈ F1 * ((diag(D1) .^ 2) .* x) + L = D1 * T1 * D1 * F2 + x = randn(6) + @test L isa Compose + @test length(L.A) == 2 + @test L * x ≈ ((diag(D1) .^ 2) .* (F2 * x)) + + redirect_stdout(original_stdout) + close(write_pipe) + log_str = read(read_pipe, String) + @test occursin("Replacing", log_str) || occursin("Combining", log_str) + finally + AbstractOperators.DEBUG_COMPOSE[] = old_debug end +end - @testset "Triple combination path (requires 3-arg combine specialization)" begin - - struct TripleCombTestOp <: LinearOperator end - LinearAlgebra.size(::TripleCombTestOp) = ((5,), (5,)) - AbstractOperators.domain_type(::TripleCombTestOp) = Float64 - AbstractOperators.codomain_type(::TripleCombTestOp) = Float64 - AbstractOperators.fun_name(::TripleCombTestOp) = "TCT" - - AbstractOperators.can_be_combined(::TripleCombTestOp, ::TripleCombTestOp, ::TripleCombTestOp) = true - AbstractOperators.combine(::TripleCombTestOp, ::TripleCombTestOp, ::TripleCombTestOp) = DiagOp(3 .* ones(5)) - AbstractOperators.can_be_combined(::FiniteDiff, ::DiagOp, ::TripleCombTestOp) = true - AbstractOperators.combine(L1::FiniteDiff, L2::DiagOp, ::TripleCombTestOp) = L1 * L2 - AbstractOperators.can_be_combined(::TripleCombTestOp, ::DiagOp, ::FiniteDiff) = true - AbstractOperators.combine(::TripleCombTestOp, L1::DiagOp, L2::FiniteDiff) = L1 * L2 - - T1, T2, T3 = TripleCombTestOp(), TripleCombTestOp(), TripleCombTestOp() - D1, D2 = DiagOp(randn(5)), DiagOp(randn(4)) - F1, F2 = FiniteDiff((5,)), FiniteDiff((6,)) - - - old_debug = AbstractOperators.DEBUG_COMPOSE[] - try - AbstractOperators.DEBUG_COMPOSE[] = true - - # Capture stdout to check logging output - original_stdout = stdout - (read_pipe, write_pipe) = redirect_stdout() - - # Compose that triggers triple combination that reduces to DiagOp - L = T1 * T2 * T3 - @test L isa DiagOp - @test diag(L) == 3 .* ones(5) - - L = F1 * T1 * T2 * T3 - @test L isa Compose - x = randn(5) - y = L * x - @test y ≈ F1 * (3 .* x) - - L = T1 * T2 * T3 * F2 - @test L isa Compose - x = randn(6) - y = L * x - @test y ≈ 3 .* (F2 * x) - - # Compose that triggers triple combination that reduces to Compose - L = F1 * D1 * T1 - @test L isa Compose - @test length(L.A) == 2 - x = randn(5) - y = L * x - @test y ≈ F1 * (D1 * x) - - L = F1 * D1 * T1 * F2 - @test L isa Compose - @test length(L.A) == 3 - x = randn(6) - y = L * x - @test y ≈ F1 * (D1 * (F2 * x)) - - L = D2 * F1 * D1 * T1 - @test L isa Compose - @test length(L.A) == 3 - x = randn(5) - y = L * x - @test y ≈ D2 * (F1 * (D1 * x)) - - L = D2 * F1 * D1 * T1 * F2 - @test L isa Compose - @test length(L.A) == 4 - x = randn(6) - y = L * x - @test y ≈ D2 * (F1 * (D1 * (F2 * x))) - - # Double combination from right - L = F1 * D1 * T1 * D1 # (F1 * D1 * T1) -> (F1 * D1), then (F1 * D1) * D1 -> F1 * (D1 * D1) - @test L isa Compose - @test length(L.A) == 2 - x = randn(5) - y = L * x - @test y ≈ F1 * ((diag(D1) .^ 2) .* x) +@testitem "Compose (GPU)" tags = [:gpu, :calculus, :Compose] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv - # Double combination from left - L = D1 * T1 * D1 * F2 # (T1 * D1 * F2) -> (D1 * F2), then D1 * (D1 * F2) -> (D1 * D1) * F2 - @test L isa Compose - @test length(L.A) == 2 - x = randn(6) - y = L * x - @test y ≈ ((diag(D1) .^ 2) .* (F2 * x)) + for backend in gpu_backends() + Random.seed!(0) - # Restore stdout and read captured output - redirect_stdout(original_stdout) - close(write_pipe) - log_str = read(read_pipe, String) + m1, m2 = 4, 7 + A1 = gpu_randn(backend, m2, m1) + opC = Compose(FiniteDiff(gpu_zeros(backend, Float64, m2)), MatrixOp(A1)) + test_op(opC, gpu_randn(backend, m1), gpu_randn(backend, m2 - 1), false) - # Verify logging occurred - @test occursin("Replacing", log_str) || occursin("Combining", log_str) - finally - AbstractOperators.DEBUG_COMPOSE[] = old_debug - end + n = 5 + opC2 = DiagOp(gpu_randn(backend, n)) * DiagOp(gpu_randn(backend, n)) + test_op(opC2, gpu_randn(backend, n), gpu_randn(backend, n), false) end end diff --git a/test/calculus/test_dcat.jl b/test/calculus/test_dcat.jl index 0d4998d..d6da91d 100644 --- a/test/calculus/test_dcat.jl +++ b/test/calculus/test_dcat.jl @@ -1,78 +1,54 @@ -@testitem "DCAT" tags = [:calculus, :DCAT] setup = [TestUtils] begin +@testitem "DCAT: basic mul" tags = [:calculus, :DCAT] setup = [TestUtils] begin using Random, AbstractOperators - import Base: size Random.seed!(0) - verb && println(" --- Testing DCAT --- ") - - m1, n1, m2, n2 = 4, 7, 5, 2 - A1 = randn(m1, n1) - A2 = randn(m2, n2) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opD = DCAT(opA1, opA2) - x1 = randn(n1) - x2 = randn(n2) - y1 = test_op(opD, ArrayPartition(x1, x2), ArrayPartition(randn(m1), randn(m2)), verb) - y2 = ArrayPartition(A1 * x1, A2 * x2) - @test norm(y1 .- y2) .<= 1.0e-12 - - # test DCAT longer m1, n1, m2, n2, m3, n3 = 4, 7, 5, 2, 5, 5 A1 = randn(m1, n1) A2 = randn(m2, n2) A3 = randn(m3, n3) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opA3 = MatrixOp(A3) - opD = DCAT(opA1, opA2, opA3) + opD = DCAT(MatrixOp(A1), MatrixOp(A2), MatrixOp(A3)) x1 = randn(n1) x2 = randn(n2) x3 = randn(n3) y1 = test_op( opD, ArrayPartition(x1, x2, x3), ArrayPartition(randn(m1), randn(m2), randn(m3)), verb ) - y2 = ArrayPartition(A1 * x1, A2 * x2, A3 * x3) - @test norm(y1 .- y2) .<= 1.0e-12 + @test norm(y1 .- ArrayPartition(A1 * x1, A2 * x2, A3 * x3)) .<= 1.0e-12 - #properties + n1, n2 = 4, 7 + opEye = Eye(ArrayPartition(randn(n1), randn(n2))) + @test is_eye(opEye) == true && is_orthogonal(opEye) == true +end + +@testitem "DCAT: properties" tags = [:calculus, :DCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + + m1, n1, m2, n2, m3, n3 = 4, 7, 5, 2, 5, 5 + opD = DCAT(MatrixOp(randn(m1, n1)), MatrixOp(randn(m2, n2)), MatrixOp(randn(m3, n3))) @test is_linear(opD) == true @test is_null(opD) == false @test is_eye(opD) == false @test is_diagonal(opD) == false @test is_AcA_diagonal(opD) == false - @test is_AAc_diagonal(opD) == false - @test is_orthogonal(opD) == false - @test is_invertible(opD) == false @test is_full_row_rank(opD) == false @test is_full_column_rank(opD) == false - # DCAT of Eye - - n1, n2 = 4, 7 - x1 = randn(n1) - x2 = randn(n2) - - opD = Eye(ArrayPartition(x1, x2)) - y1 = test_op(opD, ArrayPartition(x1, x2), ArrayPartition(randn(n1), randn(n2)), verb) - - #properties - @test is_linear(opD) == true - @test is_null(opD) == false - @test is_eye(opD) == true - @test is_diagonal(opD) == true - @test is_AcA_diagonal(opD) == true - @test is_AAc_diagonal(opD) == true - @test is_orthogonal(opD) == true - @test is_invertible(opD) == true - @test is_full_row_rank(opD) == true - @test is_full_column_rank(opD) == true - - @test diag(opD) == 1 - @test diag_AcA(opD) == 1 - @test diag_AAc(opD) == 1 + m1, n1, m2, n2 = 4, 7, 5, 2 + A1 = MatrixOp(randn(m1, n1)) + A2 = MatrixOp(randn(m2, n2)) + op = DCAT(A1, A2) + @test domain_array_type(op) !== nothing + @test codomain_array_type(op) !== nothing + @test remove_displacement(op) == op + opd = DCAT(AffineAdd(A1, randn(m1)), AffineAdd(A2, randn(m2))) + opd_removed = remove_displacement(opd) + @test remove_displacement(opd_removed) == opd_removed +end - # displacement DCAT +@testitem "DCAT: displacement" tags = [:calculus, :DCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) m1, n1, m2, n2, m3, n3 = 4, 7, 5, 2, 5, 5 A1 = randn(m1, n1) @@ -81,113 +57,57 @@ d1 = randn(m1) d2 = randn(m2) d3 = randn(m3) - opA1 = AffineAdd(MatrixOp(A1), d1) - opA2 = AffineAdd(MatrixOp(A2), d2) - opA3 = AffineAdd(MatrixOp(A3), d3) - opD = DCAT(opA1, opA2, opA3) + opD = DCAT( + AffineAdd(MatrixOp(A1), d1), AffineAdd(MatrixOp(A2), d2), AffineAdd(MatrixOp(A3), d3) + ) x1 = randn(n1) x2 = randn(n2) x3 = randn(n3) y1 = opD * ArrayPartition(x1, x2, x3) - y2 = ArrayPartition(A1 * x1 + d1, A2 * x2 + d2, A3 * x3 + d3) - @test norm(y1 .- y2) .<= 1.0e-12 + @test norm(y1 .- ArrayPartition(A1 * x1 + d1, A2 * x2 + d2, A3 * x3 + d3)) .<= 1.0e-12 @test norm(displacement(opD) .- ArrayPartition(d1, d2, d3)) .<= 1.0e-12 - y1 = remove_displacement(opD) * ArrayPartition(x1, x2, x3) - y2 = ArrayPartition(A1 * x1, A2 * x2, A3 * x3) - @test norm(y1 .- y2) .<= 1.0e-12 - - m1, n1, m2, n2 = 4, 7, 5, 2 - A1 = MatrixOp(randn(m1, n1)) - A2 = MatrixOp(randn(m2, n2)) - op = DCAT(A1, A2) - _ds = domain_storage_type(op) - _cs = codomain_storage_type(op) - @test _ds !== nothing - @test _cs !== nothing - @test is_thread_safe(op) == is_thread_safe(A1) && is_thread_safe(A2) - - @test remove_displacement(op) == op - - d1 = randn(m1) - d2 = randn(m2) - opd = DCAT(AffineAdd(A1, d1), AffineAdd(A2, d2)) - opd_removed = remove_displacement(opd) - @test remove_displacement(opd_removed) == opd_removed - - # Test show output (fun_name indirectly) - opA1 = MatrixOp([1.0 0.0; 0.0 1.0]) - opA2 = MatrixOp([2.0 0.0; 0.0 2.0]) - opD2 = DCAT(opA1, opA2) - io = IOBuffer(); show(io, opD2); shown = String(take!(io)) - @test occursin("DCAT", shown) || occursin("[", shown) - - # Test permute - m1, n1, m2, n2 = 2, 2, 2, 2 - opA1 = MatrixOp([1.0 0.0; 0.0 1.0]) - opA2 = MatrixOp([2.0 0.0; 0.0 2.0]) - opD = DCAT(opA1, opA2) - x = ArrayPartition([1.0, 2.0], [3.0, 4.0]) - p = [2, 1] # permute domain - opDp = AbstractOperators.permute(opD, p) - # The permuted operator should act on permuted input - x_perm = ArrayPartition(x.x[2], x.x[1]) - # Just check that permute returns a DCAT and doesn't error - @test typeof(opDp) <: typeof(opD) - - # Test get_normal_op and has_optimized_normalop - struct DCATDummyOp <: AbstractOperator end - AbstractOperators.has_optimized_normalop(::DCATDummyOp) = true - AbstractOperators.get_normal_op(::DCATDummyOp) = DCATDummyOp() - AbstractOperators.size(::DCATDummyOp) = ((2,), (2,)) - opD3 = DCAT(DCATDummyOp(), opA1) - @test AbstractOperators.has_optimized_normalop(opD3) == true - nrm = AbstractOperators.get_normal_op(opD3) - @test typeof(nrm) <: DCAT - - # Test opnorm and estimate_opnorm - opA3 = MatrixOp([1.0 0.0; 0.0 3.0]) - opD4 = DCAT(opA1, opA3) - @test AbstractOperators.has_fast_opnorm(opD4) == (AbstractOperators.has_fast_opnorm(opA1) && AbstractOperators.has_fast_opnorm(opA3)) - @test opnorm(opD4) ≈ estimate_opnorm(opD4) rtol = 0.05 - - # Test == for general equality - opD5a = DCAT(opA1, opA2) - opD5b = DCAT(opA1, opA2) - opD5c = DCAT(opA2, opA1) - @test opD5a == opD5b - @test opD5a != opD5c - - # Test is_thread_safe for mixed thread safety - struct NotThreadSafeOp <: AbstractOperator end - AbstractOperators.is_thread_safe(::NotThreadSafeOp) = false - Base.size(::NotThreadSafeOp) = ((2,), (2,)) - opD6 = DCAT(opA1, NotThreadSafeOp()) - @test is_thread_safe(opD6) == false - - # Explicit diag_AcA/diag_AAc for all-Eye DCAT - opEye = Eye(ArrayPartition([1.0, 2.0], [3.0, 4.0])) - @test diag_AcA(opEye) == 1 - @test diag_AAc(opEye) == 1 + y_rd = remove_displacement(opD) * ArrayPartition(x1, x2, x3) + @test norm(y_rd .- ArrayPartition(A1 * x1, A2 * x2, A3 * x3)) .<= 1.0e-12 +end - # Stacked codomain path: VCAT inside DCAT yields flattened output partition - opV = VCAT(MatrixOp(randn(3, 2)), MatrixOp(randn(4, 2))) - opB = MatrixOp(randn(4, 5)) - dc_stacked = DCAT(opV, opB) - xin = ArrayPartition(randn(2), randn(5)) - yout = dc_stacked * xin - @test length(yout.x) == 3 - @test length((dc_stacked' * yout).x) == 2 +@testitem "DCAT: nonlinear" tags = [:calculus, :DCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # Test nonlinear DCAT n, m = 4, 3 x = ArrayPartition(randn(n), randn(m)) r = ArrayPartition(randn(n), randn(m)) A = randn(n, n) B = Sigmoid(Float64, (m,), 2) op = DCAT(MatrixOp(A), B) - y, grad = test_NLop(op, x, r, verb) + @test norm(ArrayPartition(A * x.x[1], B * x.x[2]) - y) < 1.0e-8 +end - Y = ArrayPartition(A * x.x[1], B * x.x[2]) - @test norm(Y - y) < 1.0e-8 +@testitem "DCAT (GPU)" tags = [:gpu, :calculus, :DCAT] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n1, n2 = 3, 4 + opD = DCAT(DiagOp(gpu_ones(backend, Float64, n1)), DiagOp(2 .* gpu_ones(backend, Float64, n2))) + test_op( + opD, + ArrayPartition(gpu_randn(backend, n1), gpu_randn(backend, n2)), + ArrayPartition(gpu_randn(backend, n1), gpu_randn(backend, n2)), + false, + ) + + m1, n1, m2, n2 = 4, 7, 5, 2 + A1 = gpu_randn(backend, m1, n1) + A2 = gpu_randn(backend, m2, n2) + opD2 = DCAT(MatrixOp(A1), MatrixOp(A2)) + test_op( + opD2, + ArrayPartition(gpu_randn(backend, n1), gpu_randn(backend, n2)), + ArrayPartition(gpu_randn(backend, m1), gpu_randn(backend, m2)), + false, + ) + end end diff --git a/test/calculus/test_hadamardprod.jl b/test/calculus/test_hadamardprod.jl index 9ca432c..7303971 100644 --- a/test/calculus/test_hadamardprod.jl +++ b/test/calculus/test_hadamardprod.jl @@ -1,6 +1,5 @@ -@testitem "HadamardProd" tags = [:calculus, :HadamardProd] setup = [TestUtils] begin +@testitem "HadamardProd: basic mul" tags = [:calculus, :HadamardProd] setup = [TestUtils] begin using AbstractOperators - verb && println(" --- Testing HadamardProd --- ") # Basic square identity factors (Eye.*Eye) n = 3 @@ -34,6 +33,24 @@ P = HadamardProd(op1, op2) y, grad = test_NLop(P, x, r, verb) @test norm((op1 * x) .* (op2 * x) - y) < 1.0e-9 +end + +@testitem "HadamardProd: properties" tags = [:calculus, :HadamardProd] setup = [TestUtils] begin + using AbstractOperators + + # Re-create the HCAT-based P for remove_displacement and permute tests + m, n = 3, 5 + x = ArrayPartition(randn(m), randn(n)) + r = randn(m) + b = randn(m) + A1 = AffineAdd(Sin(Float64, (m,)), b) + B1 = MatrixOp(randn(m, n)) + op1 = HCAT(A1, B1) + C1 = Cos(Float64, (m,)) + D1 = MatrixOp(randn(m, n)) + op2 = HCAT(C1, D1) + P = HadamardProd(op1, op2) + y, grad = test_NLop(P, x, r, verb) # remove_displacement and its idempotence y2, grad2 = test_NLop(remove_displacement(P), x, r, verb) @@ -51,8 +68,8 @@ @test_throws Exception HadamardProd(Eye(2, 2, 2), Eye(1, 2, 2)) # Storage type / thread safety accessors - _ds = domain_storage_type(P) - _cs = codomain_storage_type(P) + _ds = domain_array_type(P) + _cs = codomain_array_type(P) @test _ds !== nothing @test _cs !== nothing @test is_thread_safe(P) == false @@ -62,7 +79,11 @@ show(io, P) str = String(take!(io)) @test occursin(".*", str) - # --- Additional coverage --- +end + +@testitem "HadamardProd: equality and permute" tags = [:calculus, :HadamardProd] setup = [TestUtils] begin + using AbstractOperators + # Equality / inequality n = 3 A = Eye(n, n) @@ -78,21 +99,26 @@ @test size(P1) == ((n, n), (n, n)) @test domain_type(P1) == domain_type(A) @test codomain_type(P1) == codomain_type(A) - @test domain_storage_type(P1) !== nothing - @test codomain_storage_type(P1) !== nothing + @test domain_array_type(P1) !== nothing + @test codomain_array_type(P1) !== nothing # fun_name direct - io = IOBuffer(); show(io, P1); sP1 = String(take!(io)) + io = IOBuffer() + show(io, P1) + sP1 = String(take!(io)) @test occursin(".*", sP1) # permute with more than 2 domains (using HCAT) - mH = 4; n1 = 2; n2 = 2 + mH = 4 + n1 = 2 + n2 = 2 A1p = MatrixOp(randn(mH, n1)) A2p = MatrixOp(randn(mH, n2)) H1 = HCAT(A1p, A2p) H2 = HCAT(A2p, A1p) P = HadamardProd(H1, H2) - x1p = randn(n1); x2p = randn(n2) + x1p = randn(n1) + x2p = randn(n2) y_orig, _ = test_NLop(P, ArrayPartition(x1p, x2p), randn(mH), verb) p = [2, 1] Pp = AbstractOperators.permute(P, p) @@ -105,3 +131,47 @@ Prd = remove_displacement(Pdisp) @test remove_displacement(Prd) == Prd end + +@testitem "HadamardProd (GPU)" tags = [:gpu, :calculus, :HadamardProd] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n = 3 + P = HadamardProd( + Eye(Float64, (n, n); array_type = gpu_wrapper(backend, Float64, n, n)), + Eye(Float64, (n, n); array_type = gpu_wrapper(backend, Float64, n, n)), + ) + x = gpu_randn(backend, n, n) + r = gpu_randn(backend, n, n) + test_NLop_gpu(P, x, r, false) + + n2, l = 3, 2 + P2 = HadamardProd(Sin(gpu_zeros(backend, Float64, n2, l)), Cos(gpu_zeros(backend, Float64, n2, l))) + x2 = gpu_randn(backend, n2, l) + r2 = gpu_randn(backend, n2, l) + test_NLop_gpu(P2, x2, r2, false) + end +end + +@testitem "HadamardProd: copy_operator" tags = [:calculus, :HadamardProd] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(3) + + n = 6 + P = HadamardProd(Sin((n,)), Cos((n,))) + P2 = copy_operator(P) + @test P2 isa HadamardProd + x = randn(n) + y1 = zeros(n) + y2 = zeros(n) + mul!(y1, P, x) + mul!(y2, P2, x) + @test y1 ≈ y2 + # Verify buffer independence — a second forward should not cross-contaminate + x2 = randn(n) + y3 = zeros(n) + mul!(y3, P2, x2) + @test y3 ≈ P * x2 +end diff --git a/test/calculus/test_hcat.jl b/test/calculus/test_hcat.jl index 7ba9e5d..c23e6e5 100644 --- a/test/calculus/test_hcat.jl +++ b/test/calculus/test_hcat.jl @@ -1,7 +1,6 @@ -@testitem "HCAT" tags = [:calculus, :HCAT] setup = [TestUtils] begin +@testitem "HCAT: basic mul" tags = [:calculus, :HCAT] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing HCAT --- ") m, n1, n2 = 4, 7, 5 A1 = randn(m, n1) @@ -21,8 +20,6 @@ y1 = test_op(opHp, ArrayPartition(x2, x1), randn(m), verb) @test norm(y1 - y2) <= 1.0e-12 - # test HCAT longer - m, n1, n2, n3 = 4, 7, 5, 6 A1 = randn(m, n1) A2 = randn(m, n2) @@ -35,49 +32,22 @@ x2 = randn(n2) x3 = randn(n3) y1 = test_op(opH, ArrayPartition(x1, x2, x3), randn(m), verb) - y2 = A1 * x1 + A2 * x2 + A3 * x3 - @test norm(y1 - y2) <= 1.0e-12 + @test norm(y1 - (A1 * x1 + A2 * x2 + A3 * x3)) <= 1.0e-12 - # test HCAT of HCAT + # HCAT of HCAT (flattening) opHH = HCAT(opH, opA2, opA3) y1 = test_op(opHH, ArrayPartition(x1, x2, x3, x2, x3), randn(m), verb) - y2 = A1 * x1 + A2 * x2 + A3 * x3 + A2 * x2 + A3 * x3 - @test norm(y1 - y2) <= 1.0e-12 - - opHH = HCAT(opH, opH, opA3) - x = ArrayPartition(x1, x2, x3, x1, x2, x3, x3) - y1 = test_op(opHH, x, randn(m), verb) - y2 = A1 * x1 + A2 * x2 + A3 * x3 + A1 * x1 + A2 * x2 + A3 * x3 + A3 * x3 - @test norm(y1 - y2) <= 1.0e-12 - - opA3 = MatrixOp(randn(n1, n1)) - @test_throws Exception HCAT(opA1, opA2, opA3) - opF = MatrixOp(randn(ComplexF64, m, m)) - @test_throws Exception HCAT(opA1, opF, opA2) - - # test utilities - - # permutation - p = randperm(ndoms(opHH, 2)) - opHP = AbstractOperators.permute(opHH, p) - - xp = ArrayPartition(x.x[p]...) - - y1 = test_op(opHP, xp, randn(m), verb) + @test norm(y1 - (A1 * x1 + A2 * x2 + A3 * x3 + A2 * x2 + A3 * x3)) <= 1.0e-12 +end - pp = randperm(ndoms(opHH, 2)) - opHPP = AbstractOperators.permute(opHH, pp) - xpp = ArrayPartition(x.x[pp]...) - y1 = test_op(opHPP, xpp, randn(m), verb) +@testitem "HCAT: properties" tags = [:calculus, :HCAT] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) - #properties m, n1, n2, n3 = 4, 7, 5, 6 - A1 = randn(m, n1) - A2 = randn(m, n2) - A3 = randn(m, n3) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opA3 = MatrixOp(A3) + opA1 = MatrixOp(randn(m, n1)) + opA2 = MatrixOp(randn(m, n2)) + opA3 = MatrixOp(randn(m, n3)) op = HCAT(opA1, opA2, opA3) @test is_linear(op) == true @test is_null(op) == false @@ -92,23 +62,24 @@ d1 = randn(n1) .+ im .* randn(n1) d2 = randn(n1) .+ im .* randn(n1) - op = HCAT(DiagOp(d1), DiagOp(d2)) - @test is_null(op) == false - @test is_eye(op) == false - @test is_diagonal(op) == false - @test is_AcA_diagonal(op) == false - @test is_AAc_diagonal(op) == true - @test is_orthogonal(op) == false - @test is_invertible(op) == false - @test is_full_row_rank(op) == true - @test is_full_column_rank(op) == false - - @test diag_AAc(op) == d1 .* conj(d1) .+ d2 .* conj(d2) - + op2 = HCAT(DiagOp(d1), DiagOp(d2)) + @test is_AAc_diagonal(op2) == true + @test diag_AAc(op2) == d1 .* conj(d1) .+ d2 .* conj(d2) y1 = randn(n1) .+ im .* randn(n1) - @test norm(op * (op' * y1) .- diag_AAc(op) .* y1) < 1.0e-12 + @test norm(op2 * (op2' * y1) .- diag_AAc(op2) .* y1) < 1.0e-12 + + # storage type and thread safety + A1 = MatrixOp(randn(m, n1)) + A2 = MatrixOp(randn(m, n2)) + op3 = HCAT(A1, A2) + @test domain_array_type(op3) !== nothing + @test codomain_array_type(op3) !== nothing + @test is_thread_safe(op3) == false +end - #test displacement +@testitem "HCAT: displacement" tags = [:calculus, :HCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) m, n1, n2 = 4, 7, 5 A1 = randn(m, n1) @@ -121,52 +92,24 @@ x1 = randn(n1) x2 = randn(n2) y1 = opH * ArrayPartition(x1, x2) - y2 = A1 * x1 + d1 + A2 * x2 + d2 - @test norm(y1 - y2) <= 1.0e-12 - y1 = remove_displacement(opH) * ArrayPartition(x1, x2) - y2 = A1 * x1 + A2 * x2 - @test norm(y1 - y2) <= 1.0e-12 - - m, n1, n2 = 4, 7, 5 - A1 = MatrixOp(randn(m, n1)) - A2 = MatrixOp(randn(m, n2)) - op = HCAT(A1, A2) - # storage type accessors (results not asserted for value, just exercise paths) - _ds = domain_storage_type(op) - _cs = codomain_storage_type(op) - @test _ds !== nothing - @test _cs !== nothing - # thread safety (expected false because of shared output accumulation) - @test is_thread_safe(op) == false - - # remove_displacement idempotence (no displacement present) + @test norm(y1 - (A1 * x1 + d1 + A2 * x2 + d2)) <= 1.0e-12 + y2 = remove_displacement(opH) * ArrayPartition(x1, x2) + @test norm(y2 - (A1 * x1 + A2 * x2)) <= 1.0e-12 + + # remove_displacement idempotence + A1b = MatrixOp(randn(m, n1)) + A2b = MatrixOp(randn(m, n2)) + op = HCAT(A1b, A2b) @test remove_displacement(op) == op - - # with displacement once - d1 = randn(m) - d2 = randn(m) - opd = HCAT(AffineAdd(A1, d1), AffineAdd(A2, d2)) + opd = HCAT(AffineAdd(A1b, d1), AffineAdd(A2b, d2)) opd_removed = remove_displacement(opd) - # applying again should be a no-op @test remove_displacement(opd_removed) == opd_removed +end - # Show output / fun_name indirect check (two-operator bracket form) - E = Eye(5) - D = DiagOp(2 .* ones(5)) - H2 = HCAT(E, D) - io = IOBuffer(); show(io, H2); shown = String(take!(io)) - @test occursin("HCAT", shown) || occursin("[", shown) - - # Equality and inequality - Aeq = MatrixOp(randn(4, 3)) - Beq = MatrixOp(randn(4, 2)) - H1a = HCAT(Aeq, Beq) - H1b = HCAT(Aeq, Beq) - H1c = HCAT(Beq, Aeq) - @test H1a == H1b - @test H1a != H1c +@testitem "HCAT: slicing and permute utilities" tags = [:calculus, :HCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # Slicing related helpers using GetIndex n = 8 op1 = GetIndex(Float64, (n,), (1:4,)) op2 = GetIndex(Float64, (n,), (5:8,)) @@ -174,136 +117,122 @@ @test is_sliced(Hs) == true exprs = AbstractOperators.get_slicing_expr(Hs) @test length(exprs) == 2 - @test exprs[1] == (1:4,) - @test exprs[2] == (5:8,) + @test exprs[1] == (1:4,) && exprs[2] == (5:8,) masks = AbstractOperators.get_slicing_mask(Hs) @test length(masks) == 2 @test sum(masks[1]) == 4 && sum(masks[2]) == 4 - Hs_removed = AbstractOperators.remove_slicing(Hs) - @test is_sliced(Hs_removed) == false + @test !is_sliced(AbstractOperators.remove_slicing(Hs)) - # Slicing expressions for HCAT of sliced Compose operators d1, d2 = randn(5), randn(5) D1 = DiagOp(d1) * GetIndex((10,), 1:5) D2 = DiagOp(d2) * GetIndex((10,), 6:10) Hs_comp = HCAT(D1, D2) @test is_sliced(Hs_comp) @test AbstractOperators.get_slicing_expr(Hs_comp) == ((1:5,), (6:10,)) - @test length(collect(AbstractOperators.get_slicing_mask(Hs_comp))) == 2 @test !is_sliced(AbstractOperators.remove_slicing(Hs_comp)) - # Show path with reversed domain indices - Hrev = HCAT(opA2, opA1) - Hrev_perm = Hrev[[2, 1]] - io2 = IOBuffer(); show(io2, Hrev_perm); s2 = String(take!(io2)) - @test occursin("[", s2) - - # Stacked second operator to exercise generated mul! stacked branch - Hflat = MatrixOp(randn(5, 2)) - Hstacked_part = HCAT(MatrixOp(randn(5, 3)), MatrixOp(randn(5, 4))) - Hmix = HCAT((Hflat, Hstacked_part), zeros(5)) - xmix = ArrayPartition(randn(2), randn(3), randn(4)) - @test Hmix * xmix ≈ (Hflat * xmix.x[1]) + (Hstacked_part * ArrayPartition(xmix.x[2], xmix.x[3])) - - - # Permute domain ordering and ensure type stability + m, n1, n2 = 4, 3, 2 + Aeq = MatrixOp(randn(m, n1)) + Beq = MatrixOp(randn(m, n2)) + H1a = HCAT(Aeq, Beq) p2 = collect(Iterators.reverse(1:ndoms(H1a, 2))) Hp = AbstractOperators.permute(H1a, p2) @test typeof(Hp) <: HCAT - xA = randn(size(Aeq, 2)); xB = randn(size(Beq, 2)) + xA = randn(size(Aeq, 2)) + xB = randn(size(Beq, 2)) y_orig = H1a * ArrayPartition(xA, xB) xin = p2 == [2, 1] ? ArrayPartition(xB, xA) : ArrayPartition(xA, xB) - y_perm = Hp * xin - @test y_orig ≈ y_perm + @test y_orig ≈ Hp * xin +end - # diag_AAc accumulation explicit check on known simple ops (DiagOp and Eye) - dvals = randn(5) .+ im .* randn(5) - Hdiag = HCAT(DiagOp(dvals), Eye(ComplexF64, 5)) - @test diag_AAc(Hdiag) == dvals .* conj(dvals) .+ 1 +@testitem "HCAT: nonlinear operators" tags = [:calculus, :HCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # nonlinear operator test (Sigmoid + MatrixOp) n, m = 4, 3 x = ArrayPartition(randn(n), randn(m)) r = randn(m) A = randn(m, n) B = Sigmoid(Float64, (m,), 2) op = HCAT(MatrixOp(A), B) - y, grad = test_NLop(op, x, r, verb) + @test norm(A * x.x[1] + B * x.x[2] - y) < 1.0e-8 - Y = A * x.x[1] + B * x.x[2] - @test norm(Y - y) < 1.0e-8 - - m, n = 3, 5 + n, m = 5, 3 x = ArrayPartition(randn(m), randn(n)) r = randn(m) - A = Sin(Float64, (m,)) + A_sin = Sin(Float64, (m,)) M = randn(m, n) - B = MatrixOp(M) - op = HCAT(A, B) + op2 = HCAT(A_sin, MatrixOp(M)) + y2, grad2 = test_NLop(op2, x, r, verb) + @test norm(A_sin * x.x[1] + M * x.x[2] - y2) < 1.0e-8 +end - y, grad = test_NLop(op, x, r, verb) +@testitem "HCAT constructor errors" tags = [:calculus, :HCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - Y = A * x.x[1] + M * x.x[2] - @test norm(Y - y) < 1.0e-8 - - p = [2, 1] - opP = AbstractOperators.permute(op, p) - xp = ArrayPartition(x.x[p]...) - J = Jacobian(opP, xp)' - verb && println(size(J, 1)) - y, grad = test_NLop(opP, xp, r, verb) - - # Test HCAT constructor error paths - @testset "HCAT constructor errors" begin - # DimensionMismatch: operators with different codomain dimensions - A1 = MatrixOp(randn(4, 3)) - A2 = MatrixOp(randn(5, 2)) # Different codomain dimension (5 vs 4) - @test_throws DimensionMismatch HCAT(A1, A2) - - # codomain_type mismatch: real vs complex - A1 = MatrixOp(randn(4, 3)) - A2 = MatrixOp(randn(ComplexF64, 4, 2)) - @test_throws Exception HCAT(A1, A2) - end + A1 = MatrixOp(randn(4, 3)) + A2 = MatrixOp(randn(5, 2)) + @test_throws DimensionMismatch HCAT(A1, A2) - # Test HCAT with nested HCAT (flattening behavior) - @testset "HCAT flattening" begin - m, n1, n2, n3 = 4, 3, 2, 5 - A1 = MatrixOp(randn(m, n1)) - A2 = MatrixOp(randn(m, n2)) - A3 = MatrixOp(randn(m, n3)) - - # Create nested HCAT - H1 = HCAT(A1, A2) - H2 = HCAT(H1, A3) # This should flatten - - # Test that it works correctly - x1, x2, x3 = randn(n1), randn(n2), randn(n3) - y = H2 * ArrayPartition(x1, x2, x3) - y_expected = A1 * x1 + A2 * x2 + A3 * x3 - @test norm(y - y_expected) < 1.0e-12 - - # Test adjoint also works - y_test = randn(m) - x_adj = H2' * y_test - @test length(x_adj.x) == 3 # ArrayPartition has .x field - - # More complex nesting: HCAT(HCAT(...), HCAT(...)) - H3 = HCAT(A1, A2) - H4 = HCAT(A2, A3) - H5 = HCAT(H3, H4) # Should flatten all - - x_full = ArrayPartition(x1, x2, x2, x3) - y2 = H5 * x_full - y2_expected = A1 * x1 + A2 * x2 + A2 * x2 + A3 * x3 - @test norm(y2 - y2_expected) < 1.0e-12 - end + A1 = MatrixOp(randn(4, 3)) + A2 = MatrixOp(randn(ComplexF64, 4, 2)) + @test_throws Exception HCAT(A1, A2) +end + +@testitem "HCAT flattening" tags = [:calculus, :HCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + + m, n1, n2, n3 = 4, 3, 2, 5 + A1 = MatrixOp(randn(m, n1)) + A2 = MatrixOp(randn(m, n2)) + A3 = MatrixOp(randn(m, n3)) + H1 = HCAT(A1, A2) + H2 = HCAT(H1, A3) + + x1, x2, x3 = randn(n1), randn(n2), randn(n3) + y = H2 * ArrayPartition(x1, x2, x3) + y_expected = A1 * x1 + A2 * x2 + A3 * x3 + @test norm(y - y_expected) < 1.0e-12 + + y_test = randn(m) + x_adj = H2' * y_test + @test length(x_adj.x) == 3 + + H3 = HCAT(A1, A2) + H4 = HCAT(A2, A3) + H5 = HCAT(H3, H4) + x_full = ArrayPartition(x1, x2, x2, x3) + y2 = H5 * x_full + y2_expected = A1 * x1 + A2 * x2 + A2 * x2 + A3 * x3 + @test norm(y2 - y2_expected) < 1.0e-12 +end + +@testitem "HCAT single operator" tags = [:calculus, :HCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + + A = MatrixOp(randn(4, 3)) + H_single = HCAT(A) + @test H_single === A +end + +@testitem "HCAT (GPU)" tags = [:gpu, :calculus, :HCAT] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n = 4 + opH = HCAT(DiagOp(gpu_ones(backend, Float64, n)), DiagOp(2 .* gpu_ones(backend, Float64, n))) + test_op(opH, ArrayPartition(gpu_randn(backend, n), gpu_randn(backend, n)), gpu_randn(backend, n), false) - # Test single operator HCAT (should return the operator itself) - @testset "HCAT single operator" begin - A = MatrixOp(randn(4, 3)) - H_single = HCAT(A) - @test H_single === A # Should return the same operator + m, n1, n2 = 4, 7, 5 + A1 = gpu_randn(backend, m, n1) + A2 = gpu_randn(backend, m, n2) + opH2 = HCAT(MatrixOp(A1), MatrixOp(A2)) + test_op(opH2, ArrayPartition(gpu_randn(backend, n1), gpu_randn(backend, n2)), gpu_randn(backend, m), false) end end diff --git a/test/calculus/test_operatorwrapper.jl b/test/calculus/test_operatorwrapper.jl new file mode 100644 index 0000000..b197ece --- /dev/null +++ b/test/calculus/test_operatorwrapper.jl @@ -0,0 +1,88 @@ +@testitem "OperatorWrapper: CPU mul! forward and adjoint" tags = [ + :calculus, :OperatorWrapper, +] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(42) + + n = 8 + cpu_op = FiniteDiff(Float64, (n,), 1) # n-1 output + wrapper = OperatorWrapper(cpu_op) + + x = randn(n) + y = zeros(n - 1) + mul!(y, wrapper, x) + y_ref = zeros(n - 1) + mul!(y_ref, cpu_op, x) + @test y ≈ y_ref + + b = randn(n - 1) + z = zeros(n) + mul!(z, wrapper', b) + z_ref = zeros(n) + mul!(z_ref, cpu_op', b) + @test z ≈ z_ref + + @test size(wrapper) == size(cpu_op) + @test domain_type(wrapper) == domain_type(cpu_op) + @test codomain_type(wrapper) == codomain_type(cpu_op) + @test is_linear(wrapper) == is_linear(cpu_op) + @test domain_array_type(wrapper) == Array{Float64} + @test codomain_array_type(wrapper) == Array{Float64} + + io = IOBuffer() + show(io, wrapper) + @test occursin("CPU[", String(take!(io))) +end + +@testitem "OperatorWrapper: properties forwarded" tags = [:calculus, :OperatorWrapper] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(42) + + n = 6 + op = MatrixOp(randn(n, n)) + wrapper = OperatorWrapper(op) + @test is_full_row_rank(wrapper) == is_full_row_rank(op) + @test is_full_column_rank(wrapper) == is_full_column_rank(op) + @test is_thread_safe(wrapper) == false # always false regardless of inner op +end + +@testitem "OperatorWrapper: CPU mul! via test_op" tags = [:calculus, :OperatorWrapper] setup = [ + TestUtils, +] begin + using Random, AbstractOperators + Random.seed!(42) + + n = 6 + op = MatrixOp(randn(n, n)) + wrapper = OperatorWrapper(op) + + test_op(wrapper, randn(n), randn(n), verb) +end + +@testitem "OperatorWrapper (GPU)" tags = [:gpu, :calculus, :OperatorWrapper] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(42) + n = 32 + op = FiniteDiff(Float32, (n,), 1) + array_type = gpu_wrapper(backend, Float32, n) + wrapper = OperatorWrapper(op; array_type = array_type) + @test domain_array_type(wrapper) <: backend.array_type + @test codomain_array_type(wrapper) <: backend.array_type + + x = gpu_randn(backend, Float32, n) + y = gpu_zeros(backend, Float32, n - 1) + mul!(y, wrapper, x) + @test y isa array_type + @test collect(y) ≈ op * collect(x) + + r = gpu_randn(backend, Float32, n - 1) + z = gpu_zeros(backend, Float32, n) + mul!(z, wrapper', r) + @test z isa array_type + ref = zeros(Float32, n) + mul!(ref, op', collect(r)) + @test collect(z) ≈ ref + end +end diff --git a/test/calculus/test_reshape.jl b/test/calculus/test_reshape.jl index 4d921a3..942458a 100644 --- a/test/calculus/test_reshape.jl +++ b/test/calculus/test_reshape.jl @@ -1,7 +1,6 @@ -@testitem "Reshape" tags = [:calculus, :Reshape] setup = [TestUtils] begin +@testitem "Reshape: basic 1D->2D" tags = [:calculus, :Reshape] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Reshape --- ") m, n = 8, 4 dim_out = (2, 2, 2) @@ -25,6 +24,11 @@ @test is_invertible(opR) == is_invertible(opA1) @test is_full_row_rank(opR) == is_full_row_rank(opA1) @test is_full_column_rank(opR) == is_full_column_rank(opA1) +end + +@testitem "Reshape: displacement and storage" tags = [:calculus, :Reshape] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) # testing displacement m, n = 8, 4 @@ -42,19 +46,22 @@ @test norm(y1 - y2) <= 1.0e-12 # fun_name / storage / thread safety / idempotent remove - io = IOBuffer(); show(io, opR); s = String(take!(io)) + io = IOBuffer() + show(io, opR) + s = String(take!(io)) @test length(s) > 0 - _dst = domain_storage_type(opR) - _cst = codomain_storage_type(opR) + _dst = domain_array_type(opR) + _cst = codomain_array_type(opR) @test _dst !== nothing && _cst !== nothing @test is_thread_safe(opR) == is_thread_safe(opA1) rd1 = remove_displacement(opR) rd2 = remove_displacement(rd1) @test rd1 * x1 == rd2 * x1 +end - ####################### - ## test Scale ####### - ####################### +@testitem "Reshape: Scale mul" tags = [:calculus, :Reshape] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) m, n = 8, 4 coeff = pi @@ -78,8 +85,18 @@ y1 = test_op(opS, x1, diff(randn(m)), verb) y2 = coeff * (diff(x1)) @test norm(y1 - y2) <= 1.0e-12 +end +@testitem "Reshape: Scale properties" tags = [:calculus, :Reshape] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + + m, n = 8, 4 + coeff = pi + A1 = randn(m, n) + opA1 = MatrixOp(A1) opS = Scale(coeff, opA1) + @test is_null(opS) == is_null(opA1) @test is_eye(opS) == is_eye(opA1) @test is_diagonal(opS) == is_diagonal(opA1) @@ -125,6 +142,11 @@ y1 = remove_displacement(opS) * x1 y2 = coeff * (A1 * x1) @test norm(y1 - y2) <= 1.0e-12 +end + +@testitem "Reshape: equality and adjoint" tags = [:calculus, :Reshape] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) # Equality / inequality m, n = 8, 4 @@ -145,11 +167,14 @@ @test lhs ≈ rhs # fun_name via show should start with paragraph symbol - io = IOBuffer(); show(io, Rf); sR = String(take!(io)) + io = IOBuffer() + show(io, Rf) + sR = String(take!(io)) @test occursin("¶", sR) # has_optimized_normalop + get_normal_op passthrough (using GetIndex which has optimized normal) - nGI = 10; kGI = 7 + nGI = 10 + kGI = 7 GI = GetIndex(Float64, (nGI,), (1:kGI,)) # sliced operator returning size kGI RG = Reshape(GI, (kGI, 1)) normal_RG = AbstractOperators.get_normal_op(RG) @@ -161,14 +186,22 @@ @test exprRG == (1:kGI,) maskRG = AbstractOperators.get_slicing_mask(GI) @test sum(maskRG) == kGI +end + +@testitem "Reshape: permute and nonlinear" tags = [:calculus, :Reshape] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) # permute domain ordering (wrap HCAT to get multi-domain) and ensure same behavior when inputs permuted - mH = 6; n1 = 3; n2 = 5 + mH = 6 + n1 = 3 + n2 = 5 A1p = MatrixOp(randn(mH, n1)) A2p = MatrixOp(randn(mH, n2)) H = HCAT(A1p, A2p) RH = Reshape(H, (2, 3)) # 6 = 2*3 - x1p = randn(n1); x2p = randn(n2) + x1p = randn(n1) + x2p = randn(n2) y_orig = RH * ArrayPartition(x1p, x2p) p = [2, 1] RHp = AbstractOperators.permute(RH, p) @@ -176,6 +209,10 @@ @test y_orig ≈ y_perm # opnorm passthrough + m, n = 8, 4 + dim_out = (2, 2, 2) + Aeq = MatrixOp(randn(m, n)) + R1 = Reshape(Aeq, dim_out) @test AbstractOperators.has_fast_opnorm(R1) == AbstractOperators.has_fast_opnorm(Aeq) @test opnorm(R1) ≈ opnorm(Aeq) @@ -197,3 +234,21 @@ Y = reshape(opS * x, 2, 2) @test norm(Y - y) < 1.0e-8 end + +@testitem "Reshape (GPU)" tags = [:gpu, :calculus, :Reshape] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + m, n = 8, 4 + dim_out = (2, 2, 2) + A1 = gpu_randn(backend, m, n) + opR = Reshape(MatrixOp(A1), dim_out) + test_op(opR, gpu_randn(backend, n), gpu_randn(backend, dim_out...), false) + + n2 = 6 + opR2 = Reshape(DiagOp(gpu_randn(backend, n2)), (2, 3)) + test_op(opR2, gpu_randn(backend, n2), gpu_randn(backend, 2, 3), false) + end +end diff --git a/test/calculus/test_scale.jl b/test/calculus/test_scale.jl index 62fc7bd..0841910 100644 --- a/test/calculus/test_scale.jl +++ b/test/calculus/test_scale.jl @@ -1,7 +1,6 @@ @testitem "Scale" tags = [:calculus, :Scale] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Scale --- ") m, n = 8, 4 coeff = pi @@ -102,11 +101,14 @@ @test S1 != S3 # fun_name via show should start with α - io = IOBuffer(); show(io, S1); sS = String(take!(io)) + io = IOBuffer() + show(io, S1) + sS = String(take!(io)) @test occursin("α", sS) # has_optimized_normalop + get_normal_op passthrough (using GetIndex which has optimized normal) - nGI = 8; kGI = 5 + nGI = 8 + kGI = 5 GI = GetIndex(Float64, (nGI,), (1:kGI,)) SG = Scale(2.0, GI) @test AbstractOperators.has_optimized_normalop(SG) == true @@ -121,12 +123,15 @@ @test sum(maskSG) == kGI # permute domain ordering (wrap HCAT to get multi-domain) and ensure same behavior when inputs permuted - mH = 6; n1 = 2; n2 = 1 + mH = 6 + n1 = 2 + n2 = 1 A1p = MatrixOp(randn(mH, n1)) A2p = MatrixOp(randn(mH, n2)) H = HCAT(A1p, A2p) SH = Scale(2.0, H) - x1p = randn(n1); x2p = randn(n2) + x1p = randn(n1) + x2p = randn(n2) y_orig = SH * ArrayPartition(x1p, x2p) p = [2, 1] SHp = AbstractOperators.permute(SH, p) @@ -167,154 +172,163 @@ Y = -A * x @test norm(Y - y) < 1.0e-8 - @testset "Scale constructors and basic mapping" begin - # Base operator - A = MatrixOp(randn(6, 4)) - α = 2.5 - S = Scale(α, A) - @test size(S) == size(A) - x = randn(4) - y_ref = α * (A * x) - y = S * x - @test y ≈ y_ref - # Adjoint mapping - y2_ref = α * (A * x) - @test S * x ≈ y2_ref - z = randn(6) - # Adjoint operator action (uses coeff_conj though real here) - @test S' * z ≈ α * (A' * z) - end +end - @testset "Scale special constructor scale-of-scale" begin - A = FiniteDiff((10, 2)) - α = 3.0 - β = -2.0 - S1 = Scale(α, A) - S2 = Scale(β, S1) # should multiply coefficients and unwrap - x = randn(10, 2) - # FiniteDiff maps (10,2)->(9,2), ensure shape works - v1 = S1 * x - v2 = S2 * x - @test v1 ≈ α * (A * x) - @test v2 ≈ (β * α) * (A * x) - end +@testitem "Scale constructors and basic mapping" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = MatrixOp(randn(6, 4)) + α = 2.5 + S = Scale(α, A) + @test size(S) == size(A) + x = randn(4) + @test S * x ≈ α * (A * x) + z = randn(6) + @test S' * z ≈ α * (A' * z) +end - @testset "Scale coeff==1 returns original" begin - A = MatrixOp(randn(5, 5)) - S = Scale(1.0, A) - @test S === A || S == A - # Ensure behavior matches - x = randn(5) - @test (S * x) ≈ (A * x) - end +@testitem "Scale special constructor scale-of-scale" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = FiniteDiff((10, 2)) + α = 3.0 + β = -2.0 + S1 = Scale(α, A) + S2 = Scale(β, S1) + x = randn(10, 2) + @test S1 * x ≈ α * (A * x) + @test S2 * x ≈ (β * α) * (A * x) +end - @testset "Scale properties delegation" begin - A = Eye(7) - α = 4.2 - S = Scale(α, A) - @test domain_type(S) == domain_type(A) - @test codomain_type(S) == codomain_type(A) - @test domain_storage_type(S) == domain_storage_type(A) - @test codomain_storage_type(S) == codomain_storage_type(A) - @test is_thread_safe(S) == is_thread_safe(A) - @test is_linear(S) - @test !is_null(S) - @test is_diagonal(S) == is_diagonal(A) - @test is_invertible(S) == is_invertible(A) - @test AbstractOperators.fun_name(S) == "α" * AbstractOperators.fun_name(A) - # diag and diag_AcA / diag_AAc for Eye should be scalar scaling (Eye diag returns ones vector) - @test AbstractOperators.diag(S) == α * AbstractOperators.diag(A) - @test AbstractOperators.diag_AcA(S) == α^2 * AbstractOperators.diag_AcA(A) - @test AbstractOperators.diag_AAc(S) == α^2 * AbstractOperators.diag_AAc(A) - @test is_full_row_rank(S) == is_full_row_rank(A) - @test is_full_column_rank(S) == is_full_column_rank(A) - end +@testitem "Scale coeff==1 returns original" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = MatrixOp(randn(5, 5)) + S = Scale(1.0, A) + @test S === A || S == A + x = randn(5) + @test S * x ≈ A * x +end - @testset "Scale real vs complex coefficient error path" begin - A = Eye(5) # real codomain - αc = 1.0 + 2.0im - @test_throws ErrorException Scale(αc, A) # triggers codomain real + complex coeff error - end +@testitem "Scale properties delegation" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = Eye(7) + α = 4.2 + S = Scale(α, A) + @test domain_type(S) == domain_type(A) + @test codomain_type(S) == codomain_type(A) + @test domain_array_type(S) == domain_array_type(A) + @test codomain_array_type(S) == codomain_array_type(A) + @test is_thread_safe(S) == is_thread_safe(A) + @test is_linear(S) && !is_null(S) + @test is_diagonal(S) == is_diagonal(A) + @test is_invertible(S) == is_invertible(A) + @test AbstractOperators.fun_name(S) == "α" * AbstractOperators.fun_name(A) + @test AbstractOperators.diag(S) == α * AbstractOperators.diag(A) + @test AbstractOperators.diag_AcA(S) == α^2 * AbstractOperators.diag_AcA(A) + @test AbstractOperators.diag_AAc(S) == α^2 * AbstractOperators.diag_AAc(A) + @test is_full_row_rank(S) == is_full_row_rank(A) + @test is_full_column_rank(S) == is_full_column_rank(A) +end - @testset "Scale get_normal_op paths" begin - # Linear operator case: Eye has optimized normal op - A = Eye(8) - α = 2.0 - S = Scale(α, A) - N = AbstractOperators.get_normal_op(S) - # For linear A: should be Scale(|α|^2, |α|^2, get_normal_op(A)) - @test N isa Scale - @test N.coeff ≈ α * α - @test N.A == AbstractOperators.get_normal_op(A) - - NL = Sigmoid((5,)) # 1D nonlinear operator - @test !AbstractOperators.is_linear(NL) - S2 = Scale(α, NL) - # Nonlinear path attempts L' * L which is not allowed; ensure it throws - @test_throws ErrorException AbstractOperators.get_normal_op(S2) - end +@testitem "Scale real vs complex coefficient error path" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + @test_throws ErrorException Scale(1.0 + 2.0im, Eye(5)) +end - @testset "Scale equality and remove_displacement" begin - A = AffineAdd(Eye(6), randn(6)) # has displacement - α = 1.5 - S1 = Scale(α, A) - S2 = Scale(α, A) - @test S1 == S2 - Srd = AbstractOperators.remove_displacement(S1) - # remove_displacement(AffineAdd) removes the displacement leaving the inner operator - @test Srd.A == AbstractOperators.remove_displacement(A) - @test Srd.coeff == α - end +@testitem "Scale get_normal_op paths" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = Eye(8) + α = 2.0 + S = Scale(α, A) + N = AbstractOperators.get_normal_op(S) + @test N isa Scale + @test N.coeff ≈ α * α + @test N.A == AbstractOperators.get_normal_op(A) + NL = Sigmoid((5,)) + S2 = Scale(α, NL) + @test !AbstractOperators.is_linear(NL) + @test_throws ErrorException AbstractOperators.get_normal_op(S2) +end - @testset "Scale slicing and remove_slicing" begin - # Use Eye so that slicing puts GetIndex first in the Compose chain (required by remove_slicing) - A = Eye(11) - α = 2.0 - As = A[1:9] # sliced codomain, Compose(GetIndex, Eye) - S = Scale(α, As) - @test AbstractOperators.is_sliced(S) - expr = AbstractOperators.get_slicing_expr(S) - @test expr !== nothing - # remove_slicing should succeed and strip the GetIndex - Sr = AbstractOperators.remove_slicing(S) - @test Sr == Scale(α, AbstractOperators.remove_slicing(As)) - end +@testitem "Scale equality and remove_displacement" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = AffineAdd(Eye(6), randn(6)) + α = 1.5 + S1 = Scale(α, A) + S2 = Scale(α, A) + @test S1 == S2 + Srd = AbstractOperators.remove_displacement(S1) + @test Srd.A == AbstractOperators.remove_displacement(A) + @test Srd.coeff == α +end - @testset "Scale opnorm and estimate_opnorm" begin - A = MatrixOp(randn(7, 4)) - α = -1.2 - S = Scale(α, A) - @test AbstractOperators.has_fast_opnorm(S) == AbstractOperators.has_fast_opnorm(A) - @test opnorm(S) ≈ abs(α) * opnorm(A) - @test estimate_opnorm(S) ≈ opnorm(S) rtol = 0.05 - end +@testitem "Scale slicing and remove_slicing" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = Eye(11) + α = 2.0 + As = A[1:9] + S = Scale(α, As) + @test AbstractOperators.is_sliced(S) + @test AbstractOperators.get_slicing_expr(S) !== nothing + Sr = AbstractOperators.remove_slicing(S) + @test Sr == Scale(α, AbstractOperators.remove_slicing(As)) +end - @testset "Scale permute utility" begin - # Use HCAT of two Eyes so underlying operator supports permute - A = HCAT(Eye(5), Eye(5)) - α = 2.2 - S = Scale(α, A) - # Permutation over domain blocks (two blocks) swaps them - p = [2, 1] - Spr = AbstractOperators.permute(S, p) - @test Spr.A == AbstractOperators.permute(S.A, p) - # Build domain input as ArrayPartition matching HCAT domain ordering - x1 = randn(5); x2 = randn(5) - xp = ArrayPartition(x1, x2) - y_original = S * xp - # After permuting, operator expects swapped domain order - xp_swapped = ArrayPartition(x2, x1) - y_permuted = Spr * xp_swapped - @test y_original ≈ y_permuted - end +@testitem "Scale opnorm and estimate_opnorm" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = MatrixOp(randn(7, 4)) + α = -1.2 + S = Scale(α, A) + @test AbstractOperators.has_fast_opnorm(S) == AbstractOperators.has_fast_opnorm(A) + @test opnorm(S) ≈ abs(α) * opnorm(A) + @test estimate_opnorm(S) ≈ opnorm(S) rtol = 0.05 +end + +@testitem "Scale permute utility" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = HCAT(Eye(5), Eye(5)) + α = 2.2 + S = Scale(α, A) + p = [2, 1] + Spr = AbstractOperators.permute(S, p) + @test Spr.A == AbstractOperators.permute(S.A, p) + x1 = randn(5) + x2 = randn(5) + y_original = S * ArrayPartition(x1, x2) + y_permuted = Spr * ArrayPartition(x2, x1) + @test y_original ≈ y_permuted +end + +@testitem "Scale threaded behavior" tags = [:calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + A = MatrixOp(randn(20000, 5)) + α = 0.75 + S = Scale(α, A; threaded = true) + x = randn(5) + @test S * x ≈ α * (A * x) +end + +@testitem "Scale (GPU)" tags = [:gpu, :calculus, :Scale] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + m = 8 + op = Scale(pi, FiniteDiff(gpu_zeros(backend, Float64, m))) + test_op(op, gpu_randn(backend, m), gpu_randn(backend, m - 1), false) - @testset "Scale threaded behavior" begin - # Force threading decision by large output length - A = MatrixOp(randn(20000, 5)) - α = 0.75 - S = Scale(α, A; threaded = true) - x = randn(5) - y = S * x - @test y ≈ α * (A * x) + n = 4 + op = Scale(2.0, DiagOp(gpu_randn(backend, n))) + test_op(op, gpu_randn(backend, n), gpu_randn(backend, n), false) end end diff --git a/test/calculus/test_sum.jl b/test/calculus/test_sum.jl index 48e1f2f..6253d76 100644 --- a/test/calculus/test_sum.jl +++ b/test/calculus/test_sum.jl @@ -1,20 +1,7 @@ -@testitem "Sum" tags = [:calculus, :Sum] setup = [TestUtils] begin +@testitem "Sum: basic mul" tags = [:calculus, :Sum] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Sum --- ") - m, n = 5, 7 - A1 = randn(m, n) - A2 = randn(m, n) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opS = Sum(opA1, opA2) - x1 = randn(n) - y1 = test_op(opS, x1, randn(m), verb) - y2 = A1 * x1 + A2 * x1 - @test norm(y1 - y2) <= 1.0e-12 - - #test Sum longer m, n = 5, 7 A1 = randn(m, n) A2 = randn(m, n) @@ -25,30 +12,17 @@ opS = Sum(opA1, opA2, opA3) x1 = randn(n) y1 = test_op(opS, x1, randn(m), verb) - y2 = A1 * x1 + A2 * x1 + A3 * x1 - @test norm(y1 - y2) <= 1.0e-12 + @test norm(y1 - (A1 * x1 + A2 * x1 + A3 * x1)) <= 1.0e-12 - opA3 = MatrixOp(randn(m, m)) - @test_throws Exception Sum(opA1, opA3) - opF = DiagOp(Float64, (m,), 2.0 + im * 3.0) - @test_throws Exception Sum(opF, opA3) - - @test is_null(opS) == false - @test is_eye(opS) == false - @test is_diagonal(opS) == false - @test is_AcA_diagonal(opS) == false - @test is_AAc_diagonal(opS) == false - @test is_orthogonal(opS) == false - @test is_invertible(opS) == false + @test_throws Exception Sum(opA1, MatrixOp(randn(m, m))) @test is_full_row_rank(opS) == true @test is_full_column_rank(opS) == false +end - d = randn(10) - op = Sum(Scale(-3.1, Eye(10)), DiagOp(d)) - @test is_diagonal(op) == true - @test norm(diag(op) - (d .- 3.1)) < 1.0e-12 +@testitem "Sum: displacement" tags = [:calculus, :Sum] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - #test displacement of sum m, n = 5, 7 A1 = randn(m, n) A2 = randn(m, n) @@ -61,100 +35,66 @@ opA3 = AffineAdd(MatrixOp(A3), d3) opS = Sum(opA1, opA2, opA3) x1 = randn(n) - y2 = A1 * x1 + A2 * x1 + A3 * x1 + d1 .+ d2 + d3 - @test norm(opS * x1 - y2) <= 1.0e-12 + @test norm(opS * x1 - (A1 * x1 + A2 * x1 + A3 * x1 + d1 .+ d2 + d3)) <= 1.0e-12 @test norm(displacement(opS) - (d1 .+ d2 + d3)) <= 1.0e-12 - y2 = A1 * x1 + A2 * x1 + A3 * x1 - @test norm(remove_displacement(opS) * x1 - y2) <= 1.0e-12 - - # fun_name formatting (2-term vs multi-term) and thread safety / storage types - op2 = Sum(MatrixOp(randn(m, n)), MatrixOp(randn(m, n))) - io = IOBuffer() - show(io, op2) - str2 = String(take!(io)) - @test occursin("+", str2) || length(str2) > 0 # lightweight check without depending on exact internal symbol - - op3 = Sum(MatrixOp(randn(m, n)), MatrixOp(randn(m, n)), MatrixOp(randn(m, n))) - io = IOBuffer() - show(io, op3) - str3 = String(take!(io)) - @test length(str3) > 0 # invokes Σ path internally - - # is_thread_safe expected false - @test is_thread_safe(op3) == false + @test norm(remove_displacement(opS) * x1 - (A1 * x1 + A2 * x1 + A3 * x1)) <= 1.0e-12 +end - # storage type queries (execute for coverage) - _dst = domain_storage_type(op3) - _cst = codomain_storage_type(op3) - @test _dst !== nothing && _cst !== nothing +@testitem "Sum: properties" tags = [:calculus, :Sum] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # remove_displacement idempotence - dX = randn(m) - op_disp = Sum(AffineAdd(MatrixOp(randn(m, n)), dX), MatrixOp(randn(m, n))) - rd1 = remove_displacement(op_disp) - rd2 = remove_displacement(rd1) - @test rd1 * x1 == rd2 * x1 - # --- Additional coverage --- - # Equality / inequality m, n = 5, 7 Aeq = MatrixOp(randn(m, n)) Beq = MatrixOp(randn(m, n)) S1 = Sum(Aeq, Beq) S2 = Sum(Aeq, Beq) S3 = Sum(Beq, Aeq) - @test S1 == S2 - @test S1 != S3 + @test S1 == S2 && S1 != S3 + @test Sum(Aeq) === Aeq - # Single-operator constructor returns the operator itself - single = MatrixOp(randn(m, n)) - @test Sum(single) === single - - # Zeros are filtered out by constructor simplification z = Zeros(Float64, (n,), Float64, (m,)) - @test Sum(single, z) === single + @test Sum(Aeq, z) === Aeq - # diag aggregator for diagonal underlying - d1 = randn(m) - d2 = randn(m) - Sdiag = Sum(DiagOp(d1), DiagOp(d2)) - @test is_diagonal(Sdiag) - @test diag(Sdiag) == d1 + d2 - - # permute utility (wrap HCAT to get multi-domain) and ensure same behavior when inputs permuted - mH = 6; n1 = 2; n2 = 1 - A1p = MatrixOp(randn(mH, n1)) - A2p = MatrixOp(randn(mH, n2)) - H = HCAT(A1p, A2p) - S = Sum(H, H) - x1p = randn(n1); x2p = randn(n2) - y_orig = S * ArrayPartition(x1p, x2p) - p = [2, 1] - Sp = AbstractOperators.permute(S, p) - y_perm = Sp * ArrayPartition(x2p, x1p) - @test y_orig ≈ y_perm + d = randn(10) + op = Sum(Scale(-3.1, Eye(10)), DiagOp(d)) + @test is_diagonal(op) == true + @test norm(diag(op) - (d .- 3.1)) < 1.0e-12 - # estimate_opnorm aggregator - opnorm_S1 = opnorm(S1) - estimated_opnorm_S1 = estimate_opnorm(S1) - @test estimated_opnorm_S1 ≈ opnorm_S1 rtol = 0.05 + @test domain_array_type(S1) !== nothing + @test codomain_array_type(S1) !== nothing + @test is_thread_safe(S1) == false +end - # remove_displacement idempotence with displacement underlying - dA = randn(m) - SA = Sum(AffineAdd(MatrixOp(randn(m, n)), dA), MatrixOp(randn(m, n))) - SA_rd = remove_displacement(SA) - @test remove_displacement(SA_rd) == SA_rd +@testitem "Sum: nonlinear" tags = [:calculus, :Sum] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # Testing nonlinear Sum of MatrixOp and Sigmoid m = 5 x = randn(m) r = randn(m) A = randn(m, m) - opA = MatrixOp(A) opB = Sigmoid(Float64, (m,), 2) - op = Sum(opA, opB) - + op = Sum(MatrixOp(A), opB) y, grad = test_NLop(op, x, r, verb) + @test norm(A * x + opB * x - y) < 1.0e-8 +end + +@testitem "Sum (GPU)" tags = [:gpu, :calculus, :Sum] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + + n = 5 + x = gpu_ones(backend, Float64, n) + opS = Sum(Eye(x), Scale(2.0, Eye(x))) + test_op(opS, gpu_randn(backend, n), gpu_randn(backend, n), false) - Y = A * x + opB * x - @test norm(Y - y) < 1.0e-8 + m, n2 = 5, 7 + A1 = gpu_randn(backend, m, n2) + A2 = gpu_randn(backend, m, n2) + opS2 = Sum(MatrixOp(A1), MatrixOp(A2)) + test_op(opS2, gpu_randn(backend, n2), gpu_randn(backend, m), false) + end end diff --git a/test/calculus/test_vcat.jl b/test/calculus/test_vcat.jl index 3a62f7e..b3cac06 100644 --- a/test/calculus/test_vcat.jl +++ b/test/calculus/test_vcat.jl @@ -1,7 +1,6 @@ -@testitem "VCAT" tags = [:calculus, :VCAT] setup = [TestUtils] begin +@testitem "VCAT: basic mul" tags = [:calculus, :VCAT] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing VCAT --- ") m1, m2, n = 4, 7, 5 A1 = randn(m1, n) @@ -11,20 +10,8 @@ opV = VCAT(opA1, opA2) x1 = randn(n) y1 = test_op(opV, x1, ArrayPartition(randn(m1), randn(m2)), verb) - y2 = ArrayPartition(A1 * x1, A2 * x1) - @test norm(y1 - y2) .<= 1.0e-12 + @test norm(y1 - ArrayPartition(A1 * x1, A2 * x1)) .<= 1.0e-12 - m1, n = 4, 5 - A1 = randn(m1, n) + im * randn(m1, n) - opA1 = MatrixOp(A1) - opA2 = DiagOp(Float64, (n,), 2.0 + im * 3.0)' - opV = VCAT(opA1, opA2) - x1 = randn(n) + im * randn(n) - y1 = test_op(opV, x1, ArrayPartition(randn(m1) + im * randn(m1), randn(n)), verb) - y2 = ArrayPartition(A1 * x1, opA2 * x1) - @test norm(y1 - y2) .<= 1.0e-12 - - #test VCAT longer m1, m2, m3, n = 4, 7, 3, 5 A1 = randn(m1, n) A2 = randn(m2, n) @@ -35,36 +22,20 @@ opV = VCAT(opA1, opA2, opA3) x1 = randn(n) y1 = test_op(opV, x1, ArrayPartition(randn(m1), randn(m2), randn(m3)), verb) - y2 = ArrayPartition(A1 * x1, A2 * x1, A3 * x1) - @test norm(y1 - y2) .<= 1.0e-12 + @test norm(y1 - ArrayPartition(A1 * x1, A2 * x1, A3 * x1)) .<= 1.0e-12 - #test VCAT of VCAT + # VCAT of VCAT (flattening) opVV = VCAT(opV, opA3) y1 = test_op(opVV, x1, ArrayPartition(randn(m1), randn(m2), randn(m3), randn(m3)), verb) - y2 = ArrayPartition(A1 * x1, A2 * x1, A3 * x1, A3 * x1) - @test norm(y1 .- y2) <= 1.0e-12 - - opVV = VCAT(opA1, opV, opA3) - y1 = test_op( - opVV, x1, ArrayPartition(randn(m1), randn(m1), randn(m2), randn(m3), randn(m3)), verb - ) - y2 = ArrayPartition(A1 * x1, A1 * x1, A2 * x1, A3 * x1, A3 * x1) - @test norm(y1 - y2) <= 1.0e-12 + @test norm(y1 .- ArrayPartition(A1 * x1, A2 * x1, A3 * x1, A3 * x1)) <= 1.0e-12 +end - opA3 = MatrixOp(randn(m1, m1)) - @test_throws Exception VCAT(opA1, opA2, opA3) - opD = DiagOp(Float64, (m1,), 2.0 + im * 3.0) - @test_throws Exception VCAT(opA1, opD, opA2) +@testitem "VCAT: properties" tags = [:calculus, :VCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - ###properties m1, m2, m3, n = 4, 7, 3, 5 - A1 = randn(m1, n) - A2 = randn(m2, n) - A3 = randn(m3, n) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) - opA3 = MatrixOp(A3) - op = VCAT(opA1, opA2, opA3) + op = VCAT(MatrixOp(randn(m1, n)), MatrixOp(randn(m2, n)), MatrixOp(randn(m3, n))) @test is_linear(op) == true @test is_null(op) == false @test is_eye(op) == false @@ -76,219 +47,117 @@ @test is_full_row_rank(op) == false @test is_full_column_rank(op) == true - d = randn(n) .+ im .* randn(n) - op = VCAT(DiagOp(d), Eye(ComplexF64, n)) - @test is_AcA_diagonal(op) == true - @test diag_AcA(op) == d .* conj(d) .+ 1 + d = randn(5) .+ im .* randn(5) + op2 = VCAT(DiagOp(d), Eye(ComplexF64, 5)) + @test is_AcA_diagonal(op2) == true + @test diag_AcA(op2) == d .* conj(d) .+ 1 + + m1, m2, n = 4, 7, 5 + A1 = MatrixOp(randn(m1, n)) + A2 = MatrixOp(randn(m2, n)) + op3 = VCAT(A1, A2) + @test domain_array_type(op3) !== nothing + @test codomain_array_type(op3) !== nothing + @test is_thread_safe(op3) == false +end + +@testitem "VCAT: displacement" tags = [:calculus, :VCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - ##test displacement m1, m2, n = 4, 7, 5 A1 = randn(m1, n) A2 = randn(m2, n) - opA1 = MatrixOp(A1) - opA2 = MatrixOp(A2) d1 = randn(m1) d2 = randn(m2) - opV = VCAT(AffineAdd(opA1, d1), AffineAdd(opA2, d2)) + opV = VCAT(AffineAdd(MatrixOp(A1), d1), AffineAdd(MatrixOp(A2), d2)) x1 = randn(n) - y1 = opV * x1 - y2 = ArrayPartition(A1 * x1 + d1, A2 * x1 + d2) - @test norm(y1 - y2) <= 1.0e-12 - y1 = remove_displacement(opV) * x1 - y2 = ArrayPartition(A1 * x1, A2 * x1) - @test norm(y1 - y2) <= 1.0e-12 - - m1, m2, n = 4, 7, 5 - A1 = MatrixOp(randn(m1, n)) - A2 = MatrixOp(randn(m2, n)) - op = VCAT(A1, A2) - _ds = domain_storage_type(op) - _cs = codomain_storage_type(op) - @test _ds !== nothing - @test _cs !== nothing - @test is_thread_safe(op) == false + @test norm(opV * x1 - ArrayPartition(A1 * x1 + d1, A2 * x1 + d2)) <= 1.0e-12 + @test norm(remove_displacement(opV) * x1 - ArrayPartition(A1 * x1, A2 * x1)) <= 1.0e-12 + A1b = MatrixOp(randn(m1, n)) + A2b = MatrixOp(randn(m2, n)) + op = VCAT(A1b, A2b) @test remove_displacement(op) == op - - d1 = randn(m1) - d2 = randn(m2) - opd = VCAT(AffineAdd(A1, d1), AffineAdd(A2, d2)) + opd = VCAT(AffineAdd(A1b, d1), AffineAdd(A2b, d2)) opd_removed = remove_displacement(opd) @test remove_displacement(opd_removed) == opd_removed +end - # equality / inequality - Aeq1 = MatrixOp(randn(3, 4)) - Aeq2 = MatrixOp(randn(5, 4)) - Veq1 = VCAT(Aeq1, Aeq2) - Veq2 = VCAT(Aeq1, Aeq2) - Veq3 = VCAT(Aeq2, Aeq1) - @test Veq1 == Veq2 - @test Veq1 != Veq3 - - # single operator constructor returns operator itself - single = MatrixOp(randn(6, 6)) - @test VCAT(single) === single - - # show output (fun_name indirect) for 2-op VCAT - io = IOBuffer(); show(io, VCAT(Eye(4), Eye(4))); s = String(take!(io)) - @test occursin("VCAT", s) || occursin("[", s) - - # sliced helpers using GetIndex - n = 10 - g1 = GetIndex(Float64, (n,), (1:5,)) - g2 = GetIndex(Float64, (n,), (6:10,)) - Vs = VCAT(g1, g2) - @test is_sliced(Vs) == true - exprs = AbstractOperators.get_slicing_expr(Vs) - @test length(exprs) == 2 - @test exprs[1] == (1:5,) - @test exprs[2] == (6:10,) - Vs_removed = AbstractOperators.remove_slicing(Vs) - @test is_sliced(Vs_removed) == false - - # diag_AAc aggregator and is_AAc_diagonal true case - Veye = VCAT(Eye(5), Eye(5)) - @test is_AAc_diagonal(Veye) == true - @test diag_AAc(Veye) == (1, 1) +@testitem "VCAT: nonlinear operators" tags = [:calculus, :VCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # VCAT of nonlinear operator n, m = 4, 3 x = randn(m) r = ArrayPartition(randn(n), randn(m)) A = randn(n, m) B = Sigmoid(Float64, (m,), 2) op = VCAT(MatrixOp(A), B) - y, grad = test_NLop(op, x, r, verb) + @test norm(ArrayPartition(A * x, B * x) - y) < 1.0e-8 +end - Y = ArrayPartition(A * x, B * x) - @test norm(Y - y) < 1.0e-8 - - # Extended coverage tests - - # Test: domain_type mismatch error - A_real = MatrixOp(randn(3, 4)) - A_complex = MatrixOp(randn(ComplexF64, 3, 4)) - @test_throws Exception VCAT(A_real, A_complex) - - # Test: VCAT with stacked operators (operators with multiple codomains) - verb && println("Testing stacked operators with multiple codomains") - n = 5 - m1, m2 = 3, 4 - A1 = MatrixOp(randn(m1, n)) - A2 = MatrixOp(randn(m2, n)) - V1 = VCAT(A1, A2) # V1 has multiple codomains - B1 = MatrixOp(randn(m1, n)) - B2 = MatrixOp(randn(m2, n)) - V2 = VCAT(B1, B2) # V2 has multiple codomains - VV = VCAT(V1, V2) # VCAT of VCATs - creates stacked structure - x = randn(n) - y = VV * x - @test y isa ArrayPartition - @test length(y.x) == 4 # Should have 4 outputs - y_test = ArrayPartition(randn(m1), randn(m2), randn(m1), randn(m2)) - x_adj = VV' * y_test - @test length(x_adj) == n - - # Test: VCAT of HCAT with Zeros (remove_slicing edge case) - verb && println("Testing VCAT of HCAT with Zeros") - n, m = 5, 6 - A1 = MatrixOp(randn(n, m)) - Z1 = Zeros(Float64, (m,), Float64, (n,)) - H1 = HCAT(A1, Z1) - A2 = MatrixOp(randn(n, m)) - A3 = MatrixOp(randn(n, m)) - H2 = HCAT(A2, A3) - V = VCAT(H1, H2) - G1 = GetIndex(Float64, (2 * m,), (1:m,)) - G2 = GetIndex(Float64, (2 * m,), ((m + 1):(2 * m),)) - V_sliced = V * VCAT(G1, G2) - @test is_sliced(V_sliced) - V_removed = AbstractOperators.remove_slicing(V_sliced) - @test V == V_removed # Should be equal after removing slicing - - H1 = HCAT(A1 * G1, Z1 * G2) - H2 = HCAT(A2 * G1, A3 * G2) - V2 = VCAT(H1, H2) - @test is_sliced(V2) - V2_removed = AbstractOperators.remove_slicing(V2) - @test V == V2_removed # Should be equal after removing slicing - - # Test: permute function - verb && println("Testing permute function") - m1, m2, n = 3, 4, 5 - A1 = MatrixOp(randn(m1, n)) - A2 = MatrixOp(randn(m2, n)) - V = VCAT(A1, A2) - p = [2, 1] - V_perm = AbstractOperators.permute(V, p) - @test size(V, 1) == size(V_perm, 1)[p] +@testitem "VCAT: slicing utilities" tags = [:calculus, :VCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) - # Test: Deeply nested VCAT structures - verb && println("Testing deeply nested VCAT") - n = 4 - A1 = Eye(n) - A2 = DiagOp(randn(n)) - A3 = MatrixOp(randn(n, n)) - V1 = VCAT(A1, A2) - V2 = VCAT(V1, A3) # VCAT of VCAT - V3 = VCAT(A1, V2, A2) # Mix of regular and VCAT operators - x = randn(n) - y = V3 * x - @test y isa ArrayPartition - @test length(y.x) == 5 # A1, A1, A2, A3, A2 - y_adj = ArrayPartition([randn(n) for _ in 1:5]...) - x_back = V3' * y_adj - @test length(x_back) == n + n = 10 + g1 = GetIndex(Float64, (n,), (1:5,)) + g2 = GetIndex(Float64, (n,), (6:10,)) + Vs = VCAT(g1, g2) + @test is_sliced(Vs) == true + exprs = AbstractOperators.get_slicing_expr(Vs) + @test exprs[1] == (1:5,) && exprs[2] == (6:10,) + @test !is_sliced(AbstractOperators.remove_slicing(Vs)) - # Test: fun_name for VCAT with > 2 operators + # fun_name and equality A1 = Eye(3) A2 = Eye(3) A3 = Eye(3) V2 = VCAT(A1, A2) V3 = VCAT(A1, A2, A3) name2 = AbstractOperators.fun_name(V2) - name3 = AbstractOperators.fun_name(V3) @test occursin("[", name2) || occursin("]", name2) - @test name3 == "VCAT" + @test AbstractOperators.fun_name(V3) == "VCAT" + # Use different operators for equality/inequality test + Aeq1 = MatrixOp(randn(3, 4)) + Aeq2 = MatrixOp(randn(5, 4)) + @test VCAT(Aeq1, Aeq2) == VCAT(Aeq1, Aeq2) + @test VCAT(Aeq1, Aeq2) != VCAT(Aeq2, Aeq1) +end - # Test: VCAT equality with different internal structures (flattened vs nested) - A1 = Eye(4) - A2 = DiagOp(randn(4)) - A3 = MatrixOp(randn(4, 4)) - V1 = VCAT(A1, A2, A3) - V_temp = VCAT(A1, A2) - V2 = VCAT(V_temp, A3) # Should flatten - @test V1 == V2 +@testitem "VCAT (GPU)" tags = [:gpu, :calculus, :VCAT] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv - # Test: is_full_column_rank with mixed operators - n, m = 5, 3 - A1 = MatrixOp(randn(n, m)) - A2 = Zeros(Float64, (m,), Float64, (n,)) - V = VCAT(A1, A2) - @test is_full_column_rank(V) isa Bool + for backend in gpu_backends() + Random.seed!(0) - # Test: Adjoint with complex nested structure - verb && println("Testing adjoint with complex nested structure") - n = 4 - m1, m2 = 3, 3 - A1 = MatrixOp(randn(m1, n)) - A2 = MatrixOp(randn(m2, n)) - V1 = 2 * VCAT(A1, A2) - B1 = MatrixOp(randn(m1, n)) - V2 = VCAT(V1, B1) - y = ArrayPartition(randn(m1), randn(m2), randn(m1)) - x = V2' * y - @test length(x) == n - @test x isa AbstractArray - y2 = V2 * x - @test size.(y.x) == size.(y2.x) + n = 4 + opV = VCAT(DiagOp(gpu_ones(backend, Float64, n)), DiagOp(2 .* gpu_ones(backend, Float64, n))) + test_op(opV, gpu_randn(backend, n), ArrayPartition(gpu_randn(backend, n), gpu_randn(backend, n)), false) - V2 = VCAT(B1, V1) - y = ArrayPartition(randn(m1), randn(m1), randn(m2)) - x = V2' * y - @test length(x) == n - @test x isa AbstractArray - y2 = V2 * x - @test size.(y.x) == size.(y2.x) + m1, m2, n = 4, 7, 5 + A1 = gpu_randn(backend, m1, n) + A2 = gpu_randn(backend, m2, n) + opV2 = VCAT(MatrixOp(A1), MatrixOp(A2)) + test_op(opV2, gpu_randn(backend, n), ArrayPartition(gpu_randn(backend, m1), gpu_randn(backend, m2)), false) + end +end + +@testitem "VCAT: copy_operator" tags = [:calculus, :VCAT] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(4) + + n, m1, m2 = 8, 5, 6 + opV = VCAT(MatrixOp(randn(m1, n)), MatrixOp(randn(m2, n))) + opV2 = copy_operator(opV) + @test opV2 isa VCAT + x = randn(n) + y1 = opV * x + y2 = opV2 * x + @test collect(y1) ≈ collect(y2) + # Verify independence: forward into opV2 alone + x2 = randn(n) + @test collect(opV2 * x2) ≈ collect(opV * x2) end diff --git a/test/dsp/test_dsp_operators.jl b/test/dsp/test_dsp_operators.jl index 648f2d5..65013a3 100644 --- a/test/dsp/test_dsp_operators.jl +++ b/test/dsp/test_dsp_operators.jl @@ -31,9 +31,8 @@ @test is_full_column_rank(op) == true end -@testitem "Filt" tags = [:dsp, :Filt] setup = [TestUtils] begin +@testitem "Filt: IIR and FIR mappings" tags = [:dsp, :Filt] setup = [TestUtils] begin using DSPOperators, DSP, LinearAlgebra, Random - using AbstractOperators: is_linear, is_null, is_eye, is_diagonal, is_AcA_diagonal, is_AAc_diagonal, is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank n, m = 15, 2 b, a = [1.0; 0.0; 1.0; 0.0; 0.0], [1.0; 1.0; 1.0] @@ -51,6 +50,27 @@ end y2 = filt(h, [1.0], x1) @test all(norm.(y1 .- y2) .<= 1.0e-12) +end + +@testitem "Filt: constructors and properties" tags = [:dsp, :Filt] setup = [TestUtils] begin + using DSPOperators + using AbstractOperators: + is_linear, + is_null, + is_eye, + is_diagonal, + is_AcA_diagonal, + is_AAc_diagonal, + is_orthogonal, + is_invertible, + is_full_row_rank, + is_full_column_rank + + n, m = 15, 2 + b, a = [1.0; 0.0; 1.0; 0.0; 0.0], [1.0; 1.0; 1.0] + h = randn(10) + x1 = randn(n, m) + op = Filt(Float64, (n, m), h) # other constructors Filt(n, b, a) @@ -73,9 +93,8 @@ end @test is_full_column_rank(op) == true end -@testitem "MIMOFilt" tags = [:dsp, :MIMOFilt] setup = [TestUtils] begin +@testitem "MIMOFilt: 2-input 1-output mapping" tags = [:dsp, :MIMOFilt] setup = [TestUtils] begin using DSPOperators, DSP, LinearAlgebra, Random - using AbstractOperators: is_linear, is_null, is_eye, is_diagonal, is_AcA_diagonal, is_AAc_diagonal, is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank m, n = 10, 2 b = [[1.0; 0.0; 1.0; 0.0; 0.0], [1.0; 0.0; 1.0; 0.0; 0.0]] @@ -87,6 +106,10 @@ end y2 = filt(b[1], a[1], x1[:, 1]) + filt(b[2], a[2], x1[:, 2]) @test all(norm.(y1 .- y2) .<= 1.0e-12) +end + +@testitem "MIMOFilt: 3-input 2-output mapping" tags = [:dsp, :MIMOFilt] setup = [TestUtils] begin + using DSPOperators, DSP, LinearAlgebra, Random m, n = 10, 3 #time samples, number of inputs b = [ @@ -107,6 +130,10 @@ end y2 = [col1 col2] @test all(norm.(y1 .- y2) .<= 1.0e-12) +end + +@testitem "MIMOFilt: FIR mapping" tags = [:dsp, :MIMOFilt] setup = [TestUtils] begin + using DSPOperators, DSP, LinearAlgebra, Random m, n = 10, 3 b = [randn(10), randn(5), randn(10), randn(2), randn(10), randn(10)] @@ -120,6 +147,15 @@ end y2 = [col1 col2] @test all(norm.(y1 .- y2) .<= 1.0e-12) +end + +@testitem "MIMOFilt: constructors and errors" tags = [:dsp, :MIMOFilt] setup = [TestUtils] begin + using DSPOperators + + m, n = 10, 3 + b = [randn(10), randn(5), randn(10), randn(2), randn(10), randn(10)] + a = [[1.0], [1.0], [1.0], [1.0], [1.0], [1.0]] + x1 = randn(m, n) ## other constructors MIMOFilt((10, 3), b, a) @@ -138,7 +174,23 @@ end @test_throws ErrorException MIMOFilt(Float32, (m, n), b2, a2) a[1][1] = 0.0 @test_throws ErrorException MIMOFilt(Float64, (m, n), b, a) +end +@testitem "MIMOFilt: properties" tags = [:dsp, :MIMOFilt] setup = [TestUtils] begin + using DSPOperators + using AbstractOperators: + is_linear, + is_null, + is_eye, + is_diagonal, + is_AcA_diagonal, + is_AAc_diagonal, + is_orthogonal, + is_invertible, + is_full_row_rank, + is_full_column_rank + + m, n = 10, 3 b = [randn(10), randn(5), randn(10), randn(2), randn(10), randn(10)] a = [[1.0], [1.0], [1.0], [1.0], [1.0], [1.0]] op = MIMOFilt(Float64, (m, n), b, a) @@ -155,3 +207,119 @@ end @test is_full_row_rank(op) == true @test is_full_column_rank(op) == true end + +@testitem "Conv (GPU)" tags = [:gpu, :dsp, :Conv] setup = [TestUtils] begin + using DSPOperators, DSP, GPUEnv, LinearAlgebra, Random + for backend in gpu_backends(supports_fftw = true) + Random.seed!(0) + n, m = 20, 6 + h_cpu = randn(m) + x_cpu = randn(n) + h = to_gpu(backend, h_cpu) + op = Conv(Float64, (n,), h) + x = to_gpu(backend, x_cpu) + y = gpu_zeros(backend, Float64, n + m - 1) + mul!(y, op, x) + y_alloc = op * x + @test y_alloc isa typeof(y) + op_cpu = Conv(Float64, (n,), h_cpu) + y_cpu = zeros(n + m - 1) + mul!(y_cpu, op_cpu, x_cpu) + @test collect(y) ≈ y_cpu atol = 1.0e-10 + @test collect(y_alloc) ≈ y_cpu atol = 1.0e-10 + b_cpu = randn(n + m - 1) + b = to_gpu(backend, b_cpu) + z = gpu_zeros(backend, Float64, n) + z_cpu = zeros(n) + mul!(z_cpu, op_cpu', b_cpu) + mul!(z, op', b) + @test collect(z) ≈ z_cpu atol = 1.0e-10 + end +end + +@testitem "Xcorr (GPU)" tags = [:gpu, :dsp, :Xcorr] setup = [TestUtils] begin + using DSPOperators, DSP, GPUEnv, LinearAlgebra, Random + for backend in gpu_backends(supports_fftw = true) + Random.seed!(0) + n, m = 15, 5 + h_cpu = randn(m) + x_cpu = randn(n) + h = to_gpu(backend, h_cpu) + outlen = 2 * max(n, m) - 1 + op = Xcorr(Float64, (n,), h) + x = to_gpu(backend, x_cpu) + y = gpu_zeros(backend, Float64, outlen) + mul!(y, op, x) + y_alloc = op * x + @test y_alloc isa typeof(y) + op_cpu = Xcorr(Float64, (n,), h_cpu) + y_cpu = zeros(outlen) + mul!(y_cpu, op_cpu, x_cpu) + @test collect(y) ≈ y_cpu atol = 1.0e-10 + @test collect(y_alloc) ≈ y_cpu atol = 1.0e-10 + b_cpu = randn(outlen) + b = to_gpu(backend, b_cpu) + z = gpu_zeros(backend, Float64, n) + z_cpu = zeros(n) + mul!(z_cpu, op_cpu', b_cpu) + mul!(z, op', b) + @test collect(z) ≈ z_cpu atol = 1.0e-10 + end +end + +@testitem "Filt (GPU, FIR)" tags = [:gpu, :dsp, :Filt] setup = [TestUtils] begin + using DSPOperators, DSP, GPUEnv, LinearAlgebra, Random + for backend in gpu_backends(supports_fftw = true) + Random.seed!(42) + n = 20 + b = randn(5) + x_cpu = randn(n) + op_cpu = Filt(Float64, (n,), b) + y_cpu = zeros(n) + mul!(y_cpu, op_cpu, x_cpu) + x = to_gpu(backend, x_cpu) + op = Filt(x, b) + y = gpu_zeros(backend, Float64, n) + mul!(y, op, x) + y_alloc = op * x + @test y_alloc isa typeof(y) + @test collect(y) ≈ y_cpu atol = 1.0e-10 + @test collect(y_alloc) ≈ y_cpu atol = 1.0e-10 + r_cpu = randn(n) + z_cpu = zeros(n) + mul!(z_cpu, op_cpu', r_cpu) + r = to_gpu(backend, r_cpu) + z = gpu_zeros(backend, Float64, n) + mul!(z, op', r) + @test collect(z) ≈ z_cpu atol = 1.0e-10 + end +end + +@testitem "MIMOFilt (GPU, FIR)" tags = [:gpu, :dsp, :MIMOFilt] setup = [TestUtils] begin + using DSPOperators, DSP, GPUEnv, LinearAlgebra, Random + for backend in gpu_backends(supports_fftw = true) + Random.seed!(7) + m, n = 10, 3 + b = [randn(5), randn(3), randn(4), randn(5), randn(3), randn(4)] + a = [[1.0] for _ in b] + op_cpu = MIMOFilt(Float64, (m, n), b, a) + x_cpu = randn(m, n) + y_cpu = zeros(m, 2) + mul!(y_cpu, op_cpu, x_cpu) + x = to_gpu(backend, x_cpu) + op = MIMOFilt(x, b, a) + y = gpu_zeros(backend, Float64, m, 2) + mul!(y, op, x) + y_alloc = op * x + @test y_alloc isa typeof(y) + @test collect(y) ≈ y_cpu atol = 1.0e-10 + @test collect(y_alloc) ≈ y_cpu atol = 1.0e-10 + r_cpu = randn(m, 2) + z_cpu = zeros(m, n) + mul!(z_cpu, op_cpu', r_cpu) + r = to_gpu(backend, r_cpu) + z = gpu_zeros(backend, Float64, m, n) + mul!(z, op', r) + @test collect(z) ≈ z_cpu atol = 1.0e-10 + end +end diff --git a/test/dsp/test_quality.jl b/test/dsp/test_quality.jl index 95e0767..4f99d53 100644 --- a/test/dsp/test_quality.jl +++ b/test/dsp/test_quality.jl @@ -1,4 +1,4 @@ @testitem "Aqua" tags = [:quality, :dsp] begin using DSPOperators, Aqua - Aqua.test_all(DSPOperators) + Aqua.test_all(DSPOperators, persistent_tasks = VERSION >= v"1.11") end diff --git a/test/fftw/test_combinations.jl b/test/fftw/test_combinations.jl index 2d8143f..019b074 100644 --- a/test/fftw/test_combinations.jl +++ b/test/fftw/test_combinations.jl @@ -1,5 +1,6 @@ @testitem "Transform Combinations" tags = [:fftw, :CombinationRules] begin using FFTWOperators + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 8 # Power of 2 for DCT diff --git a/test/fftw/test_fftw_operators.jl b/test/fftw/test_fftw_operators.jl index 971b59c..dec66fe 100644 --- a/test/fftw/test_fftw_operators.jl +++ b/test/fftw/test_fftw_operators.jl @@ -1,4 +1,5 @@ @testitem "DCT" tags = [:fftw, :DCT] setup = [TestUtils] begin + using AbstractOperators using FFTW, LinearAlgebra, Random, FFTWOperators using AbstractOperators: is_linear, is_null, is_eye, is_diagonal, is_AcA_diagonal, is_AAc_diagonal, is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank @@ -37,6 +38,7 @@ end @testitem "IDCT" tags = [:fftw, :IDCT] setup = [TestUtils] begin + using AbstractOperators using FFTW, LinearAlgebra, Random, FFTWOperators using AbstractOperators: is_linear, is_null, is_eye, is_diagonal, is_AcA_diagonal, is_AAc_diagonal, is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank @@ -75,6 +77,7 @@ end end @testitem "DFT" tags = [:fftw, :DFT] setup = [TestUtils] begin + using AbstractOperators using FFTW, LinearAlgebra, Random, FFTWOperators using AbstractOperators: is_linear, is_null, is_eye, is_diagonal, is_AcA_diagonal, is_AAc_diagonal, is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank @@ -175,6 +178,7 @@ end end @testitem "IDFT" tags = [:fftw, :IDFT] setup = [TestUtils] begin + using AbstractOperators using FFTW, LinearAlgebra, Random, FFTWOperators using AbstractOperators: is_linear, is_null, is_eye, is_diagonal, is_AcA_diagonal, is_AAc_diagonal, is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank @@ -244,6 +248,7 @@ end end @testitem "RDFT" tags = [:fftw, :RDFT] setup = [TestUtils] begin + using AbstractOperators using FFTW, LinearAlgebra, Random, FFTWOperators using AbstractOperators: is_linear, is_null, is_eye, is_diagonal, is_AcA_diagonal, is_AAc_diagonal, is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank @@ -281,6 +286,7 @@ end end @testitem "IRDFT" tags = [:fftw, :IRDFT] setup = [TestUtils] begin + using AbstractOperators using FFTW, LinearAlgebra, Random, FFTWOperators using AbstractOperators: is_linear, is_null, is_eye, is_diagonal, is_AcA_diagonal, is_AAc_diagonal, is_orthogonal, is_invertible, is_full_row_rank, is_full_column_rank @@ -339,3 +345,72 @@ end @test is_full_row_rank(op) == true @test is_full_column_rank(op) == false end + +@testitem "DCT/IDCT (GPU)" tags = [:gpu, :fftw, :DCT, :IDCT] setup = [TestUtils] begin + using AbstractOperators, FFTW, FFTWOperators, GPUEnv, Random + has_accelerated_dcts = Base.find_package("AcceleratedDCTs") !== nothing + + for backend in gpu_backends(; include_jlarrays = false, supports_fftw = true) + Random.seed!(0) + + has_accelerated_dcts || continue + import AcceleratedDCTs + + dims = (8, 10) + x_cpu = randn(Float32, dims...) + x_gpu = to_gpu(backend, x_cpu) + + dct_op = DCT(x_gpu) + y_gpu = similar(x_gpu) + mul!(y_gpu, dct_op, x_gpu) + @test y_gpu isa typeof(x_gpu) + @test collect(y_gpu) ≈ dct(x_cpu) rtol = 1.0e-4 atol = 1.0e-4 + + idct_op = IDCT(x_gpu) + x_rec_gpu = similar(x_gpu) + mul!(x_rec_gpu, idct_op, y_gpu) + @test collect(x_rec_gpu) ≈ x_cpu rtol = 1.0e-4 atol = 1.0e-4 + + x_adj_gpu = similar(x_gpu) + mul!(x_adj_gpu, dct_op', y_gpu) + @test collect(x_adj_gpu) ≈ x_cpu rtol = 1.0e-4 atol = 1.0e-4 + end +end + +@testitem "DFT/RDFT/IRDFT (GPU)" tags = [:gpu, :fftw, :DFT, :RDFT, :IRDFT] setup = [TestUtils] begin + using FFTW, FFTWOperators, GPUEnv, LinearAlgebra, Random, AbstractOperators + + for backend in gpu_backends(supports_fftw = true) + Random.seed!(0) + n, m = 8, 6 + + x_cpu = randn(Float64, n, m) + op = DFT(to_gpu(backend, x_cpu)) + @test op isa DFT + x_gpu = to_gpu(backend, x_cpu) + y_gpu = similar(x_gpu, Complex{Float64}) + mul!(y_gpu, op, x_gpu) + @test Array(y_gpu) ≈ fft(x_cpu) + r_gpu = similar(x_gpu, Float64) + mul!(r_gpu, op', y_gpu) + @test Array(r_gpu) ≈ real.(bfft(Array(y_gpu))) + + xc_cpu = randn(ComplexF64, n, m) + opc = DFT(to_gpu(backend, xc_cpu)) + @test opc isa DFT + xc_gpu = to_gpu(backend, xc_cpu) + yc_gpu = similar(xc_gpu) + mul!(yc_gpu, opc, xc_gpu) + @test Array(yc_gpu) ≈ fft(xc_cpu) + + opr = RDFT(to_gpu(backend, x_cpu)) + yr_gpu = similar(x_gpu, Complex{Float64}, n ÷ 2 + 1, m) + mul!(yr_gpu, opr, x_gpu) + @test Array(yr_gpu) ≈ rfft(x_cpu, 1) + + opir = IRDFT(to_gpu(backend, Array(yr_gpu)), n) + out_gpu = similar(x_gpu) + mul!(out_gpu, opir, yr_gpu) + @test Array(out_gpu) ≈ x_cpu rtol = 1.0e-12 + end +end diff --git a/test/fftw/test_quality.jl b/test/fftw/test_quality.jl index 78fd6bd..09c7d41 100644 --- a/test/fftw/test_quality.jl +++ b/test/fftw/test_quality.jl @@ -1,5 +1,8 @@ @testitem "Aqua" tags = [:quality, :fftw] begin - using Aqua, FFTW, FFTWOperators - Aqua.test_all(FFTWOperators; piracies = false) + using Aqua, AbstractOperators, FFTW, FFTWOperators + Aqua.test_all( + FFTWOperators; + piracies = false, persistent_tasks = VERSION >= v"1.11" + ) Aqua.test_piracies(FFTWOperators; treat_as_own = [FFTW, AbstractOperators]) end diff --git a/test/fftw/test_shift_operators.jl b/test/fftw/test_shift_operators.jl index 9e91bbe..4903555 100644 --- a/test/fftw/test_shift_operators.jl +++ b/test/fftw/test_shift_operators.jl @@ -1,4 +1,5 @@ @testitem "FFTShift/IFFTShift Operators" tags = [:fftw, :FFTShift] setup = [TestUtils] begin + using AbstractOperators using LinearAlgebra, Random, FFTWOperators # 1D even length n = 4 @@ -22,7 +23,7 @@ # 2D, both dims n, m = 2, 3 - X = reshape(1.0:(n * m), n, m) + X = collect(reshape(1.0:(n * m), n, m)) A = FFTShift((n, m), (1, 2)) B = IFFTShift((n, m), (1, 2)) Y = A * X @@ -39,6 +40,7 @@ end @testitem "SignAlternation Operator" tags = [:fftw, :SignAlternation] setup = [TestUtils] begin + using AbstractOperators using LinearAlgebra, Random, FFTWOperators # 1D n = 6 @@ -73,13 +75,13 @@ end @test v == [1.0, -2.0, 3.0, -4.0] # out-of-place with dest - x = reshape(1.0:4.0, 2, 2) + x = collect(reshape(1.0:4.0, 2, 2)) y = similar(x) alternate_sign!(y, x, 1, 2) @test y == [1.0 -3.0; -2.0 4.0] # functional copy - u = alternate_sign(reshape(1.0:4.0, 2, 2), 1, 2) + u = alternate_sign(collect(reshape(1.0:4.0, 2, 2)), 1, 2) @test u == [1.0 -3.0; -2.0 4.0] # no dirs → identity @@ -133,6 +135,7 @@ end end @testitem "Combination rules: FFTShift/IFFTShift with DFT/IDFT" tags = [:fftw, :CombinationRules] setup = [TestUtils] begin + using AbstractOperators using LinearAlgebra, Random, FFTWOperators using AbstractOperators: can_be_combined, combine @@ -183,3 +186,40 @@ end X = randn(ComplexF64, n, m) @test (c5 * X) ≈ ((dft2 * sh2) * X) end + +@testitem "FFTShift/IFFTShift (GPU)" tags = [:gpu, :fftw, :FFTShift] setup = [TestUtils] begin + using FFTWOperators, GPUEnv, LinearAlgebra + + for backend in gpu_backends() + n = 4 + x = to_gpu(backend, collect(1.0:n)) + A = FFTShift((n,), (1,); array_type = typeof(x)) + y = A * x + @test collect(y) == [3.0, 4.0, 1.0, 2.0] + @test collect(A' * y) ≈ collect(x) + + n2, m2 = 2, 4 + X = gpu_ones(backend, Float64, n2, m2) + A2 = FFTShift((n2, m2), (1, 2); array_type = typeof(X)) + Y = A2 * X + @test collect(Y) ≈ ones(n2, m2) + end +end + +@testitem "SignAlternation (GPU)" tags = [:gpu, :fftw, :SignAlternation] setup = [TestUtils] begin + using FFTWOperators, GPUEnv, LinearAlgebra + + for backend in gpu_backends() + n = 6 + x = gpu_ones(backend, Float64, n) + S = SignAlternation((n,), (1,); array_type = typeof(x)) + y = S * x + @test collect(y) == [1.0, -1.0, 1.0, -1.0, 1.0, -1.0] + @test collect(S * y) ≈ collect(x) + + X2 = gpu_ones(backend, Float64, 2, 2) + S2 = SignAlternation((2, 2), (1, 2); array_type = typeof(X2)) + Y2 = S2 * X2 + @test collect(Y2) == [1.0 -1.0; -1.0 1.0] + end +end diff --git a/test/linearoperators/test_diagop.jl b/test/linearoperators/test_diagop.jl index 2f1c173..e50eb88 100644 --- a/test/linearoperators/test_diagop.jl +++ b/test/linearoperators/test_diagop.jl @@ -1,56 +1,58 @@ -@testitem "DiagOp" tags = [:linearoperator, :DiagOp] setup = [TestUtils] begin +@testmodule DiagOpTestHelper begin + using Test, AbstractOperators, LinearAlgebra + + export test_diagop_mul + + function test_diagop_mul(conv, verb, test_op, to_cpu, norm) + n = 4 + d = conv(randn(n)) + op = DiagOp(d) + x1 = conv(randn(n)) + y1 = test_op(op, x1, conv(randn(n)), verb) + y2 = d .* x1 + @test norm(to_cpu(y1) .- to_cpu(y2)) <= 1.0e-12 + + d = conv(randn(n) + im * randn(n)) + op = DiagOp(d) + x1 = conv(randn(n) .+ im * randn(n)) + y1 = test_op(op, x1, conv(randn(n) .+ im * randn(n)), verb) + y2 = d .* x1 + @test norm(to_cpu(y1) .- to_cpu(y2)) <= 1.0e-12 + end + +end # @testmodule DiagOpTestHelper + +@testitem "DiagOp: mul and constructors" tags = [:linearoperator, :DiagOp] setup = [TestUtils, DiagOpTestHelper] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing DiagOp --- ") - n = 4 - d = randn(n) - op = DiagOp(Float64, (n,), d) - x1 = randn(n) - y1 = test_op(op, x1, randn(n), verb) - y2 = d .* x1 - - @test all(norm.(y1 .- y2) .<= 1.0e-12) + test_diagop_mul(identity, verb, test_op, to_cpu, norm) + # scalar diagonal tests (CPU-only, not array-dependent) n = 4 - d = randn(n) + im * randn(n) - op = DiagOp(Float64, (n,), d) x1 = randn(n) - y1 = test_op(op, x1, randn(n) .+ im * randn(n), verb) - y2 = d .* x1 - - @test all(norm.(y1 .- y2) .<= 1.0e-12) - - n = 4 d = pi op = DiagOp(Float64, (n,), d) - x1 = randn(n) y1 = test_op(op, x1, randn(n), verb) - y2 = d .* x1 - - @test all(norm.(y1 .- y2) .<= 1.0e-12) + @test all(norm.(y1 .- d .* x1) .<= 1.0e-12) - n = 4 d = im op = DiagOp(Float64, (n,), d) - x1 = randn(n) y1 = test_op(op, x1, randn(n) + im * randn(n), verb) @test domain_type(op) == Float64 @test codomain_type(op) == Complex{Float64} - y2 = d .* x1 - - @test all(norm.(y1 .- y2) .<= 1.0e-12) # other constructors - d = randn(4) - op = DiagOp(d) - - d = randn(4) .+ im - op = DiagOp(d) + op = DiagOp(randn(4)) + op = DiagOp(randn(4) .+ im) + op = DiagOp((n,), pi) +end +@testitem "DiagOp: properties and helpers" tags = [:linearoperator, :DiagOp] setup = [TestUtils, DiagOpTestHelper] begin + using Random, AbstractOperators + Random.seed!(0) n = 4 - d = pi - op = DiagOp((n,), d) + op = DiagOp((n,), pi) #properties @test is_linear(op) == true @@ -67,25 +69,19 @@ @test is_full_column_rank(op) == true @test is_full_column_rank(DiagOp([ones(5); 0])) == false - @test diag(op) == d - @test norm(op' * (op * x1) .- diag_AcA(op) .* x1) <= 1.0e-12 - @test norm(op * (op' * x1) .- diag_AAc(op) .* x1) <= 1.0e-12 - - n = 4 + @test diag(op) == pi d = pi op = DiagOp((n,), d) x1 = randn(n) - - @test diag(op) == d @test norm(op' * (op * x1) .- diag_AcA(op) .* x1) <= 1.0e-12 @test norm(op * (op' * x1) .- diag_AAc(op) .* x1) <= 1.0e-12 - # Scale: create scaled operator and verify diagonal scaling and properties + # Scale op_scaled = Scale(3.0, op) @test diag(op_scaled) == 3.0 .* diag(op) @test size(op_scaled) == size(op) - # get_normal_op: should produce diagonal with abs2 of original diag + # get_normal_op normal_op = AbstractOperators.get_normal_op(op) @test diag(normal_op) == abs2.(diag(op)) @test is_diagonal(normal_op) == true @@ -98,12 +94,24 @@ @test AbstractOperators.has_optimized_normalop(op) == true @test AbstractOperators.has_optimized_normalop(op') == true - # invertibility false path (contains a zero) + # invertibility false path op_sing = DiagOp([1.0, 0.0, 2.0, 3.0]) @test is_invertible(op_sing) == false @test is_full_row_rank(op_sing) == false @test is_full_column_rank(op_sing) == false - # size returns (domain_dim, domain_dim) @test size(op) == ((n,), (n,)) end + +@testitem "DiagOp (GPU)" tags = [:gpu, :linearoperator, :DiagOp] setup = [TestUtils, DiagOpTestHelper] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + test_diagop_mul(x -> to_gpu(backend, x), false, test_op, to_cpu, norm) + x = gpu_randn(backend, 4) + op = DiagOp(x) + @test domain_array_type(op) <: backend.array_type + @test codomain_array_type(op) <: backend.array_type + end +end diff --git a/test/linearoperators/test_eye.jl b/test/linearoperators/test_eye.jl index 5263d5f..d37a3bb 100644 --- a/test/linearoperators/test_eye.jl +++ b/test/linearoperators/test_eye.jl @@ -1,7 +1,27 @@ -@testitem "Eye" tags = [:linearoperator, :Eye] setup = [TestUtils] begin +@testmodule EyeTestHelper begin + using Test, AbstractOperators, LinearAlgebra + + export test_eye_mul + + function test_eye_mul(conv, verb, test_op, to_cpu, norm) + n = 4 + x1 = conv(randn(n)) + op = Eye(x1) + y1 = test_op(op, x1, conv(randn(n)), verb) + @test norm(to_cpu(y1) .- to_cpu(x1)) <= 1.0e-12 + + x2 = conv(randn(n, n)) + op2 = Eye(x2) + test_op(op2, x2, conv(randn(n, n)), verb) + end + +end # @testmodule EyeTestHelper + +@testitem "Eye" tags = [:linearoperator, :Eye] setup = [TestUtils, EyeTestHelper] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Eye --- ") + + test_eye_mul(identity, verb, test_op, to_cpu, norm) n = 4 op = Eye(Float64, (n,)) @@ -50,8 +70,8 @@ # Storage type helpers @test domain_type(op) == Float64 @test codomain_type(op) == Float64 - @test domain_storage_type(op) == Array{Float64} - @test codomain_storage_type(op) == Array{Float64} + @test domain_array_type(op) == Array{Float64} + @test codomain_array_type(op) == Array{Float64} @test is_thread_safe(op) == true # Adjoint, opnorm, get_normal_op @@ -68,3 +88,12 @@ mul!(y, op, x1) @test all(y .== x1) end + +@testitem "Eye (GPU)" tags = [:gpu, :linearoperator, :Eye] setup = [TestUtils, EyeTestHelper] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + test_eye_mul(x -> to_gpu(backend, x), false, test_op, to_cpu, norm) + end +end diff --git a/test/linearoperators/test_finitediff.jl b/test/linearoperators/test_finitediff.jl index cac1629..cbdb97b 100644 --- a/test/linearoperators/test_finitediff.jl +++ b/test/linearoperators/test_finitediff.jl @@ -1,7 +1,29 @@ -@testitem "FiniteDiff" tags = [:linearoperator, :FiniteDiff] setup = [TestUtils] begin +@testmodule FiniteDiffTestHelper begin + using Test, AbstractOperators, LinearAlgebra + + export test_finitediff_mul + + function test_finitediff_mul(conv, verb, test_op) + n = 10 + op = FiniteDiff(conv(zeros(Float64, n))) + test_op(op, conv(randn(n)), conv(randn(n - 1)), verb) + + n, m = 10, 5 + op = FiniteDiff(conv(zeros(Float64, n, m))) + test_op(op, conv(randn(n, m)), conv(randn(n - 1, m)), verb) + + n, m = 10, 5 + op = FiniteDiff(conv(zeros(Float64, n, m)), 2) + test_op(op, conv(randn(n, m)), conv(randn(n, m - 1)), verb) + end + +end # @testmodule FiniteDiffTestHelper + +@testitem "FiniteDiff: basic mul" tags = [:linearoperator, :FiniteDiff] setup = [TestUtils, FiniteDiffTestHelper] begin using Random, SparseArrays, AbstractOperators Random.seed!(0) - verb && println(" --- Testing FiniteDiff --- ") + + test_finitediff_mul(identity, verb, test_op) n = 10 op = FiniteDiff(Float64, (n,)) @@ -13,16 +35,8 @@ I1, J1, V1 = SparseArrays.spdiagm_internal(0 => ones(n - 1)) I2, J2, V2 = SparseArrays.spdiagm_internal(1 => ones(n - 1)) B = -sparse(I1, J1, V1, n - 1, n) + sparse(I2, J2, V2, n - 1, n) - @test norm(B * x1 - op * x1) <= 1.0e-8 - n, m = 10, 5 - op = FiniteDiff(Float64, (n, m)) - x1 = randn(n, m) - y1 = test_op(op, x1, randn(n - 1, m), verb) - y1 = op * repeat(collect(range(0; stop = 1, length = n)), 1, m) - @test all(norm.(y1 .- 1 / 9) .<= 1.0e-12) - n, m = 10, 5 op = FiniteDiff(Float64, (n, m), 2) x1 = randn(n, m) @@ -30,48 +44,17 @@ y1 = op * repeat(collect(range(0; stop = 1, length = n)), 1, m) @test all(norm.(y1) .<= 1.0e-12) - n, m, l = 10, 5, 7 - op = FiniteDiff(Float64, (n, m, l)) - x1 = randn(n, m, l) - y1 = test_op(op, x1, randn(n - 1, m, l), verb) - y1 = op * reshape(repeat(collect(range(0; stop = 1, length = n)), 1, m * l), n, m, l) - @test all(norm.(y1 .- 1 / 9) .<= 1.0e-12) - - n, m, l = 10, 5, 7 - op = FiniteDiff(Float64, (n, m, l), 2) - x1 = randn(n, m, l) - y1 = test_op(op, x1, randn(n, m - 1, l), verb) - y1 = op * reshape(repeat(collect(range(0; stop = 1, length = n)), 1, m * l), n, m, l) - @test all(norm.(y1) .<= 1.0e-12) - - n, m, l = 10, 5, 7 - op = FiniteDiff(Float64, (n, m, l), 3) - x1 = randn(n, m, l) - y1 = test_op(op, x1, randn(n, m, l - 1), verb) - y1 = op * reshape(repeat(collect(range(0; stop = 1, length = n)), 1, m * l), n, m, l) - @test all(norm.(y1) .<= 1.0e-12) + @test_throws ErrorException FiniteDiff(Float64, (n, m), 4) + FiniteDiff((n, m)) + FiniteDiff(x1) +end - n, m, l, i = 5, 6, 2, 3 - op = FiniteDiff(Float64, (n, m, l, i), 1) - x1 = randn(n, m, l, i) - y1 = test_op(op, x1, randn(n - 1, m, l, i), verb) - y1 = op * reshape(repeat(collect(range(0; stop = 1, length = n)), 1, m * l * i), n, m, l, i) - @test all(norm.(y1 .- 1 / (n - 1)) .<= 1.0e-12) +@testitem "FiniteDiff: properties" tags = [:linearoperator, :FiniteDiff] setup = [TestUtils, FiniteDiffTestHelper] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) n, m, l, i = 5, 6, 2, 3 op = FiniteDiff(Float64, (n, m, l, i), 4) - x1 = randn(n, m, l, i) - y1 = test_op(op, x1, randn(n, m, l, i - 1), verb) - y1 = op * reshape(repeat(collect(range(0; stop = 1, length = n)), 1, m * l * i), n, m, l, i) - @test norm(y1) <= 1.0e-12 - - @test_throws ErrorException FiniteDiff(Float64, (n, m, l), 4) - - ## other constructors - FiniteDiff((n, m)) - FiniteDiff(x1) - - #properties @test is_linear(op) == true @test is_null(op) == false @test is_eye(op) == false @@ -83,47 +66,31 @@ @test is_full_row_rank(op) == true @test is_full_column_rank(op) == false - # Invalid direction > number of dims on construction (simple 2D example) - @test_throws ErrorException FiniteDiff(Float64, (4, 5), 3) - - # Minimal dimension case (size 2 along diff axis) gives single difference - op_min = FiniteDiff(Float64, (2,)) - x_min = randn(2) - @test (op_min * x_min)[1] ≈ x_min[2] - x_min[1] - @test size(op_min) == ((1,), (2,)) - - # Complex input - ensure same behavior and real/imag handled n = 6 - op_c = FiniteDiff(ComplexF64, (n,)) - x_c = randn(n) .+ im * randn(n) - y_c = similar(x_c, n - 1) - y_c .= op_c * x_c - @test all(y_c .≈ x_c[2:end] .- x_c[1:(end - 1)]) - - # Adjoint consistency: == - x = randn(n) F = FiniteDiff(Float64, (n,)) g = randn(n - 1) + x = randn(n) lhs = dot(F * x, g) tmp = zeros(n) mul!(tmp, F', g) rhs = dot(x, tmp) @test lhs ≈ rhs atol = 1.0e-10 - # In-place mul! forward and adjoint - y_store = zeros(n - 1) - mul!(y_store, F, x) - @test all(y_store .== (F * x)) - x_store = zeros(n) - mul!(x_store, F', y_store) - @test dot(F * x, y_store) ≈ dot(x, x_store) atol = 1.0e-10 - - # show() output should contain the symbolic direction labels (instead of testing non-exported fun_name) - io = IOBuffer(); show(io, F); strF = String(take!(io)); @test occursin("δx", strF) + io = IOBuffer() + show(io, F) + @test occursin("δx", String(take!(io))) Fy = FiniteDiff(Float64, (3, 4), 2) - io = IOBuffer(); show(io, Fy); strFy = String(take!(io)); @test occursin("δy", strFy) + io = IOBuffer() + show(io, Fy) + @test occursin("δy", String(take!(io))) + @test size(FiniteDiff(Float64, (3, 4, 5), 2)) == ((3, 3, 5), (3, 4, 5)) +end + +@testitem "FiniteDiff (GPU)" tags = [:gpu, :linearoperator, :FiniteDiff] setup = [TestUtils, FiniteDiffTestHelper] begin + using Random, AbstractOperators, GPUEnv - # size correctness for higher-dim - F3 = FiniteDiff(Float64, (3, 4, 5), 2) - @test size(F3) == ((3, 3, 5), (3, 4, 5)) + for backend in gpu_backends() + Random.seed!(0) + test_finitediff_mul(x -> to_gpu(backend, x), false, test_op) + end end diff --git a/test/linearoperators/test_getindex.jl b/test/linearoperators/test_getindex.jl index 1fec84c..4645acf 100644 --- a/test/linearoperators/test_getindex.jl +++ b/test/linearoperators/test_getindex.jl @@ -1,55 +1,32 @@ -@testitem "GetIndex" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin +@testitem "GetIndex: basic mul" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin using Random, LinearAlgebra, AbstractOperators Random.seed!(0) - verb && println(" --- Testing GetIndex --- ") + n, k = 5, 3 + test_op(GetIndex(zeros(Float64, n), (1:k,)), randn(n), randn(k), verb) n, m = 5, 4 k = 3 - op = GetIndex(Float64, (n,), (1:k,)) - x1 = randn(n) - y1 = test_op(op, x1, randn(k), verb) + test_op(GetIndex(zeros(Float64, n, m), (1:k, :)), randn(n, m), randn(k, m), verb) - @test all(norm.(y1 .- x1[1:k]) .<= 1.0e-12) + op = GetIndex(Float64, (n,), (1:k,); array_type = Array{ComplexF32, 2}) + @test domain_array_type(op) == Array{Float64} + @test codomain_array_type(op) == Array{Float64} - n, m = 5, 4 - k = 3 - op = GetIndex(Float64, (n, m), (1:k, :)) + op2d = GetIndex(Float64, (n, m), (:, 2)) x1 = randn(n, m) - y1 = test_op(op, x1, randn(k, m), verb) - - @test all(norm.(y1 .- x1[1:k, :]) .<= 1.0e-12) - - n, m = 5, 4 - op = GetIndex(Float64, (n, m), (:, 2)) - x1 = randn(n, m) - y1 = test_op(op, x1, randn(n), verb) - + y1 = test_op(op2d, x1, randn(n), verb) @test all(norm.(y1 .- x1[:, 2]) .<= 1.0e-12) - n, m, l = 5, 4, 3 - op = GetIndex(Float64, (n, m, l), (1:3, 2, :)) - x1 = randn(n, m, l) - y1 = test_op(op, x1, randn(3, 3), verb) - - @test all(norm.(y1 .- x1[1:3, 2, :]) .<= 1.0e-12) - - # other constructors - GetIndex((n, m), (1:k, :)) - GetIndex(x1, (1:k, :, :)) - - # SubArray constructor path - xv = randn(5, 5) - sv = @view xv[1:3, :] - g_sub = GetIndex(sv, 1:2) - @test size(g_sub, 1) == (2,) - @test_throws BoundsError GetIndex(Float64, (n, m), (1:k, :, :)) - op = GetIndex(Float64, (n, m), (1:n, 1:m)) - @test typeof(op) <: Eye + @test typeof(GetIndex(Float64, (n, m), (1:n, 1:m))) <: Eye +end - op = GetIndex(Float64, (n,), (1:k,)) +@testitem "GetIndex: properties" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) - ##properties + n, m, k = 5, 4, 3 + op = GetIndex(Float64, (n,), (1:k,)) @test is_sliced(op) == true @test is_linear(op) == true @test is_null(op) == false @@ -61,191 +38,247 @@ @test is_invertible(op) == false @test is_full_row_rank(op) == true @test is_full_column_rank(op) == false - @test diag_AAc(op) == 1 + @test AbstractOperators.has_fast_opnorm(op) == true + @test opnorm(op) == estimate_opnorm(op) +end + +@testitem "GetIndex: boolean mask and CartesianIndex" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) - # Boolean mask partial selection + n, m = 5, 4 mask = falses(n, m) mask[1:2, 2] .= true op_mask = GetIndex(Float64, (n, m), mask) xmask = randn(n, m) ymask = op_mask * xmask - @test length(ymask) == sum(mask) - @test all(ymask .== xmask[mask]) + @test length(ymask) == sum(mask) && all(ymask .== xmask[mask]) - # Boolean mask selecting all => should behave like Eye reshaped fullmask = trues(n, m) op_full = GetIndex(Float64, (n, m), fullmask) xfull = randn(n, m) yfull = op_full * xfull - @test prod(size(xfull)) == length(yfull) - # Applying adjoint should scatter back back = zeros(size(xfull)) mul!(back, op_full', yfull) @test back ≈ xfull - # Vector of CartesianIndex selection + mask_vec = [isodd(i) for i in 1:n] + op_mask_vec = GetIndex(Float64, (n,), mask_vec) + xmask_vec = randn(n) + @test op_mask_vec * xmask_vec == xmask_vec[mask_vec] + + op_mask_vec_tuple = GetIndex(Float64, (n,), (mask_vec,)) + @test op_mask_vec_tuple * xmask_vec == xmask_vec[mask_vec] + cart = collect(CartesianIndices((n, m))[1:4]) op_cart = GetIndex(Float64, (n, m), cart) xcart = randn(n, m) - ycart = op_cart * xcart - @test ycart == xcart[cart] + @test op_cart * xcart == xcart[cart] +end - # Normal operator A'A (should be diagonal identity restricted) +@testitem "GetIndex: normal op and slicing helpers" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) + + n, m = 5, 4 A = GetIndex(Float64, (n, m), (1:3, :)) xA = randn(n, m) - yA = A * xA normal = AbstractOperators.get_normal_op(A) - # normal maps full shape -> full shape tmp = similar(xA) mul!(tmp, normal, xA) - # For GetIndex selecting subset rows, A'A should be a projector with ones on selected entries proj = zeros(size(xA)) proj[1:3, :] .= xA[1:3, :] @test tmp == proj + @test typeof(AbstractOperators.get_normal_op(A')) <: Eye - # normal of adjoint is Eye on domain - normal_adj = AbstractOperators.get_normal_op(A') - @test typeof(normal_adj) <: Eye - - # Slicing helpers @test AbstractOperators.get_slicing_expr(A) == (1:3, :) - mask_expr = AbstractOperators.get_slicing_mask(op_mask) - @test sum(mask_expr) == sum(mask) + mask = falses(n, m) + mask[1:2, 2] .= true + op_mask = GetIndex(Float64, (n, m), mask) + @test sum(AbstractOperators.get_slicing_mask(op_mask)) == sum(mask) - # remove_slicing returns Eye of original domain base_eye = AbstractOperators.remove_slicing(A) @test typeof(base_eye) <: Eye @test size(base_eye) == (size(A, 1), size(A, 1)) - # opnorm vs estimate_opnorm - @test AbstractOperators.has_fast_opnorm(A) == true - @test opnorm(A) == estimate_opnorm(A) - - # show output should contain arrow-like symbol for GetIndex - io = IOBuffer(); show(io, A); strA = String(take!(io)); @test occursin("↓", strA) - - # Dimension mismatch errors in mul! - bad_y = zeros(size(A, 1)..., 2) # deliberately wrong extra dim - @test_throws DimensionMismatch mul!(bad_y, A, xA) - - # Additional coverage-focused tests for uncovered GetIndex / NormalGetIndex branches + io = IOBuffer() + show(io, A) + @test occursin("↓", String(take!(io))) +end - @testset "GetIndex scalar Int... specialized path" begin - n, m = 6, 7 - op = GetIndex(Float64, (n, m), (2, 3)) # both indices scalar Int → get_dim_out(::Dims, Int...) method - x = randn(n, m) - y = op * x - @test size(op) == ((1,), (n, m)) - @test length(y) == 1 - @test y[1] == x[2, 3] +@testitem "GetIndex (GPU)" tags = [:gpu, :linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + n, m, k = 5, 4, 3 + backend_type_wrapper = gpu_wrapper(backend, randn(1)) + + test_op(GetIndex(gpu_zeros(backend, Float64, n), (1:k,)), gpu_randn(backend, n), gpu_randn(backend, k), false) + test_op( + GetIndex(gpu_zeros(backend, Float64, n, m), (1:k, :)), + gpu_randn(backend, n, m), + gpu_randn(backend, k, m), + false, + ) + test_op( + GetIndex(gpu_zeros(backend, Float64, n, m), (:, 2)), + gpu_randn(backend, n, m), + gpu_randn(backend, n), + false, + ) + test_op( + GetIndex(gpu_zeros(backend, Float64, n, m), (2:k, 2:m)), + gpu_randn(backend, n, m), + gpu_randn(backend, k - 1, m - 1), + false, + ) + test_op( + GetIndex(gpu_zeros(backend, Float64, n, m), (1:k, 2)), + gpu_randn(backend, n, m), + gpu_randn(backend, k), + false, + ) + + idx_vec = collect(1:2:n) + op_vec = GetIndex(gpu_zeros(backend, Float64, n), idx_vec) + @test Base.typename(typeof(op_vec.idx[1])).wrapper == backend_type_wrapper + test_op(op_vec, gpu_randn(backend, n), gpu_randn(backend, length(idx_vec)), false) + op_vec_kw = GetIndex(Float64, (n,), idx_vec; array_type = gpu_wrapper(backend, Float64, 1)) + @test Base.typename(typeof(op_vec_kw.idx[1])).wrapper == backend_type_wrapper + + mask = falses(n) + mask[2:2:n] .= true + op_mask = GetIndex(gpu_zeros(backend, Float64, n), mask) + @test Base.typename(typeof(op_mask.idx[1])).wrapper == backend_type_wrapper + @test length(op_mask.idx[1]) == sum(mask) + test_op(op_mask, gpu_randn(backend, n), gpu_randn(backend, sum(mask)), false) + op_mask_kw = GetIndex(Float64, (n,), mask; array_type = gpu_wrapper(backend, Float64, 1)) + @test Base.typename(typeof(op_mask_kw.idx[1])).wrapper == backend_type_wrapper end +end - @testset "GetIndex boolean mask inside tuple branch" begin - n, m, l = 5, 4, 3 - mask_first = falses(n); mask_first[2:4] .= true # Boolean vector for first dimension - op = GetIndex(Float64, (n, m, l), (mask_first, :, 2)) # mask used as one index among others - x = randn(n, m, l) - y = op * x - @test length(y) == sum(mask_first) * m # last dim fixed → product of kept rows and full second dim - expected = x[mask_first, :, 2] - @test y == expected - end +@testitem "GetIndex scalar Int specialized path" tags = [:linearoperator, :GetIndex] setup = [ + TestUtils, +] begin + using Random, AbstractOperators + Random.seed!(0) + n, m = 6, 7 + op = GetIndex(Float64, (n, m), (2, 3)) + x = randn(n, m) + y = op * x + @test size(op) == ((1,), (n, m)) + @test length(y) == 1 + @test y[1] == x[2, 3] +end - @testset "GetIndex AbstractArray of Int indices (multi-dim)" begin - n = 10 - arr_idx = reshape([1, 2, 3, 4], 2, 2) # proper reshape with dimensions - op = GetIndex(Float64, (n,), (arr_idx,)) # index array inside tuple triggers AbstractArray branch - x = randn(n) - y = op * x - @test size(y) == (2, 2) - @test y == x[arr_idx] - end +@testitem "GetIndex boolean mask inside tuple branch" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n, m, l = 5, 4, 3 + mask_first = falses(n) + mask_first[2:4] .= true + op = GetIndex(Float64, (n, m, l), (mask_first, :, 2)) + x = randn(n, m, l) + y = op * x + @test length(y) == sum(mask_first) * m + @test y == x[mask_first, :, 2] +end - @testset "GetIndex unsupported index type error path" begin - n, m = 5, 4 - @test_throws ArgumentError GetIndex(Float64, (n, m), (1:2, "bad")) - end +@testitem "GetIndex AbstractArray of Int indices (multi-dim)" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + arr_idx = reshape([1, 2, 3, 4], 2, 2) + op = GetIndex(Float64, (10,), (arr_idx,)) + x = randn(10) + y = op * x + @test size(y) == (2, 2) + @test y == x[arr_idx] +end - @testset "NormalGetIndex non-tuple idx conversion" begin - n = 8 - N1 = AbstractOperators.NormalGetIndex(Float64, Array{Float64}, (n,), 3:3) - @test size(N1) == ((n,), (n,)) - @test AbstractOperators.domain_type(N1) == Float64 - @test AbstractOperators.domain_storage_type(N1) == Array{Float64} - # Use diag to avoid 0-dim view assignment in mul! for scalar idx - d = AbstractOperators.diag(N1) - expected = zeros(n); expected[3] = 1 - @test d == expected - end +@testitem "GetIndex unsupported index type error path" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + @test_throws ArgumentError GetIndex(Float64, (5, 4), (1:2, "bad")) +end - @testset "NormalGetIndex vector idx conversion" begin - n = 9 - idx_vec = [2, 4, 5] - N2 = AbstractOperators.NormalGetIndex(Float64, Array{Float64}, (n,), idx_vec) - @test size(N2) == ((n,), (n,)) - d = AbstractOperators.diag(N2) - @test d[idx_vec] == ones(length(idx_vec)) - @test sum(d) == length(idx_vec) - end +@testitem "NormalGetIndex non-tuple idx conversion" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n = 8 + N1 = AbstractOperators.NormalGetIndex(Float64, Array{Float64}, (n,), 3:3) + @test size(N1) == ((n,), (n,)) + @test AbstractOperators.domain_type(N1) == Float64 + @test AbstractOperators.domain_array_type(N1) == Array{Float64} + expected = zeros(n) + expected[3] = 1 + @test AbstractOperators.diag(N1) == expected +end - @testset "GetIndex get_idx accessor" begin - n = 7 - op = GetIndex(Float64, (n,), (2:5,)) - idx_back = AbstractOperators.get_idx(op) - @test idx_back == (2:5,) - end +@testitem "NormalGetIndex vector idx conversion" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + idx_vec = [2, 4, 5] + N2 = AbstractOperators.NormalGetIndex(Float64, Array{Float64}, (9,), idx_vec) + d = AbstractOperators.diag(N2) + @test d[idx_vec] == ones(length(idx_vec)) + @test sum(d) == length(idx_vec) +end - # Cover x::AbstractArray overloads and BitArray-specific slicing mask - @testset "GetIndex array-first overloads and BitArray mask" begin - # Tuple indices equal to full dims -> Eye path in x::Tuple overload - x2 = randn(4, 3) - op_eye_x = GetIndex(x2, (:, :)) - @test op_eye_x isa Eye - @test op_eye_x * x2 == x2 - - # Vector-of-Int indices equal to full length on 1D -> Eye path in x::Vector overload - xv = randn(6) - op_eye_vec = GetIndex(xv, collect(1:length(xv))) - @test op_eye_vec isa Eye - @test op_eye_vec * xv == xv - - # BitArray mask-specific get_slicing_mask specialization - n, m = 3, 5 - bmask = trues(n, m) # BitMatrix <: BitArray - op_bmask = GetIndex(Float64, (n, m), bmask) - # Only call get_slicing_mask if result is actually a GetIndex; otherwise it may be a Reshape(Eye) - if op_bmask isa GetIndex - m2 = AbstractOperators.get_slicing_mask(op_bmask) - @test m2 === bmask - else - @test size(op_bmask) == ((n * m,), (n, m)) # reshaped Eye path - end - end +@testitem "GetIndex get_idx accessor" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + @test AbstractOperators.get_idx(GetIndex(Float64, (7,), (2:5,))) == (2:5,) +end - # Small utility error path in get_dim_out(::Dims) - @testset "get_dim_out missing indices error" begin - @test_throws ErrorException AbstractOperators.get_dim_out((2, 3)) +@testitem "GetIndex array-first overloads and BitArray mask" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + x2 = randn(4, 3) + op_eye_x = GetIndex(x2, (:, :)) + @test op_eye_x isa Eye + @test op_eye_x * x2 == x2 + xv = randn(6) + op_eye_vec = GetIndex(xv, collect(1:length(xv))) + @test op_eye_vec isa Eye + @test op_eye_vec * xv == xv + n, m = 3, 5 + bmask = trues(n, m) + op_bmask = GetIndex(Float64, (n, m), bmask) + if op_bmask isa GetIndex + @test AbstractOperators.get_slicing_mask(op_bmask) === bmask + else + @test size(op_bmask) == ((n * m,), (n, m)) end - # Specialized adjoint for NormalGetIndex returns itself - @testset "AdjointOperator(NormalGetIndex) identity" begin - n = 5 - N = AbstractOperators.NormalGetIndex(Float64, Array{Float64}, (n,), 2:3) - @test AbstractOperators.AdjointOperator(N) === N - @test AbstractOperators.codomain_storage_type(N) == Array{Float64} - end + bool_vec = [i % 2 == 0 for i in 1:length(xv)] + op_bool_vec = GetIndex(xv, bool_vec) + @test op_bool_vec * xv == xv[bool_vec] +end - @testset "GetIndex vector of CartesianIndex via tuple (hits get_dim_out branch)" begin - n, m = 6, 5 - cart = collect(CartesianIndices((n, m))[1:6]) - # Wrap as a tuple to force the GetIndex(domain_type, dim_in, idx::Tuple) constructor, - # which calls get_dim_out and should take the AbstractVector{CartesianIndex} branch - op = GetIndex(Float64, (n, m), (cart,)) - x = randn(n, m) - y = op * x - @test size(y) == (length(cart),) - @test y == x[cart] - end +@testitem "get_dim_out missing indices error" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + @test_throws ErrorException AbstractOperators.get_dim_out((2, 3)) +end +@testitem "AdjointOperator(NormalGetIndex) identity" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + N = AbstractOperators.NormalGetIndex(Float64, Array{Float64}, (5,), 2:3) + @test AbstractOperators.AdjointOperator(N) === N + @test AbstractOperators.codomain_array_type(N) == Array{Float64} +end + +@testitem "GetIndex vector of CartesianIndex via tuple" tags = [:linearoperator, :GetIndex] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n, m = 6, 5 + cart = collect(CartesianIndices((n, m))[1:6]) + op = GetIndex(Float64, (n, m), (cart,)) + x = randn(n, m) + y = op * x + @test size(y) == (length(cart),) + @test y == x[cart] end diff --git a/test/linearoperators/test_lbfgs.jl b/test/linearoperators/test_lbfgs.jl index 100d389..0594e73 100644 --- a/test/linearoperators/test_lbfgs.jl +++ b/test/linearoperators/test_lbfgs.jl @@ -1,7 +1,30 @@ -@testitem "L-BFGS" tags = [:linearoperator, :LBFGS] setup = [TestUtils] begin +@testitem "L-BFGS: construction and basic mul" tags = [:linearoperator, :LBFGS] setup = [TestUtils] begin + using AbstractOperators + using AbstractOperators: LBFGS, update!, mul!, reset! + + mem = 3 + x = zeros(10) + + H = LBFGS(x, mem) + # Basic properties after construction (no memory yet) + @test size(H) == (size(x), size(x)) + @test domain_type(H) == eltype(x) + @test codomain_type(H) == eltype(x) + @test is_thread_safe(H) == false + # Initial operator should act like identity (H.H = 1, empty memory) + @test H * x == x # stochastic but value not stored; run separately below + dir = zeros(10) + verb && println(H) + + HH = LBFGS(ArrayPartition(x, x), mem) + @test size(HH) == (size.(ArrayPartition(x, x).x), size.(ArrayPartition(x, x).x)) + dirdir = ArrayPartition(zeros(10), zeros(10)) + verb && println(HH) +end + +@testitem "L-BFGS: update and two-loop recursion" tags = [:linearoperator, :LBFGS] setup = [TestUtils] begin using AbstractOperators using AbstractOperators: LBFGS, update!, mul!, reset! - verb && println(" --- Testing L-BFGS --- ") Q = [ 32.0 13.1 -4.9 -3.0 6.0 2.2 2.6 3.4 -1.9 -7.5 @@ -48,20 +71,9 @@ x = zeros(10) H = LBFGS(x, mem) - # Basic properties after construction (no memory yet) - @test size(H) == (size(x), size(x)) - @test domain_type(H) == eltype(x) - @test codomain_type(H) == eltype(x) - @test is_thread_safe(H) == false - # Initial operator should act like identity (H.H = 1, empty memory) - @test H * x == x # stochastic but value not stored; run separately below - dir = zeros(10) - verb && println(H) - HH = LBFGS(ArrayPartition(x, x), mem) - @test size(HH) == (size.(ArrayPartition(x, x).x), size.(ArrayPartition(x, x).x)) + dir = zeros(10) dirdir = ArrayPartition(zeros(10), zeros(10)) - verb && println(HH) let x_old = [], grad_old = [] for i in 1:5 @@ -113,6 +125,61 @@ grad_old = grad end end # let x_old, grad_old +end + +@testitem "L-BFGS: memory limit and reset" tags = [:linearoperator, :LBFGS] setup = [TestUtils] begin + using AbstractOperators + using AbstractOperators: LBFGS, update!, mul!, reset! + + Q = [ + 32.0 13.1 -4.9 -3.0 6.0 2.2 2.6 3.4 -1.9 -7.5 + 13.1 18.3 -5.3 -9.5 3.0 2.1 3.9 3.0 -3.6 -4.4 + -4.9 -5.3 7.7 2.1 -0.4 -3.4 -0.8 -3.0 5.3 5.5 + -3.0 -9.5 2.1 20.1 1.1 0.8 -12.4 -2.5 5.5 2.1 + 6.0 3.0 -0.4 1.1 3.8 0.6 0.5 0.9 -0.4 -2.0 + 2.2 2.1 -3.4 0.8 0.6 7.8 2.9 -1.3 -4.3 -5.1 + 2.6 3.9 -0.8 -12.4 0.5 2.9 14.5 1.7 -4.9 1.2 + 3.4 3.0 -3.0 -2.5 0.9 -1.3 1.7 6.6 -0.8 2.7 + -1.9 -3.6 5.3 5.5 -0.4 -4.3 -4.9 -0.8 7.9 5.7 + -7.5 -4.4 5.5 2.1 -2.0 -5.1 1.2 2.7 5.7 16.1 + ] + + q = [ + 2.9, 0.8, 1.3, -1.1, -0.5, -0.3, 1.0, -0.3, 0.7, -2.1, + ] + + xs = + [ + 1.0 0.01 0.02 0.03 0.04 0.05 0.06 0.07 0.08 0.09 + 0.09 1.0 0.01 0.02 0.03 0.04 0.05 0.06 0.07 0.08 + 0.08 0.09 1.0 0.01 0.02 0.03 0.04 0.05 0.06 0.07 + 0.07 0.08 0.09 1.0 0.01 0.02 0.03 0.04 0.05 0.06 + 0.06 0.07 0.08 0.09 1.0 0.01 0.02 0.03 0.04 0.05 + ]' + + mem = 3 + x = zeros(10) + + H = LBFGS(x, mem) + HH = LBFGS(ArrayPartition(x, x), mem) + + # Run all updates to put H and HH in post-loop state + let x_old = [], grad_old = [] + for i in 1:5 + x = xs[:, i] + grad = Q * x + q + if i > 1 + xx = ArrayPartition(x, copy(x)) + xx_old = ArrayPartition(x_old, copy(x_old)) + gradgrad = ArrayPartition(grad, copy(grad)) + gradgrad_old = ArrayPartition(grad_old, copy(grad_old)) + update!(H, x, x_old, grad, grad_old) + update!(HH, xx, xx_old, gradgrad, gradgrad_old) + end + x_old = x + grad_old = grad + end + end # Memory limit: ensure no more than mem updates stored @test H.currmem <= mem @@ -127,7 +194,10 @@ @test H.currmem == prev_mem # Show output symbol - io = IOBuffer(); show(io, H); s = String(take!(io)); @test occursin("LBFGS", s) + io = IOBuffer() + show(io, H) + s = String(take!(io)) + @test occursin("LBFGS", s) # Testing reset @@ -140,3 +210,43 @@ @test ones(size(H, 1)) == H * ones(size(H, 1)) @test ArrayPartition(ones.(size(HH, 1))) == HH * ArrayPartition(ones.(size(HH, 1))...) end + +@testitem "L-BFGS (GPU)" tags = [:gpu, :linearoperator, :LBFGS] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(42) + + n = 64 + mem = 5 + x_prev = gpu_randn(backend, Float32, n) + grad_prev = gpu_randn(backend, Float32, n) + H = LBFGS(x_prev, mem) + @test domain_array_type(H) <: backend.array_type + @test codomain_array_type(H) <: backend.array_type + + x = gpu_randn(backend, Float32, n) + grad = gpu_randn(backend, Float32, n) + update!(H, x, x_prev, grad, grad_prev) + + y = H * grad + @test y isa typeof(x_prev) + y2 = similar(grad) + mul!(y2, H, grad) + @test collect(y) ≈ collect(y2) + + g = gpu_randn(backend, Float32, n) + lhs = collect(H' * g) + rhs = similar(g) + mul!(rhs, H', g) + @test lhs ≈ collect(rhs) + + xp = ArrayPartition(gpu_randn(backend, Float32, n), gpu_randn(backend, Float32, n)) + gp = ArrayPartition(gpu_randn(backend, Float32, n), gpu_randn(backend, Float32, n)) + Hp = LBFGS(xp, mem) + @test domain_array_type(Hp) == typeof(xp) + @test codomain_array_type(Hp) == typeof(xp) + yp = Hp * gp + @test yp isa typeof(xp) + end +end diff --git a/test/linearoperators/test_lmatrixop.jl b/test/linearoperators/test_lmatrixop.jl index ac887af..a3c7499 100644 --- a/test/linearoperators/test_lmatrixop.jl +++ b/test/linearoperators/test_lmatrixop.jl @@ -1,15 +1,22 @@ -@testitem "LMatrixOp" tags = [:linearoperator, :LMatrixOp] setup = [TestUtils] begin +@testitem "LMatrixOp: basic mul" tags = [:linearoperator, :LMatrixOp] setup = [TestUtils] begin using Random, AbstractOperators + Random.seed!(0) - verb && println(" --- Testing LMatrixOp --- ") n, m = 5, 6 b = randn(m) op = LMatrixOp(Float64, (n, m), b) + test_op(op, randn(n, m), randn(n), verb) + + n, m = 5, 6 + b = randn(m) + op = LMatrixOp(Float64, (n, m), b) + op_array_type = LMatrixOp(Float64, (n, m), b; array_type = Array{ComplexF32, 2}) + @test domain_array_type(op_array_type) == Array{Float64} + @test codomain_array_type(op_array_type) == Array{Float64} x1 = randn(n, m) y1 = test_op(op, x1, randn(n), verb) y2 = x1 * b - @test all(norm.(y1 .- y2) .<= 1.0e-12) # size (codomain, domain) @test size(op) == ((n,), (n, m)) @@ -22,7 +29,6 @@ x1 = randn(n, m) + im * randn(n, m) y1 = test_op(op, x1, randn(n) + im * randn(n), verb) y2 = x1 * b - @test all(norm.(y1 .- y2) .<= 1.0e-12) @test size(op) == ((n,), (n, m)) @@ -32,7 +38,6 @@ x1 = randn(n, m) y1 = test_op(op, x1, randn(n, l), verb) y2 = x1 * b - @test all(norm.(y1 .- y2) .<= 1.0e-12) @test size(op) == ((n, l), (n, m)) @@ -42,9 +47,15 @@ x1 = randn(n, m) + im * randn(n, m) y1 = test_op(op, x1, randn(n, l) + im * randn(n, l), verb) y2 = x1 * b - @test all(norm.(y1 .- y2) .<= 1.0e-12) @test size(op) == ((n, l), (n, m)) +end + +@testitem "LMatrixOp: other constructors" tags = [:linearoperator, :LMatrixOp] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + + n, m, l = 5, 6, 7 ## other constructors (vector b and matrix b) bvec = randn(m) @@ -79,6 +90,16 @@ Zout2 = zeros(n, m) mul!(Zout2, op_mat2', Zrhs) @test Zout2 ≈ Zrhs * b_mat' +end + +@testitem "LMatrixOp: scale and properties" tags = [:linearoperator, :LMatrixOp] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + + n, m, l = 5, 6, 7 + bvec = randn(m) + op_vec = LMatrixOp(bvec, n) + X = randn(n, m) # Scaling Sop = Scale(2.5, op_vec) @@ -92,6 +113,9 @@ Xbad = randn(n, m + 1) @test_throws DimensionMismatch (op_vec * Xbad) + b = randn(m, l) + im * randn(m, l) + op = LMatrixOp(Complex{Float64}, (n, m), b) + ##properties @test is_linear(op) == true @test is_null(op) == false @@ -105,5 +129,19 @@ @test is_full_column_rank(op) == false # Show output symbol - io = IOBuffer(); show(io, op_vec); s = String(take!(io)); @test occursin("(⋅)b", s) + io = IOBuffer() + show(io, op_vec) + s = String(take!(io)) + @test occursin("(⋅)b", s) +end + +@testitem "LMatrixOp (GPU)" tags = [:gpu, :linearoperator, :LMatrixOp] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + for backend in gpu_backends() + Random.seed!(0) + n, m = 5, 6 + b = gpu_randn(backend, m) + op = LMatrixOp(Float64, (n, m), b) + test_op(op, gpu_randn(backend, n, m), gpu_randn(backend, n), false) + end end diff --git a/test/linearoperators/test_matrixop.jl b/test/linearoperators/test_matrixop.jl index 3a34167..4a9e3ba 100644 --- a/test/linearoperators/test_matrixop.jl +++ b/test/linearoperators/test_matrixop.jl @@ -1,62 +1,58 @@ -@testitem "MatrixOp" tags = [:linearoperator, :MatrixOp] setup = [TestUtils] begin +@testitem "MatrixOp: basic mul" tags = [:linearoperator, :MatrixOp] setup = [TestUtils] begin using Random, SparseArrays, LinearAlgebra, AbstractOperators Random.seed!(0) - verb && println(" --- Testing MatrixOp --- ") - # real matrix, real input n, m = 5, 4 A = randn(n, m) - op = MatrixOp(Float64, (m,), A) - x1 = randn(m) - y1 = test_op(op, x1, randn(n), verb) - y2 = A * x1 - - # real matrix, complex input - n, m = 5, 4 - A = randn(n, m) - op = MatrixOp(Complex{Float64}, (m,), A) - x1 = randn(m) + im .* randn(m) - y1 = test_op(op, x1, randn(n) + im * randn(n), verb) - y2 = A * x1 - - # complex matrix, complex input - n, m = 5, 4 - A = randn(n, m) + im * randn(n, m) - op = MatrixOp(Complex{Float64}, (m,), A) - x1 = randn(m) + im .* randn(m) - y1 = test_op(op, x1, randn(n) + im * randn(n), verb) - y2 = A * x1 + op = MatrixOp(A) + test_op(op, randn(m), randn(n), verb) - # complex matrix, real input - n, m = 5, 4 A = randn(n, m) + im * randn(n, m) - op = MatrixOp(Float64, (m,), A) - x1 = randn(m) - y1 = test_op(op, x1, randn(n) + im * randn(n), verb) - y2 = A * x1 + op = MatrixOp(A) + test_op(op, randn(m) + im * randn(m), randn(n) + im * randn(n), verb) - @test all(norm.(y1 .- y2) .<= 1.0e-12) + # array_type keyword is ignored for element-type selection + A = randn(n, m) + op = MatrixOp(Float64, (m,), A; array_type = Array{ComplexF32, 2}) + @test domain_array_type(op) == Array{Float64} + @test codomain_array_type(op) == Array{Float64} # complex matrix, real matrix input + A = randn(n, m) + im * randn(n, m) c = 3 op = MatrixOp(Float64, (m, c), A) - @test_throws ErrorException MatrixOp(Float64, (m, c, 3), A) - @test_throws MethodError MatrixOp(Float64, (m, c), randn(n, m, 2)) x1 = randn(m, c) - y1 = test_op(op, x1, randn(n, c) .+ im .* randn(n, c), verb) # codomain is ComplexF64 - y2 = A * x1 + y1 = test_op(op, x1, randn(n, c) .+ im .* randn(n, c), verb) + @test all(norm.(y1 .- A * x1) .<= 1.0e-12) +end + +@testitem "MatrixOp: constructors" tags = [:linearoperator, :MatrixOp] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) - # other constructors + n, m = 5, 4 + A = randn(n, m) + im * randn(n, m) + c = 3 op = MatrixOp(A) op = MatrixOp(Float64, A) op = MatrixOp(A, c) op = MatrixOp(Float64, A, c) - op = convert(LinearOperator, A) op = convert(LinearOperator, A, c) - op = convert(LinearOperator, Complex{Float64}, size(x1), A) + op = convert(LinearOperator, Complex{Float64}, (m, c), A) + + @test_throws ErrorException MatrixOp(Float64, (m, c, 3), A) + @test_throws MethodError MatrixOp(Float64, (m, c), randn(n, m, 2)) +end - #properties +@testitem "MatrixOp: properties" tags = [:linearoperator, :MatrixOp] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) + + n, m = 5, 4 + A = randn(n, m) + im * randn(n, m) + c = 3 + op = MatrixOp(Float64, (m, c), A) @test is_sliced(op) == false @test is_linear(op) == true @test is_null(op) == false @@ -69,67 +65,63 @@ @test is_full_row_rank(MatrixOp(randn(Random.seed!(0), 3, 4))) == true @test is_full_column_rank(MatrixOp(randn(Random.seed!(0), 3, 4))) == false - # Square invertible matrix B = randn(4, 4) invop = MatrixOp(Float64, (4,), B) @test is_invertible(invop) == (det(B) != 0) - # Orthogonal case (approx) Q, _ = qr(randn(5, 5)) Qmat = Matrix(Q) - Qop = MatrixOp(Float64, (5,), Qmat) - @test is_orthogonal(Qop) == true + @test is_orthogonal(MatrixOp(Float64, (5,), Qmat)) == true - # Diagonal detection path D = diagm(0 => randn(5)) Dop = MatrixOp(Float64, (5,), D) @test is_diagonal(Dop) == true @test is_AcA_diagonal(Dop) == true @test is_AAc_diagonal(Dop) == true +end - # Scale: real codomain, complex coeff should error - @test_throws ErrorException Scale(1.0 + 2im, invop) - Sop = Scale(2.0, invop) - xS = randn(4) - @test Sop * xS ≈ 2.0 * (invop * xS) - - # Normal operator A'A - Nop = AbstractOperators.get_normal_op(invop) - @test size(Nop) == ((4,), (4,)) - xN = randn(4) - @test Nop * xN ≈ invop' * (invop * xN) - - # Adjoint mapping with complex matrix and real input (special method) - Cmat = randn(3, 3) + im * randn(3, 3) - Cop = MatrixOp(Float64, (3,), Cmat) - xr = randn(3) - yr = Cop' * xr - # Compare with manual real of complex multiplication - @test yr ≈ real.(Cmat' * xr) +@testitem "MatrixOp: adjoint and in-place" tags = [:linearoperator, :MatrixOp] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) - # In-place mul! + Q, _ = qr(randn(5, 5)) + Qmat = Matrix(Q) + Qop = MatrixOp(Float64, (5,), Qmat) xvec = randn(5) yvec = zeros(5) mul!(yvec, Qop, xvec) @test yvec ≈ Qmat * xvec - - # Batched (matrix input) in-place multiplication Xmat = randn(5, 3) Ymat = zeros(5, 3) mul!(Ymat, Qop, Xmat) @test Ymat ≈ Qmat * Xmat - # opnorm vs estimate_opnorm - @test opnorm(invop) ≈ estimate_opnorm(invop) - @test opnorm(invop) ≈ opnorm(B) + Cmat = randn(3, 3) + im * randn(3, 3) + Cop = MatrixOp(Float64, (3,), Cmat) + xr = randn(3) + @test Cop' * xr ≈ real.(Cmat' * xr) - # Size variations: single vs multi-column - Tall = randn(8, 3) - TallOp = MatrixOp(Float64, (3,), Tall) - TallOpMulti = MatrixOp(Float64, (3, 2), Tall) - @test size(TallOp) == ((8,), (3,)) - @test size(TallOpMulti) == ((8, 2), (3, 2)) + Nop = AbstractOperators.get_normal_op(Qop) + v = randn(5) + @test Nop * v ≈ Qop' * (Qop * v) + + @test opnorm(Qop) ≈ estimate_opnorm(Qop) rtol = 0.02 +end - # Show output symbol - io = IOBuffer(); show(io, TallOp); str = String(take!(io)); @test occursin("▒", str) +@testitem "MatrixOp (GPU)" tags = [:gpu, :linearoperator, :MatrixOp] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + n, m = 5, 4 + A = gpu_randn(backend, n, m) + test_op(MatrixOp(A), gpu_randn(backend, m), gpu_randn(backend, n), false) + Ac = gpu_randn(backend, ComplexF64, n, m) + test_op( + MatrixOp(Ac), + gpu_randn(backend, ComplexF64, m), + gpu_randn(backend, ComplexF64, n), + false, + ) + end end diff --git a/test/linearoperators/test_mylinop.jl b/test/linearoperators/test_mylinop.jl index 2b68b21..9555932 100644 --- a/test/linearoperators/test_mylinop.jl +++ b/test/linearoperators/test_mylinop.jl @@ -1,11 +1,15 @@ -@testitem "MyLinOp" tags = [:linearoperator, :MyLinOp] setup = [TestUtils] begin +@testitem "MyLinOp basic" tags = [:linearoperator, :MyLinOp] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing MyLinOp --- ") n, m = 5, 4 A = randn(n, m) op = MyLinOp(Float64, (m,), (n,), (y, x) -> mul!(y, A, x), (y, x) -> mul!(y, A', x)) + op_array_type = MyLinOp( + Float64, (m,), (n,), (y, x) -> mul!(y, A, x), (y, x) -> mul!(y, A', x); array_type = Array{ComplexF32, 2} + ) + @test domain_array_type(op_array_type) == Array{Float64} + @test codomain_array_type(op_array_type) == Array{Float64} x1 = randn(m) y1 = test_op(op, x1, randn(n), verb) y2 = A * x1 @@ -15,36 +19,61 @@ @test size(op) == ((n,), (m,)) @test domain_type(op) == Float64 @test codomain_type(op) == Float64 +end - # adjoint application +@testitem "MyLinOp mul and scaling" tags = [:linearoperator, :MyLinOp] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n, m = 5, 4 + A = randn(n, m) + op = MyLinOp(Float64, (m,), (n,), (y, x) -> mul!(y, A, x), (y, x) -> mul!(y, A', x)) + x1 = randn(m) z = randn(n) out_adj = similar(x1) mul!(out_adj, op', z) @test out_adj ≈ A' * z - - # In-place forward mul! on matrix input via broadcasting columns X = randn(m, 3) Y = zeros(n, 3) for j in 1:3 mul!(view(Y, :, j), op, view(X, :, j)) end @test Y ≈ A * X - - # Scaling operator (real scale ok, complex should fail for real op) Sop = Scale(3.0, op) @test Sop * x1 ≈ 3.0 * (op * x1) @test_throws ErrorException Scale(1 + 2im, op) +end - # Error path: dimension mismatch input length +@testitem "MyLinOp constructors and errors" tags = [:linearoperator, :MyLinOp] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + n, m = 5, 4 + A = randn(n, m) + op = MyLinOp(Float64, (m,), (n,), (y, x) -> mul!(y, A, x), (y, x) -> mul!(y, A', x)) + x1 = randn(m) @test_throws DimensionMismatch op * randn(m + 1) + io = IOBuffer() + show(io, op) + s = String(take!(io)) + @test occursin("A", s) + op2 = MyLinOp(Float64, (m,), Float64, (n,), (y, x) -> mul!(y, A, x), (y, x) -> mul!(y, A', x)) + @test size(op2) == ((n,), (m,)) + @test op2 * x1 ≈ A * x1 +end - # Show output - io = IOBuffer(); show(io, op); s = String(take!(io)); @test occursin("A", s) +@testitem "MyLinOp (GPU)" tags = [:gpu, :linearoperator, :MyLinOp] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv - # other constructors - op = MyLinOp( - Float64, (m,), Float64, (n,), (y, x) -> mul!(y, A, x), (y, x) -> mul!(y, A', x) - ) - @test size(op) == ((n,), (m,)) - @test op * x1 ≈ A * x1 + for backend in gpu_backends() + Random.seed!(0) + n, m = 5, 4 + A = gpu_randn(backend, n, m) + AT = gpu_wrapper(backend, randn(1)) + op = MyLinOp( + Float64, (m,), (n,), (y, x) -> mul!(y, A, x), (y, x) -> mul!(y, A', x); array_type = AT + ) + x = gpu_randn(backend, m) + y = gpu_randn(backend, n) + test_op(op, x, y, false) + @test domain_array_type(op) <: AT + end end diff --git a/test/linearoperators/test_variation.jl b/test/linearoperators/test_variation.jl index 48707e4..d436bcf 100644 --- a/test/linearoperators/test_variation.jl +++ b/test/linearoperators/test_variation.jl @@ -1,12 +1,14 @@ -@testitem "Variation" tags = [:linearoperator, :Variation] setup = [TestUtils] begin +@testitem "Variation: basic mul" tags = [:linearoperator, :Variation] setup = [TestUtils] begin using Random, SparseArrays, LinearAlgebra, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Variation --- ") for threaded in (false, true) n, m = 10, 5 verb && println(" - threaded = $threaded") op = Variation(Float64, (n, m); threaded) + op_array_type = Variation(Float64, (n, m); threaded, array_type = Array{ComplexF32, 2}) + @test domain_array_type(op_array_type) == Array{Float64} + @test codomain_array_type(op_array_type) == Array{Float64} x1 = randn(n, m) y1 = test_op(op, x1, randn(m * n, 2), verb) # size & types @@ -34,7 +36,14 @@ x1 = randn(n, m) @test norm(op * x1 - reshape(TV * (x1[:]), n * m, 2)) < 1.0e-12 + end +end +@testitem "Variation: 3D mul and constructors" tags = [:linearoperator, :Variation] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(0) + + for threaded in (false, true) n, m, l = 100, 50, 30 verb && println(" - threaded = $threaded") op = Variation(Float64, (n, m, l); threaded) @@ -59,13 +68,23 @@ @test_throws ErrorException Variation(Float64, (n,)) badX = randn(n, m + 1) @test_throws DimensionMismatch op * badX + end +end + +@testitem "Variation: adjoint and properties" tags = [:linearoperator, :Variation] setup = [TestUtils] begin + using Random, LinearAlgebra, AbstractOperators + Random.seed!(0) + + for threaded in (false, true) + n, m, l = 100, 50, 30 + op = Variation(Float64, (n, m, l); threaded) # Adjoint consistency: == x_test = randn(n, m) verb && println(" - threaded = $threaded") V = Variation(Float64, (n, m); threaded) Y = randn(n * m, 2) - lhs = dot(V * x_test |> vec, vec(Y)) # vec(Vx) ⋅ vec(Y) + lhs = dot(vec(V * x_test), vec(Y)) # vec(Vx) ⋅ vec(Y) z = zeros(n, m) mul!(z, V', Y) rhs = dot(vec(x_test), vec(z)) @@ -85,7 +104,10 @@ @test_throws ErrorException Scale(1 + 2im, V) # Show output symbol - io = IOBuffer(); show(io, V); s = String(take!(io)); @test occursin("Ʋ", s) + io = IOBuffer() + show(io, V) + s = String(take!(io)) + @test occursin("Ʋ", s) ###properties @test is_linear(op) == true @@ -100,3 +122,14 @@ @test is_full_column_rank(op) == false end end + +@testitem "Variation (GPU)" tags = [:gpu, :linearoperator, :Variation] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + n, m = 10, 5 + op = Variation(gpu_zeros(backend, Float64, n, m); threaded = false) + test_op(op, gpu_randn(backend, n, m), gpu_randn(backend, n * m, 2), false) + end +end diff --git a/test/linearoperators/test_zeropad.jl b/test/linearoperators/test_zeropad.jl index 470348d..b6f9d4f 100644 --- a/test/linearoperators/test_zeropad.jl +++ b/test/linearoperators/test_zeropad.jl @@ -1,14 +1,46 @@ -@testitem "ZeroPad" tags = [:linearoperator, :ZeroPad] setup = [TestUtils] begin +@testmodule ZeroPadTestHelper begin + using Test, AbstractOperators, LinearAlgebra + + export test_zeropad_mul + + function test_zeropad_mul(conv, verb, test_op, to_cpu, norm) + n = (3,) + z = (5,) + op = ZeroPad(conv(zeros(Float64, n)), z) + x1 = conv(randn(n)) + y1 = test_op(op, x1, conv(randn(n .+ z)), verb) + y2 = conv([collect(x1); zeros(5)]) + @test norm(to_cpu(y1) .- to_cpu(y2)) <= 1.0e-12 + + n = (3, 2) + z = (5, 3) + op = ZeroPad(conv(zeros(Float64, n)), z) + x1 = conv(randn(n)) + y1 = test_op(op, x1, conv(randn(n .+ z)), verb) + y2c = zeros(n .+ z) + y2c[1:n[1], 1:n[2]] = collect(x1) + @test norm(to_cpu(y1) .- y2c) <= 1.0e-12 + + n = (3, 2, 2) + z = (5, 3, 1) + op = ZeroPad(conv(zeros(Float64, n)), z) + x1 = conv(randn(n)) + test_op(op, x1, conv(randn(n .+ z)), verb) + end + +end # @testmodule ZeroPadTestHelper + +@testitem "ZeroPad" tags = [:linearoperator, :ZeroPad] setup = [TestUtils, ZeroPadTestHelper] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing ZeroPad --- ") + test_zeropad_mul(identity, verb, test_op, to_cpu, norm) + + # CPU-only: type-based constructors, properties, errors n = (3,) z = (5,) op = ZeroPad(Float64, n, z) x1 = randn(n) - y1 = test_op(op, x1, randn(n .+ z), verb) - @test all(norm.(y1 .- [x1; zeros(5)]) .<= 1.0e-12) @test size(op) == (n .+ z, n) @test domain_type(op) == Float64 @test codomain_type(op) == Float64 @@ -16,44 +48,23 @@ @test AbstractOperators.has_fast_opnorm(op) == true @test opnorm(op) == 1 - n = (3, 2) - z = (5, 3) - op = ZeroPad(Float64, n, z) - x1 = randn(n) - y1 = test_op(op, x1, randn(n .+ z), verb) - y2 = zeros(n .+ z) - y2[1:n[1], 1:n[2]] = x1 - @test all(norm.(y1 .- y2) .<= 1.0e-12) - @test size(op) == (n .+ z, n) - n = (3, 2, 2) z = (5, 3, 1) op = ZeroPad(Float64, n, z) x1 = randn(n) - y1 = test_op(op, x1, randn(n .+ z), verb) - y2 = zeros(n .+ z) - y2[1:n[1], 1:n[2], 1:n[3]] = x1 - @test all(norm.(y1 .- y2) .<= 1.0e-12) - @test size(op) == (n .+ z, n) # Normal operator should be identity on input space Nop = AbstractOperators.get_normal_op(op) - xin = randn(n) - @test Nop * xin ≈ xin + @test Nop * x1 ≈ x1 - # Adjoint: crop back + # Adjoint crop ybig = zeros(n .+ z) ybig[1:n[1], 1:n[2], 1:n[3]] .= x1 xcropped = zeros(n) mul!(xcropped, op', ybig) @test xcropped ≈ x1 - # In-place forward padding - ybuf = similar(ybig) - mul!(ybuf, op, x1) - @test ybuf ≈ y2 - - # Scaling (real ok, complex rejected for real domain) + # Scaling Sop = Scale(2.0, op) @test Sop * x1 ≈ 2.0 * (op * x1) @test_throws ErrorException Scale(1 + 2im, op) @@ -61,15 +72,14 @@ # other constructors ZeroPad(n, z...) ZeroPad(Float64, n, z...) - ZeroPad(n, z...) ZeroPad(x1, z) ZeroPad(x1, z...) - #errors + # errors @test_throws ErrorException ZeroPad(Float64, n, (1, 2)) @test_throws ErrorException ZeroPad(Float64, n, (1, -2, 3)) - #properties + # properties @test is_linear(op) == true @test is_null(op) == false @test is_eye(op) == false @@ -80,9 +90,19 @@ @test is_invertible(op) == false @test is_full_row_rank(op) == false @test is_full_column_rank(op) == true - @test diag_AcA(op) == 1 - # Show output symbol - io = IOBuffer(); show(io, op); s = String(take!(io)); @test occursin("[I;0]", s) + io = IOBuffer() + show(io, op) + s = String(take!(io)) + @test occursin("[I;0]", s) +end + +@testitem "ZeroPad (GPU)" tags = [:gpu, :linearoperator, :ZeroPad] setup = [TestUtils, ZeroPadTestHelper] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + test_zeropad_mul(x -> to_gpu(backend, x), false, test_op, to_cpu, norm) + end end diff --git a/test/linearoperators/test_zerosop.jl b/test/linearoperators/test_zerosop.jl index 9c8eb75..fd49ad1 100644 --- a/test/linearoperators/test_zerosop.jl +++ b/test/linearoperators/test_zerosop.jl @@ -1,13 +1,21 @@ @testitem "Zeros" tags = [:linearoperator, :Zeros] setup = [TestUtils] begin using Random, AbstractOperators Random.seed!(0) - verb && println(" --- Testing Zeros --- ") + + n = (3, 4) + m = (5, 2) + op = Zeros(Float64, n, Float64, m) + y1 = test_op(op, randn(n), randn(m), verb) + @test y1 == zeros(Float64, m) n = (3, 4) D = Float64 m = (5, 2) C = Complex{Float64} op = Zeros(D, n, C, m) + op_array_type = Zeros(D, n, C, m; array_type = Array{ComplexF32, 2}) + @test domain_array_type(op_array_type) == Array{Float64} + @test codomain_array_type(op_array_type) == Array{ComplexF64} x1 = randn(n) y1 = test_op(op, x1, randn(m) + im * randn(m), verb) @test y1 == zeros(eltype(y1), m) @@ -56,5 +64,70 @@ @test diag_AAc(op) == 0 # Show output symbol - io = IOBuffer(); show(io, op); s = String(take!(io)); @test occursin("0", s) + io = IOBuffer() + show(io, op) + s = String(take!(io)) + @test occursin("0", s) +end + +@testitem "Zeros (GPU)" tags = [:gpu, :linearoperator, :Zeros] setup = [TestUtils] begin + using Random, AbstractOperators, GPUEnv + + for backend in gpu_backends() + Random.seed!(0) + n = (3, 4) + m = (5, 2) + op = Zeros(Float64, n, Float64, m; array_type = gpu_wrapper(backend, Float64, 1)) + y = test_op(op, gpu_randn(backend, n...), gpu_randn(backend, m...), false) + @test collect(y) == zeros(Float64, m) + end +end + +@testitem "Zeros: same-type constructor and Zeros(A) shortcut" tags = [:linearoperator, :Zeros] setup = [ + TestUtils, +] begin + using Random, AbstractOperators + Random.seed!(1) + + # Zeros(T, dim_in, dim_out) — same domain/codomain type, no storage keyword + n = (3,) + m = (5,) + op = Zeros(Float64, n, m) + @test domain_type(op) == Float64 + @test codomain_type(op) == Float64 + x = randn(n...) + y = zeros(m...) + mul!(y, op, x) + @test all(iszero, y) + mul!(x, op', y) + @test all(iszero, x) + + # Zeros(A::AbstractOperator) shortcut + ref_op = MatrixOp(randn(4, 3)) + zop = Zeros(ref_op) + @test domain_type(zop) == domain_type(ref_op) + @test codomain_type(zop) == codomain_type(ref_op) + @test size(zop) == size(ref_op) +end + +@testitem "Zeros: multi-domain (HCAT) and multi-codomain (VCAT) constructors" tags = [ + :linearoperator, :Zeros, +] setup = [TestUtils] begin + using Random, AbstractOperators + Random.seed!(2) + + # Multi-domain → returns HCAT of Zeros + z_hcat = Zeros((Float64, Float64), ((3,), (4,)), Float64, (5,)) + @test z_hcat isa HCAT + x1 = randn(3) + x2 = randn(4) + y = z_hcat * ArrayPartition(x1, x2) + @test all(iszero, y) + + # Multi-codomain → returns VCAT of Zeros + z_vcat = Zeros(Float64, (3,), (Float64, Float64), ((5,), (6,))) + @test z_vcat isa VCAT + x = randn(3) + yv = z_vcat * x + @test all(iszero, yv) end diff --git a/test/nfft/test_NfftOps.jl b/test/nfft/test_NfftOps.jl index 346819d..40c74f0 100644 --- a/test/nfft/test_NfftOps.jl +++ b/test/nfft/test_NfftOps.jl @@ -1,53 +1,57 @@ -@testitem "NFFTOp" tags = [:nfft, :NFFTOp] setup = [TestUtils] begin +@testmodule NFFTTestHelper begin + using Test + using AbstractOperators using LinearAlgebra, Random, NFFT, NFFTOperators function test_nufft_op(op, plan, image, dcf) - ksp₁ = similar(image, ComplexF64, size(op, 1)) - mul!(vec(ksp₁), plan, image) - ksp₂ = similar(ksp₁) - mul!(ksp₂, op, image) - @test ksp₂ == ksp₁ - image₂ = similar(image) + ksp1 = similar(image, ComplexF64, size(op, 1)) + mul!(vec(ksp1), plan, image) + ksp2 = similar(ksp1) + mul!(ksp2, op, image) + @test ksp2 == ksp1 + + image2 = similar(image) if dcf === nothing - mul!(image₂, plan', vec(ksp₂ .* op.dcf)) - @test norm(image₂ .- image) / norm(image) < 0.5 + mul!(image2, plan', vec(ksp2 .* op.dcf)) + @test norm(image2 .- image) / norm(image) < 0.5 else - image₁ = similar(image) - mul!(image₁, plan', vec(ksp₁ .*= dcf)) - mul!(image₂, op', ksp₂) - @test image₂ ≈ image₁ + image1 = similar(image) + mul!(image1, plan', vec(ksp1 .*= dcf)) + mul!(image2, op', ksp2) + @test image2 ≈ image1 end + normal_op = AbstractOperators.get_normal_op(op) - image₃ = similar(image) - mul!(image₃, normal_op, image) - return @test image₃ ≈ image₂ + image3 = similar(image) + mul!(image3, normal_op, image) + @test image3 ≈ image2 end - function test_2D_nufft_op(threaded) + function test_2d_nufft(threaded) trajectory = rand(2, 128, 50) .- 0.5 dcf = rand(128, 50) image_size = (128, 128) image = rand(ComplexF64, image_size) plan = plan_nfft(reshape(trajectory, 2, :), image_size) op = NFFTOp(image_size, trajectory, dcf; threaded) - return test_nufft_op(op, plan, image, dcf) + test_nufft_op(op, plan, image, dcf) end - function test_3D_nufft_op(threaded) + function test_3d_nufft(threaded) trajectory = rand(3, 128, 50) .- 0.5 dcf = rand(128, 50) image_size = (64, 64, 64) image = rand(ComplexF64, image_size) plan = plan_nfft(reshape(trajectory, 3, :), image_size) op = NFFTOp(image_size, trajectory, dcf; threaded) - return test_nufft_op(op, plan, image, dcf) + test_nufft_op(op, plan, image, dcf) end - function test_realistic_2D_nufft_op(threaded) + function test_realistic_2d_nufft(threaded) trajectory = Array{Float64}(undef, 2, 256, 201) - ϕₛₜₑₚ = 2π / 201 + ϕstep = 2π / 201 for i in 1:201 - ϕ = i * ϕₛₜₑₚ + ϕ = i * ϕstep trajectory[1, :, i] = cos(ϕ) .* ((-64:0.5:63.5) ./ 128) trajectory[2, :, i] = sin(ϕ) .* ((-64:0.5:63.5) ./ 128) end @@ -61,39 +65,8 @@ end plan = plan_nfft(reshape(trajectory, 2, :), image_size) op = NFFTOp(image_size, trajectory; threaded) - return test_nufft_op(op, plan, image, nothing) - end - - @testset "NFFTOp" begin - @testset "2D" begin - @testset "single-threaded" begin - test_2D_nufft_op(false) - end - @testset "multi-threaded" begin - test_2D_nufft_op(true) - end - end - @testset "realistic 2D" begin - @testset "single-threaded" begin - test_realistic_2D_nufft_op(false) - end - @testset "multi-threaded" begin - test_realistic_2D_nufft_op(true) - end - end - @testset "3D" begin - @testset "single-threaded" begin - test_3D_nufft_op(false) - end - @testset "multi-threaded" begin - test_3D_nufft_op(true) - end - end + test_nufft_op(op, plan, image, nothing) end -end - -@testitem "NfftNormalOp" tags = [:nfft, :NfftNormalOp] setup = [TestUtils] begin - using LinearAlgebra, Random, NFFT, NFFTOperators, AbstractOperators function test_nfft_normal_op(threaded) trajectory = rand(2, 128, 50) .- 0.5 @@ -102,23 +75,33 @@ end image = rand(ComplexF64, image_size) op = NFFTOp(image_size, trajectory, dcf; threaded) normal_op = AbstractOperators.get_normal_op(op) - # Normal op should compute A' * A * image + image_out1 = similar(image) mul!(image_out1, normal_op, image) - # Manually compute A' * A * image ksp = similar(image, ComplexF64, size(op, 1)) mul!(ksp, op, image) image_out2 = similar(image) mul!(image_out2, op', ksp) - return @test image_out1 ≈ image_out2 + @test image_out1 ≈ image_out2 end +end - @testset "NfftNormalOp" begin - @testset "single-threaded" begin - test_nfft_normal_op(false) - end - @testset "multi-threaded" begin - test_nfft_normal_op(true) - end - end +@testitem "NFFTOp 2D" tags = [:nfft, :NFFTOp] setup = [TestUtils, NFFTTestHelper] begin + NFFTTestHelper.test_2d_nufft(false) + NFFTTestHelper.test_2d_nufft(true) +end + +@testitem "NFFTOp realistic 2D" tags = [:nfft, :NFFTOp] setup = [TestUtils, NFFTTestHelper] begin + NFFTTestHelper.test_realistic_2d_nufft(false) + NFFTTestHelper.test_realistic_2d_nufft(true) +end + +@testitem "NFFTOp 3D" tags = [:nfft, :NFFTOp] setup = [TestUtils, NFFTTestHelper] begin + NFFTTestHelper.test_3d_nufft(false) + NFFTTestHelper.test_3d_nufft(true) +end + +@testitem "NfftNormalOp" tags = [:nfft, :NfftNormalOp] setup = [TestUtils, NFFTTestHelper] begin + NFFTTestHelper.test_nfft_normal_op(false) + NFFTTestHelper.test_nfft_normal_op(true) end diff --git a/test/nfft/test_quality.jl b/test/nfft/test_quality.jl index 93a457c..144fe45 100644 --- a/test/nfft/test_quality.jl +++ b/test/nfft/test_quality.jl @@ -1,4 +1,4 @@ @testitem "Aqua" tags = [:quality, :nfft] begin using Aqua, NFFTOperators - Aqua.test_all(NFFTOperators) + Aqua.test_all(NFFTOperators; persistent_tasks = VERSION >= v"1.11") end diff --git a/test/runtests.jl b/test/runtests.jl index 81a6d1e..9803796 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,13 +1,28 @@ using TestItemRunner -# Run all tests. Tags enable selective filtering for interactive development: -# -# julia --project=test -e ' -# using TestItemRunner -# TestItemRunner.run_tests(pwd(); filter = ti -> :calculus in ti.tags) -# ' -# -# Available tags: :calculus, :linearoperator, :nonlinearoperator, -# :batching, :misc, :quality, :jet +const FILTER_PARTS = if length(ARGS) > 0 + @assert length(ARGS) == 1 + split(ARGS[1], ",") +else + String[] +end +const FILTER_TAGS = map(p -> Symbol(p[2:end]), filter(x -> startswith(x, ":"), FILTER_PARTS)) +const FILTER_NAMES = filter(x -> !startswith(x, ":"), FILTER_PARTS) -@run_package_tests +const VERB = get(ENV, "ABSTRACTOPERATORS_TEST_VERBOSE", "false") == "true" +const FILTER = if length(FILTER_PARTS) > 0 + ti -> begin + run_item = any(t -> t in ti.tags, FILTER_TAGS) || any(n -> n == ti.name, FILTER_NAMES) + if VERB && run_item + println("Running @testitem: ", ti.name) + end + run_item + end +else + ti -> begin + VERB && println("Running @testitem: ", ti.name) + true + end +end + +@run_package_tests filter = FILTER diff --git a/test/test_LinearMapsExt.jl b/test/test_LinearMapsExt.jl index d106899..a957911 100644 --- a/test/test_LinearMapsExt.jl +++ b/test/test_LinearMapsExt.jl @@ -1,5 +1,5 @@ @testitem "LinearMapsExt" tags = [:misc, :LinearMapsExt] begin - using LinearMaps, LinearAlgebra + using AbstractOperators, LinearMaps, LinearAlgebra # Real Diagonal d = rand(10, 10) diff --git a/test/test_combination_rules.jl b/test/test_combination_rules.jl index d8848d1..b882756 100644 --- a/test/test_combination_rules.jl +++ b/test/test_combination_rules.jl @@ -1,5 +1,6 @@ @testitem "CR: AffineAdd Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 5 @@ -39,6 +40,7 @@ end @testitem "CR: Compose Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 3 @@ -86,6 +88,7 @@ end @testitem "CR: DCAT Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 4 @@ -105,6 +108,7 @@ end @testitem "CR: HCAT Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 3 @@ -126,6 +130,7 @@ end @testitem "CR: Scale+Matrix Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 4 @@ -180,6 +185,7 @@ end @testitem "CR: Sum Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 3 @@ -199,6 +205,7 @@ end @testitem "CR: MatrixOp Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 4 @@ -233,6 +240,7 @@ end @testitem "CR: Scale+Eye/DiagOp Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 3 @@ -267,6 +275,7 @@ end @testitem "CR: Zeros Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 4 @@ -310,6 +319,7 @@ end @testitem "CR: DiagOp Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 5 @@ -343,6 +353,7 @@ end @testitem "CR: Eye Combinations" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine n = 4 @@ -372,22 +383,19 @@ end @test combined_adj isa Eye end -@testitem "CR: Mixed Combinations" tags = [:calculus, :CombinationRules] begin +@testitem "CR: Mixed - Scale and Compose rules" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine - # Helper square matrices n = 4 A = randn(n, n) B = randn(n, n) - C = randn(n, n) - D = randn(n, n) - S = randn(n, n) α = 1.7 β = -0.9 comp1 = Compose(MatrixOp(A), FiniteDiff((n + 1,))) # A*(diff(x)) - comp2 = Compose(FiniteDiff((n,)), MatrixOp(B)) # diff(B*x) + comp2 = Compose(FiniteDiff((n,)), MatrixOp(B)) # diff(B*x) # Scale with Compose (combine(L::Scale, R::Compose)) scale_left = Scale(α, FiniteDiff((n,))) @@ -417,12 +425,26 @@ end combined_csa = combine(comp2, scale_right_adj) x = randn(n - 1) # domain of combined_csa = domain of scale_right_adj = (n-1,) @test combined_csa * x ≈ comp2 * (scale_right_adj * x) +end + +@testitem "CR: Mixed - Scale DiagOp adjoint rules" tags = [:calculus, :CombinationRules] begin + using LinearAlgebra + using AbstractOperators + using AbstractOperators: can_be_combined, combine + + n = 4 + A = randn(n, n) + α = 1.7 + β = -0.9 + + dvec = randn(n) + diag_op = DiagOp(dvec) + diag_op_adj = diag_op' + mat_op = MatrixOp(A) # Adjoint Scale with DiagOp (combine(T1::AdjointOperator{<:Scale}, T2::DiagOp)) scale_left = Scale(α, FiniteDiff((n + 1,))) scale_left_adj = scale_left' - dvec = randn(n) - diag_op = DiagOp(dvec) @test can_be_combined(scale_left_adj, diag_op) combined_asd = combine(scale_left_adj, diag_op) x = randn(n) @@ -430,7 +452,6 @@ end # DiagOp' with Scale (combine(T1::AdjointOperator{<:DiagOp}, T2::Scale)) scale_right = Scale(β, FiniteDiff((n + 1,))) - diag_op_adj = diag_op' @test can_be_combined(diag_op_adj, scale_right) combined_das = combine(diag_op_adj, scale_right) x = randn(n + 1) @@ -445,7 +466,6 @@ end @test combined_sad * x ≈ scale_left_adj * (diag_op_adj * x) # Adjoint DiagOp with MatrixOp (both orders to exercise can_be_combined variants) - mat_op = MatrixOp(A) @test can_be_combined(diag_op_adj, mat_op) combined_dm = combine(diag_op_adj, mat_op) @test combined_dm * x ≈ diag_op_adj * (mat_op * x) @@ -461,6 +481,16 @@ end combined_sma = combine(scale_left_adj, mat_op_adj) x = randn(n) @test combined_sma * x ≈ scale_left_adj * (mat_op_adj * x) +end + +@testitem "CR: Mixed - Compose combination branches" tags = [:calculus, :CombinationRules] begin + using LinearAlgebra + using AbstractOperators + using AbstractOperators: can_be_combined, combine + + n = 4 + A = randn(n, n) + α = 1.7 # Branch: combine(L, R::Compose) -> if branch (combined isa Compose) M = MatrixOp(rand(n, n - 2)) @@ -485,16 +515,23 @@ end combined_sc = combine(S, C3) x = randn(n + 2) @test combined_sc * x ≈ S * (C3 * x) +end + +@testitem "CR: Mixed - Scale Compose branch coverage" tags = [:calculus, :CombinationRules] begin + using LinearAlgebra + using AbstractOperators + using AbstractOperators: can_be_combined, combine - # Branch: combine(L::Scale, R::Compose) -> if branch (can_be_combined(L.A, R.A[end])) n2 = 5 A1 = randn(n2, n2) A2 = randn(n2, n2) - A3 = randn(n2, n2) - comp_tail = Compose(MatrixOp(A2), FiniteDiff((n2 + 1,))) # last is MatrixOp so combinable with MatrixOp(A1) + + comp_tail = Compose(MatrixOp(A2), FiniteDiff((n2 + 1,))) scl = Scale(2.3, MatrixOp(A1)) - comb_sc_if = combine(scl, comp_tail) x2 = randn(n2 + 1) + + # Branch: combine(L::Scale, R::Compose) -> if branch (can_be_combined(L.A, R.A[end])) + comb_sc_if = combine(scl, comp_tail) @test comb_sc_if * x2 ≈ scl * (comp_tail * x2) # Branch: combine(L::AdjointOperator{<:Scale}, R::Compose) -> if branch @@ -515,21 +552,28 @@ end comb_csa_if = combine(comp_head, scl_r_adj) x3 = randn(n2 - 1) @test comb_csa_if * x3 ≈ comp_head * (scl_r_adj * x3) +end + +@testitem "CR: Mixed - Scale DiagOp branch coverage" tags = [:calculus, :CombinationRules] begin + using LinearAlgebra + using AbstractOperators + using AbstractOperators: can_be_combined, combine + + n2 = 5 + A2 = randn(n2, n2) + x2 = randn(n2 + 1) + x3 = randn(n2) # Branch: combine(T1::Scale, T2::MatrixOp) -> if branch (can_be_combined) scl_mat_if = Scale(0.7, FiniteDiff((n2,))) mat2 = MatrixOp(A2) comb_scale_mat_if = combine(scl_mat_if, mat2) - x3 = randn(n2) @test comb_scale_mat_if * x3 ≈ scl_mat_if * (mat2 * x3) # Branch: combine(T1::Scale, T2::DiagOp) -> else branch (non-combinable path) - # Use FiniteDiff so scaled_diagop cannot be combined further fd = FiniteDiff((n2 + 1,)) scl_fd = Scale(1.5, fd) - dvec = randn(n2 + 1) - diag_long = DiagOp(dvec) - # Adjust x dimension for FiniteDiff input (n2+1) + diag_long = DiagOp(randn(n2 + 1)) comb_scale_diag_else = combine(scl_fd, diag_long) @test comb_scale_diag_else * x2 ≈ scl_fd * (diag_long * x2) @@ -541,7 +585,6 @@ end @test comb_asd_if * x3 ≈ sclA_adj * (diagA * x3) # Branch: combine(T1::DiagOp, T2::Scale) else branch (non-combinable) - # Make second scale wrap FiniteDiff so scaled_diagop cannot be combined scl_fd2 = Scale(-0.4, FiniteDiff((n2 + 1,))) diag_diff = DiagOp(randn(n2)) comb_diag_scale_else = combine(diag_diff, scl_fd2) @@ -575,6 +618,7 @@ end @testitem "CR: Fallback and Null cases" tags = [:calculus, :CombinationRules] begin using LinearAlgebra + using AbstractOperators using AbstractOperators: can_be_combined, combine # Fallback combine path should error for unsupported nonlinear pair @@ -586,3 +630,100 @@ end @test is_null(combine(z, M)) @test is_null(combine(M, z)) end + +@testitem "CR: DiagOp adjoint-left combinations" tags = [:calculus, :CombinationRules] begin + using LinearAlgebra + using AbstractOperators + using AbstractOperators: can_be_combined, combine + + n = 5 + d1 = randn(n) + d2 = randn(n) + diag1 = DiagOp(d1) + diag2 = DiagOp(d2) + + # DiagOp' * DiagOp + adj1 = diag1' + @test can_be_combined(adj1, diag2) + c1 = combine(adj1, diag2) + @test c1 isa DiagOp + @test c1.d ≈ conj.(d1) .* d2 + + # DiagOp' * DiagOp' + adj2 = diag2' + @test can_be_combined(adj1, adj2) + c2 = combine(adj1, adj2) + @test c2 isa DiagOp + @test c2.d ≈ conj.(d1) .* conj.(d2) +end + +@testitem "CR: can_be_combined with Eye" tags = [:calculus, :CombinationRules] begin + using AbstractOperators + using AbstractOperators: can_be_combined, combine + + n = 4 + eye = Eye(n) + fd = FiniteDiff(Float64, (n + 1,), 1) + + @test can_be_combined(fd, eye) + c = combine(fd, eye) + @test c isa typeof(fd) + x = randn(n + 1) + @test c * x ≈ fd * x +end + +@testitem "CR: MatrixOp and Scale combinations (adjoint variants)" tags = [ + :calculus, :CombinationRules, +] begin + using LinearAlgebra + using AbstractOperators + using AbstractOperators: can_be_combined, combine + + n = 4 + A = randn(n, n) + B = randn(n, n) + α = 1.5 + mat = MatrixOp(A) + mat_adj = mat' + scl = Scale(α, Eye(n)) + + # MatrixOp' * MatrixOp + @test can_be_combined(mat_adj, mat) + c1 = combine(mat_adj, mat) + @test c1 isa MatrixOp + @test c1.A ≈ A' * A + + # MatrixOp * MatrixOp' + @test can_be_combined(mat, mat_adj) + c2 = combine(mat, mat_adj) + @test c2.A ≈ A * A' + + # MatrixOp' * MatrixOp' + @test can_be_combined(mat_adj, mat_adj) + c3 = combine(mat_adj, mat_adj) + @test c3.A ≈ A' * A' + + # AdjointMatrixOp * Scale + @test can_be_combined(mat_adj, scl) + c4 = combine(mat_adj, scl) + x = randn(n) + @test c4 * x ≈ mat_adj * (scl * x) + + # Scale * AdjointMatrixOp + @test can_be_combined(scl, mat_adj) + c5 = combine(scl, mat_adj) + @test c5 * x ≈ scl * (mat_adj * x) + + # AdjointScale * MatrixOp + scl_op = Scale(α, FiniteDiff(Float64, (n + 1,), 1)) + scl_adj = scl_op' + @test can_be_combined(scl_adj, mat) + c6 = combine(scl_adj, mat) + x2 = randn(n) + @test c6 * x2 ≈ scl_adj * (mat * x2) + + # AdjointScale * AdjointMatrixOp + @test can_be_combined(scl_adj, mat_adj) + c7 = combine(scl_adj, mat_adj) + @test c7 * x ≈ scl_adj * (mat_adj * x) +end diff --git a/test/test_gpu_quality.jl b/test/test_gpu_quality.jl new file mode 100644 index 0000000..cc4cf4d --- /dev/null +++ b/test/test_gpu_quality.jl @@ -0,0 +1,61 @@ +@testitem "GpuExt Quality" tags = [:gpu, :quality] begin + using AbstractOperators, JLArrays, LinearAlgebra + + # Loading JLArrays triggers GpuExt (GPUArrays is transitive dep) + # Verify GpuExt loaded by checking GPU-specific dispatch + d = jl(randn(4)) + op = DiagOp(d) + x = jl(randn(4)) + y = jl(zeros(4)) + mul!(y, op, x) + @test collect(y) ≈ collect(d) .* collect(x) + + # Verify guards work: GPU output + CPU input → error + op_gi = GetIndex(Float64, (4,), (2:4,)) + x_cpu = randn(4) + y_gpu = jl(zeros(3)) + @test_throws ArgumentError mul!(y_gpu, op_gi, x_cpu) + + # CPU output + GPU input → error + x_gpu = jl(randn(4)) + y_cpu = zeros(3) + @test_throws ArgumentError mul!(y_cpu, op_gi, x_gpu) + + # Storage type and allocation checks for GPU operators + @test domain_array_type(op) <: JLArrays.JLArray + @test codomain_array_type(op) <: JLArrays.JLArray + y_alloc = op * x + @test y_alloc isa JLArrays.JLArray +end + +@testitem "GpuExt JET" tags = [:gpu, :jet] begin + using AbstractOperators, JET, JLArrays + + # Verify key GPU-dispatched functions are type-stable (no dynamic dispatch) + n = 8 + d = jl(randn(n)) + x = jl(randn(n)) + y = jl(zeros(n)) + + # DiagOp GPU mul! — target_modules restricts JET to AbstractOperators code only, + # avoiding false positives from KernelAbstractions (transitive dep of JLArrays). + op_diag = DiagOp(d) + @test_call target_modules = (AbstractOperators,) mul!(y, op_diag, x) + + # GetIndex GPU mul! + op_gi = GetIndex(d, (2:4,)) + y3 = jl(zeros(3)) + @test_call target_modules = (AbstractOperators,) mul!(y3, op_gi, x) + x3 = jl(randn(3)) + @test_call target_modules = (AbstractOperators,) mul!(y, op_gi', x3) + + # ZeroPad GPU mul! — zp is a flat NTuple{N,Int} giving per-dimension padding + op_zp = ZeroPad(d, (2,)) + y10 = jl(zeros(n + 2)) + @test_call target_modules = (AbstractOperators,) mul!(y10, op_zp, x) + @test_call target_modules = (AbstractOperators,) mul!(x, op_zp', y10) + + # _should_thread dispatches to false for GPU arrays + @test AbstractOperators._should_thread(typeof(d)) == false + @test AbstractOperators._should_thread(Array{Float64}) == (Threads.nthreads() > 1) +end diff --git a/test/test_nonlinear_operators.jl b/test/test_nonlinear_operators.jl index 191d803..d13c839 100644 --- a/test/test_nonlinear_operators.jl +++ b/test/test_nonlinear_operators.jl @@ -135,3 +135,129 @@ end op = Pow(Float64, (n,), 0.5) y, grad = test_NLop(op, x, r, verb) end + +# ─── GPU test items for nonlinear operators ─────────────────────────────────── + +@testitem "NonlinearOp: Sigmoid (GPU)" tags = [:gpu, :nonlinearoperator, :Sigmoid] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n = 4 + x = gpu_randn(backend, n) + op = Sigmoid(x; gamma = 2.0) # construct from GPU array to get GPU storage type + test_NLop_gpu(op, x, gpu_randn(backend, n), false) + end +end + +@testitem "NonlinearOp: SoftMax (GPU)" tags = [:gpu, :nonlinearoperator, :SoftMax] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n = 10 + x = gpu_randn(backend, n) + op = SoftMax(x) # construct from GPU array so buffer is GPU-typed + test_NLop_gpu(op, x, gpu_randn(backend, n), false) + end +end + +@testitem "NonlinearOp: SoftPlus (GPU)" tags = [:gpu, :nonlinearoperator, :SoftPlus] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n = 10 + x = gpu_randn(backend, n) + op = SoftPlus(x) + test_NLop_gpu(op, x, gpu_randn(backend, n), false) + end +end + +@testitem "NonlinearOp: Exp (GPU)" tags = [:gpu, :nonlinearoperator, :Exp] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n, m = 4, 5 + x = gpu_randn(backend, n, m) + op = Exp(x) + test_NLop_gpu(op, x, gpu_randn(backend, n, m), false) + end +end + +@testitem "NonlinearOp: Sin (GPU)" tags = [:gpu, :nonlinearoperator, :Sin] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n, m = 4, 5 + x = gpu_randn(backend, n, m) + op = Sin(x) + test_NLop_gpu(op, x, gpu_randn(backend, n, m), false) + end +end + +@testitem "NonlinearOp: Cos (GPU)" tags = [:gpu, :nonlinearoperator, :Cos] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n, m = 4, 5 + x = gpu_randn(backend, n, m) + op = Cos(x) + test_NLop_gpu(op, x, gpu_randn(backend, n, m), false) + end +end + +@testitem "NonlinearOp: Atan (GPU)" tags = [:gpu, :nonlinearoperator, :Atan] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n = 10 + x = gpu_randn(backend, n) + op = Atan(x) + test_NLop_gpu(op, x, gpu_randn(backend, n), false) + end +end + +@testitem "NonlinearOp: Tanh (GPU)" tags = [:gpu, :nonlinearoperator, :Tanh] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n = 10 + x = gpu_randn(backend, n) + op = Tanh(x) + test_NLop_gpu(op, x, gpu_randn(backend, n), false) + end +end + +@testitem "NonlinearOp: Sech (GPU)" tags = [:gpu, :nonlinearoperator, :Sech] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n = 10 + x = gpu_randn(backend, n) + op = Sech(x) + test_NLop_gpu(op, x, gpu_randn(backend, n), false) + end +end + +@testitem "NonlinearOp: Pow (GPU)" tags = [:gpu, :nonlinearoperator, :Pow] setup = [TestUtils] begin + using GPUEnv, Random, AbstractOperators + + for backend in gpu_backends() + Random.seed!(0) + n = 10 + x = gpu_randn(backend, n) + op = Pow(x, 2) + test_NLop_gpu(op, x, gpu_randn(backend, n), false) + + x2 = abs.(gpu_randn(backend, n)) + op2 = Pow(x2, 0.5) + test_NLop_gpu(op2, x2, abs.(gpu_randn(backend, n)), false) + end +end diff --git a/test/test_quality.jl b/test/test_quality.jl index 03a015a..8cf3a88 100644 --- a/test/test_quality.jl +++ b/test/test_quality.jl @@ -1,5 +1,5 @@ @testitem "Documentation" tags = [:quality] begin - using Documenter + using Documenter, AbstractOperators # Eagerly load weak-dep extensions so they're precompiled before the doctest # sandbox runs `using AbstractOperators, LinearMaps` (avoids stray precompile output). using LinearMaps @@ -13,6 +13,7 @@ end @testitem "Aqua" tags = [:quality] begin - using Aqua - Aqua.test_all(AbstractOperators) + using Aqua, AbstractOperators + + Aqua.test_all(AbstractOperators, persistent_tasks = VERSION >= v"1.11") end diff --git a/test/test_syntax.jl b/test/test_syntax.jl index fb4d498..48da474 100644 --- a/test/test_syntax.jl +++ b/test/test_syntax.jl @@ -1,22 +1,27 @@ -@testitem "Syntax" tags = [:misc, :Syntax] setup = [TestUtils] begin - verb && println(" --- Testing Syntax --- ") - ###### ' ###### +@testitem "Syntax: Adjoint" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators n, m = 5, 3 A = randn(n, m) - B = randn(n, m) - C = randn(n, m) - x1 = randn(m) x2 = randn(n) opA = MatrixOp(A) - opB = MatrixOp(B) - opC = MatrixOp(C) opt = opA' y1 = opt * x2 y2 = A' * x2 @test norm(y1 - y2) < 1.0e-9 +end + +@testitem "Syntax: Addition and Subtraction" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators + n, m = 5, 3 + A = randn(n, m) + B = randn(n, m) + C = randn(n, m) + x1 = randn(m) + opA = MatrixOp(A) + opB = MatrixOp(B) + opC = MatrixOp(C) - ######+,-###### opp = +opA y1 = opp * x1 y2 = A * x1 @@ -54,9 +59,10 @@ y1 = opss * x1 y2 = -A * x1 + B * x1 + C * x1 @test norm(y1 - y2) < 1.0e-9 +end - ###### * ###### - +@testitem "Syntax: Multiplication" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators n, m, l = 5, 3, 7 A = randn(n, m) B = randn(l, n) @@ -104,8 +110,16 @@ y1 = opc * x1 y2 = x1 @test norm(y1 - y2) < 1.0e-9 +end - ###### getindex ###### +@testitem "Syntax: Getindex basic" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators + n, m, l = 5, 3, 7 + A = randn(n, m) + B = randn(l, n) + x1 = randn(m) + opA = MatrixOp(A) + opB = MatrixOp(B) ops = opA[1:(n - 1)] y1 = ops * x1 @@ -121,7 +135,6 @@ @test norm(y1 - y2) < 1.0e-9 opF = DiagOp(d) - x3 = randn(n, m, l) ops = opF[1:(n - 1), 2:m, 1] y1 = ops * x3 y2 = (d .* x3)[1:(n - 1), 2:m, 1] @@ -149,9 +162,10 @@ sliced_diag = comp_diag[1:3] xdiag = randn(n) @test sliced_diag * xdiag ≈ (comp_diag * xdiag)[1:3] +end - #slicing HCAT - +@testitem "Syntax: Getindex HCAT" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators n, m1, m2, m3 = 5, 6, 7, 8 A = randn(n, m1) B = randn(n, m2) @@ -163,10 +177,12 @@ opB = MatrixOp(B) opC = MatrixOp(C) opH = HCAT(opA, opB, opC) + opH2 = opH[1:2] y1 = opH2 * ArrayPartition(x1, x2) y2 = A * x1 + B * x2 @test all(norm.(y1 .- y2) .<= 1.0e-12) + opH3 = opH[3] y1 = opH3 * x3 y2 = C * x3 @@ -190,44 +206,56 @@ @test norm(opHA[3] * x3 - (C * x3 + d)) < 1.0e-9 m4 = 9 - x4 = randn(m4) D = randn(n, n) E = randn(n, m4) opD = MatrixOp(D) opE = MatrixOp(E) opCH = opD * opH - opHCH = HCAT(opCH, opE) - opH4 = opHCH[4] @test opH4 == opE @test_throws ErrorException opHCH[1] @test_throws ErrorException opHCH[1:2] - # check() utility error paths - op_diag = DiagOp(rand(3)) - ygood = zeros(3) - xgood = rand(3) - @test_throws ArgumentError AbstractOperators.check(ygood, op_diag, "bad") - @test_throws ArgumentError AbstractOperators.check("bad", op_diag, xgood) - @test_throws ArgumentError AbstractOperators.check(ygood, op_diag, ArrayPartition(rand(3), rand(3))) - @test_throws ArgumentError AbstractOperators.check(ArrayPartition(zeros(3), zeros(3)), op_diag, xgood) + # HCAT getindex split-error path with stacked indices + h_joint = HCAT(MatrixOp(randn(4, 2)), MatrixOp(randn(4, 3))) + h_stacked = HCAT((h_joint, MatrixOp(randn(4, 5))), zeros(4)) + @test_throws ErrorException h_stacked[1] - #slicing VCAT + # AffineAdd{HCAT} ndoms==1 branch via low-level single-operator HCAT + Hsingle = HCAT((MatrixOp(randn(4, 4)),), (1,), zeros(4)) + AHsingle = AffineAdd(Hsingle, randn(4)) + @test AHsingle[1:2] isa AbstractOperator + # Compose getindex branch for diagonal tail (ndoms(A,2)>1) + Hbase = HCAT((MatrixOp(randn(4, 2)), MatrixOp(randn(4, 3))), zeros(4)) + Cdiag = DiagOp(randn(4)) * Hbase + xin = ArrayPartition(randn(2), randn(3)) + @test Cdiag[[2, 1]] * ArrayPartition(xin.x[2], xin.x[1]) ≈ Cdiag * xin + + # Compose getindex error path for non-diagonal tail + Cnondiag = MatrixOp(randn(4, 4)) * Hbase + @test_throws ErrorException Cnondiag[[2, 1]] +end + +@testitem "Syntax: Getindex VCAT" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators n1, n2, n3, m = 5, 6, 7, 8 A = randn(n1, m) B = randn(n2, m) C = randn(n3, m) x1 = randn(m) + x3 = randn(m) opA = MatrixOp(A) opB = MatrixOp(B) opC = MatrixOp(C) opV = VCAT(opA, opB, opC) + opV2 = opV[1:2] y1 = opV2 * x1 y2 = ArrayPartition(A * x1, B * x1) @test norm(y1 - y2) <= 1.0e-12 + opV3 = opV[3] y1 = opV3 * x3 y2 = C * x3 @@ -245,29 +273,21 @@ comp_single = FiniteDiff((m,)) * MatrixOp(randn(m, m)) x_comp = randn(m) @test comp_single[1:3] * x_comp ≈ (comp_single * x_comp)[1:3] +end - # Compose getindex branch for diagonal tail (ndoms(A,2)>1) - Hbase = HCAT((MatrixOp(randn(4, 2)), MatrixOp(randn(4, 3))), zeros(4)) - Cdiag = DiagOp(randn(4)) * Hbase - xin = ArrayPartition(randn(2), randn(3)) - @test Cdiag[[2, 1]] * ArrayPartition(xin.x[2], xin.x[1]) ≈ Cdiag * xin - - # Compose getindex error path for non-diagonal tail - Cnondiag = MatrixOp(randn(4, 4)) * Hbase - @test_throws ErrorException Cnondiag[[2, 1]] - - # HCAT getindex split-error path with stacked indices - h_joint = HCAT(MatrixOp(randn(4, 2)), MatrixOp(randn(4, 3))) - h_stacked = HCAT((h_joint, MatrixOp(randn(4, 5))), zeros(4)) - @test_throws ErrorException h_stacked[1] - - # AffineAdd{HCAT} ndoms==1 branch via low-level single-operator HCAT - Hsingle = HCAT((MatrixOp(randn(4, 4)),), (1,), zeros(4)) - AHsingle = AffineAdd(Hsingle, randn(4)) - @test AHsingle[1:2] isa AbstractOperator - - ###### hcat ###### +@testitem "Syntax: Check utility" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators + op_diag = DiagOp(rand(3)) + ygood = zeros(3) + xgood = rand(3) + @test_throws ArgumentError AbstractOperators.check(ygood, op_diag, "bad") + @test_throws ArgumentError AbstractOperators.check("bad", op_diag, xgood) + @test_throws ArgumentError AbstractOperators.check(ygood, op_diag, ArrayPartition(rand(3), rand(3))) + @test_throws ArgumentError AbstractOperators.check(ArrayPartition(zeros(3), zeros(3)), op_diag, xgood) +end +@testitem "Syntax: HCAT and VCAT construction" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators n, m1, m2 = 5, 6, 7 A = randn(n, m1) B = randn(n, m2) @@ -285,8 +305,6 @@ y2 = [A B B] * [x1; x2; x2] @test norm(y1 - y2) <= 1.0e-12 - ###### vcat ###### - n, m1, m2 = 5, 6, 7 A = randn(m1, n) B = randn(m2, n) @@ -302,8 +320,10 @@ y1 = opVV * x1 y2 = ArrayPartition(A * x1, A * x1, B * x1) @test norm(y1 - y2) <= 1.0e-12 +end - ###### reshape ###### +@testitem "Syntax: Reshape ndims ndoms" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators n, m = 10, 5 A = randn(n, m) x1 = randn(m) @@ -314,7 +334,6 @@ y2 = reshape(A * x1, 2, 5) @test norm(y1 - y2) <= 1.0e-12 - # testing ndims & ndoms L = Variation((3, 4, 5)) @test ndims(L) == (2, 3) @test ndims(L, 1) == 2 @@ -330,15 +349,16 @@ D = DCAT(L, L) @test ndims(D) == ((2, 2), (3, 3)) @test ndoms(D) == (2, 2) +end - ###### jacobian ###### +@testitem "Syntax: Jacobian convert displacement" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators n, m = 10, 5 A = MatrixOp(randn(n, m)) B = Sigmoid(Float64, (n,), 100.0) op = B * A J = jacobian(op, randn(m)) - #### convert #### L = Eye(10) LL = convert(AbstractOperator, Float64, (10,), L) @test LL == L @@ -348,7 +368,73 @@ @test_throws MethodError convert(NonLinearOperator, Float64, (10,), L) - ### displacement ### L = Eye(10) @test displacement(L) == 0.0 end + +@testitem "Syntax: VCAT and HCAT with Tuple input" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators, RecursiveArrayTools + n1, n2, m = 3, 4, 5 + x1 = randn(n1) + x2 = randn(n2) + + # HCAT * Tuple (returns codomain array, not ArrayPartition) + opH = HCAT(MatrixOp(randn(m, n1)), MatrixOp(randn(m, n2))) + yH = opH * (x1, x2) + @test yH ≈ opH * ArrayPartition(x1, x2) + + # VCAT * Tuple (multi-domain VCAT — each sub-op is an HCAT, so domain is a tuple) + # opV has domain (n1,)×(n2,) and codomain (m,)×(m,); input must be a Tuple + opH1 = HCAT(MatrixOp(randn(m, n1)), MatrixOp(randn(m, n2))) + opH2 = HCAT(MatrixOp(randn(m, n1)), MatrixOp(randn(m, n2))) + opV = VCAT(opH1, opH2) + yV = opV * (x1, x2) # calls *(L::VCAT, b::Tuple) → returns y.x as Tuple + @test yV isa Tuple + ref = opV * ArrayPartition(x1, x2) + @test all(collect(yV[i]) ≈ ref.x[i] for i in eachindex(yV)) +end + +@testitem "Syntax: L * coeff (operator-right scalar mul)" tags = [:misc, :Syntax] setup = [ + TestUtils, +] begin + using AbstractOperators + n = 4 + op = FiniteDiff(Float64, (n,), 1) + alpha = 3.0 + s1 = op * alpha + s2 = alpha * op + x = randn(n) + @test s1 * x ≈ s2 * x +end + +@testitem "Syntax: mul! fallback throws MethodError" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators + # Custom operator with no mul! defined so the fallback in syntax.jl fires + struct _NoMulOp <: AbstractOperators.LinearOperator end + AbstractOperators.size(::_NoMulOp) = ((3,), (3,)) + AbstractOperators.domain_type(::_NoMulOp) = Float64 + AbstractOperators.codomain_type(::_NoMulOp) = Float64 + op = _NoMulOp() + @test_throws MethodError mul!(zeros(3), op, zeros(3)) +end + +@testitem "Syntax: Sum getindex (multi-domain)" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators + n = 5 + A = MatrixOp(randn(n, n)) + B = DiagOp(randn(n)) + S = Sum(A, B) + sliced = S[2:(n - 1)] + x = randn(n) + @test sliced * x ≈ (S * x)[2:(n - 1)] +end + +@testitem "Syntax: Scale getindex (ndoms == 1 branch)" tags = [:misc, :Syntax] setup = [TestUtils] begin + using AbstractOperators + n = 5 + op = FiniteDiff(Float64, (n,), 1) # ndoms(Scale(coeff, op), 2) == 1 + s = Scale(2.0, op) + sliced = s[1:2] + x = randn(n) + @test sliced * x ≈ (s * x)[1:2] +end diff --git a/test/utils.jl b/test/utils.jl index 5551aa6..de05bad 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -1,25 +1,89 @@ @testmodule TestUtils begin using Test using AbstractOperators + using RecursiveArrayTools using RecursiveArrayTools: ArrayPartition using LinearAlgebra using Random + using Pkg export verb, test_op, test_NLop, gradient_fd, ArrayPartition, norm, dot, diag, opnorm + export to_cpu, test_NLop_gpu + export assert_cpu_approx, assert_adjoint_invariant + + if VERSION < v"1.11" + Pkg.develop(path = normpath(joinpath(@__DIR__, ".."))) # AbstractOperators + Pkg.develop(path = normpath(joinpath(@__DIR__, "..", "DSPOperators"))) # DSPOperators + Pkg.develop(path = normpath(joinpath(@__DIR__, "..", "FFTWOperators"))) # FFTWOperators + Pkg.develop(path = normpath(joinpath(@__DIR__, "..", "NFFTOperators"))) # NFFTOperators + Pkg.develop(path = normpath(joinpath(@__DIR__, "..", "WaveletOperators"))) # WaveletOperators + end + + using GPUEnv + GPUEnv.activate(; persist = true) + + if VERSION >= v"1.11" && Base.find_package("AcceleratedDCTs") === nothing + Pkg.add(name = "AcceleratedDCTs", version = "0.4") + end - const verb = false + const verb = get(ENV, "ABSTRACTOPERATORS_TEST_VERBOSE", "false") == "true" + + to_cpu(x::AbstractArray) = collect(x) + to_cpu(x::RecursiveArrayTools.ArrayPartition) = RecursiveArrayTools.ArrayPartition(collect.(x.x)...) + + function assert_cpu_approx(x, y; atol = 1.0e-8) + @test norm(to_cpu(x) .- to_cpu(y)) <= atol + end + + function assert_adjoint_invariant(x, Ax, y, Acy; atol = 1.0e-8) + s1 = real(dot(to_cpu(Ax), to_cpu(y))) + s2 = real(dot(to_cpu(x), to_cpu(Acy))) + @test abs(s1 - s2) < atol + end + + function assert_backend_output_type(op::AbstractOperator, x::AbstractArray, y) + expected = typeof(similar(x, codomain_type(op), size(op, 1)...)) + @test y isa expected + end + + assert_backend_output_type(::AbstractOperator, ::ArrayPartition, ::Any) = nothing + assert_backend_output_type(::AbstractOperator, ::AbstractArray, ::ArrayPartition) = nothing + + function test_NLop_gpu(A::AbstractOperator, x, y, verb::Bool = false) + verb && (println(), println(A)) + + Ax = A * x + assert_backend_output_type(A, x, Ax) + Ax2 = similar(Ax) + mul!(Ax2, A, x) + assert_cpu_approx(Ax, Ax2) + + @test_throws ErrorException A' + + J = Jacobian(A, x) + grad = J' * y + mul!(Ax2, A, x) # redo forward + grad2 = similar(grad) + mul!(grad2, J', y) + assert_cpu_approx(grad, grad2; atol = 1.0e-8) + + return Ax, grad + end ########### Test for LinearOperators function test_op(A::AbstractOperator, x, y, verb::Bool = false) verb && (println(); show(A); println()) Ax = A * x + if !(x isa Array) && !(x isa ArrayPartition) + assert_backend_output_type(A, x, Ax) + end Ax2 = similar(Ax) verb && println("forward preallocated") mul!(Ax2, A, x) #verify in-place linear operator works verb && @time mul!(Ax2, A, x) - @test norm(Ax .- Ax2) <= 1.0e-8 + assert_cpu_approx(Ax, Ax2) Acy = A' * y Acy2 = similar(Acy) @@ -28,11 +92,8 @@ mul!(Acy2, At, y) #verify in-place linear operator works verb && @time mul!(Acy2, At, y) - @test norm(Acy .- Acy2) <= 1.0e-8 - - s1 = real(dot(Ax2, y)) - s2 = real(dot(x, Acy2)) - @test abs(s1 - s2) < 1.0e-8 + assert_cpu_approx(Acy, Acy2) + assert_adjoint_invariant(x, Ax2, y, Acy2) return Ax end diff --git a/test/wavelets/test_quality.jl b/test/wavelets/test_quality.jl index a85a5e8..61b6f5b 100644 --- a/test/wavelets/test_quality.jl +++ b/test/wavelets/test_quality.jl @@ -1,4 +1,4 @@ @testitem "Aqua" tags = [:quality, :wavelet] begin using Aqua, WaveletOperators - Aqua.test_all(WaveletOperators) + Aqua.test_all(WaveletOperators; persistent_tasks = VERSION >= v"1.11") end