diff --git a/Project.toml b/Project.toml index c8358d0..7afdb14 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorNetworksNext" uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c" -version = "0.4.3" +version = "0.4.4" authors = ["ITensor developers and contributors"] [workspace] @@ -19,8 +19,10 @@ FunctionImplementations = "7c7cc465-9c6a-495f-bdd1-f42428e86d0c" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" +MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4" NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde" NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" @@ -47,11 +49,13 @@ FunctionImplementations = "0.4.1" Graphs = "1.13.1" LinearAlgebra = "1.10" MacroTools = "0.5.16" -NamedDimsArrays = "0.14.3, 0.15" +MatrixAlgebraKit = "0.6" +NamedDimsArrays = "0.15.5" NamedGraphs = "0.11" +Random = "1.10" SimpleTraits = "0.9.5" SplitApplyCombine = "1.2.3" -TensorAlgebra = "0.9.2" +TensorAlgebra = "0.9.3" TensorOperations = "5.3.1" TermInterface = "2" TypeParameterAccessors = "0.4.4" diff --git a/src/AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl b/src/AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl index a95e0e0..627a482 100644 --- a/src/AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl +++ b/src/AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl @@ -52,44 +52,6 @@ function Base.propertynames(state::NestedState) return (fieldnames(typeof(state))..., :iterate) end -# ============================ select_algorithm / default_algorithm ======================== - -# Like `MatrixAlgebraKit.select_algorithm` / `default_algorithm`, but -# selection-relevant inputs are packed into an `args` tuple so the value -# and type domains stay disjoint: `(1.2,)` vs `Tuple{Float64}`. Strategy -# types subtype `AbstractAlgorithm` so the passthrough overload is generic. -abstract type AbstractAlgorithm end - -function default_algorithm(f, ::Type{Args}; kwargs...) where {Args <: Tuple} - return throw(MethodError(default_algorithm, (f, Args))) -end -function default_algorithm(f, args::Tuple; kwargs...) - return default_algorithm(f, typeof(args); kwargs...) -end - -function select_algorithm(f, alg, args::Tuple; kwargs...) - return select_algorithm(f, alg, typeof(args); kwargs...) -end -function select_algorithm(f, ::Nothing, ::Type{Args}; kwargs...) where {Args <: Tuple} - return default_algorithm(f, Args; kwargs...) -end -function select_algorithm(f, alg::NamedTuple, ::Type{Args}; kwargs...) where {Args <: Tuple} - isempty(kwargs) || throw( - ArgumentError( - "Additional keyword arguments are not allowed when `alg` is a `NamedTuple`." - ) - ) - return default_algorithm(f, Args; alg...) -end -function select_algorithm(f, alg::AbstractAlgorithm, ::Type{<:Tuple}; kwargs...) - isempty(kwargs) || throw( - ArgumentError( - "Additional keyword arguments are not allowed when `alg` is an `AbstractAlgorithm` instance." - ) - ) - return alg -end - # ============================ StopWhenConverged =========================================== # Stopping criterion that fires once `iterate_diff(iterate, previous_iterate) < tol`. diff --git a/src/ITensorNetworksNext.jl b/src/ITensorNetworksNext.jl index 41ce78e..4b0f8c3 100644 --- a/src/ITensorNetworksNext.jl +++ b/src/ITensorNetworksNext.jl @@ -6,8 +6,10 @@ module ITensorNetworksNext # dependency by Aqua. using TensorAlgebra: TensorAlgebra +include("select_algorithm.jl") include("AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl") include("LazyNamedDimsArrays/LazyNamedDimsArrays.jl") +include("tensoralgebra.jl") include("abstracttensornetwork.jl") include("tensornetwork.jl") include("TensorNetworkGenerators/TensorNetworkGenerators.jl") @@ -15,5 +17,8 @@ include("contract_network.jl") include("beliefpropagation/messagecache.jl") include("beliefpropagation/beliefpropagation.jl") +include("beliefpropagation/normnetwork.jl") + +include("apply/apply_operators.jl") end diff --git a/src/abstracttensornetwork.jl b/src/abstracttensornetwork.jl index 121073d..ca8652f 100644 --- a/src/abstracttensornetwork.jl +++ b/src/abstracttensornetwork.jl @@ -181,8 +181,6 @@ function rand_trivial_namedunitrange( return namedunitrange(trivial_unitrange(R), randname(N)) end -dag(x) = x - function insert_trivial_link!(tn, e) add_edge!(tn, e) l = rand_trivial_namedunitrange(eltype(inds(tn[src(e)]))) diff --git a/src/apply/apply_operators.jl b/src/apply/apply_operators.jl new file mode 100644 index 0000000..d1d1fd9 --- /dev/null +++ b/src/apply/apply_operators.jl @@ -0,0 +1,266 @@ +import .AlgorithmsInterfaceExtensions as AIE +import AlgorithmsInterface as AI +import NamedDimsArrays as NDA +import TensorAlgebra as TA +using Base: @kwdef +using Graphs: dst, src, vertices +using LinearAlgebra: norm +using NamedDimsArrays: AbstractNamedDimsArray, dimnames, domainnames, nameddims, operator, + randname, replacedimnames +using NamedGraphs.GraphsExtensions: all_edges, boundary_edges +using TensorAlgebra: gram_eigh_full, gram_eigh_full_with_pinv + +# === Top-level user entry point === + +# Apply a list of operators to a state given the environments. +function apply_operators(operators, state, env; alg = nothing, kwargs...) + algorithm = select_algorithm( + apply_operators, alg, (operators, state, env); kwargs... + ) + return apply_operators(algorithm, operators, state, env) +end + +# The `apply_operators` iteration algorithm wraps the per-operator algorithm, +# which is itself resolved via `apply_operator` (overridable with `operator_alg`). +function default_algorithm( + ::typeof(apply_operators), args::Tuple; + operator_alg = nothing, environment_alg = nothing, kwargs... + ) + operators, state, env = args + # `apply_operator` acts on a single operator, so select on the operator + # element type, keeping the remaining `(state, env)` argument types. + # We use types here in case the operator list is empty. + operator_args = Tuple{eltype(operators), typeof(state), typeof(env)} + operator_algorithm = + select_algorithm(apply_operator, operator_alg, operator_args; kwargs...) + # `apply_operator_environment_preparation` signature (minus the env algorithm): + # `(operator_algorithm, operators, iteration::Int, iterate, env)`. + prepare_args = (operator_algorithm, operators, 0, state, env) + environment_algorithm = select_algorithm( + apply_operator_environment_preparation, environment_alg, prepare_args + ) + return ApplyOperatorsAlgorithm(; + operator_algorithm, + environment_algorithm, + stopping_criterion = AI.StopAfterIteration(length(operators)) + ) +end + +function apply_operators(algorithm, operators, state, env) + isempty(operators) && return copy(state), copy(env) + problem = ApplyOperatorsProblem(; operators, init = state) + return AI.solve(problem, algorithm; iterate = state, env) +end + +# === Layer 1: apply_operators iteration === + +@kwdef struct ApplyOperatorsProblem{Ops, Init} <: AI.Problem + operators::Ops + init::Init +end + +@kwdef struct ApplyOperatorsAlgorithm{ + OperatorAlgorithm, + EnvironmentAlgorithm, + StoppingCriterion <: AI.StoppingCriterion, + } <: AI.Algorithm + operator_algorithm::OperatorAlgorithm + environment_algorithm::EnvironmentAlgorithm = NoApplyOperatorEnvironmentPreparation() + stopping_criterion::StoppingCriterion = AI.StopAfterIteration(0) +end + +@kwdef mutable struct ApplyOperatorsState{ + Iterate, Env, StoppingCriterionState <: AI.StoppingCriterionState, + } <: AI.State + iterate::Iterate + env::Env + iteration::Int = 0 + stopping_criterion_state::StoppingCriterionState +end + +function AI.initialize_state( + problem::ApplyOperatorsProblem, algorithm::ApplyOperatorsAlgorithm; + iterate, env, iteration::Int = 0 + ) + stopping_criterion_state = AI.initialize_state( + problem, algorithm, algorithm.stopping_criterion; iterate + ) + return ApplyOperatorsState(; + iterate, env, iteration, stopping_criterion_state + ) +end + +function AI.step!( + problem::ApplyOperatorsProblem, algorithm::ApplyOperatorsAlgorithm, + state::ApplyOperatorsState + ) + # Prepare for the operator application, for example by updating the + # environments in a path between where the operators are being applied. + state.iterate, state.env = apply_operator_environment_preparation( + algorithm.environment_algorithm, algorithm.operator_algorithm, + problem.operators, state.iteration, state.iterate, state.env + ) + state.iterate, state.env = apply_operator( + algorithm.operator_algorithm, problem.operators[state.iteration], state.iterate, + state.env + ) + return state +end + +function AI.finalize_state!( + ::ApplyOperatorsProblem, ::ApplyOperatorsAlgorithm, state::ApplyOperatorsState + ) + return state.iterate, state.env +end + +# === Layer 2: environment-preparation strategy === + +# Update the environment (and possibly the factors) before the next operator is +# applied. The full `operators`/`iteration` and `operator_algorithm` are passed so +# a strategy can judge which messages went stale and how much to recompute; it may +# also return regauged/orthogonalized factors. Only the no-op is implemented for +# now (reconvergence policies are follow-up work). +struct NoApplyOperatorEnvironmentPreparation <: AbstractAlgorithm end + +function apply_operator_environment_preparation( + ::NoApplyOperatorEnvironmentPreparation, operator_algorithm, operators, iteration, + iterate, env + ) + return iterate, env +end + +function default_algorithm( + ::typeof(apply_operator_environment_preparation), ::Type{<:Tuple}; kwargs... + ) + return NoApplyOperatorEnvironmentPreparation() +end + +# === Layer 3: single-operator strategy === + +abstract type ApplyOperatorAlgorithm <: AbstractAlgorithm end + +# Apply a single operator to the state, given the specified environments. +# Returns an updated state along with updated environments where relevant. +# Note that it isn't expected that environments are fully recomputed, +# generally only minimal updates will be made (say to the edge where a 2-site +# operator is applied). +function apply_operator(operator, state, env; alg = nothing, kwargs...) + algorithm = select_algorithm(apply_operator, alg, (operator, state, env); kwargs...) + return apply_operator(algorithm, operator, state, env) +end + +function apply_operator(algorithm::ApplyOperatorAlgorithm, operator, state, env) + dest, env_dest = initialize_output(apply_operator!, algorithm, operator, state, env) + apply_operator!(algorithm, dest, operator, state, env_dest) + return dest, env_dest +end + +# === Default strategy: BPApplyGate === + +@kwdef struct BPApplyGate{Trunc} <: ApplyOperatorAlgorithm + trunc::Trunc = nothing + normalize::Bool = false +end + +function apply_operator!( + algorithm::BPApplyGate, dest, operator, state, env + ) + apply_gate_bp!( + dest, operator, state, env; + algorithm.trunc, algorithm.normalize + ) + return dest +end + +function initialize_output( + ::typeof(apply_operator!), ::BPApplyGate, operator, state, env + ) + return copy(state), copy(env) +end + +function default_algorithm(::typeof(apply_operator), ::Type{<:Tuple}; kwargs...) + return BPApplyGate(; kwargs...) +end + +# === BP simple-update implementation === + +function apply_gate_bp!( + dest::AbstractTensorNetwork, op::AbstractNamedDimsArray, + state::AbstractTensorNetwork, env; kwargs... + ) + op_in = domainnames(op) + vs = [v for v in vertices(state) if !isempty(intersect(op_in, sitenames(state, v)))] + isempty(vs) && throw( + ArgumentError("operator shares no indices with the tensor network") + ) + return apply_gate_bp_nsite!(Val(length(vs)), dest, op, state, env, vs; kwargs...) +end + +function apply_gate_bp_nsite!( + ::Val{N}, dest::AbstractTensorNetwork, op::AbstractNamedDimsArray, + state::AbstractTensorNetwork, env, vs; kwargs... + ) where {N} + return throw(ArgumentError("$N-site gate decomposition not implemented")) +end + +function apply_gate_bp_nsite!( + ::Val{1}, dest::AbstractTensorNetwork, op::AbstractNamedDimsArray, + state::AbstractTensorNetwork, env, vs; + normalize, kwargs... + ) + v = only(vs) + ψv = NDA.apply(op, state[v]) + if normalize + gauges = [ + gram_eigh_full(env[e]) + for e in boundary_edges(state, vs; dir = :in) + ] + ψv /= norm(prod([[ψv]; gauges])) + end + dest[v] = ψv + return dest +end + +function apply_gate_bp_nsite!( + ::Val{2}, dest::AbstractTensorNetwork, op::AbstractNamedDimsArray, + state::AbstractTensorNetwork, env, vs; + trunc, normalize + ) + v1, v2 = vs + edges_in = boundary_edges(state, vs; dir = :in) + grams_v1 = + [gram_eigh_full_with_pinv(env[e]) for e in edges_in if dst(e) == v1] + grams_v2 = + [gram_eigh_full_with_pinv(env[e]) for e in edges_in if dst(e) == v2] + gauges_v1, inv_gauges_v1 = first.(grams_v1), last.(grams_v1) + gauges_v2, inv_gauges_v2 = first.(grams_v2), last.(grams_v2) + + ψ_v1 = prod([[state[v1]]; gauges_v1]) + ψ_v2 = prod([[state[v2]]; gauges_v2]) + + Q_v1, R_v1 = TA.qr(ψ_v1, setdiff(dimnames(ψ_v1), dimnames(ψ_v2), dimnames(op))) + Q_v2, R_v2 = TA.qr(ψ_v2, setdiff(dimnames(ψ_v2), dimnames(ψ_v1), dimnames(op))) + op_R_v1v2 = NDA.apply(op, R_v1 * R_v2) + U_v1, S, U_v2 = TA.svd(op_R_v1v2, setdiff(dimnames(R_v1), dimnames(R_v2)); trunc) + if normalize + S = S / norm(S) + end + name_v1, name_v2 = dimnames(S) + sqrt_S = sqrt(S, (name_v1,), (name_v2,)) + R_v1 = replacedimnames(U_v1 * sqrt_S, name_v2 => name_v1) + R_v2 = sqrt_S * U_v2 + + dest[v1] = prod([[Q_v1 * R_v1]; inv_gauges_v1]) + dest[v2] = prod([[Q_v2 * R_v2]; inv_gauges_v2]) + + fresh_12 = randname(name_v1) + fresh_21 = randname(name_v1) + # TODO: if `replacedimnames` preserved the operator wrapper (updating the + # codomain/domain `Bijection` accordingly), we could drop the outer + # `operator(...)` wrap here. + env[v1 => v2] = + operator(replacedimnames(S, name_v2 => fresh_12), (name_v1,), (fresh_12,)) + env[v2 => v1] = + operator(replacedimnames(S, name_v2 => fresh_21), (name_v1,), (fresh_21,)) + return dest +end diff --git a/src/beliefpropagation/beliefpropagation.jl b/src/beliefpropagation/beliefpropagation.jl index d6dfabc..458bc3c 100644 --- a/src/beliefpropagation/beliefpropagation.jl +++ b/src/beliefpropagation/beliefpropagation.jl @@ -64,7 +64,7 @@ function beliefpropagation( cache = MessageCache(messages) # No concrete `edge` value here, so the args tuple uses `edgetype(factors)`. - message_update_algorithm = AIE.select_algorithm( + message_update_algorithm = select_algorithm( message_update!, message_update_algorithm, Tuple{typeof(cache), typeof(factors), edgetype(factors)} @@ -203,21 +203,21 @@ end # message is computed and written back into the message store. Plug in a # new strategy by subtyping `MessageUpdateAlgorithm` and overloading # `message_update!(strategy, cache, factors, edge)`. -abstract type MessageUpdateAlgorithm <: AIE.AbstractAlgorithm end +abstract type MessageUpdateAlgorithm <: AbstractAlgorithm end function message_update! end # `args` tuple mirrors the `message_update!(cache, factors, edge)` call shape. -function AIE.default_algorithm(::typeof(message_update!), ::Type{<:Tuple}; kwargs...) +function default_algorithm(::typeof(message_update!), ::Type{<:Tuple}; kwargs...) return SimpleMessageUpdate(; kwargs...) end -# Convenience entry: pick the strategy via `AIE.select_algorithm` +# Convenience entry: pick the strategy via `select_algorithm` # (accepts either `alg = ::MessageUpdateAlgorithm` / `::NamedTuple`, or flat # kwargs forwarded to the default algorithm), then dispatch. function message_update!(cache, factors, edge; alg = nothing, kwargs...) return message_update!( - AIE.select_algorithm(message_update!, alg, (cache, factors, edge); kwargs...), + select_algorithm(message_update!, alg, (cache, factors, edge); kwargs...), cache, factors, edge ) end diff --git a/src/beliefpropagation/normnetwork.jl b/src/beliefpropagation/normnetwork.jl new file mode 100644 index 0000000..abeb018 --- /dev/null +++ b/src/beliefpropagation/normnetwork.jl @@ -0,0 +1,179 @@ +using DataGraphs: underlying_graph +using Graphs: edges, src +using NamedDimsArrays: + codomainnames, domainnames, name, operator, randname, replacedimnames, state +using NamedGraphs.GraphsExtensions: all_edges, incident_edges +using Random: Random + +# === Norm-network environment constructors === +# +# `*_norm_message_env(tn)` builds a `MessageCache` shaped to act as the BP environment +# for the norm network ⟨tn|tn⟩, with each entry filled per the leading verb (`identity`, +# `ones`, `randn`, `rand`). The `_env` suffix is reserved for the high-level +# environment-builder interface; the low-level `MessageCache` / `messagecache(...)` +# constructors are used internally. A parallel `*_norm_ctm_env` family is planned for +# CTMRG environments. + +""" + similar_norm_message_env(tn) -> MessageCache + +Allocate a BP environment for the norm network ⟨tn|tn⟩ with **undefined** message data: +one square operator message per directed edge of `tn` (both directions on every +undirected edge). Each message's codomain is the link axes on that edge in `tn`; the +domain has dual axes with fresh `randname`-generated names. Element type and backend are +inherited from the factor tensors of `tn` via `Base.similar`. + +Used internally by [`norm_message_env`](@ref) and the filled environment constructors +([`identity_norm_message_env`](@ref), [`ones_norm_message_env`](@ref), +[`randn_norm_message_env`](@ref), [`rand_norm_message_env`](@ref)). Use it directly to +construct environments with custom message data, e.g. by mutating each entry after +allocation. +""" +function similar_norm_message_env(tn) + return messagecache(all_edges(tn)) do e + return similar_operator(tn[src(e)], linkinds(tn, e)) + end +end + +""" + norm_message_env(f, tn) -> MessageCache + +Allocate a norm-network BP environment via [`similar_norm_message_env`](@ref) and apply +`f` to each operator-message entry. Shared building block for the filled-environment +constructors. +""" +function norm_message_env(f, tn) + env = similar_norm_message_env(tn) + # TODO: replace with `map(f, env)` once `map` is defined on `MessageCache`. + foreach(e -> env[e] = f(env[e]), edges(env)) + return env +end + +""" + identity_norm_message_env(tn) -> MessageCache + +Build a norm-network BP environment with identity-operator messages on every edge — the +"uncorrelated environment" starting point for belief-propagation simple-update gauging +on ⟨tn|tn⟩. + +See also: [`ones_norm_message_env`](@ref), [`randn_norm_message_env`](@ref), +[`rand_norm_message_env`](@ref), [`similar_norm_message_env`](@ref). +""" +identity_norm_message_env(tn) = norm_message_env(one_operator, tn) + +""" + ones_norm_message_env(tn) -> MessageCache + +Build a norm-network BP environment whose per-edge messages have every entry equal to +`1` — the rank-1 outer product of all-ones vectors on each (codomain, domain) pair. + +See also: [`identity_norm_message_env`](@ref), [`randn_norm_message_env`](@ref), +[`rand_norm_message_env`](@ref). +""" +ones_norm_message_env(tn) = norm_message_env(msg -> fill!(msg, one(eltype(msg))), tn) + +randn_norm_message_env(tn) = randn_norm_message_env(Random.default_rng(), tn) + +""" + randn_norm_message_env([rng], tn) -> MessageCache + +Build a norm-network BP environment whose per-edge messages have entries drawn from a +standard normal distribution. `rng` defaults to `Random.default_rng()`. + +See also: [`rand_norm_message_env`](@ref), [`identity_norm_message_env`](@ref), +[`ones_norm_message_env`](@ref). +""" +function randn_norm_message_env(rng::Random.AbstractRNG, tn) + return norm_message_env(msg -> randn_operator!(rng, msg), tn) +end + +rand_norm_message_env(tn) = rand_norm_message_env(Random.default_rng(), tn) + +""" + rand_norm_message_env([rng], tn) -> MessageCache + +Build a norm-network BP environment whose per-edge messages have entries drawn from a +uniform distribution on `[0, 1)`. `rng` defaults to `Random.default_rng()`. + +See also: [`randn_norm_message_env`](@ref), [`identity_norm_message_env`](@ref), +[`ones_norm_message_env`](@ref). +""" +function rand_norm_message_env(rng::Random.AbstractRNG, tn) + return norm_message_env(msg -> rand_operator!(rng, msg), tn) +end + +# === Double-layer construction and BP wrapper === + +""" + normnetwork(tn) -> norm_tn, linknames_map + +Build the double-layer norm network `⟨tn|tn⟩` together with the per-edge ket→bra name +mapping used to construct it. + +Each ket link axis on every edge is paired with a fresh `randname`-generated bra link +name; the bra layer at every vertex is the ket tensor with all of its incident link +names renamed accordingly. The returned `linknames_map` is keyed by both directions of +each undirected edge (the values are shared `Dict`s, so a directed edge and its reverse +look up the same `ketname => braname` table) and is the source of truth for adapting +externally-supplied messages onto the double-layer network. + +Anticipates a future `NormNetwork(tn)` struct that bundles `norm_tn` and `linknames_map` +into a single value with `beliefpropagation` dispatch. +""" +function normnetwork(tn) + linknames_map = Dict( + e => Dict(name(ind) => randname(name(ind)) for ind in linkinds(tn, e)) + for e in edges(tn) + ) + merge!(linknames_map, Dict(reverse(e) => m for (e, m) in linknames_map)) + norm_tn = TensorNetwork(underlying_graph(tn)) do v + t = tn[v] + ket_to_bra = Dict(p for e in incident_edges(tn, v) for p in linknames_map[e]) + return t * replacedimnames(n -> get(ket_to_bra, n, n), dag(t)) + end + return norm_tn, linknames_map +end + +""" + beliefpropagation_normnetwork(tn, messages; kwargs...) -> MessageCache + +Run belief propagation on the norm network `⟨tn|tn⟩` (treating `tn` as the ket), +starting from a pre-built operator `MessageCache` `messages` (e.g. from +[`identity_norm_message_env`](@ref) or any of the other `*_norm_message_env` +constructors). + +The norm network built by [`normnetwork`](@ref) is the source of truth for bra-link +names. Each input operator message's domain (bra) axes are renamed to match the +norm-network's bra names before BP iterates; the converged messages are wrapped back as +operators using those same bra names on output. `kwargs` are forwarded to +[`beliefpropagation`](@ref). + +Anticipates a future `beliefpropagation(NormNetwork(tn), messages)` once a `NormNetwork` +wrapper type lands; until then this is the canonical entry point for BP on the norm +network. +""" +function beliefpropagation_normnetwork(tn, messages; kwargs...) + norm_tn, linknames_map = normnetwork(tn) + + # Adapt input messages onto the norm network: rename each operator's domain (bra) + # axes to the bra names `linknames_map` chose, paired via the operator's own + # codomain → domain bijection. + es = collect(keys(messages)) + raws = map(es) do e + msg, ket_to_bra = messages[e], linknames_map[e] + bra_rename = Dict( + cur => ket_to_bra[kn] for + (kn, cur) in zip(codomainnames(msg), domainnames(msg)) + ) + return replacedimnames(n -> get(bra_rename, n, n), state(msg)) + end + raw_messages = Dict(es .=> raws) + + cache = beliefpropagation(norm_tn, raw_messages; kwargs...) + + # Re-wrap each converged message as an operator with codomain = ket names and + # domain = paired bra names from the map. + return messagecache(keys(cache)) do e + return operator(cache[e], keys(linknames_map[e]), values(linknames_map[e])) + end +end diff --git a/src/select_algorithm.jl b/src/select_algorithm.jl new file mode 100644 index 0000000..35ec885 --- /dev/null +++ b/src/select_algorithm.jl @@ -0,0 +1,57 @@ +# MAK-style algorithm selection helpers (cf. `MatrixAlgebraKit.select_algorithm` +# / `default_algorithm`), but with selection-relevant inputs packed into an +# `args` tuple so the value and type domains stay disjoint: `(1.2,)` vs +# `Tuple{Float64}`. Strategy types subtype `AbstractAlgorithm` so the passthrough +# overload is generic. + +abstract type AbstractAlgorithm end + +function default_algorithm(f, ::Type{Args}; kwargs...) where {Args <: Tuple} + return throw(MethodError(default_algorithm, (f, Args))) +end +function default_algorithm(f, args::Tuple; kwargs...) + return default_algorithm(f, typeof(args); kwargs...) +end + +function select_algorithm(f, alg, args::Tuple; kwargs...) + return select_algorithm(f, alg, typeof(args); kwargs...) +end +function select_algorithm(f, ::Nothing, args::Tuple; kwargs...) + return default_algorithm(f, args; kwargs...) +end +function select_algorithm(f, alg::NamedTuple, args::Tuple; kwargs...) + isempty(kwargs) || throw( + ArgumentError( + "Additional keyword arguments are not allowed when `alg` is a `NamedTuple`." + ) + ) + return default_algorithm(f, args; alg...) +end +function select_algorithm(f, ::Nothing, ::Type{Args}; kwargs...) where {Args <: Tuple} + return default_algorithm(f, Args; kwargs...) +end +function select_algorithm(f, alg::NamedTuple, ::Type{Args}; kwargs...) where {Args <: Tuple} + isempty(kwargs) || throw( + ArgumentError( + "Additional keyword arguments are not allowed when `alg` is a `NamedTuple`." + ) + ) + return default_algorithm(f, Args; alg...) +end +function select_algorithm(f, alg::AbstractAlgorithm, ::Type{<:Tuple}; kwargs...) + isempty(kwargs) || throw( + ArgumentError( + "Additional keyword arguments are not allowed when `alg` is an `AbstractAlgorithm` instance." + ) + ) + return alg +end +function select_algorithm(f, alg::AbstractAlgorithm, args::Tuple; kwargs...) + return select_algorithm(f, alg, typeof(args); kwargs...) +end + +# Allocate the destination for an in-place call to `f`. Operations overload +# `initialize_output(::typeof(f), args...)` to control allocation. +function initialize_output(f, args...; kwargs...) + return throw(MethodError(initialize_output, (f, args...))) +end diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl new file mode 100644 index 0000000..bb4e2c5 --- /dev/null +++ b/src/tensoralgebra.jl @@ -0,0 +1,130 @@ +using MatrixAlgebraKit: MatrixAlgebraKit +using NamedDimsArrays: AbstractNamedDimsArray, AbstractNamedDimsOperator, codomainnames, + denamed, dimnames, domainnames, name, nameddims, operator, randname, setname, state +using Random: Random +using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, FusionStyle, bipermutedims, + blockedperm_indexin, blocks, matricize, trivialbiperm, unmatricize + +# Local stand-ins for what would eventually become upstream interface functions in +# `TensorAlgebra` / `NamedDimsArrays`. Naming: +# +# - `similar_operator(prototype, [T,] codomain)` — eventual +# `TensorAlgebra.similar_operator` / `NamedDimsArrays.similar_operator`. +# - `one_tensor(a, ...)` — eventual `TensorAlgebra.one` (paralleling `TA.svd`, +# `TA.eigen`). +# - `one_operator(op)` / `one_operator(na, codomain, domain)` — eventual methods +# of `Base.one` on `AbstractNamedDimsOperator` and `AbstractNamedDimsArray`. +# Held under the local name `one_operator` until then to avoid piracy on +# `NamedDimsArrays` types. +# - `randn_operator!([rng,] op)` / `rand_operator!([rng,] op)` — eventual methods +# of `Random.randn!` / `Random.rand!` on `AbstractNamedDimsOperator`. Held locally +# for the same piracy reason, plus to hide the workaround for the ITensor +# `eltype(::Type) === Any` issue (peeling to the concrete storage so the +# stdlib `randn!` / `rand!` sees the runtime eltype). +# - `dag`, `dual` — no-op stubs for the tensor and axis involutions. + +# Tensor-algebra interface no-op stubs. Currently identity; backends (graded sectors, +# complex tensors, etc.) will overload these for their semantics. +# +# `dag` is the involution on TENSORS (conjugate-transpose, sector-direction flip, …). +# `dual` is the involution on AXES (vector space → dual vector space). +dag(x) = x +dual(x) = x + +# Allocate a square operator with the given `codomain` named axes. Domain axes are +# derived as `dual.(codomain)` with fresh `randname`-generated names; backend / device +# inherited from `prototype` via `Base.similar`. +function similar_operator(prototype, ::Type{T}, codomain) where {T} + domain_names = randname.(name.(codomain)) + domain_axes = setname.(dual.(codomain), domain_names) + raw = similar(prototype, T, (codomain..., domain_axes...)) + return operator(raw, name.(codomain), domain_names) +end +function similar_operator(prototype, codomain) + return similar_operator(prototype, eltype(prototype), codomain) +end + +# === Identity operator/tensor: TA-style layered API === +# +# Mirrors `TensorAlgebra.svd` / `eigen`: a chain of dispatches accepting named +# operators, named arrays with codomain/domain names, raw arrays with labels, with +# biperms, with perms, or in canonical `(codomain..., domain...)` layout — all funnel +# into the canonical worker `one_tensor(style, a, ndims_codomain::Val)`, which +# matricizes the array, calls `MatrixAlgebraKit.one!` on the matrix, and unmatricizes +# back. +# +# All forms are out-of-place: `a` is treated as a shape prototype, not mutated. We +# rely on `matricize` returning a fresh non-aliasing array; a future view-returning +# `matricized` would be the lower-level building block for an in-place variant. + +# --- Named layers (local `one_operator`; would be `Base.one` upstream) --- + +function one_operator(op::AbstractNamedDimsOperator) + co, dom = codomainnames(op), domainnames(op) + return operator(one_operator(state(op), co, dom), co, dom) +end + +function one_operator(na::AbstractNamedDimsArray, codomain_names, domain_names) + raw = one_tensor(denamed(na), dimnames(na), codomain_names, domain_names) + return nameddims(raw, dimnames(na)) +end + +# --- Raw-array layers (`one_tensor`; would be `TensorAlgebra.one` upstream) --- + +# Label form: derive a biperm from per-axis labels. +function one_tensor(a::AbstractArray, labels_a, labels_codomain, labels_domain) + biperm = blockedperm_indexin(Tuple.((labels_a, labels_codomain, labels_domain))...) + return one_tensor(a, blocks(biperm)...) +end + +# Biperm form. +function one_tensor(a::AbstractArray, biperm::AbstractBlockPermutation{2}) + return one_tensor(a, blocks(biperm)...) +end + +# Explicit codomain/domain permutation form: physically permute axes into canonical +# layout, then dispatch to the canonical form. +function one_tensor( + a::AbstractArray, + perm_codomain::Tuple{Vararg{Int}}, + perm_domain::Tuple{Vararg{Int}} + ) + a_perm = bipermutedims(a, perm_codomain, perm_domain) + return one_tensor(a_perm, Val(length(perm_codomain))) +end + +# Canonical form: matricize → matrix-level identity → unmatricize. +function one_tensor(a::AbstractArray, ndims_codomain::Val) + return one_tensor(FusionStyle(a), a, ndims_codomain) +end +function one_tensor(style::FusionStyle, a::AbstractArray, ndims_codomain::Val) + a_mat = matricize(style, a, ndims_codomain) + MatrixAlgebraKit.one!(a_mat) + biperm = trivialbiperm(ndims_codomain, Val(ndims(a))) + axes_codomain, axes_domain = blocks(axes(a)[biperm]) + return unmatricize(style, a_mat, axes_codomain, axes_domain) +end + +# === Random fills for operators === +# +# Local helpers that would eventually become methods of `Random.randn!` and +# `Random.rand!` on `AbstractNamedDimsOperator`. They hide the workaround for the +# ITensor `eltype(typeof(::ITensor)) === Any` issue: a direct `randn!(op)` / `rand!(op)` +# dispatches on `Type{Any}` and fails, so we peel down to the concrete storage where +# the runtime eltype is honored. + +function randn_operator!(op::AbstractNamedDimsOperator) + return randn_operator!(Random.default_rng(), op) +end +function randn_operator!(rng::Random.AbstractRNG, op::AbstractNamedDimsOperator) + Random.randn!(rng, denamed(state(op))) + return op +end + +function rand_operator!(op::AbstractNamedDimsOperator) + return rand_operator!(Random.default_rng(), op) +end +function rand_operator!(rng::Random.AbstractRNG, op::AbstractNamedDimsOperator) + Random.rand!(rng, denamed(state(op))) + return op +end diff --git a/test/Project.toml b/test/Project.toml index 5fa41df..4f08271 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -11,11 +11,14 @@ ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7" ITensorNetworksNext = "302f2e75-49f0-4526-aef7-d8ba550cb06c" ITensorPkgSkeleton = "3d388ab1-018a-49f4-ae50-18094d5f71ea" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4" NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde" NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" +TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" @@ -36,11 +39,14 @@ Graphs = "1.13.1" ITensorBase = "0.5" ITensorNetworksNext = "0.4" ITensorPkgSkeleton = "0.3.42" +MatrixAlgebraKit = "0.6" NamedDimsArrays = "0.14, 0.15" NamedGraphs = "0.11" QuadGK = "2.11.2" +Random = "1.10" SafeTestsets = "0.1" Suppressor = "0.2.8" +TensorAlgebra = "0.9.2" TensorOperations = "5.3.1" TermInterface = "2" Test = "1.10" diff --git a/test/test_apply_operator.jl b/test/test_apply_operator.jl new file mode 100644 index 0000000..d436d03 --- /dev/null +++ b/test/test_apply_operator.jl @@ -0,0 +1,127 @@ +import Graphs +import NamedDimsArrays as NDA +import TensorAlgebra as TA +using ITensorBase: Index +using ITensorNetworksNext: TensorNetwork, apply_operator, apply_operators, + beliefpropagation_normnetwork, identity_norm_message_env, ones_norm_message_env, + rand_norm_message_env, randn_norm_message_env, similar_norm_message_env +using MatrixAlgebraKit: truncrank +using NamedDimsArrays: name, operator, randname, setname +using NamedGraphs.GraphsExtensions: incident_edges +using NamedGraphs.NamedGraphGenerators: named_path_graph +using Test: @test, @testset + +# The helpers below are written against the `NamedDimsArrays` interface (named +# axes, `randname`, `operator`, `randn`), so the array type is determined by the +# axes passed in. Here we use ITensor `Index`es. + +# Random tensor network on `g`: one named site axis per vertex (`site_axes`) and +# one named link axis per edge (`link_axes`). +function random_tensornetwork(g, link_axes, site_axes) + link_axis(e) = haskey(link_axes, e) ? link_axes[e] : link_axes[reverse(e)] + return TensorNetwork(g) do v + return randn((site_axes[v], (link_axis(e) for e in incident_edges(g, v))...)) + end +end + +# Random operator acting on `domain_namedaxes`, mapping them to fresh codomain +# names so that `apply` leaves the acted-on dimension names unchanged. The fresh +# names come from `randname` on the dimension *names* (not the axes), which is +# collision-free. +function randn_operator(domain_namedaxes) + codomain_namedaxes = setname.(domain_namedaxes, randname.(name.(domain_namedaxes))) + data = randn((codomain_namedaxes..., domain_namedaxes...)) + return operator(data, name.(codomain_namedaxes), name.(domain_namedaxes)) +end + +@testset "apply_operator on a path graph" begin + N, χ, d = 4, 4, 2 + g = named_path_graph(N) + + # `@testset` reseeds the global RNG on entry to every (nested) testset, so we + # build the network, environment, and gates inside each one. That keeps the + # link `Index`es as the first draws from each testset's RNG stream, so every + # later `randname` — the gate codomains here, and the rank names created + # inside the gate application — stays distinct from the link names. + @testset "untruncated gates are exact (gauge-invariant)" begin + link_axes = Dict(e => Index(χ) for e in Graphs.edges(g)) + site_axes = Dict(v => Index(d) for v in Graphs.vertices(g)) + state = random_tensornetwork(g, link_axes, site_axes) + env = beliefpropagation_normnetwork( + state, ones_norm_message_env(state); + stopping_criterion = (; maxiter = 100, tol = 1.0e-13) + ) + # Without truncation the gate is applied exactly, so the gated network + # reproduces exact contraction regardless of the gauge. + for gate in ( + randn_operator((site_axes[2],)), + randn_operator((site_axes[2], site_axes[3])), + ) + gated, _ = apply_operator(gate, state, env) + @test prod(gated) ≈ NDA.apply(gate, prod(state)) + end + end + + @testset "truncated 2-site gate matches global optimal SVD (rank $k)" for k in 1:3 + link_axes = Dict(e => Index(χ) for e in Graphs.edges(g)) + site_axes = Dict(v => Index(d) for v in Graphs.vertices(g)) + state = random_tensornetwork(g, link_axes, site_axes) + env = beliefpropagation_normnetwork( + state, ones_norm_message_env(state); + stopping_criterion = (; maxiter = 100, tol = 1.0e-13) + ) + gate = randn_operator((site_axes[2], site_axes[3])) + # Exact oracle: gate the fully contracted state, then take the globally + # optimal rank-`k` SVD truncation across the 2 | 3 cut. + gated_full = NDA.apply(gate, prod(state)) + left = [name(site_axes[v]) for v in 1:2] + U, S, Vt = TA.svd(gated_full, left; trunc = truncrank(k)) + gated, _ = apply_operator(gate, state, env; trunc = truncrank(k)) + @test prod(gated) ≈ U * S * Vt + end + + @testset "apply_operators applies a sequence" begin + link_axes = Dict(e => Index(χ) for e in Graphs.edges(g)) + site_axes = Dict(v => Index(d) for v in Graphs.vertices(g)) + state = random_tensornetwork(g, link_axes, site_axes) + env = beliefpropagation_normnetwork( + state, ones_norm_message_env(state); + stopping_criterion = (; maxiter = 100, tol = 1.0e-13) + ) + # Gates on neighboring edges sharing site 3, applied in sequence. + g1 = randn_operator((site_axes[2], site_axes[3])) + g2 = randn_operator((site_axes[3], site_axes[4])) + gated, _ = apply_operators([g1, g2], state, env) + @test prod(gated) ≈ NDA.apply(g2, NDA.apply(g1, prod(state))) + end + + @testset "norm-message-env constructors" begin + link_axes = Dict(e => Index(χ) for e in Graphs.edges(g)) + site_axes = Dict(v => Index(d) for v in Graphs.vertices(g)) + state = random_tensornetwork(g, link_axes, site_axes) + + # All three constructors build a `MessageCache` with two directed edges per + # undirected edge of the state. + n_directed = 2 * length(collect(Graphs.edges(g))) + for ctor in ( + similar_norm_message_env, identity_norm_message_env, + ones_norm_message_env, randn_norm_message_env, + rand_norm_message_env, + ) + cache = ctor(state) + @test length(collect(Graphs.edges(cache))) == n_directed + end + + # Identity env reproduces the gauge-invariant exact-gate property: an + # untruncated gate gives the exact result regardless of which valid env we + # gauge against. + env = identity_norm_message_env(state) + for gate in ( + randn_operator((site_axes[2],)), + randn_operator((site_axes[2], site_axes[3])), + ) + gated, _ = apply_operator(gate, state, env) + @test prod(gated) ≈ NDA.apply(gate, prod(state)) + end + end +end