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:.

2

I think this sounds like exactly the right way to go about it!

Is the general goal also, that your ProcessTensors can be used in place of ITensors MPS / MPOs? Then extending the ITensor interfaces to your own types is exactly the way to make generic code work on both ITensors MPSs / MPOs and your ProcessTensors.

Two smaller notes:

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

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

is maybe a bit confusing, because the names AbstractMPS and MPS are already taken in ITensors. how about AbstractProcessMPS and ProcessMPS?

And if you have a lot of code that just forwards to applying the same function to the core you can reduce the boilerplate using the @eval macro:

for f in [:siteinds, :linkinds]
    @eval $f(s::AbstractMPS, args...; kwargs...) = $f(s.core, args...; kwargs...)
end
3

Thank you very much for your comments @jlbosse!

Yes. I want my process tensors package to be used without manually importing ITensorMPS MPS, and MPOs, and also have my own features on top of it.

Good that you brought this up. Initially, I was playing with

abstract type AbstractMyMPS{S <: AbstractSpace} <: end

struct MyMPS{S <: AbstractSpace, C} <: AbstractMyMPS{S}
    core::MPS
    combiners::C
end

and avoiding the confusion of redundant names. But then later, when I started using my own package to define some problem, I realised they needed to be changed to a more appropriate name :sweat_smile:. ProcessMPS and ProcessMPO are not a good choice because the word “process” is an established term in the open quantum systems community. So a process has some inherent properties that ordinary MPS/MPOs do not. So it might confuse users when they define a ProcessMPS(s, ["Dn"]) when they simply want to define an ordinary MPS that has a down spin at one site.

I thought about some good names for the MPS and MPO that I could define and ship, but I felt keeping it the same as ITensorMPS would keep the user experience seamless in adapting to my API. The downside is, of course, that I have some slightly confusing definitions in my source code :smiling_face_with_tear:.