Add range indexing to UniformScaling (JuliaLang#24359)

* Add range indexing to UniformScaling

* Add tests

* Return sparse matrix when indexing UniformScaling

* start -> first

* Correctly handle non-integer steps

* Make getindex(::UniformScaling, ranges...) return a dense array

* Remove unneccessary (and slower) branch

* Add (range, integer) indexing to UniformScaling

* Update tests

* Fix unbound param

* Add NEWS item

* Update docstring and add compat annotation

Co-authored-by: Evey Dee <eveydee@users.noreply.github.com>

NEWS.md

@@ -50,6 +50,7 @@ Standard library changes
 
 #### LinearAlgebra
 * New method `LinearAlgebra.issuccess(::CholeskyPivoted)` for checking whether pivoted Cholesky factorization was successful ([#36002]).
+* `UniformScaling` can now be indexed into using ranges to return dense matrices and vectors ([#24359]).
 
 #### Markdown
 

stdlib/LinearAlgebra/src/uniformscaling.jl

@@ -6,8 +6,13 @@ import Base: copy, adjoint, getindex, show, transpose, one, zero, inv,
 """
  UniformScaling{T<:Number}
 
-Generically sized uniform scaling operator defined as a scalar times the
-identity operator, `λ*I`. See also [`I`](@ref).
+Generically sized uniform scaling operator defined as a scalar times
+the identity operator, `λ*I`. Although without an explicit `size`, it
+acts similarly to a matrix in many cases and includes support for some
+indexing. See also [`I`](@ref).
+
+!!! compat "Julia 1.6"
+ Indexing using ranges is available as of Julia 1.6.
 
 # Examples
 ```jldoctest
@@ -24,6 +29,11 @@ julia> J*A
 2×2 Array{Float64,2}:
  2.0 4.0
  6.0 8.0
+
+julia> J[1:2, 1:2]
+2×2 Array{Float64,2}:
+ 2.0 0.0
+ 0.0 2.0
 ```
 """
 struct UniformScaling{T<:Number}
@@ -78,6 +88,28 @@ ndims(J::UniformScaling) = 2
 Base.has_offset_axes(::UniformScaling) = false
 getindex(J::UniformScaling, i::Integer,j::Integer) = ifelse(i==j,J.λ,zero(J.λ))
 
+getindex(x::UniformScaling, n::Integer, m::AbstractRange{<:Integer}) = getindex(x, m, n)
+function getindex(x::UniformScaling{T}, n::AbstractRange{<:Integer}, m::Integer) where T
+ v = zeros(T, length(n))
+ @inbounds for (i,ii) in enumerate(n)
+ if ii == m
+ v[i] = x.λ
+ end
+ end
+ return v
+end
+
+
+function getindex(x::UniformScaling{T}, n::AbstractRange{<:Integer}, m::AbstractRange{<:Integer}) where T
+ A = zeros(T, length(n), length(m))
+ @inbounds for (j,jj) in enumerate(m), (i,ii) in enumerate(n)
+ if ii == jj
+ A[i,j] = x.λ
+ end
+ end
+ return A
+end
+
 function show(io::IO, ::MIME"text/plain", J::UniformScaling)
  s = "$(J.λ)"
  if occursin(r"\w+\s*[\+\-]\s*\w+", s)

stdlib/LinearAlgebra/test/uniformscaling.jl

@@ -11,8 +11,6 @@ using .Main.Quaternions
 Random.seed!(123)
 
 @testset "basic functions" begin
- @test I[1,1] == 1 # getindex
- @test I[1,2] == 0 # getindex
  @test I === I' # transpose
  @test ndims(I) == 2
  @test one(UniformScaling{Float32}) == UniformScaling(one(Float32))
@@ -27,6 +25,30 @@ Random.seed!(123)
  @test opnorm(UniformScaling(1+im)) ≈ sqrt(2)
 end
 
+@testset "getindex" begin
+ @test I[1,1] == 1
+ @test I[1,2] == 0
+
+ J = I(15)
+ for (a, b) in [
+ # indexing that returns a Vector
+ (1:10, 1),
+ (4, 1:10),
+ (11, 1:10),
+ # indexing that returns a Matrix
+ (1:2, 1:2),
+ (1:2:3, 1:2:3),
+ (1:2:8, 2:2:9),
+ (1:2:8, 9:-4:1),
+ (9:-4:1, 1:2:8),
+ (2:3, 1:2),
+ (2:-1:1, 1:2),
+ (1:2:9, 5:2:13),
+ ]
+ @test I[a,b] == J[a,b]
+ end
+end
+
 @testset "sqrt, exp, log, and trigonometric functions" begin
  # convert to a dense matrix with random size
  M(J) = (N = rand(1:10); Matrix(J, N, N))