TT_weingarten

 Weingarten map computation for the fixed TT-rank manifold.

 NOTE: this manifold requires the use of a modified version of TTeMPS_1.1,
 which is packaced with Manopt and can be found in 
 manopt/manopt/manifolds/ttfixedrank/TTeMPS_1.1
 
 Y = TT_weingarten(V, Z, ind)
 
 This function implements the efficient way of computing the Weingarten
 map for the Riemannian submanifold of fixed TT-rank tensors embedded in
 its natural embedding Euclidean space, as described in the paper:

   Psenka and Boumal,
   Second-order optimization for tensors with fixed tensor-train rank,
   Optimization workshop at NeurIPS 2020.
 
 Y is the output tangent vector, V is a tangent vector (which contains
 needed information of base point X), Z is a tensor (embedding space).
 
 TT_weingarten supports the following formats for Z:
 1) Full
 2) TT-format (any rank)
 3) sparse, with non-zero indices indexed by ind (see fixedTTrankfactory).

 See also: fixedTTrankfactory

0001 function Y = TT_weingarten(V, Z, ind)
0002 % Weingarten map computation for the fixed TT-rank manifold.
0003 %
0004 % NOTE: this manifold requires the use of a modified version of TTeMPS_1.1,
0005 % which is packaced with Manopt and can be found in
0006 % manopt/manopt/manifolds/ttfixedrank/TTeMPS_1.1
0007 %
0008 % Y = TT_weingarten(V, Z, ind)
0009 %
0010 % This function implements the efficient way of computing the Weingarten
0011 % map for the Riemannian submanifold of fixed TT-rank tensors embedded in
0012 % its natural embedding Euclidean space, as described in the paper:
0013 %
0014 %   Psenka and Boumal,
0015 %   Second-order optimization for tensors with fixed tensor-train rank,
0016 %   Optimization workshop at NeurIPS 2020.
0017 %
0018 % Y is the output tangent vector, V is a tangent vector (which contains
0019 % needed information of base point X), Z is a tensor (embedding space).
0020 %
0021 % TT_weingarten supports the following formats for Z:
0022 % 1) Full
0023 % 2) TT-format (any rank)
0024 % 3) sparse, with non-zero indices indexed by ind (see fixedTTrankfactory).
0025 %
0026 % See also: fixedTTrankfactory
0027 
0028 % This file is part of Manopt: www.manopt.org.
0029 % Original author: Michael Psenka, Nov. 24, 2020.
0030 % Contributors: Nicolas Boumal
0031 % Change log:
0032 
0033     % Preliminary tangent vector set-up for Y
0034     d = V.order;
0035     r = V.rank;
0036     n = V.size;
0037     
0038 
0039     Y = V; % all properties except for dU cores inherited from V
0040     normV = norm(V);
0041     V = (1 / normV) * V;
0042     Y.dU = cell(1, d);
0043 
0044 
0045 
0046     %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0047     % Begin calculating variational components
0048     %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0049 
0050     % Pre-calculate all R-components as needed for tangent space conversion
0051     R = cell(1, d-1);
0052     % Intermediary used to calculate R-terms for re-orthogonalization
0053     Xr = V.U;
0054 
0055     for k = d:-1:2
0056 
0057         sz = size(Xr{k});
0058 
0059         if length(sz) == 2
0060             sz = [sz, 1]; %#ok<AGROW>
0061         end
0062 
0063         [Q, R{k - 1}] = qr_unique(unfold(Xr{k}, 'right')');
0064         Xr{k} = reshape(Q', [size(Q, 2), sz(2), sz(3)]);
0065 
0066         Xr{k - 1} = tensorprod_ttemps(Xr{k - 1}, R{k - 1}, 3);
0067     end
0068 
0069     % Calculate V.dU under the original parametrization of tangent space
0070 
0071     dUR = cell(1, d);
0072 
0073     for k = 1:(d - 1)
0074         dUR{k} = unfold(V.dU{k}, 'left') / R{k}';
0075         dUR{k} = reshape(dUR{k}, [r(k), n(k), r(k + 1)]);
0076     end
0077 
0078     dUR{d} = V.dU{d};
0079 
0080     % Calculate ~X^tV_{\ge k} efficiently
0081     XtVg = cell(1, d);
0082 
0083     XtVg{d} = conj(unfold(V.V{d}, 'right')) * unfold(V.dU{d}, 'right').';
0084     XtVgTmp = conj(unfold(V.V{d}, 'right')) * unfold(V.U{d}, 'right').';
0085 
0086     for k = (d - 1):-1:2
0087         tmp = tensorprod_ttemps(V.V{k}, XtVg{k + 1}', 3);
0088         XtVg{k} = conj(unfold(tmp, 'right')) * unfold(V.U{k}, 'right').';
0089 
0090         tmp2 = tensorprod_ttemps(V.V{k}, XtVgTmp', 3);
0091         XtVg{k} = XtVg{k} + conj(unfold(tmp2, 'right')) * unfold(dUR{k}, 'right').';
0092         XtVgTmp = conj(unfold(tmp2, 'right')) * unfold(V.U{k}, 'right').';
0093     end
0094 
0095     %%%%%%%%%%%%%%%%%%%%%%%%%%%%% CASES FOR TYPE OF EUCLIDEAN GRADIENT Z %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0096 
0097     if isa(Z, 'TTeMPS') || isa(Z, 'TTeMPS_tangent_orth')% Z either TT-tensor or tangent vec
0098 
0099         if isa(Z, 'TTeMPS_tangent_orth')
0100             Z = tangent_to_TTeMPS(Z); % efficient to treat both as TTeMPS case
0101         end
0102 
0103         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0104         %%%%%%%%%%%%%%%               DIAGONAL TERMS                 %%%%%%%%%%%%%%%%%%
0105         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0106 
0107         ZVr = cell(1, d); % stores values of V_<^T(Z<k>)
0108         Zr = cell(1, d);  % stores values of X_<^T(Z<k>)
0109         ZVl = cell(1, d); % stores values of (Z<k>)^T(V_>)
0110         Zl = cell(1, d);  % stores values of (Z<k>)^T(X_>)
0111 
0112 
0113         ZVr{d}  = conj(unfold(Z.U{d}, 'right')) * unfold(V.dU{d}, 'right').';
0114         XtVgTmp = conj(unfold(Z.U{d}, 'right')) * unfold(V.U{d},  'right').';
0115 
0116         for k = (d - 1):-1:2
0117             tmp = tensorprod_ttemps(Z.U{k}, ZVr{k + 1}', 3);
0118             ZVr{k} = conj(unfold(tmp, 'right')) * unfold(V.U{k}, 'right').';
0119 
0120             tmp2 = tensorprod_ttemps(Z.U{k}, XtVgTmp', 3);
0121             ZVr{k} = ZVr{k} + conj(unfold(tmp2, 'right')) * unfold(dUR{k}, 'right').';
0122             XtVgTmp = conj(unfold(tmp2, 'right')) * unfold(V.U{k}, 'right').';
0123         end
0124 
0125         for k=1:d
0126             ZVr{k} = ZVr{k}';
0127         end
0128    
0129 
0130         Zr{d} = conj(unfold(V.V{d}, 'right')) * unfold(Z.U{d}, 'right').';
0131 
0132         for k = (d - 1):-1:2
0133             tmp = tensorprod_ttemps(V.V{k}, Zr{k + 1}', 3);
0134             Zr{k} = conj(unfold(tmp, 'right')) * unfold(Z.U{k}, 'right').';
0135         end
0136 
0137         % Computation of ZVl and Zl
0138 
0139         ZVl{1} = unfold(dUR{1}, 'left')' * unfold(Z.U{1}, 'left');
0140         Zl{1}  = unfold(V.U{1}, 'left')' * unfold(Z.U{1}, 'left');
0141 
0142         for k = 2:(d - 1)
0143             % (V_k-1)(U_k)
0144             tmp = tensorprod_ttemps(V.U{k}, ZVl{k - 1}', 1);
0145             ZVl{k} = unfold(tmp, 'left')' * unfold(Z.U{k}, 'left');
0146             % + (X_k-1)(dU_k)
0147             tmp = tensorprod_ttemps(dUR{k}, Zl{k - 1}', 1);
0148             ZVl{k} = ZVl{k} + unfold(tmp, 'left')' * unfold(Z.U{k}, 'left');
0149             % update Zr to keep up w/ recursive definition
0150             tmp = tensorprod_ttemps(V.U{k}, Zl{k - 1}', 1);
0151             Zl{k} = unfold(tmp, 'left')' * unfold(Z.U{k}, 'left');
0152         end
0153 
0154         ZZ = cell(1, d); % final cores for normal (I \otimes X)Z(X^T)
0155         Zv = cell(1, d); % final cores for (I \otimes X)Z(V^T)
0156         vZ = cell(1, d); % final cores for (I \otimes V)Z(X^T)
0157 
0158 
0159         % contract to first core
0160         ZZ{1} = tensorprod_ttemps(Z.U{1}, Zr{2}, 3);
0161         % contract to inner cores
0162         for k = 2:(d - 1)
0163             res = tensorprod_ttemps(Z.U{k}, Zl{k - 1}, 1);
0164             ZZ{k} = tensorprod_ttemps(res, Zr{k + 1}, 3);
0165         end
0166 
0167         % contract to last core
0168         ZZ{d} = tensorprod_ttemps(Z.U{d}, Zl{d - 1}, 1);
0169 
0170         Zv{1} = tensorprod_ttemps(Z.U{1}, ZVr{2}, 3);
0171 
0172         for k = 2:(d - 1)
0173             res = tensorprod_ttemps(Z.U{k}, Zl{k - 1}, 1);
0174             Zv{k} = tensorprod_ttemps(res, ZVr{k + 1}, 3);
0175         end
0176 
0177         Zv{d} = tensorprod_ttemps(Z.U{d}, Zl{d - 1}, 1);
0178 
0179 
0180         vZ{1} = tensorprod_ttemps(Z.U{1}, Zr{2}, 3);
0181 
0182         for k = 2:(d - 1)
0183             res = tensorprod_ttemps(Z.U{k}, ZVl{k - 1}, 1);
0184             vZ{k} = tensorprod_ttemps(res, Zr{k + 1}, 3);
0185         end
0186 
0187         vZ{d} = tensorprod_ttemps(Z.U{d}, ZVl{d - 1}, 1);
0188 
0189         % Left unfold everything to apply additional operations
0190         for k = 1:d
0191             ZZ{k} = unfold(ZZ{k}, 'left');
0192             Zv{k} = unfold(Zv{k}, 'left');
0193             vZ{k} = unfold(vZ{k}, 'left');
0194         end
0195         % %%%%%%%%%%%%%%%%%%%%%%%%% TESTING %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0196 
0197 
0198         % 1st and last cases special, so they're computed outside the loop
0199         UL  = unfold(V.U{1}, 'left');
0200         dUL = unfold(dUR{1}, 'left');
0201 
0202         % right variational term without (I - UU^T) projection
0203         rightV = (Zv{1} - ZZ{1} * XtVg{2}) / R{1};
0204 
0205         Y.dU{1} = -dUL * UL' * ZZ{1} + rightV - UL * (UL' * rightV);
0206 
0207         for k = 2:(d - 1)
0208             UL  = unfold(V.U{k}, 'left');
0209             dUL = unfold(dUR{k}, 'left');
0210 
0211             % right variational term without (I - UU^T) projection
0212             rightV = (Zv{k} - ZZ{k} * XtVg{k + 1}) / R{k};
0213 
0214             Y.dU{k} = vZ{k} - UL * (UL' * vZ{k}) - dUL * UL' * ZZ{k} ...
0215                        + rightV - UL * (UL' * rightV);
0216         end
0217 
0218         Y.dU{d} = vZ{d};
0219 
0220         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0221         %%%%%%%%%%%%%%%                  CROSS TERMS                 %%%%%%%%%%%%%%%%%%
0222         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0223 
0224         for k = 1:(d - 1)
0225             U = unfold(V.U{k}, 'left');
0226             ZZ{k} = ZZ{k} - U * (U' * ZZ{k});
0227             ZZ{k} = reshape(ZZ{k}, [r(k), n(k), r(k + 1)]);
0228         end
0229 
0230         ZZ{d} = reshape(ZZ{d}, [r(d), n(d), 1]);
0231 
0232         % Double loop to represent P_k DP_m, 1 <= k < d
0233         for k = 1:(d - 1)
0234 
0235             UL = unfold(V.U{k}, 'left');
0236             dUL = unfold(dUR{k}, 'left');
0237 
0238             for m = 1:(k - 1)
0239                 VtZ = crossTermVariational(k, m, ZZ, dUR{m}, V);
0240                 projUVtZ = VtZ - UL * (UL' * VtZ);
0241 
0242                 Y.dU{k} = Y.dU{k} - projUVtZ;
0243             end
0244 
0245             for m = (k + 1):d
0246                 ZtX = crossTermMatrixRight(k + 1, m, ZZ, V);
0247                 Y.dU{k} = Y.dU{k} + dUL * ZtX;
0248             end
0249 
0250         end
0251 
0252         % Final terms for k = d
0253         for m = 1:(d - 1)
0254             VtZ = crossTermVariational(d, m, ZZ, dUR{m}, V);
0255             Y.dU{d} = Y.dU{d} - VtZ;
0256         end
0257 
0258         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%FINAL RESHAPE%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0259 
0260         for k = 1:d
0261             Y.dU{k} = reshape(Y.dU{k}, [r(k), n(k), r(k + 1)]);
0262         end
0263 
0264     elseif ~exist('ind', 'var')% Z is a full array
0265 
0266         %{
0267         Note that you can split up the full expression for the Weingarten map in
0268         the following way:
0269 
0270         1) Diagonal terms(P_i DP_i terms, or non - cross terms)
0271         a) This can be split into(left variational) + (right variational),
0272         which correspond to two terms in product rule expansion
0273         2) Cross terms(P_i DP_j, we show mathematically equivalent expression
0274         in companion paper.)
0275         a) This likewise can be split into(i < j) + (i > j)
0276         %}
0277 
0278         % Cell that represents the LEFT VARIATIONAL SIDE
0279         Zv = cell(1, d);
0280         % Cell that represents the RIGHT VARIATIONAL SIDE
0281         vZ = cell(1, d);
0282 
0283         % Initialization step
0284         Zv{d} = Z(:);
0285         vZ{d} = Z(:);
0286 
0287         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0288         %%%%%%%%%%%%%%%               DIAGONAL TERMS                 %%%%%%%%%%%%%%%%%%
0289         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0290 
0291         %%%%%%%%%%%%%%%               LEFT VARIATIONAL               %%%%%%%%%%%%%%%%%%
0292 
0293         % First, calculate where the right side is the variational matrix
0294         % Note that the calculation here is split between VR^-1 and XX^tVR^-1
0295         % in (I - XX^t)VR^{-1} through xx and xx2 terms respectively
0296 
0297         % (d-1) term done separately to initialize the temp variable handoff
0298         zz = reshape(Z, [prod(n(1:(d - 1))), n(d)]);
0299         xx = transpose(unfold(V.dU{d}, 'right'));
0300         xxTmp = transpose(unfold(V.U{d}, 'right')); % temp var to help calc xx terms
0301         xx2 = transpose(unfold(V.V{d}, 'right'));
0302 
0303         % Z-matrix-k-variational, represents the xx calculation for Zv{k}
0304         Zmkv = zz * xx;
0305         vZ{d - 1} = zz * xx2;
0306         Zv{d - 1} = (Zmkv - vZ{d - 1} * XtVg{d}) / R{d - 1};
0307 
0308         % auxillery term for computing anything with variational matrix
0309         ZZtmp = zz * xxTmp;
0310 
0311         % Main calcualtion for right size of Zv
0312         for k = (d - 2):-1:1
0313             zz = reshape(Zmkv, [prod(n(1:k)), n(k + 1) * r(k + 2)]);
0314             xx = transpose(unfold(V.U{k + 1}, 'right'));
0315             Zmkv = zz * xx;
0316 
0317             zzTmp = reshape(ZZtmp, [prod(n(1:k)), n(k + 1) * r(k + 2)]);
0318             xxTmp = transpose(unfold(dUR{k + 1}, 'right'));
0319             Zmkv = Zmkv + zzTmp * xxTmp;
0320             ZZtmp = zzTmp * xx;
0321 
0322             % NOTE: vZ is only used here to store these values for later use
0323             % It stores the standard right side calc for the orthogonal projector
0324             zz2 = reshape(vZ{k + 1}, [prod(n(1:k)), n(k + 1) * r(k + 2)]);
0325             xx2 = transpose(unfold(V.V{k + 1}, 'right'));
0326             vZ{k} = zz2 * xx2;
0327 
0328             Zv{k} = (Zmkv - vZ{k} * XtVg{k + 1}) / R{k};
0329         end
0330 
0331         % Standard left side calculations for Zv
0332         for k = 2:d
0333 
0334             for i = 1:(k - 1)
0335                 Z_i = reshape(Zv{k}, [r(i) * n(i), prod(n(i + 1:k)) * r(k + 1)]);
0336                 X_i = unfold(V.U{i}, 'left');
0337                 Zm = X_i' * Z_i;
0338                 Zv{k} = Zm;
0339             end
0340 
0341             Zv{k} = reshape(Zv{k}, [r(k) * n(k), r(k + 1)]);
0342         end
0343 
0344         % Final U core projector mult for Zv
0345         for k = 1:(d - 1)
0346             U = unfold(V.U{k}, 'left');
0347             Zv{k} = Zv{k} - U * (U' * Zv{k});
0348         end
0349 
0350         %%%%%%%%%%%%%%%               RIGHT VARIATIONAL               %%%%%%%%%%%%%%%%%%
0351 
0352         % represents standard projection, also needed as intermediary for this section
0353         ZZ = vZ;
0354         % Start left side calculation for vZ
0355         for k = 2:d
0356 
0357             % Do first in loop separately to initialize intermediary variable
0358             Z_i = reshape(vZ{k}, [n(1), prod(n(2:k)) * r(k + 1)]);
0359             dU_i = unfold(dUR{1}, 'left');
0360             Zm = dU_i' * Z_i;
0361             vZ{k} = Zm;
0362 
0363             X_i = unfold(V.U{1}, 'left');
0364             ZZ{k} = X_i' * Z_i;
0365 
0366             for m = 2:(k - 1)
0367                 Z_i = reshape(vZ{k}, [r(m) * n(m), prod(n(m + 1:k)) * r(k + 1)]);
0368                 X_i = unfold(V.U{m}, 'left');
0369                 Zm = X_i' * Z_i;
0370 
0371                 Z_tmp = reshape(ZZ{k}, [r(m) * prod(n(m)), prod(n(m + 1:k)) * r(k + 1)]);
0372                 dU_i = unfold(dUR{m}, 'left');
0373 
0374                 vZ{k} = Zm + dU_i' * Z_tmp;
0375 
0376                 ZZ{k} = X_i' * Z_tmp;
0377             end
0378 
0379             vZ{k} = reshape(vZ{k}, [r(k) * n(k), r(k + 1)]);
0380             ZZ{k} = reshape(ZZ{k}, [r(k) * n(k), r(k + 1)]);
0381         end
0382 
0383         % Final U core projector mult for vZ
0384 
0385         % 1st core separate since vZ{1} = 0 up until now. Purely for efficiency
0386         U  = unfold(V.U{1},  'left');
0387         dU = unfold(V.dU{1}, 'left');
0388         vZ{1} = -dU * U' * ZZ{1};
0389 
0390         for k = 2:(d - 1)
0391             U  = unfold(V.U{k},  'left');
0392             dU = unfold(V.dU{k}, 'left');
0393             vZ{k} = vZ{k} - U * (U' * vZ{k});
0394 
0395             vZ{k} = vZ{k} - dU * U' * ZZ{k};
0396 
0397         end
0398 
0399         for k = 1:d - 1
0400             Y.dU{k} = Zv{k} + vZ{k};
0401         end
0402 
0403         Y.dU{d} = vZ{d};
0404 
0405 
0406         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0407         %%%%%%%%%%%%%%%                  CROSS TERMS                 %%%%%%%%%%%%%%%%%%
0408         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0409 
0410         for k = 1:(d - 1)
0411             % now only need ZZ as dU cores of normal projection of V
0412             U = unfold(V.U{k}, 'left');
0413             ZZ{k} = ZZ{k} - U * (U' * ZZ{k});
0414             ZZ{k} = reshape(ZZ{k}, [r(k), n(k), r(k + 1)]);
0415         end
0416 
0417         ZZ{d} = reshape(ZZ{d}, [r(d), n(d), 1]);
0418 
0419         % Double loop to represent P_k DP_m, 1 <= k < d
0420         for k = 1:(d - 1)
0421 
0422             UL  = unfold(V.U{k},  'left');
0423             dUL = unfold(V.dU{k}, 'left');
0424 
0425             for m = 1:(k - 1)
0426                 VtZ = crossTermVariational(k, m, ZZ, dUR{m}, V);
0427 
0428                 projUVtZ = VtZ - UL * (UL' * VtZ);
0429 
0430                 Y.dU{k} = Y.dU{k} - projUVtZ; 
0431             end
0432 
0433             for m = (k + 1):d
0434                 ZtX = crossTermMatrixRight(k + 1, m, ZZ, V);
0435                 Y.dU{k} = Y.dU{k} + dUL * ZtX;
0436             end
0437 
0438         end
0439 
0440         % Final terms for k = d
0441         for m = 1:(d - 1)
0442             VtZ = crossTermVariational(d, m, ZZ, dUR{m}, V);
0443 
0444             Y.dU{d} = Y.dU{d} - VtZ;
0445         end
0446 
0447         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% FINAL RESHAPE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0448 
0449         for k = 1:d
0450             Y.dU{k} = reshape(Y.dU{k}, [r(k), n(k), r(k + 1)]);
0451         end
0452 
0453     else
0454         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% SPARSE CASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0455 
0456         % permuted U, V, and dU cells for the efficient computation in C
0457         CU = cell(1, d);
0458         CV = cell(1, d);
0459         CdUR = cell(1, d);
0460 
0461         for k = 1:d
0462             CU{k} = permute(V.U{k}, [1 3 2]);
0463             CV{k} = permute(V.V{k}, [1 3 2]);
0464             CdUR{k} = permute(dUR{k}, [1 3 2]);
0465         end
0466 
0467         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0468         %%%%%%%%%%%%%%%               DIAGONAL TERMS                 %%%%%%%%%%%%%%%%%%
0469         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0470 
0471         [ZZ, vZ, Zv] = weingarten_omega(n, r, CU, CV, CdUR, ind.', Z);
0472 
0473         for k = 1:d
0474 
0475             ZZ{k} = reshape(ZZ{k}, [r(k), r(k + 1), n(k)]);
0476             ZZ{k} = unfold(permute(ZZ{k}, [1 3 2]), 'left');
0477 
0478             vZ{k} = reshape(vZ{k}, [r(k), r(k + 1), n(k)]);
0479             vZ{k} = unfold(permute(vZ{k}, [1 3 2]), 'left');
0480 
0481             Zv{k} = reshape(Zv{k}, [r(k), r(k + 1), n(k)]);
0482             Zv{k} = unfold(permute(Zv{k}, [1 3 2]), 'left');
0483         end
0484 
0485 
0486             
0487 
0488         % 1st and last cases special, so they're computed outside the loop
0489         UL  = unfold(V.U{1}, 'left');
0490         dUL = unfold(dUR{1}, 'left');
0491 
0492         % right variational term without (I - UU^T) projection
0493         rightV = (Zv{1} - ZZ{1} * XtVg{2}) / R{1};
0494         Y.dU{1} = -dUL * UL' * ZZ{1} + rightV - UL * (UL' * rightV);
0495 
0496         for k = 2:(d - 1)
0497             UL  = unfold(V.U{k}, 'left');
0498             dUL = unfold(dUR{k}, 'left');
0499 
0500             % right variational term without (I - UU^T) projection
0501             rightV = (Zv{k} - ZZ{k} * XtVg{k + 1}) / R{k};
0502             Y.dU{k} = vZ{k} - UL * (UL' * vZ{k}) - dUL * UL' * ZZ{k} + rightV - UL * (UL' * rightV);
0503             % norm(vZ{k} - UL * (UL' * vZ{k}) + (Zv{k} - UL * (UL' * Zv{k})) / R{k}) / norm(V)
0504             % norm(V)
0505         end
0506 
0507 
0508         Y.dU{d} = vZ{d};
0509 
0510 
0511         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0512         %%%%%%%%%%%%%%%                  CROSS TERMS                 %%%%%%%%%%%%%%%%%%
0513         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0514 
0515         for k = 1:(d - 1)
0516             U = unfold(V.U{k}, 'left');
0517             ZZ{k} = ZZ{k} - U * (U' * ZZ{k});
0518             ZZ{k} = reshape(ZZ{k}, [r(k), n(k), r(k + 1)]);
0519         end
0520 
0521         ZZ{d} = reshape(ZZ{d}, [r(d), n(d), 1]);
0522 
0523         % Double loop to represent P_k DP_m, 1 <= k < d
0524         for k = 1:(d - 1)
0525 
0526             UL  = unfold(V.U{k}, 'left');
0527             dUL = unfold(dUR{k}, 'left');
0528 
0529             for m = 1:(k - 1)
0530                 % XtZ = crossTermMatrixLeft(k, m, ZZ, V);
0531                 VtZ = crossTermVariational(k, m, ZZ, dUR{m}, V);
0532                 projUVtZ = VtZ - UL * (UL' * VtZ);
0533 
0534                 Y.dU{k} = Y.dU{k} - projUVtZ; %+ dUL * XtZ;
0535 
0536                 % Y.dU{k} = Y.dU{k} - (projUk * kron(eye(n(k)), Vllk') - dUkL * Xlek') * Zkm;
0537             end
0538 
0539             for m = (k + 1):d
0540                 ZtX = crossTermMatrixRight(k + 1, m, ZZ, V);
0541                 Y.dU{k} = Y.dU{k} + dUL * ZtX;
0542 
0543                 %         Y.dU{k} = Y.dU{k} + dUkL * Zkm' * Xg{k+1};
0544             end
0545 
0546         end
0547 
0548         % Final terms for k = d
0549         for m = 1:(d - 1)
0550             VtZ = crossTermVariational(d, m, ZZ, dUR{m}, V);
0551 
0552             Y.dU{d} = Y.dU{d} - VtZ;
0553 
0554             %     Y.dU{d} = Y.dU{d} - kron(eye(n(d)), Vlld') * Zkm;
0555         end
0556 
0557         %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%FINAL RESHAPE%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0558 
0559         for k = 1:d
0560             Y.dU{k} = reshape(Y.dU{k}, [r(k), n(k), r(k + 1)]);
0561         end
0562 
0563         Y = (normV) * Y;
0564     end
0565 
0566 end
0567 
0568 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0569 % Helper methods
0570 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
0571 
0572 % Efficient ZtX term for P_k DP_m, m > k
0573 % Zp here is the projection of Z (just dU cores)
0574 % V is only passed for access to U and V cores
0575 function ZtX = crossTermMatrixRight(k, m, Zp, V)
0576     ZtX = conj(unfold(Zp{m}, 'right')) * unfold(V.V{m}, 'right').';
0577 
0578     for p = (m - 1):-1:k
0579         tmp = tensorprod_ttemps(V.U{p}, (ZtX)', 3);
0580         ZtX = conj(unfold(tmp, 'right')) * unfold(V.V{p}, 'right').';
0581     end
0582 
0583 end
0584 
0585 % Efficient XtZ term for P_k DP_m, m < k
0586 % Zp here is the projection of Z
0587 % V is only passed for access to U and V cores
0588 function XtZ = crossTermMatrixLeft(k, m, Zp, V) %#ok<DEFNU>
0589 
0590     XtZ = unfold(V.U{m}, 'left')' * unfold(Zp{m}, 'left');
0591 
0592     for p = (m + 1):k
0593         XtZ = tensorprod_ttemps(V.U{p}, (XtZ)', 1);
0594         XtZ = unfold(XtZ, 'left')' * unfold(V.V{p}, 'left');
0595     end
0596 
0597 end
0598 
0599 % Efficient VtZ (more accurately kron(eye(n(k)), V') * Z) term for P_k DP_m, m < k
0600 % Zp here is the projection of Z, dURm is normally parametrized dU{m}
0601 % V is only passed for access to U and V cores
0602 function VtZ = crossTermVariational(k, m, Zp, dURm, V)
0603 
0604     VtZ = unfold(dURm, 'left')' * unfold(Zp{m}, 'left');
0605 
0606     for p = (m + 1):(k - 1)
0607         VtZ = tensorprod_ttemps(V.U{p}, (VtZ)', 1);
0608         VtZ = unfold(VtZ, 'left')' * unfold(V.V{p}, 'left');
0609     end
0610 
0611     VtZ = tensorprod_ttemps(V.V{k}, VtZ, 1);
0612     VtZ = unfold(VtZ, 'left');
0613 
0614 end