[ITensors] Better constructors from arrays (#1696)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensors"
 uuid = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>", "Miles Stoudenmire <mstoudenmire@flatironinstitute.org>"]
-version = "0.9.17"
+version = "0.9.18"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/itensor.jl

@@ -306,7 +306,7 @@ function emptyITensor((::Type{ElT}) = EmptyNumber) where {ElT <: Number}
     return itensor(EmptyTensor(ElT, ()))
 end
 
-using NDTensors.TypeParameterAccessors: set_eltype, type_parameters, specify_type_parameters
+using NDTensors: AllowAlias, NeverAlias
 """
     ITensor([ElT::Type, ]A::AbstractArray, inds)
     ITensor([ElT::Type, ]A::AbstractArray, inds::Index...)
@@ -350,17 +350,35 @@ T[i => 1, j => 1] == 3.3
 """
 function ITensor(
         as::AliasStyle,
-        eltype::Type{<:Number},
+        elt::Type{<:Number},
         A::AbstractArray{<:Number},
         inds::Indices;
         kwargs...,
     )
-    length(A) ≠ dim(inds) && throw(
-        DimensionMismatch(
-            "In ITensor(::AbstractArray, inds), length of AbstractArray ($(length(A))) must match total dimension of IndexSet ($(dim(inds)))",
-        ),
-    )
-    data = set_eltype(typeof(A), eltype)(as, A)
+    check_dims(A, inds)
+    # Other cases already handle when `elt ≡ eltype(A)`
+    # so here we only need to handle when they differ
+    # (which always involves a copy).
+    # This is equivalent to `copy!(similar(A, elt), A)`.
+    data = AbstractArray{elt}(A)
+    return itensor(Dense(data), inds)
+end
+
+function ITensor(
+        as::AliasStyle,
+        eltype::Type{T},
+        A::AbstractArray{T},
+        inds::Indices;
+        kwargs...,
+    ) where {T <: Number}
+    check_dims(A, inds)
+    data = if as ≡ AllowAlias()
+        A
+    elseif as ≡ NeverAlias()
+        copy(A)
+    else
+        error("Unknown AliasStyle: $as")
+    end
     return itensor(Dense(data), inds)
 end
 
@@ -381,6 +399,15 @@ function ITensor(
     return ITensor(as, eltype, Matrix(A), inds; kwargs...)
 end
 
+function check_dims(A::AbstractArray, inds::Indices)
+    length(A) ≠ dim(inds) && throw(
+        DimensionMismatch(
+            "In ITensor(::AbstractArray, inds), length of AbstractArray ($(length(A))) must match total dimension of IndexSet ($(dim(inds)))",
+        ),
+    )
+    return nothing
+end
+
 function ITensor(
         as::AliasStyle, eltype::Type{<:Number}, A::AbstractArray{<:Number}, is...; kwargs...
     )
@@ -470,13 +497,36 @@ The version `diagitensor` might output an ITensor whose storage data
 is an alias of the input vector data in order to minimize operations.
 """
 function diag_itensor(
-        as::AliasStyle, eltype::Type{<:Number}, v::AbstractVector{<:Number}, is::Indices
+        as::AliasStyle, elt::Type{<:Number}, v::AbstractVector{<:Number}, is::Indices
     )
+    check_diag_dims(v, is)
+    # Other cases already handle when `elt ≡ eltype(v)`
+    # so here we only need to handle when they differ
+    # (which always involves a copy).
+    # This is equivalent to `copy!(similar(v, elt), v)`.
+    data = AbstractVector{elt}(v)
+    return itensor(Diag(data), is)
+end
+
+function diag_itensor(
+        as::AliasStyle, elt::Type{T}, v::AbstractVector{T}, is::Indices
+    ) where {T <: Number}
+    check_diag_dims(v, is)
+    data = if as ≡ AllowAlias()
+        v
+    elseif as ≡ NeverAlias()
+        copy(v)
+    else
+        error("Unknown AliasStyle: $as")
+    end
+    return itensor(Diag(data), is)
+end
+
+function check_diag_dims(v::AbstractVector, is::Indices)
     length(v) ≠ mindim(is) && error(
         "Length of vector for diagonal must equal minimum of the dimension of the input indices",
     )
-    data = set_eltype(typeof(v), eltype)(as, v)
-    return itensor(Diag(data), is)
+    return nothing
 end
 
 function diag_itensor(
@@ -671,16 +721,6 @@ zero(T::ITensor)::ITensor = itensor(zero(tensor(T)))
 # Construct from Array
 #
 
-# Helper functions for different view behaviors
-# TODO: Move to NDTensors.jl
-function (arraytype::Type{<:AbstractArray})(::NeverAlias, A::AbstractArray)
-    return specify_type_parameters(arraytype, type_parameters(A))(A)
-end
-
-function (arraytype::Type{<:AbstractArray})(::AllowAlias, A::AbstractArray)
-    return convert(specify_type_parameters(arraytype, type_parameters(A)), A)
-end
-
 """
     Array{ElT, N}(T::ITensor, i:Index...)
     Array{ElT}(T::ITensor, i:Index...)

test/base/Project.toml

@@ -1,8 +1,10 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"

test/base/test_diagitensor.jl

@@ -106,6 +106,15 @@ using Test
             end
         end
 
+        @testset "diag_itensor constructor from Range" begin
+            t = diag_itensor(1:d, i, j)
+            @test ITensors.data(t) == 1:d
+            @test ITensors.data(t) isa Vector{Float64}
+
+            t = diag_itensor(Int, 1:d, i, j)
+            @test ITensors.data(t) ≡ 1:d
+        end
+
         @testset "reductions (sum, prod)" for elt in (
                 Float32, Float64, Complex{Float32}, Complex{Float64},
             )

test/base/test_itensor.jl

@@ -995,6 +995,21 @@ end
         @test storage(T) isa NDTensors.Dense{Float64}
     end
 
+    @testset "Construct from Range" begin
+        # https://github.com/ITensor/ITensors.jl/issues/1691
+
+        # Automatic conversion to float.
+        i, j = Index.((2, 2))
+        t = ITensor(1:4, (i, j))
+        @test ITensors.data(t) == 1:4
+        @test ITensors.data(t) isa Vector{Float64}
+
+        # Preserve integer element type.
+        i, j = Index.((2, 2))
+        t = ITensor(Int, 1:4, (i, j))
+        @test ITensors.data(t) ≡ 1:4
+    end
+
     @testset "ITensor Array constructor view behavior" begin
         d = 2
         i = Index(d)