rework tensor contructors

Author: lkdvos

src/tensors/tensor.jl

@@ -71,10 +71,22 @@ dim(t::TensorMap) = length(t.data)
 
 # General TensorMap constructors
 #--------------------------------
-# undef constructors
+# hook for mapping input types to storage types -- to be implemented in extensions
+_tensormap_storagetype(::Type{A}) where {A <: AbstractArray} = _tensormap_storagetype(eltype(A))
+_tensormap_storagetype(::Type{T}) where {T <: Number} = Vector{T}
+
+# utility type alias that makes constructors also work for type aliases that specify
+# different storage types. (i.e. CuTensorMap = _TensorMap{T, CuVector{T}, ...})
+# TODO: do we need a name for this and do we want to consider this as public?
+const _TensorMap{T, A <: DenseVector{T}, S, N₁, N₂} = TensorMap{T, S, N₁, N₂, A}
+const _Tensor{T, A <: DenseVector{T}, S, N} = Tensor{T, S, N, A}
+
+# undef constructors:
+# - dispatch start with TensorMap{T}
+# - select A and map to _TensorMap{T, A}
+# - select S, N1, N2 and map to TensorMap{T,S,N1,N2,A}
 """
-    TensorMap{T}(undef, codomain::ProductSpace{S,N₁}, domain::ProductSpace{S,N₂})
-                where {T,S,N₁,N₂}
+    TensorMap{T}(undef, codomain::ProductSpace{S,N₁}, domain::ProductSpace{S,N₂}) where {T,S,N₁,N₂}
     TensorMap{T}(undef, codomain ← domain)
     TensorMap{T}(undef, domain → codomain)
     # expert mode: select storage type `A`
@@ -83,33 +95,35 @@ dim(t::TensorMap) = length(t.data)
 
 Construct a `TensorMap` with uninitialized data.
 """
-function TensorMap{T}(::UndefInitializer, V::TensorMapSpace{S, N₁, N₂}) where {T, S, N₁, N₂}
-    return TensorMap{T, S, N₁, N₂, Vector{T}}(undef, V)
-end
-function TensorMap{T}(
-        ::UndefInitializer, codomain::TensorSpace{S}, domain::TensorSpace{S}
-    ) where {T, S}
-    return TensorMap{T}(undef, codomain ← domain)
-end
-function Tensor{T}(::UndefInitializer, V::TensorSpace{S}) where {T, S}
-    return TensorMap{T}(undef, V ← one(V))
-end
+TensorMap{T}(::UndefInitializer, V::TensorMapSpace) where {T} =
+    _TensorMap{T, _tensormap_storagetype(T)}(undef, V)
+TensorMap{T}(::UndefInitializer, codomain::TensorSpace, domain::TensorSpace) where {T} =
+    TensorMap{T}(undef, codomain ← domain)
+Tensor{T}(::UndefInitializer, V::TensorSpace) where {T} = TensorMap{T}(undef, V ← one(V))
+
+# specifying storagetype, fill in other parameters
+_TensorMap{T, A}(::UndefInitializer, V::TensorMapSpace) where {T, A} =
+    TensorMap{T, spacetype(V), numout(V), numin(V), A}(undef, V)
+_TensorMap{T, A}(::UndefInitializer, codomain::TensorSpace, domain::TensorSpace) where {T, A} =
+    _TensorMap{T, A}(undef, codomain ← domain)
+_Tensor{T, A}(::UndefInitializer, V::TensorSpace) where {T, A} = _TensorMap{T, A}(undef, V ← one(V))
 
 # constructor starting from vector = independent data (N₁ + N₂ = 1 is special cased below)
 # documentation is captured by the case where `data` is a general array
-# here, we force the `T` argument to distinguish it from the more general constructor below
-function TensorMap{T}(
-        data::A, V::TensorMapSpace{S, N₁, N₂}
-    ) where {T, S, N₁, N₂, A <: DenseVector{T}}
-    return TensorMap{T, S, N₁, N₂, A}(data, V)
-end
-function TensorMap{T}(
-        data::DenseVector{T}, codomain::TensorSpace{S}, domain::TensorSpace{S}
-    ) where {T, S}
-    return TensorMap(data, codomain ← domain)
-end
+# here, we force the `T` and/or `A` argument to distinguish it from the more general constructor below
+TensorMap{T}(data::DenseVector{T}, V::TensorMapSpace) where {T} =
+    _TensorMap{T, typeof(data)}(data, V)
+TensorMap{T}(data::DenseVector{T}, codomain::TensorSpace, domain::TensorSpace) where {T} =
+    TensorMap{T}(data, codomain ← domain)
+
+_TensorMap{T, A}(data::DenseVector{T}, V::TensorMapSpace) where {T, A} =
+    TensorMap{T, spacetype(V), numout(V), numin(V), A}(data, V)
+_TensorMap{T, A}(data::DenseVector{T}, codomain::TensorSpace, domain::TensorSpace) where {T, A} =
+    _TensorMap{T, A}(data, codomain ← domain)
 
 # constructor starting from block data
+const _BlockData{I <: Sector, A <: AbstractMatrix} = AbstractDict{I, A}
+
 """
     TensorMap(data::AbstractDict{<:Sector,<:AbstractMatrix}, codomain::ProductSpace{S,N₁},
                 domain::ProductSpace{S,N₂}) where {S<:ElementarySpace,N₁,N₂}
@@ -129,29 +143,33 @@ Construct a `TensorMap` by explicitly specifying its block data.
 Alternatively, the domain and codomain can be specified by passing a [`HomSpace`](@ref)
 using the syntax `codomain ← domain` or `domain → codomain`.
 """
-function TensorMap(
-        data::AbstractDict{<:Sector, <:AbstractMatrix}, V::TensorMapSpace{S, N₁, N₂}
-    ) where {S, N₁, N₂}
-    T = eltype(valtype(data))
-    t = TensorMap{T}(undef, V)
+function TensorMap(data::_BlockData, V::TensorMapSpace)
+    A = _tensormap_storagetype(valtype(data))
+    return _TensorMap{scalartype(A), A}(data, V)
+end
+TensorMap(data::_BlockData, codom::TensorSpace, dom::TensorSpace) =
+    TensorMap(data, codom ← dom)
+
+function _TensorMap{T, A}(data::_BlockData, V::TensorMapSpace) where {T, A}
+    t = _TensorMap{T, A}(undef, V)
+
+    # check that there aren't too many blocks
+    for (c, b) in data
+        c ∈ blocksectors(t) || isempty(b) || throw(SectorMismatch("data for block sector $c not expected"))
+    end
+
+    # fill in the blocks -- rely on conversion in copy
     for (c, b) in blocks(t)
         haskey(data, c) || throw(SectorMismatch("no data for block sector $c"))
         datac = data[c]
-        size(datac) == size(b) ||
-            throw(DimensionMismatch("wrong size of block for sector $c"))
+        size(datac) == size(b) || throw(DimensionMismatch("wrong size of block for sector $c"))
         copy!(b, datac)
     end
-    for (c, b) in data
-        c ∈ blocksectors(t) || isempty(b) ||
-            throw(SectorMismatch("data for block sector $c not expected"))
-    end
+
     return t
 end
-function TensorMap(
-        data::AbstractDict{<:Sector, <:AbstractMatrix}, codom::TensorSpace{S}, dom::TensorSpace{S}
-    ) where {S}
-    return TensorMap(data, codom ← dom)
-end
+_TensorMap{T, A}(data::_BlockData, codom::TensorSpace, dom::TensorSpace) where {T, A} =
+    _TensorMap{T, A}(data, codom ← dom)
 
 @doc """
     zeros([T=Float64,], codomain::ProductSpace{S,N₁}, domain::ProductSpace{S,N₂}) where {S,N₁,N₂,T}
@@ -317,49 +335,45 @@ cases.
     to a plain array is possible, and only in the case where the `data` actually respects
     the specified symmetry structure, up to a tolerance `tol`.
 """
-function TensorMap(
-        data::AbstractArray, V::TensorMapSpace{S, N₁, N₂};
-        tol = sqrt(eps(real(float(eltype(data)))))
-    ) where {S <: IndexSpace, N₁, N₂}
+function TensorMap(data::AbstractArray, V::TensorMapSpace; tol = sqrt(eps(real(float(eltype(data))))))
     T = eltype(data)
-    if ndims(data) == 1 && length(data) == dim(V)
-        if data isa DenseVector # refer to specific data-capturing constructor
-            return TensorMap{T}(data, V)
-        else
-            return TensorMap{T}(collect(data), V)
-        end
-    end
+    A = _tensormap_storagetype(typeof(data))
+    return _TensorMap{T, A}(data, V; tol)
+end
+function _TensorMap{T, A}(
+        data::AbstractArray, V::TensorMapSpace; tol = sqrt(eps(real(float(eltype(data)))))
+    ) where {T, A}
+    # refer to specific data-capturing constructors if input is a vector of the correct length
+    ndims(data) == 1 && length(data) == dim(V) && return _TensorMap{T, A}(data, V)
 
     # dimension check
-    codom = codomain(V)
-    dom = domain(V)
+    codom, dom = codomain(V), domain(V)
     arraysize = dims(V)
     matsize = (dim(codom), dim(dom))
+    (size(data) == arraysize || size(data) == matsize) || throw(DimensionMismatch())
 
-    if !(size(data) == arraysize || size(data) == matsize)
-        throw(DimensionMismatch())
-    end
-
-    if sectortype(S) === Trivial # refer to same method, but now with vector argument
-        return TensorMap(reshape(data, length(data)), V)
+    if sectortype(V) === Trivial # refer to same method, but now with vector argument
+        return _TensorMap{T, A}(reshape(data, length(data)), V)
     end
 
-    t = TensorMap{T}(undef, codom, dom)
+    t = _TensorMap{T, A}(undef, V)
     arraydata = reshape(collect(data), arraysize)
     t = project_symmetric!(t, arraydata)
     if !isapprox(arraydata, convert(Array, t); atol = tol)
         throw(ArgumentError("Data has non-zero elements at incompatible positions"))
     end
     return t
 end
-function TensorMap(
-        data::AbstractArray, codom::TensorSpace{S}, dom::TensorSpace{S}; kwargs...
-    ) where {S}
-    return TensorMap(data, codom ← dom; kwargs...)
-end
-function Tensor(data::AbstractArray, codom::TensorSpace, ; kwargs...)
-    return TensorMap(data, codom ← one(codom); kwargs...)
-end
+
+TensorMap(data::AbstractArray, codom::TensorSpace, dom::TensorSpace; kwargs...) =
+    TensorMap(data, codom ← dom; kwargs...)
+_TensorMap{T, A}(data::AbstractArray, codom::TensorSpace, dom::TensorSpace; kwargs...) where {T, A} =
+    _TensorMap(data, codom ← dom; kwargs...)
+
+Tensor(data::AbstractArray, codom::TensorSpace; kwargs...) =
+    TensorMap(data, codom ← one(codom); kwargs...)
+_Tensor{T, A}(data::AbstractArray, codom::TensorSpace; kwargs...) where {T, A} =
+    _TensorMap{T, A}(data, codom ← one(codom); kwargs...)
 
 """
     project_symmetric!(t::TensorMap, data::DenseArray) -> TensorMap