Fix integration tests, update gates and noises, and update deps (#53)

* fix: bump integ test python versions

* change: remove extraneous whitespace

* fix: finalizer error

* fix: fix dm1 integration test

* change: schema update

* fix: adjust to updates in AWS.jl

* bump package versions

Project.toml

@@ -1,7 +1,7 @@
 name = "Braket"
 uuid = "19504a0f-b47d-4348-9127-acc6cc69ef67"
 authors = ["Katharine Hyatt <hyatkath@amazon.com>"]
-version = "0.6.0"
+version = "0.7.0"
 
 [deps]
 AWS = "fbe9abb3-538b-5e4e-ba9e-bc94f4f92ebc"
@@ -30,7 +30,7 @@ Tar = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
-AWS = "=1.81.0"
+AWS = "=1.82.0"
 AWSS3 = "=0.10.3"
 Aqua = "=0.6"
 AxisArrays = "=0.4.6"

PyBraket/Project.toml

@@ -1,7 +1,7 @@
 name = "PyBraket"
 uuid = "e85266a6-1825-490b-a80e-9b9469c53660"
 authors = ["Katharine Hyatt <hyatkath@amazon.com>"]
-version = "0.6.0"
+version = "0.7.0"
 
 [deps]
 Braket = "19504a0f-b47d-4348-9127-acc6cc69ef67"
@@ -14,7 +14,7 @@ StructTypes = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
 
 [compat]
 Aqua = "=0.6"
-Braket = "=0.6.0"
+Braket = "=0.7.0"
 CondaPkg = "=0.2.17"
 DataStructures = "=0.18.13"
 PythonCall = "=0.9.12"

PyBraket/src/py_type_conversion.jl

@@ -0,0 +1,84 @@
+using PythonCall, Braket, Braket.IR
+import Braket.IR: IRObservable, AbstractProgramResult, AbstractProgram, AbstractIR
+import Braket: ResultTypeValue
+
+function union_convert(::Type{IRObservable}, x)
+    PythonCall.pyisinstance(x, PythonCall.pybuiltins.str) && return pyconvert(String, x)
+    x_vec = Union{String, Vector{Vector{Vector{Float64}}}}[PythonCall.pyisinstance(x_, PythonCall.pybuiltins.str) ? pyconvert(String, x_) : pyconvert(Vector{Vector{Vector{Float64}}}, x_) for x_ in x] 
+    return x_vec
+end
+
+function union_convert(union_type, x)
+    union_ts = union_type isa Union ? Type[] : [union_type]
+    union_t = union_type
+    while union_t isa Union
+        union_t.a != Nothing && push!(union_ts, union_t.a)
+        !(union_t.b isa Union) && push!(union_ts, union_t.b)
+        union_t = union_t.b
+    end
+    arg = nothing
+    for t_ in union_ts
+        try
+            if pyisinstance(x, PythonCall.pybuiltins.list) && t_ <: Vector
+                return [union_convert(Union{eltype(t_)}, attr_) for attr_ in x]
+            elseif pyisinstance(x, pybuiltins.str) && t_ <: Integer
+                return tryparse(t_, pyconvert(String, x))
+            elseif t_ == ResultTypeValue
+                if pyhasattr(x, "value")
+                    typ = jl_convert(AbstractProgramResult, pygetattr(x, "type"))
+                    val = union_convert(Union{Dict{String, ComplexF64}, Float64, Vector}, pygetattr(x, "value"))
+                    return PythonCall.pyconvert_return(ResultTypeValue(typ, val))
+                else
+                    rt = jl_convert(AbstractProgramResult, x)
+                    return PythonCall.pyconvert_return(ResultTypeValue(rt, 0.0))
+                end
+            else
+                return pyconvert(t_, x)
+            end
+        catch e
+        end
+    end
+    arg isa Vector{Nothing} && (arg = nothing)
+    return arg
+end
+
+function jl_convert_attr(n, t, attr)
+    if !(t isa Union)
+        if pyisinstance(attr, PythonCall.pybuiltins.list)
+            if eltype(t) isa Union
+                return [union_convert(eltype(t), attr_) for attr_ in attr]
+            else
+                return [pyconvert(eltype(t), attr_) for attr_ in attr]
+            end
+        else
+            return pyconvert(t, attr)
+        end
+    else
+        PythonCall.pyisnone(attr) && return nothing
+        return union_convert(t, attr)
+    end
+end
+
+function jl_convert(::Type{T}, x::Py) where {T}
+    fts = fieldtypes(T)
+    fns = fieldnames(T)
+    args = Any[]
+    for (n, t) in zip(fns, fts)
+        attr = pygetattr(x, string(n))
+        arg = jl_convert_attr(n, t, attr)
+        push!(args, arg)
+    end
+    PythonCall.pyconvert_return(T(args...))
+end
+
+function jl_convert(::Type{T}, x::Py) where {T<:AbstractIR}
+    fts = fieldtypes(T)[1:end-1]
+    fns = fieldnames(T)[1:end-1]
+    args = Any[]
+    for (n, t) in zip(fns, fts)
+        attr = pygetattr(x, string(n))
+        arg = jl_convert_attr(n, t, attr)
+        push!(args, arg)
+    end
+    PythonCall.pyconvert_return(T(args..., pyconvert(String, pygetattr(x, "type"))))
+end

src/ahs.jl

@@ -1,7 +1,5 @@
 import StructTypes
 
-
-
 export AtomArrangementItem, AtomArrangement, TimeSeriesItem, TimeSeries, Field, DrivingField
 export ShiftingField, Hamiltonian, AnalogHamiltonianSimulation, Pattern, vacant, filled, SiteType, discretize
 

src/aws_jobs.jl

@@ -29,10 +29,10 @@ Cancels the job `j`.
 cancel(j::AwsQuantumJob) = (r = BRAKET.cancel_job(HTTP.escapeuri(arn(j))); return nothing)
 
 function describe_log_streams(log_group::String, stream_prefix::String, limit::Int=-1, next_token::String="")
-    args = Dict{String, Any}("logStreamNamePrefix"=>stream_prefix, "orderBy"=>"LogStreamName")
+    args = Dict{String, Any}("logGroupName"=>log_group, "logStreamNamePrefix"=>stream_prefix, "orderBy"=>"LogStreamName")
     limit > 0 && (args["limit"] = limit)
     !isempty(next_token) && (args["nextToken"] = next_token)
-    return CLOUDWATCH_LOGS.describe_log_streams(log_group, args) 
+    return CLOUDWATCH_LOGS.describe_log_streams(args) 
 end
 
 """
@@ -55,7 +55,7 @@ function log_stream(ch::Channel, log_group::String, stream_name::String, start_t
     next_token = nothing
     event_count = 1
     while event_count > 0
-        response = CLOUDWATCH_LOGS.get_log_events(log_group, stream_name, Dict("startTime"=>start_time, "nextToken"=>next_token, "startFromHead"=>true))
+        response = CLOUDWATCH_LOGS.get_log_events(stream_name, Dict("logGroupName"=>log_group, "startTime"=>start_time, "nextToken"=>next_token, "startFromHead"=>true))
         next_token = response["nextForwardToken"]
         events = response["events"]
         event_count = length(events)

src/gates.jl

@@ -5,31 +5,28 @@ export Gate,AngledGate,DoubleAngledGate,H,I,X,Y,Z,S,Si,T,Ti,V,Vi,CNot,Swap,ISwap
 Abstract type representing a quantum gate.
 """
 abstract type Gate <: QuantumOperator end
-
 StructTypes.StructType(::Type{Gate}) = StructTypes.AbstractType()
-
 StructTypes.subtypes(::Type{Gate}) = (angledgate=AngledGate, doubleangledgate=DoubleAngledGate, h=H, i=I, x=X, y=Y, z=Z, s=S, si=Si, t=T, ti=Ti, v=V, vi=Vi, cnot=CNot, swap=Swap, iswap=ISwap, cv=CV, cy=CY, cz=CZ, ecr=ECR, ccnot=CCNot, cswap=CSwap, unitary=Unitary, rx=Rx, ry=Ry, rz=Rz, phaseshift=PhaseShift, pswap=PSwap, xy=XY, cphaseshift=CPhaseShift, cphaseshift00=CPhaseShift00, cphaseshift01=CPhaseShift01, cphaseshift10=CPhaseShift10, xx=XX, yy=YY, zz=ZZ, gpi=GPi, gpi2=GPi2, ms=MS)
+
 """
     AngledGate <: Gate
 
 Abstract type representing a quantum gate with an `angle` parameter.
 """
 abstract type AngledGate <: Gate end
-
 StructTypes.StructType(::Type{AngledGate}) = StructTypes.AbstractType()
-
 StructTypes.subtypes(::Type{AngledGate}) = (rx=Rx, ry=Ry, rz=Rz, phaseshift=PhaseShift, pswap=PSwap, xy=XY, cphaseshift=CPhaseShift, cphaseshift00=CPhaseShift00, cphaseshift01=CPhaseShift01, cphaseshift10=CPhaseShift10, xx=XX, yy=YY, zz=ZZ, gpi=GPi, gpi2=GPi2)
+
 """
     DoubleAngledGate <: Gate
 
 Abstract type representing a quantum gate with two `angle` parameters.
 """
 abstract type DoubleAngledGate <: Gate end
-
 StructTypes.StructType(::Type{DoubleAngledGate}) = StructTypes.AbstractType()
- 
-StructTypes.subtypes(::Type{DoubleAngledGate}) = (ms=MS,)
-for (G, IRG, label, qc) in zip((:Rx, :Ry, :Rz, :PhaseShift, :PSwap, :XY, :CPhaseShift, :CPhaseShift00, :CPhaseShift01, :CPhaseShift10, :XX, :YY, :ZZ, :GPi, :GPi2), (:(IR.Rx), :(IR.Ry), :(IR.Rz), :(IR.PhaseShift), :(IR.PSwap), :(IR.XY), :(IR.CPhaseShift), :(IR.CPhaseShift00), :(IR.CPhaseShift01), :(IR.CPhaseShift10), :(IR.XX), :(IR.YY), :(IR.ZZ), :(IR.GPi), :(IR.GPi2)), ("rx", "ry", "rz", "phaseshift", "pswap", "xy", "cphaseshift", "cphaseshift00", "cphaseshift01", "cphaseshift10", "xx", "yy", "zz", "gpi", "gpi2"), (:1, :1, :1, :1, :2, :2, :2, :2, :2, :2, :2, :2, :2, :1, :1))
+StructTypes.subtypes(::Type{DoubleAngledGate}) = (ms=MS)
+
+for (G, IRG, label, qc, c) in zip((:Rx, :Ry, :Rz, :PhaseShift, :PSwap, :XY, :CPhaseShift, :CPhaseShift00, :CPhaseShift01, :CPhaseShift10, :XX, :YY, :ZZ, :GPi, :GPi2), (:(IR.Rx), :(IR.Ry), :(IR.Rz), :(IR.PhaseShift), :(IR.PSwap), :(IR.XY), :(IR.CPhaseShift), :(IR.CPhaseShift00), :(IR.CPhaseShift01), :(IR.CPhaseShift10), :(IR.XX), :(IR.YY), :(IR.ZZ), :(IR.GPi), :(IR.GPi2)), ("rx", "ry", "rz", "phaseshift", "pswap", "xy", "cphaseshift", "cphaseshift00", "cphaseshift01", "cphaseshift10", "xx", "yy", "zz", "gpi", "gpi2"), (:1, :1, :1, :1, :2, :2, :2, :2, :2, :2, :2, :2, :2, :1, :1), (["Rx(ang)"], ["Ry(ang)"], ["Rz(ang)"], ["PHASE(ang)"], ["PSWAP(ang)", "PSWAP(ang)"], ["XY(ang)", "XY(ang)"], ["C", "PHASE(ang)"], ["C", "PHASE00(ang)"], ["C", "PHASE01(ang)"], ["C", "PHASE10(ang)"], ["XX(ang)", "XX(ang)"], ["YY(ang)", "YY(ang)"], ["ZZ(ang)", "ZZ(ang)"], ["GPi(ang)"], ["GPi2(ang)"]))
     @eval begin
         @doc """
             $($G) <: AngledGate
@@ -40,6 +37,7 @@ for (G, IRG, label, qc) in zip((:Rx, :Ry, :Rz, :PhaseShift, :PSwap, :XY, :CPhase
         struct $G <: AngledGate
             angle::Union{Float64, FreeParameter}
         end
+        chars(g::$G) = map(char->replace(string(char), "ang"=>string(g.angle)), $c)
         qubit_count(g::$G) = $qc
         qubit_count(::Type{$G}) = $qc
         function ir(g::$G, target::QubitSet, ::Val{:JAQCD}; kwargs...)
@@ -53,7 +51,7 @@ for (G, IRG, label, qc) in zip((:Rx, :Ry, :Rz, :PhaseShift, :PSwap, :XY, :CPhase
         end
     end
 end
-for (G, IRG, label, qc) in zip((:MS,), (:(IR.MS),), ("ms",), (:2,))
+for (G, IRG, label, qc, c) in zip((:MS,), (:(IR.MS),), ("ms",), (:2,), (["MS(ang1, ang2)", "MS(ang1, ang2)"],))
     @eval begin
         @doc """
             $($G) <: DoubleAngledGate
@@ -65,6 +63,7 @@ for (G, IRG, label, qc) in zip((:MS,), (:(IR.MS),), ("ms",), (:2,))
             angle1::Union{Float64, FreeParameter}
             angle2::Union{Float64, FreeParameter}
         end
+        chars(g::$G) = map(char->replace(string(char), "ang1"=>string(g.angle1), "ang2"=>string(g.angle2)), $c)
         qubit_count(g::$G) = $qc
         qubit_count(::Type{$G}) = $qc
         function ir(g::$G, target::QubitSet, ::Val{:JAQCD}; kwargs...)
@@ -78,7 +77,7 @@ for (G, IRG, label, qc) in zip((:MS,), (:(IR.MS),), ("ms",), (:2,))
         end
     end
 end
-for (G, IRG, label, qc) in zip((:H, :I, :X, :Y, :Z, :S, :Si, :T, :Ti, :V, :Vi, :CNot, :Swap, :ISwap, :CV, :CY, :CZ, :ECR, :CCNot, :CSwap), (:(IR.H), :(IR.I), :(IR.X), :(IR.Y), :(IR.Z), :(IR.S), :(IR.Si), :(IR.T), :(IR.Ti), :(IR.V), :(IR.Vi), :(IR.CNot), :(IR.Swap), :(IR.ISwap), :(IR.CV), :(IR.CY), :(IR.CZ), :(IR.ECR), :(IR.CCNot), :(IR.CSwap)), ("h", "i", "x", "y", "z", "s", "si", "t", "ti", "v", "vi", "cnot", "swap", "iswap", "cv", "cy", "cz", "ecr", "ccnot", "cswap"), (:1, :1, :1, :1, :1, :1, :1, :1, :1, :1, :1, :2, :2, :2, :2, :2, :2, :2, :3, :3))
+for (G, IRG, label, qc, c) in zip((:H, :I, :X, :Y, :Z, :S, :Si, :T, :Ti, :V, :Vi, :CNot, :Swap, :ISwap, :CV, :CY, :CZ, :ECR, :CCNot, :CSwap), (:(IR.H), :(IR.I), :(IR.X), :(IR.Y), :(IR.Z), :(IR.S), :(IR.Si), :(IR.T), :(IR.Ti), :(IR.V), :(IR.Vi), :(IR.CNot), :(IR.Swap), :(IR.ISwap), :(IR.CV), :(IR.CY), :(IR.CZ), :(IR.ECR), :(IR.CCNot), :(IR.CSwap)), ("h", "i", "x", "y", "z", "s", "si", "t", "ti", "v", "vi", "cnot", "swap", "iswap", "cv", "cy", "cz", "ecr", "ccnot", "cswap"), (:1, :1, :1, :1, :1, :1, :1, :1, :1, :1, :1, :2, :2, :2, :2, :2, :2, :2, :3, :3), (["H"], ["I"], ["X"], ["Y"], ["Z"], ["S"], ["Si"], ["T"], ["Ti"], ["V"], ["Vi"], ["C", "X"], ["SWAP", "SWAP"], ["ISWAP", "ISWAP"], ["C", "V"], ["C", "Y"], ["C", "Z"], ["ECR", "ECR"], ["C", "C", "X"], ["C", "SWAP", "SWAP"]))
     @eval begin
         @doc """
             $($G) <: Gate
@@ -87,6 +86,7 @@ for (G, IRG, label, qc) in zip((:H, :I, :X, :Y, :Z, :S, :Si, :T, :Ti, :V, :Vi, :
         $($G) gate.
         """
         struct $G <: Gate end
+        chars(g::$G) = $c
         qubit_count(g::$G) = $qc
         qubit_count(::Type{$G}) = $qc
         function ir(g::$G, target::QubitSet, ::Val{:JAQCD}; kwargs...)
@@ -110,9 +110,9 @@ struct Unitary <: Gate
     matrix::Matrix{ComplexF64}
 end
 Unitary(mat::Vector{Vector{Vector{Float64}}}) = Unitary(complex_matrix_from_ir(mat))
-Base.:(==)(u1::Unitary, u2::Unitary) = u1.matrix ≈ u2.matrix
+Base.:(==)(u1::Unitary, u2::Unitary) = u1.matrix == u2.matrix
 qubit_count(g::Unitary) = convert(Int, log2(size(g.matrix, 1)))
-
+chars(g::Unitary) = ntuple(i->"U", qubit_count(g))
 function ir(g::Unitary, target::QubitSet, ::Val{:JAQCD}; kwargs...)
     mat = complex_matrix_to_ir(g.matrix) 
     t_c = target[1:convert(Int, log2(size(g.matrix, 1)))]
@@ -154,39 +154,3 @@ function bind_value!(::Parametrized, g::G, params::Dict{Symbol, Number}) where {
 end
 ir(g::Gate, target::Int, args...) = ir(g, QubitSet(target), args...)
 Base.copy(g::G) where {G<:Gate} = G((copy(getproperty(g, fn)) for fn in fieldnames(G))...)
-
-for (G, c) in zip((:H, :I, :X, :Y, :Z, :S, :Si, :T, :Ti, :V, :Vi), ("H", "I", "X", "Y", "Z", "S", "Si", "T", "Ti", "V", "Vi"))
-    @eval begin
-        chars(g::$G) = ($c,)
-    end
-end
-
-for (G, c) in zip((:PhaseShift, :Rx, :Ry, :Rz), ("Phase", "Rx", "Ry", "Rz"))
-    @eval begin
-        chars(g::$G) = ($c * "(" * string(g.angle) * ")",)
-    end
-end
-
-for (G, cs) in zip((:CNot, :Swap, :ECR, :ISwap, :CV, :CY, :CZ), (("C", "X"), ("SWAP", "SWAP"), ("ECR", "ECR"), ("iSWAP", "iSWAP"), ("C", "V"), ("C", "Y"), ("C", "Z")))
-    @eval begin
-        chars(g::$G) = $cs
-    end
-end
-
-for (G, cs) in zip((:CCNot, :CSwap), (("C", "C", "X"), ("C", "SWAP", "SWAP")))
-    @eval begin
-        chars(g::$G) = $cs
-    end
-end
-
-for (G, c) in zip((:CPhaseShift, :CPhaseShift00, :CPhaseShift01, :CPhaseShift10), ("Phase", "Phase00", "Phase01", "Phase10"))
-    @eval begin
-        chars(g::$G) = ("C", $c * "(" * string(g.angle) * ")",)
-    end
-end
-
-for (G, c) in zip((:XX, :YY, :ZZ, :PSwap), ("X", "Y", "Z", "PSwap"))
-    @eval begin
-        chars(g::$G) = ($c * "(" * string(g.angle) * ")",$c * "(" * string(g.angle) * ")")
-    end
-end

src/local_jobs.jl

@@ -8,7 +8,13 @@ mutable struct LocalJobContainer
     function LocalJobContainer(image_uri::String, create_job_args; config::AWSConfig=global_aws_config(), container_name::String="", container_code_path::String="/opt/ml/code", force_update::Bool=false)
         c = new(image_uri, container_name, container_code_path, Dict{String, String}(), "", config)
         c = start_container!(c, force_update)
-        finalizer(stop_container!, c)
+        finalizer(c) do c
+            # check that the container is still running
+            c_list = read(`docker container ls -q`, String)
+            stop_flag = occursin(first(c.container_name, 10), c_list)
+            stop_flag && read(Cmd(["docker", "stop", c.container_name]), String)
+            return
+        end
         return setup_container!(c, create_job_args)
     end
 end
@@ -126,8 +132,8 @@ function run_local_job!(c::LocalJobContainer)
     entry_point_cmd = `docker exec $c_name printenv SAGEMAKER_PROGRAM`
     entry_program, err, code = capture_docker_cmd(entry_point_cmd)
     (isnothing(entry_program) || isempty(entry_program)) && throw(ErrorException("Start program not found. The specified container is not setup to run Braket Jobs. Please see setup instructions for creating your own containers."))
-    env_list  = reduce(vcat, ["-e", k*"="*v] for (k,v) in c.env)
-    cmd  = Cmd(["docker", "exec", "-w", code_path, env_list..., c_name, "python", entry_program])
+    env_list  = String.(reduce(vcat, ["-e", k*"="*v] for (k,v) in c.env))
+    cmd  = Cmd(["docker", "exec", "-w", String(code_path), env_list..., String(c_name), "python", String(entry_program)])
     proc_out, proc_err, code = capture_docker_cmd(cmd)
     if code == 0
         c.run_log *= proc_out 
@@ -193,12 +199,6 @@ function start_container!(c::LocalJobContainer, force_update::Bool)
     return c 
 end
 
-function stop_container!(c::LocalJobContainer)
-    c_name = name(c)
-    proc_out, proc_err, code = capture_docker_cmd(`docker stop $c_name`)
-    return
-end
-
 function copy_from_container!(c::LocalJobContainer, src::String, dst::String)
     c_name = c.container_name
     cmd = `docker cp $c_name:$src $dst`