From b9bb3a584facb22bedeae9ad90ec8b8182bed84f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 13 Oct 2026 07:25:49 -0400 Subject: [PATCH] Make implicit diff use the symmetry properties of C4v --- .../optimization/implicit_differentiation.jl | 103 +++++++++++++++++- test/gradients/c4v_env_tangent.jl | 10 ++ test/testsuite/TestSuite.jl | 6 + test/testsuite/gradients/c4v_env_tangent.jl | 98 +++++++++++++++++ 4 files changed, 216 insertions(+), 1 deletion(-) create mode 100644 test/gradients/c4v_env_tangent.jl create mode 100644 test/testsuite/gradients/c4v_env_tangent.jl diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index f64b1b17e..75afeca47 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -252,6 +252,107 @@ function _rrule( return (env, info), leading_boundary_fixed_pullback end +""" + C4vEnvTangent(corner, edge) + +The independent corner and edge of a C₄ᵥ-symmetric environment, used as the vector type in +the fixed-point gradient solvers. +""" +struct C4vEnvTangent{C, T} + corner::C + edge::T +end + +function VI.scalartype(::Type{C4vEnvTangent{C, T}}) where {C, T} + return promote_type(scalartype(C), scalartype(T)) +end + +VI.zerovector(x::C4vEnvTangent, ::Type{S}) where {S <: Number} = + C4vEnvTangent(zerovector(x.corner, S), zerovector(x.edge, S)) +function VI.zerovector!(x::C4vEnvTangent) + zerovector!(x.corner) + zerovector!(x.edge) + return x +end +VI.zerovector!!(x::C4vEnvTangent) = zerovector!(x) + +VI.scale(x::C4vEnvTangent, α::Number) = C4vEnvTangent(scale(x.corner, α), scale(x.edge, α)) +function VI.scale!(x::C4vEnvTangent, α::Number) + scale!(x.corner, α) + scale!(x.edge, α) + return x +end +function VI.scale!(y::C4vEnvTangent, x::C4vEnvTangent, α::Number) + scale!(y.corner, x.corner, α) + scale!(y.edge, x.edge, α) + return y +end +VI.scale!!(x::C4vEnvTangent, α::Number) = scale!(x, α) +VI.scale!!(y::C4vEnvTangent, x::C4vEnvTangent, α::Number) = scale!(y, x, α) + +function VI.add(y::C4vEnvTangent, x::C4vEnvTangent, α::Number, β::Number) + return C4vEnvTangent(add(y.corner, x.corner, α, β), add(y.edge, x.edge, α, β)) +end +function VI.add!(y::C4vEnvTangent, x::C4vEnvTangent, α::Number, β::Number) + add!(y.corner, x.corner, α, β) + add!(y.edge, x.edge, α, β) + return y +end +VI.add!!(y::C4vEnvTangent, x::C4vEnvTangent, α::Number, β::Number) = add!(y, x, α, β) + +VI.inner(x::C4vEnvTangent, y::C4vEnvTangent) = inner((x.corner, x.edge), (y.corner, y.edge)) +VI.norm(x::C4vEnvTangent) = norm((x.corner, x.edge)) + +# C₄ᵥ-symmetric version: the environment is fully determined by a single corner and edge, +# so we only need to diff the fixed-point iteration is differentiated with respect to those. +function _rrule( + gradmode::FixedPointGradient, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + envinit::CTMRGEnv, + state, + alg::C4vCTMRG, + ) + _check_algorithm_combination(alg, gradmode) + + env, = leading_boundary(envinit, state, alg) + + alg_fixed = _set_fixed_truncation(alg) # fix spaces during differentiation + alg_gauge = _scrambling_env_gauge(alg) # select appropriate gauge-fixing algorithm + env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed) + signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env_conv, env, alg_gauge) + function gauge_fixed_iteration(A, C, E) + env′ = fix_phases( + ctmrg_iteration(InfiniteSquareNetwork(A), CTMRGEnv(C, E), alg_fixed)[1], + signs, corner_phases, edge_phases, + ) + return env′.corners[NORTHWEST, 1, 1], env′.edges[NORTH, 1, 1] + end + C, E = env.corners[NORTHWEST, 1, 1], env.edges[NORTH, 1, 1] + _, iteration_vjp = rrule_via_ad(config, gauge_fixed_iteration, state, C, E) + _, c4v_env_vjp = rrule_via_ad(config, CTMRGEnv, C, E) + ∂f∂A(x)::typeof(state) = iteration_vjp((x.corner, x.edge))[2] + function ∂f∂x(x)::typeof(C4vEnvTangent(C, E)) + _, _, ΔC, ΔE = iteration_vjp((x.corner, x.edge)) + return C4vEnvTangent(ΔC, ΔE) + end + + function leading_boundary_fixed_pullback((Δenv′, Δinfo)) + Δenv = unthunk(Δenv′) + + # accumulate adjoints from all corners & edges through the C4v env constructor + _, ΔC, ΔE = c4v_env_vjp(Δenv) + Δx = C4vEnvTangent(unthunk(ΔC), unthunk(ΔE)) + + # evaluate the geometric sum + ∂F∂env = fixedpoint_gradient(Δx, ∂f∂x, ∂f∂A, Δx, gradmode.solver_alg) + + return NoTangent(), ZeroTangent(), ∂F∂env, NoTangent() + end + + return (env, info), leading_boundary_fixed_pullback +end + @doc raw""" fixedpoint_gradient(x̆, ∂ₓf, ∂ₚf, y₀, alg) @@ -310,7 +411,7 @@ function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, y₀, alg::ManualIter dx = ∂f∂A(y) ϵ = 1.0 for i in 1:(alg.maxiter) - y′ = ∂E∂x + ∂f∂x(y) + y′ = VI.add(∂f∂x(y), ∂E∂x) dxnew = ∂f∂A(y′) ϵnew = norm(dxnew - dx) diff --git a/test/gradients/c4v_env_tangent.jl b/test/gradients/c4v_env_tangent.jl new file mode 100644 index 000000000..f3917f76c --- /dev/null +++ b/test/gradients/c4v_env_tangent.jl @@ -0,0 +1,10 @@ +using PEPSKit + +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") +using .TestSuite + +is_buildkite = get(ENV, "BUILDKITE", "false") == "true" + +if !is_buildkite + TestSuite.gradients_c4v_env_tangent(Vector) +end diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index 018a68a82..39d3131ef 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -191,6 +191,12 @@ module GradientsC4vCTMRG end using .GradientsC4vCTMRG +module GradientsC4vEnvTangent + include("gradients/c4v_env_tangent.jl") + export gradients_c4v_env_tangent +end +using .GradientsC4vEnvTangent + module GradientsCTMRG include("gradients/ctmrg_gradients.jl") export gradients_asymmetric, gradients_asymmetric_276 diff --git a/test/testsuite/gradients/c4v_env_tangent.jl b/test/testsuite/gradients/c4v_env_tangent.jl new file mode 100644 index 000000000..855d335f2 --- /dev/null +++ b/test/testsuite/gradients/c4v_env_tangent.jl @@ -0,0 +1,98 @@ +using Test +using Adapt +using Random +using PEPSKit +using PEPSKit: C4vEnvTangent +using TensorKit +using VectorInterface + +Venv = ComplexSpace(4) +Vpeps = ComplexSpace(2) + +function random_corner(AT, kind) + return if kind === :diagonal + adapt(AT, DiagonalTensorMap(randn(Float64, dim(Venv)), Venv)) + else + adapt(AT, randn(Float64, Venv ← Venv)) + end +end +random_edge(AT, T) = adapt(AT, randn(T, Venv ⊗ Vpeps ⊗ Vpeps' ← Venv)) +random_tangent(AT, kind, T) = C4vEnvTangent(random_corner(AT, kind), random_edge(AT, T)) + +parts(x::C4vEnvTangent) = (x.corner, x.edge) +issame(x::C4vEnvTangent, (C, E)) = x.corner ≈ C && x.edge ≈ E + +# this file mostly tests the helper methods, because the +# main differentiation handling is already tested by the existing +# C4v gradient tests + +function gradients_c4v_env_tangent(AT) + return @testset "C4vEnvTangent vector interface ($kind corner, $T edge) ($AT)" for (kind, T) in + ((:diagonal, ComplexF64), (:dense, Float64)) + Random.seed!(1234) + x = random_tangent(AT, kind, T) + y = random_tangent(AT, kind, T) + α, β = 0.7, -1.3 + Cx, Ex = deepcopy(parts(x)) + Cy, Ey = deepcopy(parts(y)) + + @testset "scalartype" begin + @test scalartype(x) == promote_type(Float64, T) + end + + @testset "zerovector" begin + z = zerovector(x) + @test z isa C4vEnvTangent + @test scalartype(z.corner) == scalartype(z.edge) == scalartype(x) + @test iszero(norm(z)) + zc = zerovector(x, ComplexF64) + @test scalartype(zc) == ComplexF64 + @test iszero(norm(zc)) + z′ = deepcopy(x) + @test zerovector!(z′) === z′ + @test iszero(norm(z′)) + z″ = deepcopy(x) + @test zerovector!!(z″) === z″ + @test iszero(norm(z″)) + end + + # test the VI methods + @testset "scale" begin + @test issame(scale(x, α), (scale(Cx, α), scale(Ex, α))) + @test issame(x, (Cx, Ex)) # not mutated + x′ = deepcopy(x) + @test scale!(x′, α) === x′ + @test issame(x′, (scale(Cx, α), scale(Ex, α))) + x″ = deepcopy(x) + @test scale!!(x″, α) === x″ + @test issame(x″, (scale(Cx, α), scale(Ex, α))) + z = zerovector(x) + @test scale!(z, x, α) === z + @test issame(z, (scale(Cx, α), scale(Ex, α))) + z′ = zerovector(x) + @test scale!!(z′, x, α) === z′ + @test issame(z′, (scale(Cx, α), scale(Ex, α))) + @test issame(x, (Cx, Ex)) # source not mutated + end + + @testset "add" begin + expected = (add(Cy, Cx, α, β), add(Ey, Ex, α, β)) + @test issame(add(y, x, α, β), expected) + @test issame(y, (Cy, Ey)) && issame(x, (Cx, Ex)) # not mutated + y′ = deepcopy(y) + @test add!(y′, x, α, β) === y′ + @test issame(y′, expected) + y″ = deepcopy(y) + @test add!!(y″, x, α, β) === y″ + @test issame(y″, expected) + @test issame(x, (Cx, Ex)) # source not mutated + end + + @testset "inner and norm" begin + @test inner(x, y) ≈ inner(Cx, Cy) + inner(Ex, Ey) + @test inner(x, y) ≈ conj(inner(y, x)) + @test norm(x) ≈ sqrt(norm(Cx)^2 + norm(Ex)^2) + @test norm(x)^2 ≈ real(inner(x, x)) + end + end +end