SubOperator index with colon (#130)

Author: jishnub

src/Operators/Operator.jl

@@ -5,6 +5,8 @@ export ldirichlet,rdirichlet,lneumann,rneumann
 export ldiffbc,rdiffbc,diffbcs
 export domainspace,rangespace
 
+const VectorIndices = Union{AbstractRange, Colon}
+const IntOrVectorIndices = Union{Integer, VectorIndices}
 
 abstract type Operator{T} end #T is the entry type, Float64 or Complex{Float64}
 

src/Operators/SubOperator.jl

@@ -272,11 +272,13 @@ end
 size(V::SubOperator) = V.dims
 size(V::SubOperator,k::Int) = V.dims[k]
 
+axes(V::SubOperator) = map(Base.OneTo, size(V))
+axes(V::SubOperator, k::Integer) = k <= 2 ? axes(V)[k] : Base.OneTo(1)
+
 unsafe_getindex(V::SubOperator,k::Integer,j::Integer) = V.parent[reindex(V,parentindices(V),(k,j))...]
-getindex(V::SubOperator,k::Integer,j::Integer) = V.parent[reindex(V,parentindices(V),(k,j))...]
-getindex(V::SubOperator,k::Integer,j::AbstractRange) = V.parent[reindex(V,parentindices(V),(k,j))...]
-getindex(V::SubOperator,k::AbstractRange,j::Integer) = V.parent[reindex(V,parentindices(V),(k,j))...]
-getindex(V::SubOperator,k::AbstractRange,j::AbstractRange) = V.parent[reindex(V,parentindices(V),(k,j))...]
+function getindex(V::SubOperator,k::IntOrVectorIndices,j::IntOrVectorIndices)
+    V.parent[reindex(V,parentindices(V),(k,j))...]
+end
 Base.parent(S::SubOperator) = S.parent
 Base.parentindices(S::SubOperator) = S.indexes
 

test/runtests.jl

@@ -164,5 +164,35 @@ end
     end
 end
 
+@testset "operator indexing" begin
+    @testset "SubOperator" begin
+        D = Dirichlet(ConstantSpace(0..1))
+        S = D[:, :]
+        @test S[1,1] == 1
+        ax1 = axes(S, 1)
+        ax2 = axes(S, 2)
+        inds1 = Any[ax1, StepRange(ax1), :]
+        inds2 = Any[ax2, StepRange(ax2), :]
+        @testset for r2 in inds2, r1 in inds1
+            M = S[r1, r2]
+            @test M isa AbstractMatrix
+            @test size(M) == (2,1)
+            @test all(==(1), M)
+        end
+        @testset for r1 in inds1
+            V = S[r1, 1]
+            @test V isa AbstractVector
+            @test size(V) == (2,)
+            @test all(==(1), V)
+        end
+        @testset for r2 in inds2
+            V = S[1, r2]
+            @test V isa AbstractVector
+            @test size(V) == (1,)
+            @test all(==(1), V)
+        end
+    end
+end
+
 @time include("ETDRK4Test.jl")
 include("show.jl")