Add exact affine pullback operator, fix SimpleTT buffer bug, and add tests

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

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

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
 

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

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

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