-
Notifications
You must be signed in to change notification settings - Fork 28
/
grassmann.jl
161 lines (129 loc) · 4.58 KB
/
grassmann.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
# We separate some of the internals needed for implementing GradientGrassmann into a
# submodule, to keep the MPSKit module namespace cleaner.
"""
A module for functions related to treating an InfiniteMPS in left-canonical form as a bunch
of points on Grassmann manifolds, and performing things like retractions and transports
on these Grassmann manifolds.
The module exports nothing, and all references to it should be qualified, e.g.
`GrassmannMPS.fg`.
"""
module GrassmannMPS
using ..MPSKit
using TensorKit
import TensorKitManifolds.Grassmann
"""
Compute the expectation value, and its gradient with respect to the tensors in the unit
cell as tangent vectors on Grassmann manifolds.
"""
function fg(x::Tuple{T,<:Cache}) where T <: Union{<:InfiniteMPS,FiniteMPS}
(state, envs) = x
# The partial derivative with respect to AL, al_d, is the partial derivative with
# respect to AC times CR'.
al_d = Vector{MPSKit.site_type(T)}(undef,length(state));
for i in 1:length(state)
h_eff = MPSKit.∂∂AC(i,state,envs.opp,envs);
al_d[i] = (h_eff*state.AC[i]*state.CR[i]')
end
g = Grassmann.project.(al_d,state.AL)
f = real(sum(expectation_value(state, envs)))
return f, g
end
function fg(x::Tuple{<:MPSMultiline,<:Cache})
(state, envs) = x
# The partial derivative with respect to AL, al_d, is the partial derivative with
# respect to AC times CR'.
ac_d = [MPSKit.∂∂AC(v,state,envs.opp,envs)*state.AC[v] for v in CartesianIndices(state.AC)]
al_d = [d*c' for (d, c) in zip(ac_d, state.CR)]
g = [Grassmann.project(d, a) for (d, a) in zip(al_d, state.AL)]
f = expectation_value(state, envs)
fi = imag.(f); fr = real.(f);
sum(fi) > MPSKit.Defaults.tol && @warn "mpo is not hermitian $(fi)"
g = -2*g./abs.(fr);
return -log(sum(fr)^2), g[:]
end
"""
Retract a left-canonical MPSMultiline along Grassmann tangent `g` by distance `alpha`.
"""
function retract(x::Tuple{<:MPSMultiline,<:Cache}, tg, alpha)
# the idea is to re-use imps transportation
(state, envs) = x
g = reshape(tg,size(state));
transported = map(1:length(state)) do row
(tx,th) = retract((state[row],envs),g[row,:],alpha);
tx[1],th
end
nstate = MPSKit.Multiline(first.(transported));
h = reduce(hcat,last.(transported))[:]
return (nstate,envs), h
end
"""
Retract a left-canonical infinite MPS along Grassmann tangent `g` by distance `alpha`.
"""
function retract(x::Tuple{<:InfiniteMPS,<:Cache}, g, alpha)
(state, envs) = x
nal = similar(state.AL);
h = similar(g) # The tangent at the end-point
for i in 1:length(g)
(nal[i], h[i]) = Grassmann.retract(state.AL[i], g[i], alpha)
end
nstate = InfiniteMPS(nal,state.CR[end]);
return (nstate,envs), h
end
"""
Retract a left-canonical finite MPS along Grassmann tangent `g` by distance `alpha`.
"""
function retract(x::Tuple{<:FiniteMPS,<:Cache}, g, alpha)
(state, envs) = x
y = copy(state) # The end-point
h = similar(g) # The tangent at the end-point
for i in 1:length(g)
(yal, h[i]) = Grassmann.retract(state.AL[i], g[i], alpha)
y.AC[i] = (yal,state.CR[i])
end
normalize!(y)
return (y,envs), h
end
"""
Transport a tangent vector `h` along the retraction from `x` in direction `g` by distance
`alpha`. `xp` is the end-point of the retraction.
"""
function transport!(h, x, g, alpha, xp)
(state, envs) = x
for i in 1:length(h)
h[i] = Grassmann.transport!(h[i], state.AL[i], g[i], alpha, xp[1].AL[i])
end
return h
end
"""
Euclidean inner product between two Grassmann tangents of an infinite MPS.
"""
function inner(x, g1, g2)
(state, envs) = x
tot = sum(Grassmann.inner(a, d1, d2) for (a, d1, d2) in zip(state.AL, g1, g2))
return 2*real(tot)
end
"""
Scale a tangent vector by scalar `alpha`.
"""
scale!(g, alpha) = g .* alpha
"""
Add two tangents vectors, scaling the latter by `alpha`.
"""
add!(g1, g2, alpha) = g1 + g2 .* alpha
"""
Precondition a given Grassmann tangent `g` at state `x` by the Hilbert space inner product.
This requires inverting the right MPS transfer matrix. This is done using `reginv`, with a
regularisation parameter that is the norm of the tangent `g`.
"""
function precondition(x, g)
(state, envs) = x
#hacky workaround - what is eltype(state)?
delta = min(real(one(eltype(state.AL[1]))), sqrt(inner(x, g, g)))
crinvs = MPSKit.reginv.(state.CR[1:end],delta)
g_prec = similar(g);
for i in 1:length(state)
g_prec[i] = Grassmann.project(g[i][]*crinvs[i]'*crinvs[i],state.AL[i])
end
return g_prec
end
end # module GrassmannMPS