JuliaLang / julia

The Julia Programming Language
https://julialang.org/
MIT License
44.98k stars 5.42k forks source link

Support broadcasting `I` #53240

Open LilithHafner opened 1 year ago

LilithHafner commented 1 year ago

It would be nice to create this matrix:

1 0 0 -1 0 0
0 1 0 0 -1 0
0 0 1 0 0 -1

with a single allocation by broadcasting I

x = Matrix{Int}(undef, 3, 6)
x[1:3, 1:3] .= I
x[1:3, 4:6] .= -I
ygol55 commented 1 year ago

You mean being able to define a spesific pattern like the matrix given, and setting for another matrix? Or you mean the ability to connect few matrasies with each other (as defined in python)?

LilithHafner commented 1 year ago

Copying elements from one matrix to another is already supported (e.g. matrix1 .= matrix2) as is combining two matrices (e.g. hcat(matrix1, matrix2)). My example above can be constructed as hcat(Matrix{Int}(I, 3, 3), -I).

I mean supporting broadcasting the special identity matrix I so that it is convenient to instantiate the above example more efficiently.

stevengj commented 1 year ago

This doesn't seem like it should be a broadcasting assignment — it seems like an ordinary setindex! should work. i.e.

x[1:3, 1:3] = I
x[1:3, 4:6] = -I

(Bonus: this should much easier to implement than wrestling with the broadcast machinery.)

stevengj commented 1 year ago

For example, a patch like the following (untested):

diff --git a/src/uniformscaling.jl b/src/uniformscaling.jl
index 661bd28..c421dd9 100644
--- a/src/uniformscaling.jl
+++ b/src/uniformscaling.jl
@@ -87,6 +87,7 @@ eltype(::Type{UniformScaling{T}}) where {T} = T
 ndims(J::UniformScaling) = 2
 Base.has_offset_axes(::UniformScaling) = false
 getindex(J::UniformScaling, i::Integer,j::Integer) = ifelse(i==j,J.λ,zero(J.λ))
+getindex(J::UniformScaling, c::CartesianIndex{2}) = J[Tuple(c)...]

 getindex(J::UniformScaling, n::Integer, m::AbstractVector{<:Integer}) = getindex(J, m, n)
 function getindex(J::UniformScaling{T}, n::AbstractVector{<:Integer}, m::Integer) where T
@@ -109,6 +110,23 @@ function getindex(J::UniformScaling{T}, n::AbstractVector{<:Integer}, m::Abstrac
     return A
 end

+Base.setindex!(A::AbstractMatrix{<:Number}, J::UniformScaling, ::Colon) =
+    setindex!(A, J, axes(A,1), axes(A,2))
+function Base.setindex!(A::AbstractMatrix{<:Number}, J::UniformScaling, I1_, I2_)
+    I1, I2 = to_indices(A, (I1_, I2_))
+    for i2 in I2, i1 in I1
+        A[i1,i2] = J[i1,i2]
+    end
+    return J
+end
+function Base.setindex!(A::AbstractMatrix{<:Number}, J::UniformScaling, I::AbstractArray{CartesianIndex{2}})
+    for i in I
+        A[i] = J[i]
+    end
+    return J
+end
+
+
 function show(io::IO, ::MIME"text/plain", J::UniformScaling)
     s = "$(J.λ)"
     if occursin(r"\w+\s*[\+\-]\s*\w+", s)
vtjnash commented 5 months ago

Seems like normally either may work

julia> x = Matrix{Int}(undef, 3, 6)
3×6 Matrix{Int64}:
 7358993307422584690  8387230136615138145  7309475598976773232  8314035669885219394  3709670234487346254  8245940694597068590
 3539896516525191791  6879638067882520424  8233194082049016443  3205830026326201204  3205856487619313964  8389772277106828396
 4764873483910918458  8454426816625079398  3347427285159666031  4196364628411372576  7310293556787950368  7286889245055677029

julia> x[1:3, 1:3] = [1 0 0; 0 1 0; 0 0 1]
3×3 Matrix{Int64}:
 1  0  0
 0  1  0
 0  0  1

julia> x
3×6 Matrix{Int64}:
 1  0  0  8314035669885219394  3709670234487346254  8245940694597068590
 0  1  0  3205830026326201204  3205856487619313964  8389772277106828396
 0  0  1  4196364628411372576  7310293556787950368  7286889245055677029

julia> x[1:3, 4:6] .= [1 0 0; 0 1 0; 0 0 1]
3×3 view(::Matrix{Int64}, 1:3, 4:6) with eltype Int64:
 1  0  0
 0  1  0
 0  0  1

julia> x
3×6 Matrix{Int64}:
 1  0  0  1  0  0
 0  1  0  0  1  0
 0  0  1  0  0  1
stevengj commented 5 months ago

Yes, but since I doesn't have a fixed size, it may be a bit confusing in a broadcast context.

Overloading setindex! seems less ambiguous as well as a lot less work, and in any case it's a good first step that can be done regardless of how we decide I should work in broadcasting.