API design help: Wrapping ITensorMPS structs and extending it to my package

General Usage
, , ,
1

Hi Julia community!

I am currently developing a package, ProcessTensors.jl, that uses an MPS/MPO architecture to construct process tensors and simulate non-Markovian quantum dynamics. I built this package on top of the features offered in ITensors.jl and ITensorMPS.jl. Since I am a big fan of these two packages and find their user interface appealing (thank you, @mtfishman!), I wanted to have the same API experience in my package as well. However, I also need to add some more abstraction on top of whatever is currently offered in ITensors so that they can be used by my package to handle problems tailored to non-Markovianity and open-quantum systems.

In particular, I have to distinguish between Hilbert space (computational basis space) and Liouville space (operator basis space) and store metadata that allows us to move between the two spaces. This made me wrap the existing MPS/MPO structs in my own package’s abstractions:

abstract type AbstractMPS{S <: AbstractSpace} <:
    ITensorMPS.AbstractMPS end

struct MPS{S <: AbstractSpace, C} <: AbstractMPS{S}
    core::ITensorMPS.MPS
    combiners::C
end

Here, for example, MPS{Hilbert} represents an ordinary state, whereas MPS{Liouville} represents a vectorized density operator and stores the combiners needed to reconstruct its Hilbert-space MPO. Along with this, I also wanted the existing functions of ITensorMPS to be available for users to be used the same way in my package as you would in the former. So I had to export functions such as siteinds, linkinds, OpSum, apply, and tdvp in my package so that users can use them without worrying about the wrappers I introduce in my package.

After having some conversation with ChatGPT, I went with a strategy which, at that time, felt solid:

  1. Have certain functions import and re-export directly from ITensorMPS. Functions like OpSum and op fall under this category.

    using ProcessTensors
import ITensorMPS

ProcessTensors.OpSum === ITensorMPS.OpSum  # true

  2. Have inspection functions such as siteinds, linkinds extend the same ITensorMPS generic and delegate to the wrapped core.

    import ITensorMPS: siteinds, linkinds

siteinds(m::AbstractMPS; kwargs...) =
    siteinds(m.core; kwargs...)

linkinds(m::AbstractMPS; kwargs...) =
    linkinds(m.core; kwargs...)

    Consequently, these two calls reach different methods:

    using ProcessTensors

s = siteinds("S=1/2", 3)
psi = MPS(s, fill("Dn",3)) # 3-element MPS{Hilbert}

@which siteinds(psi)       # ProcessTensors wrapper method
@which siteinds(psi.core)  # native ITensorMPS method

  3. Finally, for functions that return an MPS/MPO object, I rewrap the same object in my package’s convention after forwarding the calculations to ITensorMPS. For example, in tdvp function that returns a time-evolved MPS, given an input and a Hamiltonian,

    import ITensorMPS: tdvp

tdvp(H::AbstractMPO, t::Number, ψ::AbstractMPS; kwargs...) =
    _rewrap(ψ, tdvp(H.core, t, ψ.core; kwargs...))

This gives users the familiar ITensorMPS function names while allowing operations on MPS{Hilbert} and MPS{Liouville}.

Now, I am planning to register the package in General, and I wish to keep it such that I (and others willing to contribute) can develop it comfortably in the future. So I want to hear from the experts here on whether this is a robust way of extending an existing software package’s features and building upon it to customize it to my specific set of problems, while keeping the existing package’s features intact while doing so. I am not aware of any better way of doing this, so if you guys have any comments or instances where this could go wrong, and a workaround for that, I’d be grateful for receiving your advice/suggestions :folded_hands:t2:.