Start adding support for automatic allocation (#2)

Author: mtfishman

Project.toml

@@ -1,13 +1,23 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.0"
+version = "0.1.1"
 
 [deps]
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 BroadcastMapConversion = "4a4adec5-520f-4750-bb37-d5e66b4ddeb2"
+Derive = "a07dfc7f-7d04-4eb5-84cc-a97f051f655a"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
+UnallocatedArrays = "43c9e47c-e622-40fb-bf18-a09fc8c466b6"
+UnspecifiedTypes = "42b3faec-625b-4613-8ddc-352bf9672b8d"
 
 [compat]
+Accessors = "0.1.39"
 BroadcastMapConversion = "0.1.2"
+Derive = "0.3.6"
+FillArrays = "1.13.0"
 NamedDimsArrays = "0.3.0"
+UnallocatedArrays = "0.1.1"
+UnspecifiedTypes = "0.1.1"
 julia = "1.10"

README.md

@@ -44,6 +44,21 @@ q, r = qr(a, (i,))
 @test q * r ≈ a
 ````
 
+Automatic allocation
+
+````julia
+a = ITensor(i, j)
+````
+
+Broken, need to fix:
+a[j[1], i[2]] = 1 + 2im
+
+````julia
+a[2, 1] = 1 + 2im
+eltype(a) == Complex{Int}
+@test a[i[2], j[1]] == 1 + 2im
+````
+
 ---
 
 *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*

examples/README.jl

@@ -39,3 +39,11 @@ d = a + a′
 @test a ≈ aligndims(a, (j, i))
 q, r = qr(a, (i,))
 @test q * r ≈ a
+
+# Automatic allocation
+a = ITensor(i, j)
+# Broken, need to fix:
+# a[j[1], i[2]] = 1 + 2im
+a[2, 1] = 1 + 2im
+eltype(a) == Complex{Int}
+@test a[i[2], j[1]] == 1 + 2im

src/ITensorBase.jl

@@ -8,10 +8,12 @@ using NamedDimsArrays:
   AbstractNamedInteger,
   AbstractNamedUnitRange,
   AbstractNamedVector,
+  NamedDimsArray,
   dename,
   dimnames,
   name,
   named,
+  nameddimsindices,
   unname
 
 @kwdef struct IndexName <: AbstractName
@@ -69,11 +71,65 @@ NamedDimsArrays.nameddimsarraytype(::Type{<:IndexName}) = ITensor
 
 Base.ndims(::Type{<:AbstractITensor}) = Any
 
-struct ITensor <: AbstractITensor
+using FillArrays: Zeros
+using UnallocatedArrays: UnallocatedZeros, allocate
+using UnspecifiedTypes: UnspecifiedZero
+
+# TODO: Make this more general, maybe with traits `is_unallocated`
+# and `is_eltype_unspecified`.
+function specify_eltype(a::Zeros{UnspecifiedZero}, elt::Type)
+  return Zeros{elt}(axes(a))
+end
+
+# TODO: Use `adapt` to reach down into the storage.
+function specify_eltype!(a::AbstractITensor, elt::Type)
+  setdenamed!(a, specify_eltype(dename(a), elt))
+  return a
+end
+
+# Assume it is allocated.
+allocate!(a::AbstractArray) = a
+
+# TODO: Use `adapt` to reach down into the storage.
+function allocate!(a::AbstractITensor)
+  setdenamed!(a, allocate(dename(a)))
+  return a
+end
+
+using Derive: @derive, @interface, AbstractArrayInterface
+
+abstract type AbstractAllocatableArrayInterface <: AbstractArrayInterface end
+struct AllocatableArrayInterface <: AbstractAllocatableArrayInterface end
+
+unallocatable(a::AbstractITensor) = NamedDimsArray(a)
+
+@interface ::AbstractAllocatableArrayInterface function Base.setindex!(
+  a::AbstractArray, value, I::Int...
+)
+  allocate!(specify_eltype!(a, typeof(value)))
+  # TODO: Maybe use `@interface interface(a) a[I...] = value`?
+  unallocatable(a)[I...] = value
+  return a
+end
+
+@derive AllocatableArrayInterface() (T=AbstractITensor,) begin
+  Base.setindex!(::T, ::Any, ::Int...)
+end
+
+mutable struct ITensor <: AbstractITensor
   parent::AbstractArray
   nameddimsindices
 end
 Base.parent(a::ITensor) = a.parent
 NamedDimsArrays.nameddimsindices(a::ITensor) = a.nameddimsindices
 
+using Accessors: @set
+setdenamed(a::ITensor, denamed) = (@set a.parent = denamed)
+setdenamed!(a::ITensor, denamed) = (a.parent = denamed)
+
+function ITensor(I1::Index, I_rest::Index...)
+  I = (I1, I_rest...)
+  return ITensor(Zeros{UnspecifiedZero}(length.(dename.(I))...), I)
+end
+
 end

test/test_aqua.jl

@@ -3,5 +3,5 @@ using Aqua: Aqua
 using Test: @testset
 
 @testset "Code quality (Aqua.jl)" begin
-  Aqua.test_all(ITensorBase)
+  Aqua.test_all(ITensorBase; ambiguities=false, persistent_tasks=false)
 end