add option to reweight likelihood for datapoints

Author: stecrotti

src/MatrixProductStates/derivatives.jl

@@ -66,7 +66,8 @@ Return also the loglikelihood, which is a byproduct of the computation.
 function grad_loglikelihood_two_site(p::MPS, k::Integer, X;
     prodA_left = [precompute_left_environments(p.ψ, x) for x in X],
     prodA_right = [precompute_right_environments(p.ψ, x) for x in X],
-    Aᵏᵏ⁺¹ =_merge_tensors(p[k], p[k+1]))
+    Aᵏᵏ⁺¹ =_merge_tensors(p[k], p[k+1]),
+    weights = ones(length(X)))
 
     Zprime, Z = grad_normalization_two_site_canonical(p, k; Aᵏᵏ⁺¹)
     ll = -log(Z)
@@ -78,8 +79,8 @@ function grad_loglikelihood_two_site(p::MPS, k::Integer, X;
         gr, val = grad_evaluate_two_site(p.ψ, k, x;
             Ax_left = prodA_left[n][k-1], Ax_right = prodA_right[n][k+2], Aᵏᵏ⁺¹
             )
-        gA[:,:,x[k]...,x[k+1]...] .+= 2/T * gr / val
-        ll += 1/T * log(abs2(val))
+        gA[:,:,x[k]...,x[k+1]...] .+= 2/T * gr / val * weights[n]
+        ll += 1/T * log(abs2(val)) * weights[n]
     end
     return gA, ll
 end

test/mps.jl

@@ -277,8 +277,9 @@ end
         X = [sample(p)[1] for _ in 1:10^2]
         q = rand_mps(ComplexF64, 2, length(p), 2,2)
         ll = loglikelihood(q, X)
+        weights = ones(length(X))
         two_site_dmrg!(q, X, 1;
-            η=1e-4, ndesc=10, svd_trunc=TruncBond(5))
+            η=1e-4, ndesc=10, svd_trunc=TruncBond(5), weights)
         @test loglikelihood(q, X) > ll
     end
 end