Redesign to support slicing (#10)

Author: mtfishman

Project.toml

@@ -1,10 +1,11 @@
 name = "NamedDimsArrays"
 uuid = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.2.0"
+version = "0.3.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BroadcastMapConversion = "4a4adec5-520f-4750-bb37-d5e66b4ddeb2"
 Derive = "a07dfc7f-7d04-4eb5-84cc-a97f051f655a"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -13,8 +14,16 @@ SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
+[weakdeps]
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+
+[extensions]
+NamedDimsArraysBlockArraysExt = "BlockArrays"
+
 [compat]
 Adapt = "4.1.1"
+ArrayLayouts = "1.11.0"
+BlockArrays = "1.3.0"
 BroadcastMapConversion = "0.1.2"
 Derive = "0.3.6"
 LinearAlgebra = "1.10"

README.md

@@ -32,33 +32,67 @@ julia> Pkg.add("NamedDimsArrays")
 ## Examples
 
 ````julia
-using NamedDimsArrays: aligndims, dename, dimnames, named
+using NamedDimsArrays: aligndims, dimnames, named, nameddimsindices, namedoneto, unname
 using TensorAlgebra: contract
+using Test: @test
 
 # Named dimensions
-i = named(2, "i")
-j = named(2, "j")
-k = named(2, "k")
+i = namedoneto(2, "i")
+j = namedoneto(2, "j")
+k = namedoneto(2, "k")
 
 # Arrays with named dimensions
-na1 = randn(i, j)
-na2 = randn(j, k)
+a1 = randn(i, j)
+a2 = randn(j, k)
 
-@show dimnames(na1) == ("i", "j")
+@test dimnames(a1) == ("i", "j")
+@test nameddimsindices(a1) == (i, j)
+@test axes(a1) == (named(1:2, i), named(1:2, j))
+@test size(a1) == (named(2, i), named(2, j))
 
 # Indexing
-@show na1[j => 2, i => 1] == na1[1, 2]
+@test a1[j => 2, i => 1] == a1[1, 2]
+@test a1[j[2], i[1]] == a1[1, 2]
 
 # Tensor contraction
-na_dest = contract(na1, na2)
+a_dest = contract(a1, a2)
 
-@show issetequal(dimnames(na_dest), ("i", "k"))
-# `dename` removes the names and returns an `Array`
-@show dename(na_dest, (i, k)) ≈ dename(na1) * dename(na2)
+@test issetequal(nameddimsindices(a_dest), (i, k))
+# `unname` removes the names and returns an `Array`
+@test unname(a_dest, (i, k)) ≈ unname(a1, (i, j)) * unname(a2, (j, k))
 
 # Permute dimensions (like `ITensors.permute`)
-na1 = aligndims(na1, (j, i))
-@show na1[i => 1, j => 2] == na1[2, 1]
+a1′ = aligndims(a1, (j, i))
+@test a1′[i => 1, j => 2] == a1[i => 1, j => 2]
+@test a1′[i[1], j[2]] == a1[i[1], j[2]]
+
+# Contiguous slicing
+b1 = a1[i => 1:2, j => 1:1]
+@test b1 == a1[i[1:2], j[1:1]]
+
+b2 = a2[j => 1:1, k => 1:2]
+@test b2 == a2[j[1:1], k[1:2]]
+
+@test nameddimsindices(b1) == (i[1:2], j[1:1])
+@test nameddimsindices(b2) == (j[1:1], k[1:2])
+
+b_dest = contract(b1, b2)
+
+@test issetequal(nameddimsindices(b_dest), (i, k))
+
+# Non-contiguous slicing
+c1 = a1[i[[2, 1]], j[[2, 1]]]
+@test nameddimsindices(c1) == (i[[2, 1]], j[[2, 1]])
+@test unname(c1, (i[[2, 1]], j[[2, 1]])) == unname(a1, (i, j))[[2, 1], [2, 1]]
+@test c1[i[2], j[1]] == a1[i[2], j[1]]
+@test c1[2, 1] == a1[1, 2]
+
+a1[i[[2, 1]], j[[2, 1]]] = [22 21; 12 11]
+@test a1[i[1], j[1]] == 11
+
+x = randn(i[1:2], j[2:2])
+a1[i[1:2], j[2:2]] = x
+@test a1[i[1], j[2]] == x[i[1], j[2]]
 ````
 
 ---

TODO.md

@@ -1,20 +1,33 @@
+- Define `@align`/`@aligned` such that:
+```julia
+i = namedoneto(2, "i")
+j = namedoneto(2, "j")
+a = randn(i, j)
+@align a[j, i]
+@aligned a[j, i]
+```
+aligns the dimensions (currently `a[j, i]` doesn't align the dimensions).
+It could be written in terms of `align_getindex`/`align_view`.
 - `svd`, `eigen` (including tensor versions)
-- `reshape`, `vec`
-- `swapdimnames`
-- `mapdimnames(f, a::AbstractNamedDimsArray)` (rename `replacedimnames(f, a)` to `mapdimnames(f, a)`, or have both?)
+- `reshape`, `vec`, including fused dimension names.
+- Dimension name set logic, i.e. `setdiffnameddimsindices(a::AbstractNamedDimsArray, b::AbstractNamedDimsArray)`, etc.
+- `swapnameddimsindices` (written in terms of `mapnameddimsindices`/`replacenameddimsindices`).
+- `mapnameddimsindices(f, a::AbstractNamedDimsArray)` (rename `replacenameddimsindices(f, a)` to `mapnameddimsindices(f, a)`, or have both?)
 - `cat` (define `CatName` as a combination of the input names?).
 - `canonize`/`flatten_array_wrappers` (https://github.com/mcabbott/NamedPlus.jl/blob/v0.0.5/src/permute.jl#L207)
-  - `nameddims(PermutedDimsArray(a, perm), dimnames)` -> `nameddims(a, dimnames[invperm(perm)])`
-  - `nameddims(transpose(a), dimnames)` -> `nameddims(a, reverse(dimnames))`
-  - `Transpose(nameddims(a, dimnames))` -> `nameddims(a, reverse(dimnames))`
+  - `nameddims(PermutedDimsArray(a, perm), nameddimsindices)` -> `nameddims(a, nameddimsindices[invperm(perm)])`
+  - `nameddims(transpose(a), nameddimsindices)` -> `nameddims(a, reverse(nameddimsindices))`
+  - `Transpose(nameddims(a, nameddimsindices))` -> `nameddims(a, reverse(nameddimsindices))`
   - etc.
 - `MappedName(old_name, name)`, acts like `Name(name)` but keeps track of the old name.
-  - `namedmap(a, ::Pair...)`: `namedmap(named(randn(2, 2, 2, 2), i, j, k, l), i => k, j => l)`
+  - `nameddimsmap(a, ::Pair...)`: `namedmap(named(randn(2, 2, 2, 2), i, j, k, l), i => k, j => l)`
     represents that the names map back and forth to each other for the sake of `transpose`,
     `tr`, `eigen`, etc. Operators are generally `namedmap(named(randn(2, 2), i, i'), i => i')`.
 - `prime(:i) = PrimedName(:i)`, `prime(:i, 2) = PrimedName(:i, 2)`, `prime(prime(:i)) = PrimedName(:i, 2)`,
   `Name(:i)' = prime(:i)`, etc.
-- `transpose`/`adjoint` based on `swapdimnames` and `MappedName(old_name, new_name)`.
+    - Also `prime(f, a::AbstractNamedDimsArray)` where `f` is a filter function to determine
+      which dimensions to filter.
+- `transpose`/`adjoint` based on `swapnameddimsindices` and `MappedName(old_name, new_name)`.
   - `adjoint` could make use of a lazy `ConjArray`.
   - `transpose(a, dimname1 => dimname1′, dimname2 => dimname2′)` like `https://github.com/mcabbott/NamedPlus.jl`.
     - Same as `replacedims(a, dimname1 => dimname1′, dimname1′ => dimname1, dimname2 => dimname2′, dimname2′ => dimname2)`.
@@ -23,4 +36,5 @@
 - Slicing: `nameddims(a, "i", "j")[1:2, 1:2] = nameddims(a[1:2, 1:2], Name(named(1:2, "i")), Name(named(1:2, "j")))`, i.e.
   the parent gets sliced and the new dimensions names are the named slice.
   - Should `NamedDimsArray` store the named axes rather than just the dimension names?
-  - Should `NamedDimsArray` have special axes types so that `axes(nameddims(a, "i", "j")) == axes(nameddims(a', "j", "i"))`?
+  - Should `NamedDimsArray` have special axes types so that `axes(nameddims(a, "i", "j")) == axes(nameddims(a', "j", "i"))`,
+    i.e. equality is based on `issetequal` and not dependent on the ordering of the dimensions?

examples/README.jl

@@ -37,30 +37,64 @@ julia> Pkg.add("NamedDimsArrays")
 
 # ## Examples
 
-using NamedDimsArrays: aligndims, dename, dimnames, named
+using NamedDimsArrays: aligndims, dimnames, named, nameddimsindices, namedoneto, unname
 using TensorAlgebra: contract
+using Test: @test
 
 ## Named dimensions
-i = named(2, "i")
-j = named(2, "j")
-k = named(2, "k")
+i = namedoneto(2, "i")
+j = namedoneto(2, "j")
+k = namedoneto(2, "k")
 
 ## Arrays with named dimensions
-na1 = randn(i, j)
-na2 = randn(j, k)
+a1 = randn(i, j)
+a2 = randn(j, k)
 
-@show dimnames(na1) == ("i", "j")
+@test dimnames(a1) == ("i", "j")
+@test nameddimsindices(a1) == (i, j)
+@test axes(a1) == (named(1:2, i), named(1:2, j))
+@test size(a1) == (named(2, i), named(2, j))
 
 ## Indexing
-@show na1[j => 2, i => 1] == na1[1, 2]
+@test a1[j => 2, i => 1] == a1[1, 2]
+@test a1[j[2], i[1]] == a1[1, 2]
 
 ## Tensor contraction
-na_dest = contract(na1, na2)
+a_dest = contract(a1, a2)
 
-@show issetequal(dimnames(na_dest), ("i", "k"))
-## `dename` removes the names and returns an `Array`
-@show dename(na_dest, (i, k)) ≈ dename(na1) * dename(na2)
+@test issetequal(nameddimsindices(a_dest), (i, k))
+## `unname` removes the names and returns an `Array`
+@test unname(a_dest, (i, k)) ≈ unname(a1, (i, j)) * unname(a2, (j, k))
 
 ## Permute dimensions (like `ITensors.permute`)
-na1 = aligndims(na1, (j, i))
-@show na1[i => 1, j => 2] == na1[2, 1]
+a1′ = aligndims(a1, (j, i))
+@test a1′[i => 1, j => 2] == a1[i => 1, j => 2]
+@test a1′[i[1], j[2]] == a1[i[1], j[2]]
+
+## Contiguous slicing
+b1 = a1[i => 1:2, j => 1:1]
+@test b1 == a1[i[1:2], j[1:1]]
+
+b2 = a2[j => 1:1, k => 1:2]
+@test b2 == a2[j[1:1], k[1:2]]
+
+@test nameddimsindices(b1) == (i[1:2], j[1:1])
+@test nameddimsindices(b2) == (j[1:1], k[1:2])
+
+b_dest = contract(b1, b2)
+
+@test issetequal(nameddimsindices(b_dest), (i, k))
+
+## Non-contiguous slicing
+c1 = a1[i[[2, 1]], j[[2, 1]]]
+@test nameddimsindices(c1) == (i[[2, 1]], j[[2, 1]])
+@test unname(c1, (i[[2, 1]], j[[2, 1]])) == unname(a1, (i, j))[[2, 1], [2, 1]]
+@test c1[i[2], j[1]] == a1[i[2], j[1]]
+@test c1[2, 1] == a1[1, 2]
+
+a1[i[[2, 1]], j[[2, 1]]] = [22 21; 12 11]
+@test a1[i[1], j[1]] == 11
+
+x = randn(i[1:2], j[2:2])
+a1[i[1:2], j[2:2]] = x
+@test a1[i[1], j[2]] == x[i[1], j[2]]

ext/NamedDimsArraysBlockArraysExt/NamedDimsArraysBlockArraysExt.jl

@@ -0,0 +1,45 @@
+module NamedDimsArraysBlockArraysExt
+using ArrayLayouts: ArrayLayouts
+using BlockArrays: Block, BlockRange
+using NamedDimsArrays:
+  AbstractNamedDimsArray,
+  AbstractNamedUnitRange,
+  named_getindex,
+  nameddims_getindex,
+  nameddims_view
+
+function Base.getindex(r::AbstractNamedUnitRange{<:Integer}, I::Block{1})
+  # TODO: Use `Derive.@interface NamedArrayInterface() r[I]` instead.
+  return named_getindex(r, I)
+end
+
+function Base.getindex(r::AbstractNamedUnitRange{<:Integer}, I::BlockRange{1})
+  # TODO: Use `Derive.@interface NamedArrayInterface() r[I]` instead.
+  return named_getindex(r, I)
+end
+
+const BlockIndex{N} = Union{Block{N},BlockRange{N},AbstractVector{<:Block{N}}}
+
+function Base.view(a::AbstractNamedDimsArray, I1::Block{1}, Irest::BlockIndex{1}...)
+  # TODO: Use `Derive.@interface NamedDimsArrayInterface() r[I]` instead.
+  return nameddims_view(a, I1, Irest...)
+end
+
+function Base.view(a::AbstractNamedDimsArray, I::Block)
+  # TODO: Use `Derive.@interface NamedDimsArrayInterface() r[I]` instead.
+  return nameddims_view(a, Tuple(I)...)
+end
+
+function Base.view(a::AbstractNamedDimsArray, I1::BlockIndex{1}, Irest::BlockIndex{1}...)
+  # TODO: Use `Derive.@interface NamedDimsArrayInterface() r[I]` instead.
+  return nameddims_view(a, I1, Irest...)
+end
+
+# Fix ambiguity error.
+function Base.getindex(
+  a::AbstractNamedDimsArray, I1::BlockRange{1}, Irest::BlockRange{1}...
+)
+  return ArrayLayouts.layout_getindex(a, I1, Irest...)
+end
+
+end

src/NamedDimsArrays.jl

@@ -4,6 +4,8 @@ include("isnamed.jl")
 include("randname.jl")
 include("abstractnamedinteger.jl")
 include("namedinteger.jl")
+include("abstractnamedarray.jl")
+include("namedarray.jl")
 include("abstractnamedunitrange.jl")
 include("namedunitrange.jl")
 include("abstractnameddimsarray.jl")

src/abstractnamedarray.jl

@@ -0,0 +1,79 @@
+using TypeParameterAccessors: unspecify_type_parameters
+
+abstract type AbstractNamedArray{T,N,Value<:AbstractArray,Name} <: AbstractArray{T,N} end
+
+const AbstractNamedVector{T,Value<:AbstractVector,Name} = AbstractNamedArray{T,1,Value,Name}
+const AbstractNamedMatrix{T,Value<:AbstractVector,Name} = AbstractNamedArray{T,2,Value,Name}
+
+# Minimal interface.
+dename(a::AbstractNamedArray) = throw(MethodError(dename, Tuple{typeof(a)}))
+name(a::AbstractNamedArray) = throw(MethodError(name, Tuple{typeof(a)}))
+
+# This can be customized to output different named integer types,
+# such as `namedarray(a::AbstractArray, name::IndexName) = Index(a, name)`.
+namedarray(a::AbstractArray, name) = NamedArray(a, name)
+
+# Shorthand.
+named(a::AbstractArray, name) = namedarray(a, name)
+
+# Derived interface.
+# TODO: Use `Accessors.@set`?
+setname(a::AbstractNamedArray, name) = namedarray(dename(a), name)
+
+# TODO: Use `TypeParameterAccessors`.
+denametype(::Type{<:AbstractNamedArray{<:Any,<:Any,Value}}) where {Value} = Value
+nametype(::Type{<:AbstractNamedArray{<:Any,<:Any,<:Any,Name}}) where {Name} = Name
+
+# Traits.
+isnamed(::Type{<:AbstractNamedArray}) = true
+
+# TODO: Should they also have the same base type?
+function Base.:(==)(a1::AbstractNamedArray, a2::AbstractNamedArray)
+  return name(a1) == name(a2) && dename(a1) == dename(a2)
+end
+function Base.hash(a::AbstractNamedArray, h::UInt)
+  h = hash(Symbol(unspecify_type_parameters(typeof(a))), h)
+  # TODO: Double check how this is handling blocking/sector information.
+  h = hash(dename(a), h)
+  return hash(name(a), h)
+end
+
+named_getindex(a::AbstractArray, I...) = named(getindex(dename(a), I...), name(a))
+
+# Array funcionality.
+Base.size(a::AbstractNamedArray) = map(s -> named(s, name(a)), size(dename(a)))
+Base.axes(a::AbstractNamedArray) = map(s -> named(s, name(a)), axes(dename(a)))
+Base.eachindex(a::AbstractNamedArray) = eachindex(dename(a))
+function Base.getindex(a::AbstractNamedArray{<:Any,N}, I::Vararg{Int,N}) where {N}
+  return named_getindex(a, I...)
+end
+function Base.getindex(a::AbstractNamedArray, I::Int)
+  return named_getindex(a, I)
+end
+Base.isempty(a::AbstractNamedArray) = isempty(dename(a))
+
+## function Base.AbstractArray{Int}(a::AbstractNamedArray)
+##   return AbstractArray{Int}(dename(a))
+## end
+## 
+## Base.iterate(a::AbstractNamedArray) = isempty(a) ? nothing : (first(a), first(a))
+## function Base.iterate(a::AbstractNamedArray, i)
+##   i == last(a) && return nothing
+##   next = named(dename(i) + dename(step(a)), name(a))
+##   return (next, next)
+## end
+
+function randname(ang::AbstractRNG, a::AbstractNamedArray)
+  return named(dename(a), randname(name(a)))
+end
+
+function Base.show(io::IO, a::AbstractNamedArray)
+  print(io, "named(", dename(a), ", ", repr(name(a)), ")")
+  return nothing
+end
+function Base.show(io::IO, mime::MIME"text/plain", a::AbstractNamedArray)
+  print(io, "named(\n")
+  show(io, mime, dename(a))
+  print(io, ",\n ", repr(name(a)), ")")
+  return nothing
+end