diff --git a/README.md b/README.md index d82d1c0..123731c 100644 --- a/README.md +++ b/README.md @@ -189,6 +189,11 @@ aff = affine_operator( ones(Int64, 3); bc=[Open, Periodic, Periodic], ) +pull = affine_pullback_operator( + 3, + AffineParams([1 0; 1 1], [0, 0]); + bc=[Periodic, Periodic], +) ``` ### Interpolation Modules diff --git a/deps/build.jl b/deps/build.jl index 862c66d..b4314c0 100644 --- a/deps/build.jl +++ b/deps/build.jl @@ -12,7 +12,7 @@ using Libdl using RustToolChain: cargo # Configuration -const TENSOR4ALL_RS_FALLBACK_COMMIT = "44ddedb3ff801f80c2f8d1609bf43eb28081a9f1" +const TENSOR4ALL_RS_FALLBACK_COMMIT = "4ee57fee0a71d385576c11d42850304548c6949d" const TENSOR4ALL_RS_REPO = "https://github.com/tensor4all/tensor4all-rs.git" # Paths diff --git a/src/C_API.jl b/src/C_API.jl index d51c4d1..1dbbb8e 100644 --- a/src/C_API.jl +++ b/src/C_API.jl @@ -2326,6 +2326,32 @@ function t4a_qtransform_affine( ) end +""" + t4a_qtransform_affine_pullback(r, m, n, a_num, a_den, b_num, b_den, bc, out) -> Cint + +Create an affine pullback operator: `f(y) = g(A*y + b)`. +The affine matrix is MxN in column-major order, with rational entries encoded by +numerator and denominator arrays. +""" +function t4a_qtransform_affine_pullback( + r::Csize_t, + m::Csize_t, + n::Csize_t, + a_num, + a_den, + b_num, + b_den, + bc, + out, +) + return ccall( + _sym(:t4a_qtransform_affine_pullback), + Cint, + (Csize_t, Csize_t, Csize_t, Ptr{Int64}, Ptr{Int64}, Ptr{Int64}, Ptr{Int64}, Ptr{Cint}, Ptr{Ptr{Cvoid}}), + r, m, n, a_num, a_den, b_num, b_den, bc, out, + ) +end + """ t4a_qtransform_binaryop(r, a1, b1, a2, b2, bc1, bc2, out) -> Cint diff --git a/src/QuanticsTransform.jl b/src/QuanticsTransform.jl index 36c3ff1..1ff58ab 100644 --- a/src/QuanticsTransform.jl +++ b/src/QuanticsTransform.jl @@ -21,9 +21,10 @@ using ..C_API using ..TreeTN: TreeTensorNetwork export LinearOperator +export AffineParams export shift_operator, flip_operator, phase_rotation_operator, cumsum_operator, fourier_operator export shift_operator_multivar, flip_operator_multivar, phase_rotation_operator_multivar -export affine_operator, binaryop_operator +export affine_operator, affine_pullback_operator, binaryop_operator export apply export set_input_space!, set_output_space!, set_iospaces! export BoundaryCondition, Periodic, Open @@ -76,6 +77,32 @@ mutable struct LinearOperator end end +_as_rational64(value::Rational) = Rational{Int64}(Int64(numerator(value)), Int64(denominator(value))) +_as_rational64(value::Integer) = Rational{Int64}(Int64(value), 1) + +""" + AffineParams(a, b) + +Affine pullback parameters representing `f(y) = g(A*y + b)`. + +- `a`: source-dimension by output-dimension affine matrix +- `b`: source-dimension shift vector +""" +struct AffineParams + a::Matrix{Rational{Int64}} + b::Vector{Rational{Int64}} + + function AffineParams(a::AbstractMatrix, b::AbstractVector) + size(a, 1) == length(b) || error("Affine matrix row count must match shift length") + a_rat = _as_rational64.(a) + b_rat = _as_rational64.(b) + return new(reshape(a_rat, size(a)), collect(b_rat)) + end +end + +source_ndims(params::AffineParams) = size(params.a, 1) +output_ndims(params::AffineParams) = size(params.a, 2) + # ============================================================================ # Operator construction functions # ============================================================================ @@ -206,6 +233,40 @@ function affine_operator(r::Integer, return LinearOperator(out[]) end +""" + affine_pullback_operator(r::Integer, params::AffineParams; + bc=fill(Periodic, source_ndims(params))) -> LinearOperator + +Create an affine pullback operator implementing `f(y) = g(A*y + b)`. + +The input state has `source_ndims(params)` variables and the output state has +`output_ndims(params)` variables. Boundary conditions apply to the transformed +source coordinates `A*y + b`. +""" +function affine_pullback_operator( + r::Integer, + params::AffineParams; + bc::AbstractVector{<:BoundaryCondition}=fill(Periodic, source_ndims(params)), +) + length(bc) == source_ndims(params) || + error("Boundary condition length must match source dimension") + + out = Ref{Ptr{Cvoid}}(C_NULL) + status = C_API.t4a_qtransform_affine_pullback( + Csize_t(r), + Csize_t(source_ndims(params)), + Csize_t(output_ndims(params)), + Int64[numerator(value) for value in vec(params.a)], + Int64[denominator(value) for value in vec(params.a)], + Int64[numerator(value) for value in params.b], + Int64[denominator(value) for value in params.b], + _bc_array_cint(bc), + out, + ) + C_API.check_status(status) + return LinearOperator(out[]) +end + """ binaryop_operator(r::Integer, a1::Integer, b1::Integer, a2::Integer, b2::Integer; bc1=Periodic, bc2=Periodic) -> LinearOperator diff --git a/test/test_build_script.jl b/test/test_build_script.jl index eaaad1c..1e364ca 100644 --- a/test/test_build_script.jl +++ b/test/test_build_script.jl @@ -3,7 +3,7 @@ using Test @testset "build.jl" begin script = read(joinpath(dirname(@__DIR__), "deps", "build.jl"), String) - @test occursin("const TENSOR4ALL_RS_FALLBACK_COMMIT = \"44ddedb3ff801f80c2f8d1609bf43eb28081a9f1\"", script) + @test occursin("const TENSOR4ALL_RS_FALLBACK_COMMIT = \"4ee57fee0a71d385576c11d42850304548c6949d\"", script) @test occursin("checkout --detach", script) @test !occursin("--branch", script) end diff --git a/test/test_quanticstransform.jl b/test/test_quanticstransform.jl index ed83d8a..7778120 100644 --- a/test/test_quanticstransform.jl +++ b/test/test_quanticstransform.jl @@ -2,10 +2,12 @@ using Test using Tensor4all using Tensor4all: dim using Tensor4all.SimpleTT: SimpleTensorTrain -using Tensor4all.TreeTN: MPS, siteinds +using Tensor4all.TreeTN: MPS, siteinds, to_dense using Tensor4all.QuanticsTransform: + AffineParams, LinearOperator, affine_operator, + affine_pullback_operator, apply, binaryop_operator, flip_operator_multivar, @@ -16,6 +18,86 @@ using Tensor4all.QuanticsTransform: const CAPI = Tensor4all.C_API +function _product_mps(vectors::Vector{<:AbstractVector{ComplexF64}}) + nsites = length(vectors) + site_inds = [Tensor4all.Index(length(vector)) for vector in vectors] + arrays = Array{ComplexF64}[] + + for (site, vector) in enumerate(vectors) + if nsites == 1 + push!(arrays, reshape(collect(vector), length(vector))) + elseif site == 1 + push!(arrays, reshape(collect(vector), length(vector), 1)) + elseif site == nsites + push!(arrays, reshape(collect(vector), 1, length(vector))) + else + push!(arrays, reshape(collect(vector), 1, length(vector), 1)) + end + end + + return MPS(arrays, site_inds) +end + +_dense_state(mps) = ComplexF64.(Tensor4all.data(to_dense(mps))) + +function _decode_coords(level_digits::Vector{Int}, nvars::Int) + r = length(level_digits) + coords = zeros(Int, nvars) + for var in 1:nvars + value = 0 + for digit in level_digits + value = (value << 1) | ((digit >> (var - 1)) & 1) + end + coords[var] = value + end + return coords +end + +function _encode_coords(coords::Vector{Int}, r::Int) + digits = zeros(Int, r) + for level in 1:r + bitpos = r - level + digit = 0 + for (var, coord) in enumerate(coords) + digit |= ((coord >> bitpos) & 1) << (var - 1) + end + digits[level] = digit + end + return digits +end + +function _expected_affine_pullback(source_dense, a::AbstractMatrix{<:Integer}, + b::AbstractVector{<:Integer}, bc::Vector) + source_ndims, output_ndims = size(a) + r = ndims(source_dense) + output_dim = 1 << output_ndims + source_size = 1 << r + expected = zeros(ComplexF64, ntuple(_ -> output_dim, r)) + + for output_index in CartesianIndices(expected) + output_digits = Int[index - 1 for index in Tuple(output_index)] + output_coords = _decode_coords(output_digits, output_ndims) + source_coords = vec(a * output_coords .+ b) + + valid = true + for i in eachindex(source_coords) + if bc[i] == Tensor4all.QuanticsTransform.Periodic + source_coords[i] = mod(source_coords[i], source_size) + elseif source_coords[i] < 0 || source_coords[i] >= source_size + valid = false + break + end + end + + if valid + source_digits = _encode_coords(source_coords, r) + expected[output_index] = source_dense[CartesianIndex((source_digits .+ 1)...)] + end + end + + return expected +end + @testset "QuanticsTransform C API bindings" begin @testset "multivar constructors" begin out = Ref{Ptr{Cvoid}}(C_NULL) @@ -56,6 +138,22 @@ const CAPI = Tensor4all.C_API @test out[] != C_NULL CAPI.t4a_linop_release(out[]) + out[] = C_NULL + status = CAPI.t4a_qtransform_affine_pullback( + Csize_t(4), + Csize_t(1), + Csize_t(2), + Int64[1, 0], + Int64[1, 1], + Int64[0], + Int64[1], + Cint[1], + out, + ) + @test status == 0 + @test out[] != C_NULL + CAPI.t4a_linop_release(out[]) + out[] = C_NULL status = CAPI.t4a_qtransform_binaryop( Csize_t(4), Int8(1), Int8(1), Int8(1), Int8(-1), Cint(1), Cint(0), out) @@ -80,6 +178,7 @@ const CAPI = Tensor4all.C_API @test flip_operator_multivar(3, 2, 1; bc=Tensor4all.QuanticsTransform.Open) isa LinearOperator @test phase_rotation_operator_multivar(3, pi / 4, 2, 1) isa LinearOperator @test binaryop_operator(3, 1, 1, 1, -1) isa LinearOperator + @test affine_pullback_operator(2, AffineParams([1 0; 0 1], [0, 0])) isa LinearOperator end @testset "affine wrapper supports explicit output space" begin @@ -107,4 +206,77 @@ const CAPI = Tensor4all.C_API @test result isa Tensor4all.TreeTN.TreeTensorNetwork @test dim(siteinds(result, 1)[1]) == 8 end + + @testset "affine pullback semantics" begin + @testset "identity" begin + a = Int64[1 0; 0 1] + b = Int64[0, 0] + bc = [Tensor4all.QuanticsTransform.Periodic, Tensor4all.QuanticsTransform.Periodic] + + op = affine_pullback_operator(2, AffineParams(a, b); bc=bc) + state = _product_mps([ + ComplexF64[1, 2, 3, 4], + ComplexF64[5, 6, 7, 8], + ]) + set_iospaces!(op, state) + result = apply(op, state) + + @test _dense_state(result) ≈ _dense_state(state) + end + + @testset "2d shear" begin + a = Int64[1 0; 1 1] + b = Int64[0, 0] + bc = [Tensor4all.QuanticsTransform.Periodic, Tensor4all.QuanticsTransform.Periodic] + + op = affine_pullback_operator(2, AffineParams(a, b); bc=bc) + state = _product_mps([ + ComplexF64[1, 3, 5, 7], + ComplexF64[2, 4, 6, 8], + ]) + set_iospaces!(op, state) + result = apply(op, state) + + source_dense = _dense_state(state) + expected = _expected_affine_pullback(source_dense, a, b, bc) + @test _dense_state(result) ≈ expected + end + + @testset "embedding" begin + a = reshape(Int64[1, 0], 1, 2) + b = Int64[0] + bc = [Tensor4all.QuanticsTransform.Open] + + op = affine_pullback_operator(2, AffineParams(a, b); bc=bc) + input_state = _product_mps([ + ComplexF64[1, 2], + ComplexF64[3, 4], + ]) + output_state = MPS(SimpleTensorTrain(fill(4, 2), 0.0)) + set_iospaces!(op, input_state, output_state) + result = apply(op, input_state) + + source_dense = _dense_state(input_state) + expected = _expected_affine_pullback(source_dense, a, b, bc) + @test _dense_state(result) ≈ expected + end + + @testset "open shift" begin + a = reshape(Int64[1], 1, 1) + b = Int64[1] + bc = [Tensor4all.QuanticsTransform.Open] + + op = affine_pullback_operator(2, AffineParams(a, b); bc=bc) + state = _product_mps([ + ComplexF64[1, 2], + ComplexF64[3, 4], + ]) + set_iospaces!(op, state) + result = apply(op, state) + + source_dense = _dense_state(state) + expected = _expected_affine_pullback(source_dense, a, b, bc) + @test _dense_state(result) ≈ expected + end + end end