Refactor gradients into GradientWaveform subtype

src/components/abstract_types.jl

+module AbstracTypes
+import ..Variables: AbstractBlock, duration, variables, effective_time
+
+"""
+Super-type for all individual components that form an MRI sequence (i.e., RF pulse, gradient waveform, or readout event).
+
+RF pulses, instant gradients, and readouts are grouped together into [`EventComponent`](@ref).
+
+These all should have a [`duration`](@ref) in addition to any other relevant [`variables`](@ref).
+"""
+abstract type BaseComponent <: AbstractBlock end
+
+"""
+Super-type for all parts of a gradient waveform.
+
+N should be 1 for a 1D gradient waveform or 3 for a 3D one.
+"""
+abstract type GradientWaveform{N} <: BaseComponent end
+
+
+"""
+Super-type for all RF pulses, instant gradients and readouts that might play out during a gradient waveform.
+
+These all have an [`effective_time`](@ref), which should quantify at what single time one can approximate the RF pulse or readout to have taken place.
+"""
+abstract type EventComponent <: BaseComponent end
+
+"""
+Super type for all RF pulses.
+"""
+abstract type RFPulseComponent <: EventComponent end
+
+"""
+Super type for all readout events.
+"""
+abstract type ReadoutComponent <: EventComponent end
+
+end
\ No newline at end of file

src/components/components.jl

+module Components
+include("abstract_types.jl")
+include("gradient_waveforms/gradient_waveforms.jl")
+
+end
\ No newline at end of file

src/gradients/changing_gradient_blocks.jl → src/components/gradient_waveforms/changing_gradient_blocks.jl

 module ChangingGradientBlocks
 import StaticArrays: SVector
-import ...Variables: VariableType, variables, get_free_variable
-import ...BuildingBlocks: GradientBlock
-import ...Variables: qvec, bmat_gradient, gradient_strength, slew_rate, duration, variables, VariableType
-import ...BuildingBlocks: GradientBlock, RFPulseBlock
+import ....Variables: VariableType, duration, qvec, bmat_gradient, gradient_strength, slew_rate, get_free_variable
+import ...AbstractTypes: GradientWaveform
+
+
+abstract type ChangingGradient{N} <: GradientWaveform{N} end
 
 """
-    ChangingGradientBlock(grad1, slew_rate, duration, rotate, scale)
+    ChangingGradient1D(grad1, slew_rate, duration)
 
-Underlying type for any linearly changing part in a gradient waveform.
+Underlying type for any linearly changing part in a 1D gradient waveform.
 
-Usually, you do not want to create this object directly, use a gradient waveform instead.
+Usually, you do not want to create this object directly, use a `BuildingBlock` instead.
+"""
+struct ChangingGradient1D <: ChangingGradient{1}
+    gradient_strength_start :: VariableType
+    slew_rate :: VariableType
+    duration :: VariableType
+end
 
-## Parameters
-- `rotate`: with which user-set parameter will this gradient be rotated (e.g., :bvec). Default is no rotation.
-- `scale`: with which user-set parameter will this gradient be scaled (e.g., :bval). Default is no scaling.
 """
-struct ChangingGradientBlock <: GradientBlock
+    ChangingGradient3D(grad1, slew_rate, duration)
+
+Underlying type for any linearly changing part in a 3D gradient waveform.
+
+Usually, you do not want to create this object directly, use a `BuildingBlock` instead.
+"""
+struct ChangingGradient3D <: ChangingGradient{3}
     gradient_strength_start :: SVector{3, <:VariableType}
     slew_rate :: SVector{3, <:VariableType}
     duration :: VariableType
-    rotate :: Union{Nothing, Symbol}
-    scale :: Union{Nothing, Symbol}
 end
 
 
-duration(cgb::ChangingGradientBlock) = cgb.duration
+duration(cgb::ChangingGradient) = cgb.duration
+
+grad_start(cgb::ChangingGradient) = cgb.gradient_strength_start
+slew_rate(cgb::ChangingGradient) = cgb.slew_rate
+grad_end(cgb::ChangingGradient) = grad_start(cgb) .+ slew_rate(cgb) .* duration(cgb)
+gradient_strength(cgb::ChangingGradient) = max.(grad_start(cgb), grad_end(cgb))
+qvec(cgb::ChangingGradient) = (grad_start(cgb) .+ grad_end(cgb)) .* (duration(cgb) * π)
 
-grad_start(cgb::ChangingGradientBlock) = cgb.gradient_strength_start
-slew_rate(cgb::ChangingGradientBlock) = cgb.slew_rate
-grad_end(cgb::ChangingGradientBlock) = grad_start(cgb) .+ slew_rate(cgb) .* duration(cgb)
-gradient_strength(cgb::ChangingGradientBlock) = max.(grad_start(cgb), grad_end(cgb))
-qvec(cgb::ChangingGradientBlock) = (grad_start(cgb) .+ grad_end(cgb)) .* (duration(cgb) * π)
+_mult(g1::VariableType, g2::VariableType) = g1 * g2
+_mult(g1::AbstractVector, g2::AbstractVector) = g1 .* permutedims(g2)
 
 function bmat_gradient(cgb::ChangingGradientBlock, qstart)
     # grad = (g1 * (duration - t) + g2 * t) / duration
@@ -45,10 +56,8 @@ function bmat_gradient(cgb::ChangingGradientBlock, qstart)
     #   g1 * (g2 - g1) * duration^3 / 4 +
     #   (g2 - g1)^2 * duration^3 / 10
     return (
-        qstart .* permutedims(qstart) .* duration(cgb) .+
-        duration(cgb)^2 .* qstart .* permutedims(
-            2 .* grad_start(cgb) .+
-            grad_end(cgb)) .* 2π ./ 3 .+
+        _mult(qstart, qstart) .* duration(cgb) .+
+        duration(cgb)^2 .* _mult(qstart, 2 .* grad_start(cgb) .+ grad_end(cgb)) .* 2π ./ 3 .+
         bmat_gradient(cgb)
     )
 end
@@ -56,9 +65,9 @@ end
 function bmat_gradient(cgb::ChangingGradientBlock)
     diff = slew_rate(cgb) .* duration(cgb)
     return (2π)^2 .* (
-        grad_start(cgb) .* permutedims(grad_start(cgb)) ./ 3 .+
-        grad_start(cgb) .* permutedims(diff) ./ 4 .+
-        diff .* permutedims(diff) ./ 10
+        _mult(grad_start(cgb), grad_start(cgb)) ./ 3 .+
+        _mult(grad_start(cgb), diff) ./ 4 .+
+        _mult(diff, diff) ./ 10
     ) .* duration(cgb)^3
 end
 
@@ -73,10 +82,10 @@ Times are assumed to be in increasing order and between 0 and the duration of th
 
 For N times this returns a vector with the N+1 replacement [`ConstantGradientBlock`](@ref) or [`ChangingGradientBlock`](@ref) objects.
 """
-function split_gradient(cgb::ChangingGradientBlock, times::VariableType...)
+function split_gradient(cgb::ChangingGradient, times::VariableType...)
     all_times = [0., times...]
     durations = [times[1], [t[2] - t[1] for t in zip(times[1:end-1], times[2:end])]..., duration(cgb) - times[end]]
-    return [ChangingGradientBlock(cgb.gradient_strength .+ cgb.slew_rate .* t, cgb.slew_rate, d, cgb.rotate, cgb.scale) for (t, d) in zip(all_times, durations)]
+    return [cls(cgb)(cgb.gradient_strength .+ cgb.slew_rate .* t, cgb.slew_rate, d) for (t, d) in zip(all_times, durations)]
 end
 
 

src/components/gradient_waveforms/constant_gradient_blocks.jl

+module ConstantGradientBlocks
+import StaticArrays: SVector
+import ....Variables: VariableType, duration, qvec, bmat_gradient, gradient_strength, slew_rate, get_free_variable
+import ...AbstractTypes: GradientWaveform
+
+abstract type ConstantGradient{N} <: GradientWaveform{N} end
+
+"""
+    ConstantGradient1D(gradient_strength, duration)
+
+Underlying type for any flat part in a 1D gradient waveform.
+
+Usually, you do not want to create this object directly, use a `BuildingBlock` instead.
+"""
+struct ConstantGradient1D <: ConstantGradient{1}
+    gradient_strength :: VariableType
+    duration :: VariableType
+end
+
+"""
+    ConstantGradient1D(gradient_strength, duration)
+
+Underlying type for any flat part in a 3D gradient waveform.
+
+Usually, you do not want to create this object directly, use a `BuildingBlock` instead.
+"""
+struct ConstantGradient3D <: ConstanGradient{3}
+    gradient_strength :: SVector{3, <:VariableType}
+    duration :: VariableType
+end
+
+duration(cgb::ConstantGradient) = cgb.duration
+gradient_strength(cgb::ConstantGradient) = cgb.gradient_strength
+slew_rate(::ConstantGradient1D) = 0.
+slew_rate(::ConstantGradient3D) = zero(SVector{3, Float64})
+qvec(cgb::ConstantGradient1D) = duration(cgb) * gradient_strength(cgb) * 2π
+qvec(cgb::ConstantGradient3D) = @. duration(cgb) * gradient_strength(cgb) * 2π
+
+_mult(g1::VariableType, g2::VariableType) = g1 * g2
+_mult(g1::AbstractVector, g2::AbstractVector) = g1 .* permutedims(g2)
+
+function bmat_gradient(cgb::ConstantGradient)
+    grad = 2π .* gradient_strength(cgb)
+    return _mult(grad, grad) .* duration(cgb)^3 ./3
+end
+
+function bmat_gradient(cgb::ConstantGradient, qstart)
+    # \int dt (qstart + t * grad)^2 = 
+    #   qstart^2 * duration +
+    #   qstart * grad * duration^2 +
+    #   grad * grad * duration^3 / 3 +
+    grad = 2π .* gradient_strength(cgb)
+    return (
+        _mult(qstart, qstart) .* duration(cgb) .+
+        _mult(qstart, grad) .* duration(cgb)^2 .+
+        bmat_gradient(cgb)
+    )
+end
+
+function split_gradient(cgb::ConstantGradient, times::VariableType...)
+    durations = [times[1], [t[2] - t[1] for t in zip(times[1:end-1], times[2:end])]..., duration(cgb) - times[end]]
+    @assert all(durations >= 0.)
+    return [typeof(cgb)(cgb.gradient_strength, d) for d in durations]
+end
+
+end

src/gradients/gradients.jl → src/components/gradient_waveforms/gradient_waveforms.jl

 """
-Module defining sub-types of the [`GradientBlock`](@ref), i.e., any [`BuildingBlock`](@ref) that only defines a gradient profile.
+Module defining sub-types of the [`GradientWaveform`](@ref).
 
 There are only three types of [`GradientBlock`] objects:
 - [`ChangingGradientBlock`](@ref): any gradient changing linearly in strength.
 - [`ConstantGradientBlock`](@ref): any gradient staying constant in strength. These can overlap with a pulse (`SliceSelectPulse`).
-- [`InstantGradientBlock`](@ref): an infinitely short gradient pulse.
+- [`NoGradientBlock`](@ref): any part of the gradient waveform when no gradient is active.
 
-Combinations of these should sub-type from [`SpecificWaveform`](@ref).
+These parts are combined into a full gradient waveform in a `BuildingBlock`.
+
+Each part of this gradient waveform can compute:
+- [`gradient_strength`]: maximum gradient strength in each dimension.
+- [`slew_rate`]: maximum slew rate in each dimension.
+- [`qvec`]: area under curve in each dimension
+- [`bmat_gradient`]: diffusion weighting (scalar in 1D or matrix in 3D).
 """
 module Gradients
 
 include("changing_gradient_blocks.jl")
 include("constant_gradient_blocks.jl")
-include("instant_gradients.jl")
+include("no_gradient_blocks.jl")
 
 
-import ..BuildingBlocks: GradientBlock
+import ..AbstractTypes: GradientWaveform
+import .NoGradientBlocks: NoGradientBlock
 import .ChangingGradientBlocks: ChangingGradientBlock, split_gradient
 import .ConstantGradientBlocks: ConstantGradientBlock
-import .InstantGradients: InstantGradientBlock
+
 end
\ No newline at end of file

src/components/gradient_waveforms/no_gradient_blocks.jl

+module NoGradientBlocks
+import StaticArrays: SVector, SMatrix
+import ....Variables: VariableType, duration, qvec, bmat_gradient, gradient_strength, slew_rate, get_free_variable
+import ...AbstractTypes: GradientWaveform
+import ..ChangingGradientBlocks: split_gradient
+
+"""
+    NoGradientBlock(duration)
+
+Part of a gradient waveform when there is no gradient active.
+
+Usually, you do not want to create this object directly, use a `BuildingBlock` instead.
+"""
+struct NoGradientBlock{N} <: GradientWaveform{N}
+    duration :: VariableType
+    function NoGradientBlock{N}(duration)
+        if !(N in (1, 3))
+            error("Dimensionality of the gradient should be 1 or 3, not $N")
+        end
+        new(duration)
+    end
+end
+
+duration(ngb::NoGradientBlock) = duration(ngb)
+for func in (:qvec, :gradient_strength, :slew_rate)
+    @eval $func(::NoGradientBlock{1}) = 0.
+    @eval $func(::NoGradientBlock{3}) = zero(SVector{3, Float64})
+end
+
+bmat_gradient(::NoGradientBlock{1}) = 0.
+bmat_gradient(::NoGradientBlock{3}) = zero(SMatrix{3, 3, Float64, 9})
+
+bmat_gradient(::NoGradientBlock, ) = 0.
+bmat_gradient(ngb::NoGradientBlock{1}, qstart::VariableType) = qstart^2 * duration(ngb)
+bmat_gradient(ngb::NoGradientBlock{3}, qstart::AbstractVector{<:VariableType}) = @. qstart * permutedims(qstart) * duration(ngb)
+
+function split_gradient(ngb::NoGradientBlock, times::VariableType...)
+    durations = [times[1], [t[2] - t[1] for t in zip(times[1:end-1], times[2:end])]..., duration(ngb) - times[end]]
+    @assert all(durations >= 0.)
+    return [NoGradientBlock(d) for d in durations]
+end
+
+end
\ No newline at end of file

src/gradients/constant_gradient_blocks.jl

-module ConstantGradientBlocks
-import StaticArrays: SVector
-import ...Variables: VariableType, variables
-import ...BuildingBlocks: GradientBlock
-import ...Variables: qvec, bmat_gradient, gradient_strength, slew_rate, duration, variables, VariableType
-import ...BuildingBlocks: GradientBlock, RFPulseBlock
-import ..ChangingGradientBlocks: split_gradient
-
-"""
-    ConstantGradientBlock(gradient_strength, duration, rotate, scale)
-
-Underlying type for any flat part in a gradient waveform.
-It can overlap with an [`RFPulseBlock`], in which case it will be a `SliceSelectPulse`
-
-Usually, you do not want to create this object directly, use a gradient waveform instead.
-
-## Parameters
-- `rotate`: with which user-set parameter will this gradient be rotated (e.g., :bvec). Default is no rotation.
-- `scale`: with which user-set parameter will this gradient be scaled (e.g., :bval). Default is no scaling.
-"""
-struct ConstantGradientBlock <: GradientBlock
-    gradient_strength :: SVector{3, <:VariableType}
-    duration :: VariableType
-    rotate :: Union{Nothing, Symbol}
-    scale :: Union{Nothing, Symbol}
-end
-
-duration(cgb::ConstantGradientBlock) = cgb.duration
-gradient_strength(cgb::ConstantGradientBlock) = cgb.gradient_strength
-slew_rate(::ConstantGradientBlock) = zero(SVector{3, Float64})
-qvec(cgb::ConstantGradientBlock) = duration(cgb) .* gradient_strength(cgb) .* 2π
-
-function bmat_gradient(cgb::ConstantGradientBlock)
-    grad = 2π .* gradient_strength(cgb)
-    return (grad .* permutedims(grad)) .* duration(cgb)^3 ./3
-end
-
-function bmat_gradient(cgb::ConstantGradientBlock, qstart)
-    # \int dt (qstart + t * grad)^2 = 
-    #   qstart^2 * duration +
-    #   qstart * grad * duration^2 +
-    #   grad * grad * duration^3 / 3 +
-    grad = 2π .* gradient_strength(cgb)
-    return (
-        qstart .* permutedims(qstart) .* duration(cgb) .+
-        qstart .* permutedims(grad) .* duration(cgb)^2 .+
-        bmat_gradient(cgb)
-    )
-end
-
-function split_gradient(cgb::ConstantGradientBlock, times::VariableType...)
-    durations = [times[1], [t[2] - t[1] for t in zip(times[1:end-1], times[2:end])]..., duration(cgb) - times[end]]
-    return [ConstantGradientBlock(cgb.gradient_strength, d, cgb.rotate, cgb.scale) for d in durations]
-end
-
-end

src/variables.jl

+"""
+Defines the functions that can be called on parts of an MRI sequence to query or constrain any variables.
+
+In addition this defines:
+- [`variables`](@ref): dictionary containing all variables.
+- [`VariableType`](@ref): parent type for any variables (whether number or JuMP variable).
+- [`get_free_variable`](@ref): helper function to create new JuMP variables.
+"""
 module Variables
 import JuMP: @variable, Model, @objective, objective_function, value, AbstractJuMPScalar
 import ..Scanners: gradient_strength, slew_rate
 import ..BuildSequences: global_model
 
+"""
+Parent type of all components, building block, and sequences that form an MRI sequence.
+"""
+abstract type AbstractBlock end
+
 all_variables_symbols = [
     :block => [
         :duration => "duration of the building block in ms.",
@@ -10,7 +23,6 @@ all_variables_symbols = [
     :sequence => [
         :TR => "Time on which an MRI sequence repeats itself in ms.",
     ],
-
     :pulse => [
         :flip_angle => "The flip angle of the RF pulse in degrees",
         :amplitude => "The maximum amplitude of an RF pulse in kHz",
@@ -23,8 +35,6 @@ all_variables_symbols = [
         :slice_thickness => "Slice thickness of an RF pulse that is active during a gradient in mm. To set constraints it is often better to use [`inverse_slice_thickness`](@ref).",
         :inverse_slice_thickness => "Inverse of slice thickness of an RF pulse that is active during a gradient in 1/mm. Also, see [`slice_thickness`](@ref).",
     ],
-
-    # gradients
     :gradient => [
         :qvec => "The spatial range and orientation on which the displacements can be detected due to this gradient in rad/um.",
         :qval => "The spatial range on which the displacements can be detected due to this gradient in rad/um (i.e., norm of [`qvec`](@ref)). To set constraints it is often better to use [`qvec`](@ref) or [`qval_square`](@ref).",
@@ -48,10 +58,12 @@ all_variables_symbols = [
     ]
 ]
 
+"""
+Collection of all functions that return variables that can be used to query or constrain their values.
+"""
 variables = Dict{Symbol, Function}()
 
 
-
 for (block_symbol, all_functions) in all_variables_symbols
     for (func_symbol, description) in all_functions
         as_string = "    $func_symbol($block_symbol)\n\n$description\n\nThis represents a variable within the sequence. Variables can be set during the construction of a [`BuildingBlock`](@ref) or used to create constraints after the fact."
@@ -65,6 +77,11 @@ for (block_symbol, all_functions) in all_variables_symbols
 end
 
 
+"""
+Dictionary with alternative versions of specific function.
+    
+Setting constraints on these alternative functions can be helpful as it avoids some operations, which the optimiser might struggle with.
+"""
 alternative_variables = Dict(
     qval => (qval_square, n->n^2, sqrt, false),
     slice_thickness => (inverse_slice_thickness, inv, inv, true),
@@ -74,12 +91,16 @@ alternative_variables = Dict(
 )
 
 
-# These functions are more fully defined in building_blocks.jl
-function start_time end
-function end_time end
-function effective_time end
+"""
+Parent type for any variable in the MRI sequence.
 
+Each variable can be one of:
+- a new JuMP variable
+- an expression linking this variable to other JuMP variable
+- a number
 
+Create these using [`get_free_variable`](@ref).
+"""
 const VariableType = Union{Number, AbstractJuMPScalar}
 
 
@@ -87,6 +108,8 @@ const VariableType = Union{Number, AbstractJuMPScalar}
     get_free_variable(value; integer=false)
 
 Get a representation of a given `variable` given a user-defined constraint.
+
+The result is guaranteed to be a [`VariableType`](@ref).
 """
 get_free_variable(value::Number; integer=false) = integer ? Int(value) : Float64(value)
 get_free_variable(value::VariableType; integer=false) = value