Better definition of equality, hash (#124)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "NamedDimsArrays"
 uuid = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.8.3"
+version = "0.8.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -19,16 +19,19 @@ TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [weakdeps]
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208"
 
 [extensions]
+NamedDimsArraysAbstractTreesExt = "AbstractTrees"
 NamedDimsArraysBlockArraysExt = "BlockArrays"
 NamedDimsArraysGradedArraysExt = "GradedArrays"
 NamedDimsArraysSparseArraysBaseExt = "SparseArraysBase"
 
 [compat]
+AbstractTrees = "0.4.5"
 Adapt = "4.1.1"
 ArrayLayouts = "1.11"
 BlockArrays = "1.3"

ext/NamedDimsArraysAbstractTreesExt/NamedDimsArraysAbstractTreesExt.jl

@@ -0,0 +1,12 @@
+module NamedDimsArraysAbstractTreesExt
+
+using AbstractTrees: AbstractTrees
+using NamedDimsArrays: AbstractNamedDimsArray, dimnames
+
+# Only print the dimension names when printing with `AbstractTrees.print_tree`.
+function AbstractTrees.printnode(io::IO, a::AbstractNamedDimsArray)
+    show(IOContext(io, :compact => true, :limit => true), dimnames(a))
+    return nothing
+end
+
+end

src/abstractnameddimsarray.jl

@@ -400,19 +400,31 @@ end
 # Base version ignores dimension names.
 # TODO: Use `mapreduce(isequal, &&, a1, a2)`?
 function Base.isequal(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray)
-    return all(eachindex(a1, a2)) do I
-        isequal(a1[I], a2[I])
-    end
+    issetequal(inds(a1), inds(a2)) || return false
+    return isequal(unname(a1), unnamed(a2, inds(a1)))
 end
 
 # Base version ignores dimension names.
 # TODO: Use `mapreduce(==, &&, a1, a2)`?
 # TODO: Handle `missing` values properly.
 function Base.:(==)(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray)
     issetequal(inds(a1), inds(a2)) || return false
-    return all(eachindex(a1, a2)) do I
-        a1[I] == a2[I]
+    return unname(a1) == unnamed(a2, inds(a1))
+end
+
+# Generalization of `Base.sort` to Tuples for Julia v1.10 compatibility.
+# TODO: Remove when we drop support for Julia v1.10.
+_sort(x; kwargs...) = sort(x; kwargs...)
+_sort(x::NTuple{N}; kwargs...) where {N} = NTuple{N}(sort(collect(x); kwargs...))
+
+function Base.hash(a::AbstractNamedDimsArray, h::UInt64)
+    h = hash(:NamedDimsArray, h)
+    a′ = aligneddims(a, _sort(dimnames(a)))
+    h = hash(dename(a′), h)
+    for i in inds(a′)
+        h = hash(i, h)
     end
+    return h
 end
 
 # Indexing.
@@ -674,6 +686,7 @@ function aligndims(a::AbstractArray, dims)
     return constructorof(typeof(a))(permutedims(dename(a), perm), new_inds)
 end
 
+using DerivableInterfaces: permuteddims
 function aligneddims(a::AbstractArray, dims)
     new_inds = to_inds(a, dims)
     perm = getperm(inds(a), new_inds)
@@ -683,7 +696,7 @@ function aligneddims(a::AbstractArray, dims)
         ),
     )
     return constructorof_nameddims(typeof(a))(
-        PermutedDimsArray(dename(a), perm), new_inds
+        permuteddims(dename(a), perm), new_inds
     )
 end
 

src/naiveorderedset.jl

@@ -24,6 +24,10 @@ Base.Broadcast.BroadcastStyle(s1::AbstractArrayStyle, s2::Style{NaiveOrderedSet}
 Base.Broadcast.broadcastable(s::NaiveOrderedSet) = s
 Base.to_shape(s::NaiveOrderedSet) = s
 
+# Needed for functionality such as `CartesianIndices(::AbstractNamedDimsArray)`,
+# `pairs(::AbstractNamedDimsArray)`, etc.
+Base.CartesianIndices(s::NaiveOrderedSet) = CartesianIndices(values(s))
+
 function Base.copy(
         bc::Broadcasted{Style{NaiveOrderedSet}, <:Any, <:Any, <:Tuple{<:NaiveOrderedSet}}
     )

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
@@ -18,6 +19,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [compat]
+AbstractTrees = "0.4.5"
 Adapt = "4"
 Aqua = "0.8.9"
 BlockArrays = "1"

test/test_abstracttreesext.jl

@@ -0,0 +1,9 @@
+using AbstractTrees: printnode
+using NamedDimsArrays: nameddims
+using Test: @test, @testset
+
+@testset "AbstractTreesExt" begin
+    a = randn(3, 4)
+    na = nameddims(a, ("i", "j"))
+    @test sprint(printnode, na) == """("i", "j")"""
+end

test/test_basics.jl

@@ -1,39 +1,9 @@
 using Combinatorics: Combinatorics
-using NamedDimsArrays:
-    NamedDimsArrays,
-    AbstractNamedDimsArray,
-    AbstractNamedDimsMatrix,
-    NaiveOrderedSet,
-    Name,
-    NameMismatch,
-    NamedDimsCartesianIndex,
-    NamedDimsCartesianIndices,
-    NamedDimsArray,
-    NamedDimsMatrix,
-    aligndims,
-    aligneddims,
-    apply,
-    dename,
-    denamed,
-    dim,
-    dimnames,
-    dims,
-    fusednames,
-    isnamed,
-    mapinds,
-    name,
-    named,
-    nameddims,
-    inds,
-    namedoneto,
-    operator,
-    product,
-    replaceinds,
-    setinds,
-    state,
-    unname,
-    unnamed,
-    @names
+using NamedDimsArrays: NamedDimsArrays, AbstractNamedDimsArray, AbstractNamedDimsMatrix,
+    NaiveOrderedSet, Name, NameMismatch, NamedDimsCartesianIndex, NamedDimsCartesianIndices,
+    NamedDimsArray, NamedDimsMatrix, aligndims, aligneddims, apply, dename, denamed, dim,
+    dimnames, dims, fusednames, isnamed, mapinds, name, named, nameddims, inds, namedoneto,
+    operator, product, replaceinds, setinds, state, unname, unnamed, @names
 using Test: @test, @test_throws, @testset
 using VectorInterface: scalartype
 
@@ -74,6 +44,26 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
         @test dims(na, ("j", "i")) == (2, 1)
         @test na[1, 1] == a[1, 1]
 
+        # equals (==)/isequal
+        a = randn(elt, 3, 4)
+        na = nameddims(a, ("i", "j"))
+        @test na == na
+        @test na == aligndims(na, ("j", "i"))
+        @test isequal(na, na)
+        @test isequal(na, aligndims(na, ("j", "i")))
+        @test hash(na) == hash(aligndims(na, ("j", "i")))
+        # Regression test that NamedDimsArrays
+        # with different names are not equal (as opposed to
+        # erroring).
+        @test na ≠ nameddims(a, ("j", "k"))
+        @test !isequal(na, nameddims(a, ("j", "k")))
+        @test hash(na) ≠ hash(nameddims(a, ("j", "k")))
+
+        a = randn(elt, 2, 2)
+        na = nameddims(a, ("i", "j"))
+        @test CartesianIndices(na) == CartesianIndices(a)
+        @test collect(pairs(na)) == (CartesianIndices(a) .=> a)
+
         @test_throws ErrorException NamedDimsArray(randn(4), namedoneto.((2, 2), ("i", "j")))
         @test_throws ErrorException NamedDimsArray(randn(2, 2), namedoneto.((2, 3), ("i", "j")))