Add output support for pulseq extensions

src/sequence_io/pulseq_io/components.jl

@@ -15,13 +15,11 @@ struct PulseqComponents
     pulses:: Dict{Int, PulseqRFPulse}
     grads:: Dict{Int, AnyPulseqGradient}
     adc:: Dict{Int, PulseqADC}
-    extensions:: Dict{Int, Any}
-    PulseqComponents(shapes, pulses, grads, adc, extensions) = new(
+    PulseqComponents(shapes, pulses, grads, adc) = new(
         _convert_to_dict(shapes, PulseqShape),
         _convert_to_dict(pulses, PulseqRFPulse),
         _convert_to_dict(grads, AnyPulseqGradient),
         _convert_to_dict(adc, PulseqADC),
-        _convert_to_dict(extensions, Any),
     )
 end
 
@@ -30,21 +28,11 @@ _convert_to_dict(d::Dict{Int, T}, ::Type{T}) where {T} = d
 _convert_to_dict(vec::Vector, ::Type{T}) where {T} = Dict{Int, T}(i => v for (i, v) in enumerate(vec))
 
 
-"""
-    PulseqComponents(sequence::PulseqSequence)
-
-Indentifies and lists all the unique components in the sequence.
-"""
-function PulseqComponents(sequence::PulseqSequence)
-
-end
-
 PulseqComponents() = PulseqComponents(
     PulseqShape[],
     PulseqRFPulse[],
     AnyPulseqGradient[],
     PulseqADC[],
-    PulseqExtension[],
 )
 
 """

src/sequence_io/pulseq_io/extensions.jl

@@ -38,7 +38,7 @@ Get the name under which the given `obj` should be stored in a Pulseq extension.
 To write an object to a Pulseq file extension, 
 one needs to define both this function and [`add_extension_definition`](@ref).
 """
-get_extension_name(ext::PulseqExtension{N}) where {N} = N
+get_extension_name(::PulseqExtension{N}) where {N} = N
 
 """
     add_extension_definition!(content::Vector{String}, obj)

src/sequence_io/pulseq_io/parse_sections.jl

@@ -33,9 +33,11 @@ function gen_all_sections(seq:: PulseqSequence)
     sections = Dict{Symbol, PulseqSection}()
     sections[:version] = gen_section(seq, Val(:version))
     sections[:definitions] = gen_section(seq, Val(:definitions))
+    (section, extension_mapping) = gen_section(seq, Val(:extensions))
+    sections[:extensions] = section
 
     comp = PulseqComponents()
-    sections[:blocks] = gen_section(seq, comp, Val(:blocks))
+    sections[:blocks] = gen_section(seq, comp, Val(:blocks), extension_mapping)
     for symbol in [:rf, :gradients, :trap, :adc, :shapes]
         sections[symbol] = gen_section(comp, Val(symbol))
     end

src/sequence_io/pulseq_io/parsers/blocks.jl

@@ -25,7 +25,7 @@ function parse_section(section::PulseqSection{:blocks}; version, rf=Dict(), grad
     return res
 end
 
-function gen_section(seq::PulseqSequence, comp:: PulseqComponents, ::Val{:blocks})
+function gen_section(seq::PulseqSequence, comp:: PulseqComponents, ::Val{:blocks}, extension_mapping::Dict{Vector, Int})
     res = PulseqSection{:blocks}(String[])
 
     for (i, block) in enumerate(seq.blocks)
@@ -39,9 +39,10 @@ function gen_section(seq::PulseqSequence, comp:: PulseqComponents, ::Val{:blocks
         ]
             push!(values, add_components!(comp, search_vec, part))
         end
-        push!(values, 0)
         if length(block.ext) > 0
-            error("Cannot write extensions yet.")
+            push!(values, extension_mapping[block.ext])
+        else
+            push!(values, 0)
         end
         push!(res.content, join(string.(values), " "))
     end

src/sequence_io/pulseq_io/parsers/extensions.jl

@@ -34,10 +34,12 @@ function parse_section(section::PulseqSection{:extensions}; kwargs...)
     return Dict(key => get_extension_list(key) for key in keys(linked_list))
 end
 
-function gen_section(comp:: PulseqComponents, ::Val{:extensions})
+function gen_section(seq:: PulseqSequence, ::Val{:extensions})
     definitions = Dict{Symbol, PulseqExtensionDefinition}()
     extensions_ref_id = Dict{Any, Int}()
-    for ext_vec in comp.extensions
+
+    all_ext_vec = [block.ext for block in seq.blocks if length(block.ext) > 0]
+    for ext_vec in all_ext_vec
         for ext in ext_vec
             name = get_extension_name(ext)
             if !(name in keys(definitions))
@@ -58,17 +60,17 @@ function gen_section(comp:: PulseqComponents, ::Val{:extensions})
         end
         next_id = add_extension_vec!(ext_vec[2:end])
         ext = ext_vec[1]
-        type_id = findfirst(get_extension_name(ext), definitions_order)
+        type_id = findfirst(isequal(get_extension_name(ext)), definitions_order)
         ref_id = extensions_ref_id[ext]
         all_ids = (type_id, ref_id, next_id)
         if all_ids in labels
-            extension_vec_id[ext_vec] = findfirst(lables, all_ids)
+            extension_vec_id[ext_vec] = findfirst(isequal(all_ids), labels)
         else
             push!(labels, all_ids)
-            extension_vec_id[ext_vec] = length(lables)
+            extension_vec_id[ext_vec] = length(labels)
         end
     end
-    add_extension_definition!.(values(comp.extensions))
+    add_extension_vec!.(all_ext_vec)
 
     content = String[]
     for (id0, (id1, id2, id3)) in enumerate(labels)
@@ -83,5 +85,5 @@ function gen_section(comp:: PulseqComponents, ::Val{:extensions})
         push!(content, "")
     end
 
-    return PulseqSection{:extensions}(content)
+    return PulseqSection{:extensions}(content), extension_vec_id
 end
\ No newline at end of file

src/sequence_io/pulseq_io/pulseq_io.jl

@@ -38,7 +38,7 @@ function write_pulseq(io::IO, sequence::Types.PulseqSequence)
         error("Can only write to pulseq version 1.4 or later.")
     end
     sections = ParseSections.gen_all_sections(sequence)
-    for key in [:version, :definitions, :blocks, :rf, :gradients, :trap, :adc, :shapes]
+    for key in [:version, :definitions, :blocks, :rf, :gradients, :trap, :adc, :extensions, :shapes]
         if length(sections[key].content) == 0
             continue
         end

src/sequence_io/pulseq_io/types.jl

@@ -125,7 +125,7 @@ struct PulseqBlock
     gy :: Union{Nothing, AnyPulseqGradient}
     gz :: Union{Nothing, AnyPulseqGradient}
     adc :: Union{Nothing, PulseqADC}
-    ext :: Vector{<:Tuple{PulseqExtension, Int}}
+    ext :: Vector
 end
 
 """