Use AbstractImageReconstruction as interface package

[nHackel]

Hello everyone! This PR adds AbstractImageReconstruction (AIR) as a backend for MRIReco which I've mentioned in some MRIReco discussions and the Julia health slack. Unfortuantely, this is quite a large PR and while I've tried my best to make the migration transparent to users, it does change how reconstructions are implemented. So we should only go ahead if we see a clear benefit to this approach and people who are interested in developing new or extending existing algorithms feel capable of doing so. It would be great if some of you could spend some time playing around with this setup and give me some feedback @cncastillo, @aTrotier, @atsanda, @tknopp (or if you have questions to the wall of text that is about to follow and which I can hopefuly clarify in the comments)

I have not yet touched the Documenter docs, since that is an amount of effort I only want to do if we go ahead with the new backend. But I'll provide some context in the comments.

Why did MPIReco switch away from a dict based approach

Our reco MPIReco hit some limits of a dictionary/keyword argument based interface. To give some examples (transformed into MRIReco context):

We could dispatch to different types of arguments based on the input data, imagine acqusition data for multi-coil reconstructions had a type we could distinguish from a multi-echo file:

reconstruction(::MultiCoilAcqData, recoParams) -> reconstruction_multiCoil
reconstruction(::MultiEchoAcqData, recoParams) -> reconstruction_multiEcho

In MPI this approach to algorithm selection worked fairly well up to a point. However, when we have multiple algorithms for the same data, we can't dispatch on just the type anymore. So we need to look at either the content of a file or at a user supplied keyword. Either way we used some form of a switch/if statement:

function reconstruction(acqData::AcqData, recoParams)
	method = get(recoParams, :method, "multiCoil")
	if method == "multiCoil"
		reconstruction_multiCoil(acqData, recoParams)
	elseif method == "multiEcho"
		reconstruction_multiCoilMultiEcho(acqData, recoParams)
	else
		error("Unknown method $method")
	end
end

This works but everytime we wanted to add a new variant, we need to update the if-statement. And if other researchers wanted to adapt a method, they'd have to fork the package and edit the source code. So extendabilty was always constrainted to either editing the code or having some arguments which you can overwrite with values such as encodingOps.

A second issue we faced is that our algorithms are inherently stateful. In particular, our encoding operator is not a matrix-free function, but instead we usuaslly construct it from large measurement files. That meant 50% of our reconstruction time is spent on building up the operator. This resulted in a split user-base. People who knew the internals of the package, could build the operator once and reuse it for quick reconstruction experiments, while people who relied on the most high-level user friendly interface always paid the cost. We also have a GUI for our measuring systems and since we are recording volumes multiple times per second, we also want to be able to do quick reconstructions in the GUI. In the end this meant that we hardcoded one algorithm for the quick path and had no access to all the other variants.

And lastly, we also passed kwargs... down long function chains which also partially mutated the given set of parameters (either the kwargs or the dict). It was really difficult to reason about parameters, if they were set by a user or as a default from another function, if they should be passed on after being used or dropped. We even had cases where the same parameter needed to be used twice in an algorithm, but with different values. Which is not possible in a flat dict.

Just to highlight this in MRIReco we currently have:

function reconstruction(acqData, recoParams)
	# ...
	recoParams = merge(defaultRecoParams(), recoParams)
	# ...
	setupIterativeReco!(acqData, recoParams)
end

function defaultRecoParams()
  # ...
  params[:reg] = L1Regularization(0.0)
  params[:solver] = ADMM
  # ...
end

function setupIterativeReco!(acqData, recoParams)
  reg = get(recoParams,:reg,L1Regularization(zero(T))) # this defaults is never used
  solver = get(recoParams, :solver, FISTA) # neither is this and it even differs from the "default" ADMM
  # ...
end

We also pass user-defined kwargs into function from MRIOperators and RegularizedLeastSquares with no context or validation and it's just not apparent what parameters an algorithm actually has, when they are set correctly and what effect they will have on the overall reconstruction. Overall the dicts we use for reconstruction parameters in MRIReco have "state" distributed across many functions and some state is even spread across different packages.

A great example of this is #263. The ADMM and SplitBregman solvers are a bit special in the way they handle transformation to other domain. While most solver want to get reg = TransformedRegularization(L1Regularization, WaveletOp), those two solvers want to get reg = L1Regularization, regTrafo = WaveletOp. This means if a user uses those solvers and supplies the sparseTrafo as usualy, we currently generate the arguments to the solver incorrectly. If a user supplies sparseTrafo = "Wavelet", regTrafo = WaveletOp, the penalty suddenly is in the wavelet domain of the wavelet domain (might be interesting 🙈).

AbstractImageReconstruction

To solve these issues, we switched from a pure-function based setup to a struct based setup and developed AbstractImageReconstruction (AIR). Note that AIR is not a finished package and the interface might change down the line, but so far we are happy with the features and would always strive to not lose any features. In the case of breaking interface changes, I would of course update MRIReco again.

AIR has three important concepts:

• Algorithms are responsible for the (stateful) runtime behaviour of an image reconstruction
• Parameters are responsible for processing data during an reconstruction according to their parametrization
• RecoPlans are serializable blueprints for both algorithms and parameters

Parameters are implemented as callable structs (so parameters are the active part in defining processing), while algorithms are implemented as a FIFO queue (which is not super important for MRIReco, but is relevant for measurement systems). I've recently update AIRs documentation if you want more details.

The basic idea is that we generate templates for common algorithms:

julia> MRIRecoPlan("multiCoil")
RecoPlan{MultiCoilReconstruction}
└─ parameter::RecoPlan{ThreadedIterativeMRIRecoContextParameter}
   ├─ reconSize
   ├─ scheduler
   ├─ arrayType
   └─ parameter::RecoPlan{MultiCoilIterativeParameters}
      ├─ weightingParams::RecoPlan{DensityWeightingParameters}
      ├─ encodingParams::RecoPlan{EncodingParameters}
      │  ├─ method
      │  ├─ K
      │  ├─ K_tol
      │  ├─ correctionMap
      │  ├─ kernelSize
      │  ├─ oversamplingFactor
      │  └─ toeplitz
      ├─ coilParams::RecoPlan{CoilParameters}
      │  ├─ noiseData
      │  └─ senseMaps
      └─ solverParams::RecoPlan{LeastSquaresSolverParameter}
         ├─ rho
         ├─ absTol
         ├─ normalizeReg
         ├─ relTol
         ├─ tolInner
         ├─ solver
         ├─ regParams::RecoPlan{RegularizationParameters}
         │  ├─ reg
         │  └─ sparsityParams::RecoPlan{SimpleSparsityParameters}
         │     └─ sparseTrafo
         ├─ iterationsInner
         ├─ restart
         ├─ verbose
         ├─ iterationsCG
         ├─ iterations
         └─ vary_rho

which usually have a nested structure. Within this tree of parameters, all usable parameters are discoverable, we can validate them:

@parameter constructor = false struct LeastSquaresSolverParameter{SL <: AbstractLinearSolver, R <: RegularizationParameters, T <: AbstractFloat} <: SolverParameters
  solver::Type{SL} 
  regParams::R 
  normalizeReg::AbstractRegularizationNormalization 
  iterations::Int 
  rho::T
  # ...
  
  @validate begin
    @assert iterations > 0 "iterations must be positive"
    @assert rho >= 0 "rho must be positive"
    @assert isnothing(iterationsInner) || iterationsInner > 0 "iterationsInner must be positive"
    @assert isnothing(iterationsCG) || iterationsCG > 0 "iterationsCG must be positive"
    @assert vary_rho in (:none, :balance, :PnP) "vary_rho must be :none, :balance, or :PnP"
  end
end

And we can implement small functions which are somewhat self-contained, testable and have more context than just a series of kwargs. The collective of all "parameter" functions of an specific algorithm implement the reconstruction.

For example, I've defined a parameter struct for regularization terms which contains a set of nested parameters for the sparsity transformations to be used (so our previous sparseTrafo parameter).

julia> RecoPlan(RegularizationParameters; sparsityParams = RecoPlan(SimpleSparsityParameters))
RecoPlan{RegularizationParameters}
├─ reg
└─ sparsityParams::RecoPlan{SimpleSparsityParameters}
   └─ sparseTrafo

And this parameter implements a function: (param::RegularizationParameters)(algo, solver) which returns either reg= ... or reg = ..., regTrafo = ... according to the given solver.

If we ever want to create a new variant of an existing algorithm we can just swap out for example the SparsityParameters for something else. I've already prepared a CustomSparsityOperator where someone can just directly add Operators, but they could also define MyDomainOperators in some other package with some other dependencies, which they can then plug into the existing components of MRIReco.

Changelog

• Changed reconstruction "backend" from functions with dicts to AIRs struct-based approach
• Added pre-configured templates to our existing reconstruction algorithms: direct, standard, multiCoil, multiEcho, multiCoilMultiEcho and multiCoilMultiEchoSubspace. These just contain the "structure" that corresponds to the previous algorithms
• Added a mechanism for other packages to "register" new templates with MRIReco, s.t. they can make their algorithms available via the new interface
• The reconstruction(acqData, params) function now internally uses the pre-configured templates and does some parameter pre-processing to (hopefully) achieve transparent changes
• The reconstruction_x functions are now marked as deprecated and are only reachable if a user explicitly calls them. Internally, they are unchanged
• Reg (and regTrafo) are now generated per solver and not once per function
• Reduced number of vcat(weights...) in reconstruction loop
• ADMM and SplitBregman are now getting correct sparsity proximal maps without users having to specify regTrafo. This means that reconstructions with these now apply a transformation twice instead of once, so there is a performance cost, but according to the solver docstring this is the correct way
• The loop body of the iterative solvers is conceptually split from the loop itself, this allows us to swap out providers for multi-threading and potentially offers a way for multi-processing/distributed reconstructions (see (DaggerImageReconstruction)[https://github.com/JuliaImageRecon/DaggerImageReconstruction.jl] where I've played around with this concept in a generic setting)

Implementation details

I'll try to give a bottom up overview of the new interface. All iterative reconstruction shared a very similar structure of essentially:

function reconstruction_x(acqData; kwargs...)
 # Validate dimensions
 # Prepare weights, regularization and image Ireco
 for index in indices()
   for nested in 1:numX
     E = encodingOps()
     EFull = ∘(W, E), EFullᴴEFull = normalOperator(EFull; ...)
     solv = createLinearSolver(...)
     Ireco[...] = solve!(solv, kdata)
   end
 end
 # Postprocessing
end

At the very core of this loop is the linear solver provided by RegularizedLeastSquares (which we could also switch out now). This solver is essentially independant from most reco parameters. It just needs as inputs the kdata, EFull and EFullᴴEFull. And while previously we constructed the regularization terms outside of the solver, I moved this inside the solver parameter.

This solver core is parameterized by the solver arguments of RegularizedLeastSquares, which includes the regularization terms. The regularization terms are parameterized by the "usual" reg terms of L1Regularization and by sparsity transformations. Regularization terms need for their construction a solver.

The next part of our loop are the encoding parameters. These are essentially just direct function calls to MRIOperators encodingOps_x where x depends on the algorithm in question. So here we only need to have struct listing the common keyword arguments and dispatch to the correct call based on a given algorithm type. To finish constructing EFull we also need weighting.

At the moment MRIReco supports two weighting strategies: density weighting or:

    numContr = numContrasts(acqData)
    weights = Array{S}(undef,numContr)
    for contr=1:numContr
      numNodes = size(acqData.kdata[contr],1)
      weights[contr] = S([1.0/sqrt(prod(reconSize)) for node=1:numNodes])
    end

which I've represented by two weighting parameters DensityWeightingParameters and UniformWeightingParameters.
Both of these just depend on the acqData.

The multi coil algorithms additionally also depend on coil parameters, which in our case are the senseMaps and the noiseData.

With all these things in place, we have everything to construct EFull and EFullᴴEFull and so the core of the loop is done.
That means to fully describe our loop we need encodingParams, weightingParams, solverParams and (optionally) coilParams. These describe our existing algorithms.

I went a step further and also seperated the loop-body from the loop-setup/context. While all iterative algorithms have a loop, the loop structure and the loop "inputs" differ per algorithm. For that I've created "Context" parameters, whose job it is to setup a loop and iterate over it (in parallel) and internally call our algorithm parameters with the loop index. This context is responsible for managing parallel computing. The context parameters essentially do:

# Prepare loop
Ireco, indices, extra... = ctx.parameter(algo, reconSize) 
# Loop over indices (in parallel)
for index in indices
  ctx.parameter(algo, Ireco, index, extra...)
end    
# Finalization - call algorithm parameter
return ctx.parameter(algo, Ireco)

And extra here contains things like the weights or any other information we need to prepare before we loop over everything.

There were some parameters we needed in basically every function inside this "context", which are the reconSize, the acqData itself, the arrayType that should be used (CPU vs GPU) and related the storage_type for the LinearOperators. Instead of passing these arguments to every function, I chose to use a ScopedValue to provide these to every function inside the "context core". For those who haven't used ScopedValue yet, you can imagine this working like every reconstruction having access to its very own "global" variable for those four values. All of the ScopedValue setup is also handled by the context parameters.

Overall this gives us the following type hierarchy:

AbstractMRIRecoParameters
├── AbstractIterativeRecoParameters
│   ├── AbstractStandardParameters -> StandardIterativeParameters
│   │   ├── encodingParams::AbstractMRIRecoEncodingParameters
│   │   ├── weightingParams::AbstractMRIRecoWeightingParameters
│   │   └── solverParams::LeastSquaresSolverParameter
│   ├── AbstractMultiEchoParameters -> MultiEchoIterativeParameters
│   │   ├── encodingParams::AbstractMRIRecoEncodingParameters
│   │   ├── weightingParams::AbstractMRIRecoWeightingParameters
│   │   └── solverParams::LeastSquaresSolverParameter
│   └── AbstractMultiCoilParameters -> MultiCoilIterativeParameters
│       ├── encodingParams::AbstractMRIRecoEncodingParameters
│       ├── weightingParams::AbstractMRIRecoWeightingParameters
│       ├── solverParams::LeastSquaresSolverParameter
│       └── coilParams::AbstractCoilParameters
│
├── Context Parameters (execution framework)
│   ├── SerialIterativeMRIRecoContextParameter{P}
│   │   └── parameter::AbstractIterativeRecoParameters
│   └── ThreadedIterativeMRIRecoContextParameter{P}
│       └── parameter::AbstractIterativeRecoParameters
│
├── Configuration Parameters (composable)
│   ├── AbstractMRIRecoEncodingParameters
│   │   ├── EncodingParameters
│   │   ├── CustomEncodingParameters{E}
│   │   └── SubspaceEncodingParameters{E,B}
│   │       └── inner::AbstractMRIRecoEncodingParameters
│   │
│   ├── AbstractMRIRecoWeightingParameters
│   │   ├── DensityWeightingParameters
│   │   ├── UniformWeightingParameters
│   │   └── CustomWeightedParameters{W}
│   │
│   ├── AbstractCoilParameters
│   │   └── CoilParameters
│   │
│   ├── AbstractSparsityParameters
│   │   ├── SimpleSparsityParameters
│   │   └── CustomSparsityParameters
│   │
│   ├── RegularizationParameters
│   │   └── sparsityParams::AbstractSparsityParameters
│   │
│   └── LeastSquaresSolverParameter
│       └── regParams::RegularizationParameters

I know this is a lot, so feel free to ask any questions about this setup and if you have better ideas for structuring the parameters let me know. This is what I did based on the code, but there might be more intuitive ways based on the way MRI researchers think about these things.

Next Steps

If we approve these changes, the next steps would be:

• Updating documentation with sections on the new interface + examples of AIR specific features like parameter sweeps
• While implementing the parameters I had the most conceptual issues with the EncodingOps parameter. At the moment this essentially just calls encodingOps_x from MRIOperators with kwargs.... My gut feeling at the moment is to let MRIOperators just define the operator themselves and then have MRIReco actually construct them. I'm not sure what the most suitable type-hierarchy for this would be. We can leave it as is for now until someone has some need/benefit from that
• Connect KomaMRI reco GUI with the RecoPlan concept. This can be done generically and would allow KomaMRI to support new variations of algorithms without (or with few) changes in the code
• We can also theoretically serialize fully configured reconstructions, so far I haven't setup the MRIRecoStyle that would be needed here
• Improve usability with custom-constructors. The plan/struct interface is of course a bit more involved to setup custom algorithms. As long as we still offer a pure kwarg based constructor we can also add some defaults like:

function Regularization(reg, sparseTrafo::Union{Nothing, String, Vector{String}}) 
 return Regularization(; reg = reg, sparsityParams = SimpleSparsityParameters(sparseTrafo = sparseTrafo))
end

Here we just need to see how and if people interact with reconstruction on the struct/plan level or on the highlevel dict approaches, either plan or reconstruction based.

Third-party MRI reconstruction packages could also wrap their algorithms in very simple structs in a package extension to MRIReco and become available to KomaMRI this way.

Disclaimer

I've used some AI tools while implementing these features, however I had to manually change most of the output from that and the untouched things are mostly docstrings and some of the more boilerplate-y tests. Before switchout the backend of the dict-based reconstructions, I've compared results against the new recos and results where identially with isapprox for recos which had unchanged behaviour (see ADMM + sparseTrafos for changed behaviour). I did not visually check any examples from the docs yet

[codecov]

Codecov Report

❌ Patch coverage is 0% with 90 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.77%. Comparing base (509ded2) to head (0a5b1d5).
⚠️ Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
src/Algorithms/Storage.jl	0.00%	55 Missing ⚠️
src/Algorithms/Direct.jl	0.00%	35 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #277      +/-   ##
==========================================
- Coverage   74.99%   73.77%   -1.23%     
==========================================
  Files         103      105       +2     
  Lines        5419     5509      +90     
==========================================
  Hits         4064     4064              
- Misses       1355     1445      +90     

Flag	Coverage Δ
MRIBase	50.19% <ø> (ø)
MRICoilSensitivities	93.01% <ø> (ø)
MRIFiles	78.59% <ø> (ø)
MRIOperators	40.56% <ø> (ø)
MRIReco	41.69% <0.00%> (-0.72%)	⬇️
MRISampling	0.35% <ø> (ø)
MRISimulation	63.07% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cncastillo]

I really like this! We will try it out as soon as possible!

[aTrotier]

I will give it a look and try it for subspace reconstruction for this package : https://github.com/CRMSB/PAPER_subspace_MESE

Nice work !

[nHackel]

Feel free to ping me here or in the Julia slack if you need help or stumble over something with the new interface.

I think it would be really cool if not only existing recos work but if we can also find an example where we create a new parameter which results in functionality we didnt have before