Refactor LockedAmplitude and ShapedAmplitude

Both of these amplitudes could be made more flexible by getting rid of
the type hierarchy that distinguished between, e.g.,
`ShapedContinousAmplitude` and `ShapedPulseAmplitude`. The shape and
control can now be vectors or functions independently of each other.
This enables substituting optimized pulses into existing amplitudes,
among other benefits.

Author: goerz

src/amplitudes.jl

@@ -15,29 +15,34 @@ using ..Controls: t_mid
 ampl = LockedAmplitude(shape)
 ```
 
-wraps around `shape`, which must be either a vector of values defined on the
-midpoints of a time grid or a callable `shape(t)`.
+wraps around `shape`, which must be either a `Vector{Float64}` of values defined
+on the midpoints of a time grid or a callable `shape(t)`.
 
 ```julia
 ampl = LockedAmplitude(shape, tlist)
 ```
 
 discretizes `shape` to the midpoints of `tlist`.
 """
-abstract type LockedAmplitude end
-
-
-function LockedAmplitude(shape)
-    if shape isa Vector{Float64}
-        return LockedPulseAmplitude(shape)
-    else
-        return LockedContinuousAmplitude(shape)
+struct LockedAmplitude{ST}
+    shape::ST
+
+    function LockedAmplitude(shape::ST; check = true) where {ST}
+        if check && !(shape isa Vector{Float64})
+            try
+                shape(0.0)
+            catch
+                msg = "A LockedAmplitude shape must either be a Vector{Float64} or a callable"
+                error(msg)
+            end
+        end
+        return new{ST}(shape)
     end
 end
 
 
 function LockedAmplitude(shape, tlist)
-    return LockedPulseAmplitude(discretize_on_midpoints(shape, tlist))
+    return LockedAmplitude(discretize_on_midpoints(shape, tlist); check = false)
 end
 
 
@@ -46,53 +51,29 @@ function Base.show(io::IO, ampl::LockedAmplitude)
 end
 
 
-struct LockedPulseAmplitude <: LockedAmplitude
-    shape::Vector{Float64}
-end
-
-
-Base.Array(ampl::LockedPulseAmplitude) = ampl.shape
-
-
-struct LockedContinuousAmplitude <: LockedAmplitude
-
-    shape
-
-    function LockedContinuousAmplitude(shape)
-        try
-            S_t = shape(0.0)
-        catch
-            error("A LockedAmplitude shape must either be a vector of values or a callable")
-        end
-        return new(shape)
-    end
-
-end
-
-(ampl::LockedContinuousAmplitude)(t::Float64) = ampl.shape(t)
+Base.Array(ampl::LockedAmplitude{Vector{Float64}}) = ampl.shape
 
 get_controls(ampl::LockedAmplitude) = ()
 
 function substitute(ampl::LockedAmplitude, replacements)
     return get(replacements, ampl, ampl)
 end
 
-function evaluate(ampl::LockedPulseAmplitude, tlist, n; _...)
+function evaluate(ampl::LockedAmplitude{Vector{Float64}}, tlist, n::Int; _...)
     return ampl.shape[n]
 end
 
-function evaluate(ampl::LockedPulseAmplitude, t; _...)
-    error(
-        "A LockedAmplitude initialized on a `tlist` can only be evaluated with arguments `tlist, n`."
-    )
+function evaluate(ampl::LockedAmplitude{Vector{Float64}}, t::Float64; _...)
+    msg = "A LockedAmplitude initialized from a vector can only be evaluated with (tlist, n)."
+    error(msg)
 end
 
-function evaluate(ampl::LockedContinuousAmplitude, tlist, n; _...)
-    return ampl(t_mid(tlist, n))
+function evaluate(ampl::LockedAmplitude, tlist, n::Int; _...)
+    return ampl.shape(t_mid(tlist, n))
 end
 
-function evaluate(ampl::LockedContinuousAmplitude, t; _...)
-    return ampl(t)
+function evaluate(ampl::LockedAmplitude, t::Float64; _...)
+    return ampl.shape(t)
 end
 
 
@@ -124,63 +105,122 @@ ampl = ShapedAmplitude(control; shape=shape)
 ```
 
 produces an amplitude ``a(t) = S(t) ϵ(t)``, where ``S(t)`` corresponds to
-`shape` and ``ϵ(t)`` corresponds to `control`. Both `control` and `shape`
-should be either a vector of values defined on the midpoints of a time grid or
-a callable `control(t)`, respectively `shape(t)`. In the latter case, `ampl`
-will also be callable.
+`shape` and ``ϵ(t)`` corresponds to `control`. Each of `control` and `shape`
+must be either a `Vector{Float64}` of values defined on the midpoints of a time
+grid or a callable `control(t)`, respectively `shape(t)`. If both are callables,
+`ampl` will also be callable. If both are vectors, they must have the same length.
 
 ```julia
 ampl = ShapedAmplitude(control, tlist; shape=shape)
 ```
 
 discretizes `control` and `shape` to the midpoints of `tlist`.
 """
-abstract type ShapedAmplitude <: ControlAmplitude end
+struct ShapedAmplitude{CT,ST} <: ControlAmplitude
+    control::CT
+    shape::ST
+end
 
-function ShapedAmplitude(control; shape)
-    if (control isa Vector{Float64}) && (shape isa Vector{Float64})
-        return ShapedPulseAmplitude(control, shape)
-    else
-        try
-            ϵ_t = control(0.0)
-        catch
-            error(
-                "A ShapedAmplitude control must either be a vector of values or a callable"
-            )
+
+function ShapedAmplitude(control; shape, check = true)
+    if check
+        if !(control isa Vector{Float64})
+            try
+                control(0.0)
+            catch
+                msg = "A ShapedAmplitude control must either be a Vector{Float64} or a callable"
+                error(msg)
+            end
         end
-        try
-            S_t = shape(0.0)
-        catch
-            error("A ShapedAmplitude shape must either be a vector of values or a callable")
+        if !(shape isa Vector{Float64})
+            try
+                shape(0.0)
+            catch
+                msg = "A ShapedAmplitude shape must either be a Vector{Float64} or a callable"
+                error(msg)
+            end
+        end
+        if (control isa Vector{Float64}) && (shape isa Vector{Float64})
+            if length(control) ≠ length(shape)
+                msg = "ShapedAmplitude control and shape vectors must have the same length"
+                error(msg)
+            end
         end
-        return ShapedContinuousAmplitude(control, shape)
     end
+    return ShapedAmplitude{typeof(control),typeof(shape)}(control, shape)
 end
 
+
 function Base.show(io::IO, ampl::ShapedAmplitude)
     print(io, "ShapedAmplitude(::$(typeof(ampl.control)); shape::$(typeof(ampl.shape)))")
 end
 
-function ShapedAmplitude(control, tlist; shape)
+
+function ShapedAmplitude(control, tlist::Vector{Float64}; shape)
     control = discretize_on_midpoints(control, tlist)
     shape = discretize_on_midpoints(shape, tlist)
-    return ShapedPulseAmplitude(control, shape)
+    return ShapedAmplitude{typeof(control),typeof(shape)}(control, shape)
+end
+
+
+function substitute(ampl::ShapedAmplitude, replacements)
+    ampl in keys(replacements) && return replacements[ampl]
+    control = substitute(ampl.control, replacements)
+    return ShapedAmplitude(control; shape = ampl.shape, check = true)
+end
+
+
+Base.Array(ampl::ShapedAmplitude{Vector{Float64},Vector{Float64}}) =
+    ampl.control .* ampl.shape
+
+(ampl::ShapedAmplitude)(t::Float64) = ampl.shape(t) * ampl.control(t)
+
+
+# Vector shape: index directly; control may be vector or callable (evaluated via Controls.evaluate)
+function evaluate(
+    ampl::ShapedAmplitude{CT,Vector{Float64}},
+    tlist::Vector{Float64},
+    n::Int;
+    vals_dict = IdDict()
+) where {CT}
+    S_t = ampl.shape[n]
+    ϵ_t = evaluate(ampl.control, tlist, n; vals_dict)
+    return S_t * ϵ_t
 end
 
-struct ShapedPulseAmplitude <: ShapedAmplitude
-    control::Vector{Float64}
-    shape::Vector{Float64}
+# Callable shape: call directly with t_mid; control may be vector or callable
+function evaluate(
+    ampl::ShapedAmplitude,
+    tlist::Vector{Float64},
+    n::Int;
+    vals_dict = IdDict()
+)
+    S_t = ampl.shape(t_mid(tlist, n))
+    ϵ_t = evaluate(ampl.control, tlist, n; vals_dict)
+    return S_t * ϵ_t
 end
 
-Base.Array(ampl::ShapedPulseAmplitude) = ampl.control .* ampl.shape
+# Callable shape and callable control: call both directly at t
+function evaluate(ampl::ShapedAmplitude, t::Float64; vals_dict = IdDict())
+    S_t = ampl.shape(t)
+    ϵ_t = evaluate(ampl.control, t; vals_dict)
+    return S_t * ϵ_t
+end
 
+function evaluate(ampl::ShapedAmplitude{Vector{Float64},Vector{Float64}}, t::Float64; _...)
+    msg = "A ShapedAmplitude with vector control and shape can only be evaluated with (tlist, n)."
+    error(msg)
+end
 
-struct ShapedContinuousAmplitude <: ShapedAmplitude
-    control
-    shape
+function evaluate(ampl::ShapedAmplitude{Vector{Float64},ST}, t::Float64; _...) where {ST}
+    msg = "A ShapedAmplitude with a vector control can only be evaluated with (tlist, n)."
+    error(msg)
 end
 
-(ampl::ShapedContinuousAmplitude)(t::Float64) = ampl.shape(t) * ampl.control(t)
+function evaluate(ampl::ShapedAmplitude{CT,Vector{Float64}}, t::Float64; _...) where {CT}
+    msg = "A ShapedAmplitude with a vector shape can only be evaluated with (tlist, n)."
+    error(msg)
+end
 
 
 function evaluate(ampl::ShapedAmplitude, args...; vals_dict = IdDict())

test/test_amplitudes.jl

@@ -2,8 +2,8 @@ using Test
 using QuantumPropagators.Interfaces: check_amplitude, check_control
 using QuantumPropagators.Shapes: flattop
 using QuantumPropagators.Amplitudes: LockedAmplitude, ShapedAmplitude
-using QuantumPropagators.Controls: evaluate, get_controls, t_mid
-using QuantumPropagators.Controls: discretize_on_midpoints # DEBUG
+using QuantumPropagators.Controls: evaluate, get_controls, substitute, t_mid
+using QuantumPropagators.Controls: discretize_on_midpoints
 
 @testset "LockedAmplitude" begin
 
@@ -25,6 +25,7 @@ using QuantumPropagators.Controls: discretize_on_midpoints # DEBUG
     @test length(get_controls(ampl)) == 0
     t = t_mid(tlist, 20)
     @test evaluate(ampl, tlist, 20) ≈ S(t)
+    @test_throws Exception evaluate(ampl, t)
 
 end
 
@@ -57,5 +58,54 @@ end
     @test check_amplitude(ampl; tlist)
     t = t_mid(tlist, 20)
     @test evaluate(ampl, tlist, 20) ≈ ϵ(t) * S(t)
+    @test_throws Exception evaluate(ampl, t)
+
+end
+
+
+@testset "ShapedAmplitude mixed" begin
+    tlist = collect(range(0, 10, length = 101))
+    S(t) = flattop(t, T = 10, t_rise = 2, func = :blackman)
+    ϵ(t) = 1.0
+    S_vec = discretize_on_midpoints(S, tlist)
+    ϵ_vec = discretize_on_midpoints(ϵ, tlist)
+
+    # callable control, vector shape
+    ampl = ShapedAmplitude(ϵ; shape = S_vec)
+    @test startswith("$ampl", "ShapedAmplitude(")
+    controls = get_controls(ampl)
+    @test length(controls) == 1
+    @test check_amplitude(ampl; tlist)
+    t = t_mid(tlist, 20)
+    @test evaluate(ampl, tlist, 20) ≈ ϵ(t) * S(t)
+    @test_throws Exception evaluate(ampl, t)
+
+    # vector control, callable shape
+    ampl = ShapedAmplitude(ϵ_vec; shape = S)
+    @test startswith("$ampl", "ShapedAmplitude(")
+    controls = get_controls(ampl)
+    @test length(controls) == 1
+    @test check_amplitude(ampl; tlist)
+    t = t_mid(tlist, 20)
+    @test evaluate(ampl, tlist, 20) ≈ ϵ(t) * S(t)
+    @test_throws Exception evaluate(ampl, t)
+
+end
+
+
+@testset "ShapedAmplitude substitute" begin
+    tlist = collect(range(0, 10, length = 101))
+    S(t) = flattop(t, T = 10, t_rise = 2, func = :blackman)
+    ϵ(t) = 1.0
+    ϵ_vec = discretize_on_midpoints(ϵ, tlist)
+
+    ampl = ShapedAmplitude(ϵ; shape = S)
+    ampl2 = substitute(ampl, IdDict(ϵ => ϵ_vec))
+    @test ampl2 isa ShapedAmplitude
+    @test get_controls(ampl2)[1] === ϵ_vec
+    @test check_amplitude(ampl2; tlist)
+    t = t_mid(tlist, 20)
+    @test evaluate(ampl2, tlist, 20) ≈ ϵ(t) * S(t)
+    @test_throws Exception evaluate(ampl2, t)
 
 end