Mapped transforms (#140)

* Mapped transforms

* Move Mul here, support inv plans

Author: dlfivefifty

Project.toml

@@ -1,9 +1,9 @@
 name = "ContinuumArrays"
 uuid = "7ae1f121-cc2c-504b-ac30-9b923412ae5c"
-version = "0.12.2"
-
+version = "0.12.3"
 
 [deps]
+AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
@@ -19,6 +19,7 @@ RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
+AbstractFFTs = "1"
 ArrayLayouts = "0.7.7, 0.8"
 BandedMatrices = "0.16, 0.17"
 BlockArrays = "0.16"

src/ContinuumArrays.jl

@@ -9,7 +9,7 @@ import LazyArrays: MemoryLayout, Applied, ApplyStyle, flatten, _flatten, colsupp
                         adjointlayout, arguments, _mul_arguments, call, broadcastlayout, layout_getindex, UnknownLayout,
                         sublayout, sub_materialize, ApplyLayout, BroadcastLayout, combine_mul_styles, applylayout,
                         simplifiable, _simplify, AbstractLazyLayout, PaddedLayout
-import LinearAlgebra: pinv, dot, norm2, ldiv!, mul!
+import LinearAlgebra: pinv, inv, dot, norm2, ldiv!, mul!
 import BandedMatrices: AbstractBandedLayout, _BandedMatrix
 import BlockArrays: block, blockindex, unblock, blockedrange, _BlockedUnitRange, _BlockArray
 import FillArrays: AbstractFill, getindex_value, SquareEye
@@ -20,6 +20,7 @@ import QuasiArrays: cardinality, checkindex, QuasiAdjoint, QuasiTranspose, Inclu
                     LazyQuasiArray, LazyQuasiVector, LazyQuasiMatrix, LazyLayout, LazyQuasiArrayStyle, _factorize,
                     AbstractQuasiFill, UnionDomain, __sum, _cumsum, __cumsum, applylayout, _equals, layout_broadcasted, PolynomialLayout
 import InfiniteArrays: Infinity, InfAxes
+import AbstractFFTs: Plan
 
 export Spline, LinearSpline, HeavisideSpline, DiracDelta, Derivative, ℵ₁, Inclusion, Basis, grid, plotgrid, affine, .., transform, expand, plan_transform, basis, coefficients 
 
@@ -91,6 +92,7 @@ checkpoints(A::AbstractQuasiMatrix) = checkpoints(axes(A,1))
 
 
 include("operators.jl")
+include("plans.jl")
 include("bases/bases.jl")
 
 include("plotting.jl")

src/bases/bases.jl

@@ -173,93 +173,18 @@ grid(P, n...) = _grid(MemoryLayout(P), P, n...)
 # values(f) = 
 
 
-struct TransformFactorization{T,Grid,Plan} <: Factorization{T}
-    grid::Grid
-    plan::Plan
-end
-
-TransformFactorization{T}(grid, plan) where T = TransformFactorization{T,typeof(grid),typeof(plan)}(grid, plan)
-
-"""
-    TransformFactorization(grid, plan)
-
-associates a planned transform with a grid. That is, if `F` is a `TransformFactorization`, then
-`F \\ f` is equivalent to `F.plan * f[F.grid]`.
-"""
-TransformFactorization(grid, plan) = TransformFactorization{promote_type(eltype(eltype(grid)),eltype(plan))}(grid, plan)
-
-
-
-grid(T::TransformFactorization) = T.grid
-function size(T::TransformFactorization, k)
-    @assert k == 2 # TODO: make consistent
-    size(T.plan,1)
-end
-
-
-\(a::TransformFactorization, b::AbstractQuasiVector) = a.plan * convert(Array, b[a.grid])
-\(a::TransformFactorization, b::AbstractQuasiMatrix) = a.plan * convert(Array, b[a.grid,:])
-
-"""
-    InvPlan(factorization, dims)
-
-Takes a factorization and supports it applied to different dimensions.
-"""
-struct InvPlan{T, Fact, Dims} # <: Plan{T} We don't depend on AbstractFFTs
-    factorization::Fact
-    dims::Dims
-end
 
-InvPlan(fact, dims) = InvPlan{eltype(fact), typeof(fact), typeof(dims)}(fact, dims)
-
-size(F::InvPlan, k...) = size(F.factorization, k...)
-
-
-function *(P::InvPlan{<:Any,<:Any,Int}, x::AbstractVector)
-    @assert P.dims == 1
-    P.factorization \ x
-end
-
-function *(P::InvPlan{<:Any,<:Any,Int}, X::AbstractMatrix)
-    if P.dims == 1
-        P.factorization \ X
-    else
-        @assert P.dims == 2
-        permutedims(P.factorization \ permutedims(X))
-    end
-end
-
-function *(P::InvPlan{<:Any,<:Any,Int}, X::AbstractArray{<:Any,3})
-    Y = similar(X)
-    if P.dims == 1
-        for j in axes(X,3)
-            Y[:,:,j] = P.factorization \ X[:,:,j]
-        end
-    elseif P.dims == 2
-        for k in axes(X,1)
-            Y[k,:,:] = P.factorization \ X[k,:,:]
-        end
-    else
-        @assert P.dims == 3
-        for k in axes(X,1), j in axes(X,2)
-            Y[k,j,:] = P.factorization \ X[k,j,:]
-        end
-    end
-    Y
+function plan_grid_transform(lay, L, szs::NTuple{N,Int}, dims=1:N) where N
+    p = grid(L)
+    p, InvPlan(factorize(L[p,:]), dims)
 end
 
-function *(P::InvPlan, X::AbstractArray)
-    for d in P.dims
-        X = InvPlan(P.factorization, d) * X
-    end
-    X
+function plan_grid_transform(::MappedBasisLayout, L, szs::NTuple{N,Int}, dims=1:N) where N
+    x,F = plan_grid_transform(demap(L), szs, dims)
+    invmap(parentindices(L)[1])[x], F
 end
 
-
-function plan_grid_transform(L, szs::NTuple{N,Int}, dims=1:N) where N
-    p = grid(L)
-    p, InvPlan(factorize(L[p,:]), dims)
-end
+plan_grid_transform(L, szs::NTuple{N,Int}, dims=1:N) where N = plan_grid_transform(MemoryLayout(L), L, szs, dims)
 
 plan_grid_transform(L, arr::AbstractArray{<:Any,N}, dims=1:N) where N = 
     plan_grid_transform(L, size(arr), dims)

src/plans.jl

@@ -0,0 +1,140 @@
+
+struct TransformFactorization{T,Grid,Plan} <: Factorization{T}
+    grid::Grid
+    plan::Plan
+end
+
+TransformFactorization{T}(grid, plan) where T = TransformFactorization{T,typeof(grid),typeof(plan)}(grid, plan)
+
+"""
+    TransformFactorization(grid, plan)
+
+associates a planned transform with a grid. That is, if `F` is a `TransformFactorization`, then
+`F \\ f` is equivalent to `F.plan * f[F.grid]`.
+"""
+TransformFactorization(grid, plan) = TransformFactorization{promote_type(eltype(eltype(grid)),eltype(plan))}(grid, plan)
+
+
+
+grid(T::TransformFactorization) = T.grid
+function size(T::TransformFactorization, k)
+    @assert k == 2 # TODO: make consistent
+    size(T.plan,1)
+end
+
+
+\(a::TransformFactorization, b::AbstractQuasiVector) = a.plan * convert(Array, b[a.grid])
+\(a::TransformFactorization, b::AbstractQuasiMatrix) = a.plan * convert(Array, b[a.grid,:])
+
+"""
+    InvPlan(factorization, dims)
+
+Takes a factorization and supports it applied to different dimensions.
+"""
+struct InvPlan{T, Fact, Dims} <: Plan{T}
+    factorization::Fact
+    dims::Dims
+end
+
+InvPlan(fact, dims) = InvPlan{eltype(fact), typeof(fact), typeof(dims)}(fact, dims)
+
+size(F::InvPlan, k...) = size(F.factorization, k...)
+
+
+function *(P::InvPlan{<:Any,<:Any,Int}, x::AbstractVector)
+    @assert P.dims == 1
+    P.factorization \ x
+end
+
+function *(P::InvPlan{<:Any,<:Any,Int}, X::AbstractMatrix)
+    if P.dims == 1
+        P.factorization \ X
+    else
+        @assert P.dims == 2
+        permutedims(P.factorization \ permutedims(X))
+    end
+end
+
+function *(P::InvPlan{<:Any,<:Any,Int}, X::AbstractArray{<:Any,3})
+    Y = similar(X)
+    if P.dims == 1
+        for j in axes(X,3)
+            Y[:,:,j] = P.factorization \ X[:,:,j]
+        end
+    elseif P.dims == 2
+        for k in axes(X,1)
+            Y[k,:,:] = P.factorization \ X[k,:,:]
+        end
+    else
+        @assert P.dims == 3
+        for k in axes(X,1), j in axes(X,2)
+            Y[k,j,:] = P.factorization \ X[k,j,:]
+        end
+    end
+    Y
+end
+
+function *(P::InvPlan, X::AbstractArray)
+    for d in P.dims
+        X = InvPlan(P.factorization, d) * X
+    end
+    X
+end
+
+
+"""
+    MulPlan(matrix, dims)
+
+Takes a matrix and supports it applied to different dimensions.
+"""
+struct MulPlan{T, Fact, Dims} <: Plan{T}
+    matrix::Fact
+    dims::Dims
+end
+
+MulPlan(fact, dims) = MulPlan{eltype(fact), typeof(fact), typeof(dims)}(fact, dims)
+
+function *(P::MulPlan{<:Any,<:Any,Int}, x::AbstractVector)
+    @assert P.dims == 1
+    P.matrix * x
+end
+
+function *(P::MulPlan{<:Any,<:Any,Int}, X::AbstractMatrix)
+    if P.dims == 1
+        P.matrix * X
+    else
+        @assert P.dims == 2
+        permutedims(P.matrix * permutedims(X))
+    end
+end
+
+function *(P::MulPlan{<:Any,<:Any,Int}, X::AbstractArray{<:Any,3})
+    Y = similar(X)
+    if P.dims == 1
+        for j in axes(X,3)
+            Y[:,:,j] = P.matrix * X[:,:,j]
+        end
+    elseif P.dims == 2
+        for k in axes(X,1)
+            Y[k,:,:] = P.matrix * X[k,:,:]
+        end
+    else
+        @assert P.dims == 3
+        for k in axes(X,1), j in axes(X,2)
+            Y[k,j,:] = P.matrix * X[k,j,:]
+        end
+    end
+    Y
+end
+
+function *(P::MulPlan, X::AbstractArray)
+    for d in P.dims
+        X = MulPlan(P.matrix, d) * X
+    end
+    X
+end
+
+*(A::AbstractMatrix, P::MulPlan) = MulPlan(A*P.matrix, P.dims)
+
+inv(P::MulPlan) = InvPlan(factorize(P.matrix), P.dims)
+inv(P::InvPlan) = MulPlan(P.factorization, P.dims)

test/test_splines.jl

@@ -1,7 +1,7 @@
 using ContinuumArrays, LinearAlgebra, Base64, FillArrays, QuasiArrays, BandedMatrices, Test
 using QuasiArrays: ApplyQuasiArray, ApplyStyle, MemoryLayout, mul, MulQuasiMatrix, Vec
 import LazyArrays: MulStyle, LdivStyle, arguments, applied, apply
-import ContinuumArrays: basis, AdjointBasisLayout, ExpansionLayout, BasisLayout, SubBasisLayout, AdjointMappedBasisLayout, MappedBasisLayout
+import ContinuumArrays: basis, AdjointBasisLayout, ExpansionLayout, BasisLayout, SubBasisLayout, AdjointMappedBasisLayout, MappedBasisLayout, plan_grid_transform
 
 @testset "Splines" begin
     @testset "HeavisideSpline" begin
@@ -436,6 +436,25 @@ import ContinuumArrays: basis, AdjointBasisLayout, ExpansionLayout, BasisLayout,
             @test L[y,:] \ (y .+ y) ≈ L[y,:] \ (2y)
             @test L[y,:] \ (y .- y) ≈ zeros(10)
         end
+
+        @testset "transform" begin
+            x = Inclusion(0..1)
+            y = 2x .- 1
+            L = LinearSpline(range(-1,stop=1,length=10))
+            g,P = plan_grid_transform(L[y,:], (10,))
+            X = cos.(g)
+            @test L[y,:][g,:] * (P * X) ≈ X
+            @test P \ (P * X) ≈ P * (P \ X) ≈ X
+
+            g,P = plan_grid_transform(L[y,:], (10,10))
+            X = cos.(g .+ g')
+            @test L[y,:][g,:]*(P * X)*L[y,:][g,:]' ≈ X
+            @test P \ (P * X) ≈ P * (P \ X) ≈ X
+
+            g,P = plan_grid_transform(L[y,:], (10,10,10))
+            X = randn(10,10,10)
+            @test P \ (P * X) ≈ P * (P \ X) ≈ X
+        end
     end
 
     @testset "diff" begin