Add more tests

Author: mtfishman

src/KroneckerArrays.jl

@@ -302,6 +302,9 @@ struct KroneckerStyle{N,A,B} <: AbstractArrayStyle{N} end
 function KroneckerStyle{N}(a::BroadcastStyle, b::BroadcastStyle) where {N}
   return KroneckerStyle{N,a,b}()
 end
+function KroneckerStyle(a::AbstractArrayStyle{N}, b::AbstractArrayStyle{N}) where {N}
+  return KroneckerStyle{N}(a, b)
+end
 function KroneckerStyle{N,A,B}(v::Val{M}) where {N,A,B,M}
   return KroneckerStyle{M,typeof(A)(v),typeof(B)(v)}()
 end
@@ -316,8 +319,8 @@ end
 function Base.similar(bc::Broadcasted{<:KroneckerStyle{N,A,B}}, elt::Type) where {N,A,B}
   ax_a = map(ax -> ax.product.a, axes(bc))
   ax_b = map(ax -> ax.product.b, axes(bc))
-  bc_a = Broadcasted(A, ax_a)
-  bc_b = Broadcasted(B, ax_b)
+  bc_a = Broadcasted(A, nothing, (), ax_a)
+  bc_b = Broadcasted(B, nothing, (), ax_b)
   a = similar(bc_a, elt)
   b = similar(bc_b, elt)
   return a ⊗ b

test/test_basics.jl

@@ -1,6 +1,9 @@
+using Base.Broadcast: BroadcastStyle, Broadcasted, broadcasted
 using FillArrays: Eye
 using KroneckerArrays:
   KroneckerArrays,
+  KroneckerArray,
+  KroneckerStyle,
   CartesianProductUnitRange,
   ⊗,
   ×,
@@ -9,7 +12,7 @@ using KroneckerArrays:
   diagonal,
   kron_nd,
   unproduct
-using LinearAlgebra: Diagonal, I, det, eigen, eigvals, lq, pinv, qr, svd, svdvals, tr
+using LinearAlgebra: Diagonal, I, det, eigen, eigvals, lq, norm, pinv, qr, svd, svdvals, tr
 using StableRNGs: StableRNG
 using Test: @test, @test_broken, @test_throws, @testset
 
@@ -41,8 +44,10 @@ const elts = (Float32, Float64, ComplexF32, ComplexF64)
   a = randn(elt, 2, 2) ⊗ randn(elt, 3, 3)
   b = randn(elt, 2, 2) ⊗ randn(elt, 3, 3)
   c = a.a ⊗ b.b
+  @test a isa KroneckerArray{elt,2,typeof(a.a),typeof(a.b)}
   @test similar(typeof(a), (2, 3)) isa Matrix{elt}
   @test size(similar(typeof(a), (2, 3))) == (2, 3)
+  @test isreal(a) == (elt <: Real)
   @test a[1 × 1, 1 × 1] == a.a[1, 1] * a.b[1, 1]
   @test a[1 × 3, 2 × 1] == a.a[1, 2] * a.b[3, 1]
   @test a[1 × (2:3), 2 × 1] == a.a[1, 2] * a.b[2:3, 1]
@@ -60,6 +65,7 @@ const elts = (Float32, Float64, ComplexF32, ComplexF64)
   @test collect(-a) == -collect(a)
   @test collect(3 * a) ≈ 3 * collect(a)
   @test collect(a * 3) ≈ collect(a) * 3
+  @test collect(a / 3) ≈ collect(a) / 3
   @test a + a == 2a
   @test iszero(a - a)
   @test collect(a + c) ≈ collect(a) + collect(c)
@@ -68,6 +74,61 @@ const elts = (Float32, Float64, ComplexF32, ComplexF64)
     @test collect(f(a)) ≈ f(collect(a))
   end
   @test tr(a) ≈ tr(collect(a))
+  @test norm(a) ≈ norm(collect(a))
+
+  # Broadcasting
+  style = KroneckerStyle(BroadcastStyle(typeof(a.a)), BroadcastStyle(typeof(a.b)))
+  @test BroadcastStyle(typeof(a)) === style
+  @test_throws "not supported" sin.(a)
+  a′ = similar(a)
+  @test_throws "not supported" a′ .= sin.(a)
+  a′ = similar(a)
+  @test_broken a′ .= 2 .* a
+  bc = broadcasted(+, a, a)
+  @test bc.style === style
+  @test similar(bc, elt) isa KroneckerArray{elt,2,typeof(a.a),typeof(a.b)}
+  @test_broken copy(bc)
+  bc = broadcasted(*, 2, a)
+  @test bc.style === style
+  @test_broken copy(bc)
+
+  # Mapping
+  @test_throws "not supported" map(sin, a)
+  @test_broken map(Base.Fix1(*, 2), a)
+  a′ = similar(a)
+  @test_throws "not supported" map!(sin, a′, a)
+  a′ = similar(a)
+  map!(identity, a′, a)
+  @test collect(a′) ≈ collect(a)
+  a′ = similar(a)
+  map!(+, a′, a, a)
+  @test collect(a′) ≈ 2 * collect(a)
+  a′ = similar(a)
+  map!(-, a′, a, a)
+  @test norm(collect(a′)) ≈ 0
+  a′ = similar(a)
+  map!(Base.Fix1(*, 2), a′, a)
+  @test collect(a′) ≈ 2 * collect(a)
+  a′ = similar(a)
+  map!(Base.Fix2(*, 2), a′, a)
+  @test collect(a′) ≈ 2 * collect(a)
+  a′ = similar(a)
+  map!(Base.Fix2(/, 2), a′, a)
+  @test collect(a′) ≈ collect(a) / 2
+  a′ = similar(a)
+  map!(conj, a′, a)
+  @test collect(a′) ≈ conj(collect(a))
+
+  if elt <: Real
+    @test real(a) == a
+  else
+    @test_throws ArgumentError real(a)
+  end
+  if elt <: Real
+    @test iszero(imag(a))
+  else
+    @test_throws ArgumentError imag(a)
+  end
 
   a = randn(elt, 2, 2, 2) ⊗ randn(elt, 3, 3, 3)
   @test collect(a) ≈ kron_nd(a.a, a.b)