Feedback on how to make this bit of code more Julian

Hello there,

I am relatively new to Julia but I have developed this piece of code (I haven’t posted the full code yet) that is working as intended. However, I’d really appreciate feedbacks on how to improve it and make it more Julian. The main question is that I pass around a struct and various functions modify it and/or create new struct members. As far as I’ve understood, you don’t really create classes in Julia but here this concept might be adequate as in the end, I would like to bind it to Python.

module mps_encoder

using BitBasis
using Distributions
using ITensors
using LinearAlgebra
using Plots
using Polynomials


struct MPSEncoder
    n_qubits::Int
    degree::Int
    distr::UnivariateDistribution
end

function construct_site_inds(mps_encoder::MPSEncoder)
    siteinds("S=1/2", mps_encoder.n_qubits)
end

function quadratic_coefficients(mps_encoder::MPSEncoder)
    # Construct the x uniform grid ∈ [0,1] 
    x = range(0, 1, length = 2^mps_encoder.n_qubits + 1)
    y = pdf(mps_encoder.distr, x)

    coefficients = Vector{Polynomial}()

    for odd ∈ [x for x ∈ 1:2^mps_encoder.n_qubits if isodd(x)]
        quad_fit = Polynomials.fit(x[odd:odd+2], y[odd:odd+2], 2)
        push!(coefficients, quad_fit)
    end

    return coefficients

end

let

    n_qubits::Int = 3
    poly_degree::Int = 2
    distr::UnivariateDistribution = Normal(0.5, 0.1)

    mps_encoder::MPSEncoder = MPSEncoder(n_qubits, poly_degree, distr)

    sites_mps = construct_site_inds(mps_encoder)

    @show sites_mps

    coefficients = quadratic_coefficients(mps_encoder)

    @show coefficients

    nothing
end # let 

Any help appreciated ! Thanks a lot.

Roland

Vector{Polynomial}() is an abstract type (it can hold any type of polynomial). You should use a concrete type based on the specific type of coefficient that you need, which here is based on the element type of x:

coefficients = Vector{Polynomial{eltype{x},:x}()

No need to allocate an array with the [...] comprehension here. Just put the if statement in your loop, or better yet just loop over the odd elements directly:

for odd ∈ 1:2:2^mps_encoder.n_qubits
    @views quad_fit = Polynomials.fit(x[odd:odd+2], y[odd:odd+2], 2)
    push!(coefficients, quad_fit)
end

(Using @views to avoid allocating temporary arrays for the slices.)

Or, more simply:

@views coefficients = [Polynomials.fit(x[odd:odd+2], y[odd:odd+2], 2) for odd ∈ 1:2:2^mps_encoder.n_qubits]
1 Like

These type declarations accomplish nothing and are usually omitted.

1 Like

UnivariateDistribution is an abstract type. In performance-sensitive code, you want to use concrete types in structs. The way to do this while preserving flexibility is to define a family of struct types parameterized by a type parameter:

struct MPSEncoder{T<:UnivariateDistribution}
    n_qubits::Int
    degree::Int
    distr::T
end

and you can dod MPSEncoder(3,4,Normal(0.5, 0.1)) and it will instantiate the appropriate subtype.

2 Likes

Thanks @stevengj . It looks that I wasn’t too far away after all (playing the optimistic). What about the next part that is more convoluted I’m afraid (I know there is a lot of unecessary commented lines but I’ll tidy things up later promise):

function construct_mps_indices(degree::Int, N::Int)

    kappa = degree + 1
    d = 2

    alphas = Dict()
    s = Vector{Index}()  # Better splattable structure for bitstring creation.

    alphas[0] = Index(kappa)
    for i ∈ 1:N
        alphas[i] = Index(kappa)
        idx = Index(d)
        push!(s, idx)
    end
    alphas[N+1] = Index(kappa)

    return alphas, s

end

function piecewise_polynomial_approx(degree::Int, N::Int, coefficients, alphas, s)

    kappa = degree + 1
    d = 2

    Q = Dict()

    # Initial tensor
    Q[0] = ITensor(alphas[0])
    Q[0][alphas[0]=>1] = 1.0

    # Define and fill core tensors with ones on the diagonals.
    for i ∈ 1:N
        # push!(s, Index(d))
        # @show i
        Q[i] = ITensor(alphas[i-1], alphas[i], s[i])
        for j ∈ 1:d
            for k ∈ 1:kappa
                Q[i][alphas[i-1]=>k, alphas[i]=>k, s[i]=>j] = 1.0
                # @show k
            end
        end
    end

    # Fill the rest of the tensors elements.
    for i ∈ 1:N
        for j ∈ 1:kappa
            for k ∈ j+1:kappa
                Q[i][alphas[i-1]=>j, alphas[i]=>k, s[i]=>2] =
                    binomial(k - 1, j - 1) * 2.0^(-(k - j) * i)
            end
        end
    end

    # Final tensor filled with polynomial coefficients.
    Q[N+1] = ITensor(alphas[N])
    for i ∈ 1:kappa
        Q[N+1][alphas[N]=>i] = coefficients[i-1]
    end

    # Construct the MPS network.
    MPS_TT = Q[0]
    for mps_index ∈ 1:N+1
        MPS_TT *= Q[mps_index]
    end

    # Return the dict of indices necessary to build the
    # input bitstring.
    return MPS_TT, Q

end


function construct_mps_for_pp_approx(
    poly_degree::Int,
    nb_qubits::Int,
    coefficients,
    s_mps,
    subdomain = nothing,
)

    domain_inds = fill([1 2], nb_qubits)

    if !isnothing(subdomain)
        for i ∈ 1:length(subdomain)
            domain_inds[i] = fill(subdomain[i], 2)'
        end
    end

    # @show domain_inds

    alphas, s = construct_mps_indices(poly_degree, nb_qubits)
    MPS_TT, Q = piecewise_polynomial_approx(poly_degree, nb_qubits, coefficients, alphas, s)

    kappa = poly_degree + 1  # Physical bond dimension
    d = 2                    # Virtual bond dimension

    ψ = emptyMPS(s_mps, kappa)

    # @show inds(ψ[1])
    # @show inds(ψ[2])
    # @show inds(ψ[3])
    # @show inds(ψ[4])
    # @show inds(ψ[5])

    # Initial and final tensors site as the resulting contraction
    MPSI = Q[0] * Q[1]
    MPSF = Q[nb_qubits] * Q[nb_qubits+1]

    # @show MPSI
    # @show MPSF

    # Elementwise assignment to preserve indexing
    for i ∈ 1:kappa
        for j ∈ domain_inds[1]
            ψ[1][CartesianIndex(i, j)] = MPSI[CartesianIndex(i, j)]
        end
    end

    for i ∈ 2:nb_qubits-1
        for j ∈ domain_inds[i]
            for k ∈ 1:kappa
                for l ∈ 1:kappa
                    ψ[i][CartesianIndex(k, j, l)] = Q[i][CartesianIndex(k, l, j)]
                end
            end
        end
    end

    for i ∈ 1:kappa
        for j ∈ domain_inds[nb_qubits]
            ψ[nb_qubits][CartesianIndex(i, j)] = MPSF[CartesianIndex(i, j)]
        end
    end

    return ψ

end



    ψ = Vector()

    for (index, coeffs) in enumerate(coefficients)
         subdomain_bitstr = create_bit_array(index - 1, 2)
         mps = construct_mps_for_pp_approx(
            poly_degree,
             n_qubits,
             coeffs,
             sites_mps,
             subdomain_bitstr,
        )
         push!(ψ, mps)
     end

OK cool. I was more or less suspecting some type checking here.

OK ! Got it thanks !

Same as above. This is an abstractly typed container, which idiomatic Julia code usually avoids (if you care about performance), as explained in the performance tips. Read the performance tips.

In this case again, the easiest solution is a comprehension:

 ψ = [ construct_mps_for_pp_approx(
            poly_degree,
             n_qubits,
             coeffs,
             sites_mps,
             create_bit_array(index - 1, 2),
        ) for (index, coeffs) in enumerate(coefficients) ]

Yes, sure thing. Got it.

Why are you using a Dict here? It looks like there is an entry Q[i] for every i ∈ 1:N, in which case a Vector would be more natural?