@@ -4,23 +4,23 @@ using BlockArrays: Block
44using BlockSparseArrays: BlockIndexVector, GenericBlockIndex
55using KroneckerArrays: CartesianPair, CartesianProduct
66function Base. getindex (
7- b:: Block ,
8- I1:: Union{CartesianPair,CartesianProduct} ,
9- Irest:: Union{CartesianPair,CartesianProduct} ...,
10- )
11- return GenericBlockIndex (b, (I1, Irest... ))
7+ b:: Block ,
8+ I1:: Union{CartesianPair, CartesianProduct} ,
9+ Irest:: Union{CartesianPair, CartesianProduct} ...,
10+ )
11+ return GenericBlockIndex (b, (I1, Irest... ))
1212end
1313function Base. getindex (b:: Block , I1:: CartesianProduct , Irest:: CartesianProduct... )
14- return BlockIndexVector (b, (I1, Irest... ))
14+ return BlockIndexVector (b, (I1, Irest... ))
1515end
1616
1717using BlockSparseArrays: BlockSparseArrays, blockrange
1818using KroneckerArrays: CartesianPair, CartesianProduct, cartesianrange
1919function BlockSparseArrays. blockrange (bs:: Vector{<:CartesianPair} )
20- return blockrange (map (cartesianrange, bs))
20+ return blockrange (map (cartesianrange, bs))
2121end
2222function BlockSparseArrays. blockrange (bs:: Vector{<:CartesianProduct} )
23- return blockrange (map (cartesianrange, bs))
23+ return blockrange (map (cartesianrange, bs))
2424end
2525
2626using BlockArrays: BlockArrays, mortar
@@ -31,56 +31,56 @@ using KroneckerArrays: CartesianProductUnitRange
3131# This is helpful when indexing `BlockUnitRange`, for example:
3232# https://github.com/JuliaArrays/BlockArrays.jl/blob/v1.7.1/src/blockaxis.jl#L540-L547
3333function BlockArrays. mortar (blocks:: AbstractVector{<:CartesianProductUnitRange} )
34- return mortar (blocks, (blockrange (map (Base. axes1, blocks)),))
34+ return mortar (blocks, (blockrange (map (Base. axes1, blocks)),))
3535end
3636
3737using BlockArrays: AbstractBlockedUnitRange
3838using BlockSparseArrays: Block, ZeroBlocks, eachblockaxis, mortar_axis
3939using KroneckerArrays: KroneckerArrays, KroneckerArray, ⊗ , arg1, arg2, isactive
4040
4141function KroneckerArrays. arg1 (r:: AbstractBlockedUnitRange )
42- return mortar_axis (arg1 .(eachblockaxis (r)))
42+ return mortar_axis (arg1 .(eachblockaxis (r)))
4343end
4444function KroneckerArrays. arg2 (r:: AbstractBlockedUnitRange )
45- return mortar_axis (arg2 .(eachblockaxis (r)))
45+ return mortar_axis (arg2 .(eachblockaxis (r)))
4646end
4747
4848function block_axes (
49- ax:: NTuple{N,AbstractUnitRange{<:Integer}} , I:: Vararg{Block{1},N}
50- ) where {N}
51- return ntuple (N) do d
52- return only (axes (ax[d][I[d]]))
53- end
49+ ax:: NTuple{N, AbstractUnitRange{<:Integer}} , I:: Vararg{Block{1}, N}
50+ ) where {N}
51+ return ntuple (N) do d
52+ return only (axes (ax[d][I[d]]))
53+ end
5454end
55- function block_axes (ax:: NTuple{N,AbstractUnitRange{<:Integer}} , I:: Block{N} ) where {N}
56- return block_axes (ax, Tuple (I)... )
55+ function block_axes (ax:: NTuple{N, AbstractUnitRange{<:Integer}} , I:: Block{N} ) where {N}
56+ return block_axes (ax, Tuple (I)... )
5757end
5858
5959using DiagonalArrays: ShapeInitializer
6060
6161# # TODO : Is this needed?
6262function Base. getindex (
63- a:: ZeroBlocks{N,KroneckerArray{T,N, A1,A2}} , I:: Vararg{Int,N}
64- ) where {T,N,A1 <: AbstractArray{T,N} ,A2 <: AbstractArray{T,N} }
65- ax_a1 = map (arg1, a. parentaxes)
66- ax_a2 = map (arg2, a. parentaxes)
67- block_ax_a1 = arg1 .(block_axes (a. parentaxes, Block (I)))
68- block_ax_a2 = arg2 .(block_axes (a. parentaxes, Block (I)))
69- # TODO : Is this a good definition? It is similar to
70- # the definition of `similar` and `adapt_structure`.
71- return if isactive (A1) == isactive (A2)
72- ZeroBlocks {N,A1} (ax_a1)[I... ] ⊗ ZeroBlocks {N,A2} (ax_a2)[I... ]
73- elseif isactive (A1)
74- ZeroBlocks {N,A1} (ax_a1)[I... ] ⊗ A2 (ShapeInitializer (), block_ax_a2)
75- elseif isactive (A2)
76- A1 (ShapeInitializer (), block_ax_a1) ⊗ ZeroBlocks {N,A2} (ax_a2)[I... ]
77- end
63+ a:: ZeroBlocks{N, KroneckerArray{T, N, A1, A2}} , I:: Vararg{Int, N}
64+ ) where {T, N, A1 <: AbstractArray{T, N} , A2 <: AbstractArray{T, N} }
65+ ax_a1 = map (arg1, a. parentaxes)
66+ ax_a2 = map (arg2, a. parentaxes)
67+ block_ax_a1 = arg1 .(block_axes (a. parentaxes, Block (I)))
68+ block_ax_a2 = arg2 .(block_axes (a. parentaxes, Block (I)))
69+ # TODO : Is this a good definition? It is similar to
70+ # the definition of `similar` and `adapt_structure`.
71+ return if isactive (A1) == isactive (A2)
72+ ZeroBlocks {N, A1} (ax_a1)[I... ] ⊗ ZeroBlocks {N, A2} (ax_a2)[I... ]
73+ elseif isactive (A1)
74+ ZeroBlocks {N, A1} (ax_a1)[I... ] ⊗ A2 (ShapeInitializer (), block_ax_a2)
75+ elseif isactive (A2)
76+ A1 (ShapeInitializer (), block_ax_a1) ⊗ ZeroBlocks {N, A2} (ax_a2)[I... ]
77+ end
7878end
7979
8080using BlockSparseArrays: BlockSparseArrays
8181using KroneckerArrays: KroneckerArrays, KroneckerVector
8282function BlockSparseArrays. to_truncated_indices (values:: KroneckerVector , I)
83- return KroneckerArrays. to_truncated_indices (values, I)
83+ return KroneckerArrays. to_truncated_indices (values, I)
8484end
8585
8686end
0 commit comments