Move FiniteDifferences support to package extension (#132)

* Move FiniteDifferences support to package extension

* Refactor implementation

Project.toml

@@ -16,9 +16,11 @@ WignerSymbols = "9f57e263-0b3d-5e2e-b1be-24f2bb48858b"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 
 [extensions]
 TensorKitChainRulesCoreExt = "ChainRulesCore"
+TensorKitFiniteDifferencesExt = "FiniteDifferences"
 
 [compat]
 Aqua = "0.6, 0.7, 0.8"
@@ -27,8 +29,8 @@ ChainRulesTestUtils = "1"
 Combinatorics = "1"
 FiniteDifferences = "0.12"
 HalfIntegers = "1"
-LinearAlgebra = "1"
 LRUCache = "1.0.2"
+LinearAlgebra = "1"
 PackageExtensionCompat = "1"
 Random = "1"
 Strided = "2"

ext/TensorKitFiniteDifferencesExt.jl

@@ -0,0 +1,30 @@
+module TensorKitFiniteDifferencesExt
+
+using TensorKit
+using TensorKit: sqrtdim, isqrtdim
+using VectorInterface: scale!
+using FiniteDifferences
+
+function FiniteDifferences.to_vec(t::T) where {T<:TensorKit.TrivialTensorMap}
+    vec, from_vec = to_vec(t.data)
+    return vec, x -> T(from_vec(x), codomain(t), domain(t))
+end
+function FiniteDifferences.to_vec(t::AbstractTensorMap)
+    # convert to vector of vectors to make use of existing functionality
+    vec_of_vecs = [b * sqrtdim(c) for (c, b) in blocks(t)]
+    vec, back = FiniteDifferences.to_vec(vec_of_vecs)
+
+    function from_vec(x)
+        t′ = similar(t)
+        xvec_of_vecs = back(x)
+        for (i, (c, b)) in enumerate(blocks(t′))
+            scale!(b, xvec_of_vecs[i], isqrtdim(c))
+        end
+        return t′
+    end
+
+    return vec, from_vec
+end
+FiniteDifferences.to_vec(t::TensorKit.AdjointTensorMap) = to_vec(copy(t))
+
+end

test/ad.jl

@@ -1,7 +1,6 @@
 using ChainRulesCore
 using ChainRulesTestUtils
 using Random
-using FiniteDifferences
 using LinearAlgebra
 
 const _repartition = @static if isdefined(Base, :get_extension)
@@ -21,30 +20,6 @@ function ChainRulesTestUtils.test_approx(actual::AbstractTensorMap,
         ChainRulesTestUtils.@test_msg msg isapprox(b, block(expected, c); kwargs...)
     end
 end
-function FiniteDifferences.to_vec(t::T) where {T<:TensorKit.TrivialTensorMap}
-    vec, from_vec = to_vec(t.data)
-    return vec, x -> T(from_vec(x), codomain(t), domain(t))
-end
-function FiniteDifferences.to_vec(t::AbstractTensorMap)
-    vec = mapreduce(vcat, blocks(t); init=scalartype(t)[]) do (c, b)
-        return reshape(b, :) .* sqrt(dim(c))
-    end
-    vec_real = scalartype(t) <: Real ? vec : collect(reinterpret(real(scalartype(t)), vec))
-
-    function from_vec(x_real)
-        x = scalartype(t) <: Real ? x_real : reinterpret(scalartype(t), x_real)
-        t′ = similar(t)
-        ctr = 0
-        for (c, b) in blocks(t′)
-            n = length(b)
-            copyto!(b, reshape(view(x, ctr .+ (1:n)), size(b)) ./ sqrt(dim(c)))
-            ctr += n
-        end
-        return t′
-    end
-    return vec_real, from_vec
-end
-FiniteDifferences.to_vec(t::TensorKit.AdjointTensorMap) = to_vec(copy(t))
 
 function _randomize!(a::TensorMap)
     for b in values(blocks(a))