update

Author: GiggleLiu

src/map.jl

@@ -53,15 +53,19 @@ $(TYPEDSIGNATURES)
 Returns the largest log-probability and the most probable configuration.
 """
 function most_probable_config(tn::TensorNetworkModel; usecuda = false)::Tuple{Real, Vector}
-    ixs = OMEinsum.getixsv(tn.code)
-    unity_labels = ixs[tn.unity_tensors_idx]
-    indices = [findfirst(==([l]), unity_labels) for l in get_vars(tn)]
-    @assert !any(isnothing, indices) "To get the the most probable configuration, the unity tensors labels must include all variables"
     vars = get_vars(tn)
+    tensor_indices = check_queryvars(tn, [[v] for v in vars])
     tensors = map(t -> Tropical.(log.(t)), adapt_tensors(tn; usecuda, rescale = false))
     logp, grads = cost_and_gradient(tn.code, tensors)
     # use Array to convert CuArray to CPU arrays
-    return content(Array(logp)[]), map(k -> haskey(tn.evidence, vars[k]) ? tn.evidence[vars[k]] : argmax(grads[tn.unity_tensors_idx[indices[k]]]) - 1, 1:length(vars))
+    return content(Array(logp)[]), map(k -> haskey(tn.evidence, vars[k]) ? tn.evidence[vars[k]] : argmax(grads[tensor_indices[k]]) - 1, 1:length(vars))
+end
+# check if the queryvars are included in the unity tensors labels, if yes, return the indices of the unity tensors
+function check_queryvars(tn::TensorNetworkModel, queryvars::AbstractVector{Vector{Int}})
+    ixs = OMEinsum.getixsv(tn.code)
+    indices = [findfirst(==(l), ixs[tn.unity_tensors_idx]) for l in queryvars]
+    @assert !any(isnothing, indices) "To get the the most probable configuration, the unity tensors labels must include all variables. Query variables: $queryvars, Unity tensors labels: $(ixs[tn.unity_tensors_idx])"
+    return tn.unity_tensors_idx[indices]
 end
 
 """

src/mar.jl

@@ -136,8 +136,10 @@ tensor network.
 
 ### Arguments
 - `tn`: The [`TensorNetworkModel`](@ref) to query.
-- `usecuda`: Specifies whether to use CUDA for tensor contraction.
-- `rescale`: Specifies whether to rescale the tensors during contraction.
+
+### Keyword Arguments
+- `usecuda::Bool`: Specifies whether to use CUDA for tensor contraction.
+- `rescale::Bool`: Specifies whether to rescale the tensors during contraction.
 
 ### Example
 The following example is taken from [`examples/asia-network/main.jl`](https://tensorbfs.github.io/TensorInference.jl/dev/generated/asia-network/main/).
@@ -187,9 +189,10 @@ function marginals(tn::TensorNetworkModel; usecuda = false, rescale = true)::Dic
     cost, grads = cost_and_gradient(tn.code, adapt_tensors(tn; usecuda, rescale))
     @debug "cost = $cost"
     ixs = OMEinsum.getixsv(tn.code)
+    queryvars = ixs[tn.unity_tensors_idx]
     if rescale
-        return Dict(zip(ixs[tn.unity_tensors_idx], LinearAlgebra.normalize!.(getfield.(grads[1:length(tn.unity_tensors_idx)], :normalized_value), 1)))
+        return Dict(zip(queryvars, LinearAlgebra.normalize!.(getfield.(grads[tn.unity_tensors_idx], :normalized_value), 1)))
     else
-        return Dict(zip(ixs[tn.unity_tensors_idx], LinearAlgebra.normalize!.(grads[1:length(tn.unity_tensors_idx)], 1)))
+        return Dict(zip(queryvars, LinearAlgebra.normalize!.(grads[tn.unity_tensors_idx], 1)))
     end
-end
+end