diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
new file mode 100644
index 0000000..a3ef83e
--- /dev/null
+++ b/.github/workflows/CI.yml
@@ -0,0 +1,94 @@
+name: CI
+# Run on main, any tag or any pull request
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+    tags: '*'
+
+concurrency:
+  # cancel in progress builds for pull requests
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startswith(github.ref, 'refs/pull/') }}
+
+jobs:
+  os-test:
+    runs-on: ${{ matrix.os }}
+    # don't run on draft PRs
+    if: ${{ github.event_name == 'push' || !github.event.pull_request.draft }}
+    # allow windows to fail
+    strategy:
+      fail-fast: true
+      max-parallel: 2
+      matrix:
+        os: [macOS-latest]
+        arch: ['x64']
+    steps:
+      - uses: actions/checkout@v1
+      - name: "Set up Julia"
+        uses: julia-actions/setup-julia@v1
+        with:
+          version: 1.9
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/julia-buildpkg@latest
+      - run: |
+          git config --global user.name Tester
+          git config --global user.email te@st.er
+      - name: "Run tests"
+        uses: julia-actions/julia-runtest@v1
+  version-test:
+    runs-on: ubuntu-latest 
+    # don't run on draft PRs
+    if: ${{ github.event_name == 'push' || !github.event.pull_request.draft }}
+    # allow failures on nightly or beta Julia
+    continue-on-error: ${{ matrix.version == 'nightly'}}
+    strategy:
+      fail-fast: true
+      max-parallel: 2
+      matrix:
+        version: ['1.9', '1.10', 'nightly']
+        arch: ['x64']
+    steps:
+      - uses: actions/checkout@v1
+      - name: "Set up Julia"
+        uses: julia-actions/setup-julia@v1
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/julia-buildpkg@latest
+      - run: |
+          git config --global user.name Tester
+          git config --global user.email te@st.er
+      - name: "Run tests"
+        uses: julia-actions/julia-runtest@v1
+      - name: "Process coverage"
+        uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: src
+      - name: "Upload coverage"
+        uses: codecov/codecov-action@v2
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: lcov.info
+  docs:
+    name: Documentation
+    runs-on: ubuntu-latest
+    if: ${{ github.event_name == 'push' || !github.event.pull_request.draft }}
+    needs: [os-test, version-test]
+    steps:
+      - uses: actions/checkout@v2
+      - uses: julia-actions/setup-julia@v1
+        with:
+          version: '1'
+      - run: |
+          julia --project=docs -e '
+            using Pkg
+            Pkg.develop(PackageSpec(path=pwd()))
+            Pkg.instantiate()
+            include("docs/make.jl")'
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }}
diff --git a/Project.toml b/Project.toml
index 1110ead..9af8192 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,3 +1,9 @@
+name = "BraketAHS"
+uuid = "aad1d28c-1a3f-47da-806d-2ed76430c1d6"
+authors = ["Yaroslav Kharkov", "Katharine Hyatt"]
+version = "0.0.1"
+
+
 [deps]
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
@@ -6,9 +12,14 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Missings = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-Braket = "19504a0f-b47d-4348-9127-acc6cc69ef67"
\ No newline at end of file
+Braket = "19504a0f-b47d-4348-9127-acc6cc69ef67"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
\ No newline at end of file
diff --git a/README.md b/README.md
index d6e156e..79622db 100644
--- a/README.md
+++ b/README.md
@@ -30,7 +30,7 @@ julia --project=. -e 'using Pkg; Pkg.instantiate()'
 ### Running MPS-based simulation
 1. To run the Braket-AHS program saved as a .json file run:
 ```
-julia mps_runner.jl --program-path=<path_to_ahs.json>
+julia src/mps_runner.jl --program-path=<path_to_ahs.json>
 ```
 
 In order to generate .json configuration file there are 2 options:
@@ -54,7 +54,7 @@ with open(filename, "w") as json_file:
 Example .json config files are provided in `examples/ahs_program*.json`.
 
 #### The required input argument is: 
-* `--program-path`: specifies the path to AHS program .json file (example file is provided `./ahs_program.json`)
+* `--program-path`: specifies the path to AHS program .json file (example file is provided `examples/ahs_program.json`)
 
 #### The optional arguments are:
 * `--experiment-path`: directory where all experiment results are saved (default is `examples/experiment_braket`). 
@@ -86,6 +86,6 @@ Output files from MPS simulaiton are stored in `<experiment_path>`:
 
 To run visualization script run:
 ```
-julia plotter.jl <path_to_experiment_results>
+julia src/plotter.jl <path_to_experiment_results>
 ```
 It will generate plots from the .csv files saved in <path_to_experiment_results>.
\ No newline at end of file
diff --git a/docs/Project.toml b/docs/Project.toml
new file mode 100644
index 0000000..7b7f6b6
--- /dev/null
+++ b/docs/Project.toml
@@ -0,0 +1,3 @@
+[deps]
+BraketAHS = "aad1d28c-1a3f-47da-806d-2ed76430c1d6"
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
diff --git a/docs/make.jl b/docs/make.jl
new file mode 100644
index 0000000..7425ea6
--- /dev/null
+++ b/docs/make.jl
@@ -0,0 +1,23 @@
+using BraketAHS
+using Documenter
+
+DocMeta.setdocmeta!(BraketAHS, :DocTestSetup, :(using BraketAHS); recursive=true)
+
+makedocs(;
+    modules=[BraketAHS],
+    authors="Yaroslav Kharkov, Katharine Hyatt",
+    sitename="BraketAHS.jl",
+    format=Documenter.HTML(;
+        canonical="https://ykharkov.github.io/BraketAHS.jl",
+        edit_link="main",
+        assets=String[],
+    ),
+    pages=[
+        "Home" => "index.md",
+    ],
+)
+
+deploydocs(;
+    repo="github.com/ykharkov/BraketAHS.jl",
+    devbranch="main",
+)
diff --git a/docs/src/index.md b/docs/src/index.md
new file mode 100644
index 0000000..9aa93c3
--- /dev/null
+++ b/docs/src/index.md
@@ -0,0 +1,14 @@
+```@meta
+CurrentModule = BraketAHS
+```
+
+# BraketAHS
+
+Documentation for [BraketAHS](https://github.com/ykharkov/BraketAHS.jl).
+
+```@index
+```
+
+```@autodocs
+Modules = [BraketAHS]
+```
diff --git a/mps_runner.jl b/mps_runner.jl
deleted file mode 100644
index 11127d5..0000000
--- a/mps_runner.jl
+++ /dev/null
@@ -1,180 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-using ArgParse
-using ITensors
-using CSV, DataFrames
-using Dates
-using Missings
-using Random
-using Logging
-using JSON
-
-include("src/mps_utils.jl")
-include("src/plotter.jl")
-
-function parse_commandline()
-    s = ArgParseSettings()
-    @add_arg_table! s begin
-        "--program-path"
-            help = "the path to the AHS program JSON file"
-            arg_type = String
-            required = true
-        "--interaction-radius"
-            help = "the interaction radius in meters"
-            arg_type = Float64
-            default = 7e-6
-        "--experiment-path"
-            help = "the directory in which to store all experiment data"
-            arg_type = String
-            default = joinpath("examples", "experiment_braket")
-        "--cutoff"
-            help = "cutoff for SVD values in MPS evolution"
-            arg_type = Float64
-            default = 1e-7
-        "--shots"
-            help = "number of shots for sampling"
-            arg_type = Int
-            default = 1000
-        "--max-bond-dim"
-            help = "maximum bond dimension for MPS"
-            arg_type = Int
-            default = 16
-        "--compute-truncation-error"
-            help = "whether to compute the error induced by truncation at each step (computationally expensive)"
-            action = :store_true # default without this flag is false
-        "--tau"
-            help = "time evolution step size in seconds"
-            arg_type = Float64
-            default = 0.01e-6
-        "--n-tau-steps"
-            help = "number of time evolution steps to simulate"
-            arg_type = Int
-            default = 400
-        "--C6"
-            help = "C6 constant for van der Waals interaction between atoms in Rydberg state (Hz*m^6)"
-            arg_type = Float64
-            default = 5.42e-24
-        "--compute-correlators"
-            help = "Compute ZZ correlators at the end of the evolution (t=T)"
-            action = :store_true
-        "--compute-energies"
-            help = "Compute energies from samples at the end of the evolution (t=T)"
-            action = :store_true            
-        "--generate-plots"
-            help = "Generate plots after experiment is finished"
-            action = :store_true
-    end
-    return parse_args(s)
-end
-
-function run(ahs_json, args)
-
-    experiment_path = args["experiment-path"]
-    τ = args["tau"]
-    n_τ_steps = args["n-tau-steps"]
-    C6 = args["C6"]
-    interaction_R = args["interaction-radius"]
-    n_shots = args["shots"]    
-    Vij, protocol, N = parse_ahs_program(ahs_json, args)
-
-    @info "Preparing initial ψ MPS"
-    s = siteinds("S=1/2", N; conserve_qns=false)
-
-    # Initialize ψ to be a product state: Down state (Ground state)
-    ψ = MPS(s, n -> "Dn")
-    @info "Generating Trotterized circuit"
-    circuit = get_trotterized_circuit_2d(s, τ, n_τ_steps, N, Vij, protocol)
-
-    max_bond_dim = args["max-bond-dim"]
-    cutoff = args["cutoff"]
-    compute_truncation_error = args["compute-truncation-error"]
-
-    @info "Starting MPS evolution"
-    res = @timed begin
-        density, err_array, ψ = compute_MPS_evolution(ψ, circuit, max_bond_dim, cutoff, compute_truncation_error=compute_truncation_error)
-    end
-    summary_array = ["time: $(res.time)",
-                    "n_atoms: $N",
-                    "Trotter steps: $n_τ_steps",
-                    "interaction_R: $interaction_R",
-                    "MPS cutoff: $cutoff",
-                    "max_bond_dim: $max_bond_dim",
-                    "total truncation err: $(sum(err_array))"]
-    summary_string = join(summary_array, ", ")
-    write(joinpath(experiment_path, "summary.txt"), summary_string)
-
-    @info "Simulation complete. Elapsed time and memory used: $(res.time)."
-    @info "Number of atoms: $N, MPS cutoff: $cutoff, max_bond_dim: $max_bond_dim, Trotter steps: $n_τ_steps"
-
-    # Bitstring samples
-    @info "Sampling from final MPS state"
-    samples = Matrix{Int}(undef, N, n_shots)
-    for shot in 1:n_shots
-        sample_i = sample!(ψ) # Sampling bitstrings from a final psi(T)
-        # iTensor MPS sample outputs values [1, 2] for 2-level system
-        # Converting [1,2] -> [0,1]
-        @views samples[:, shot] = [(2 - x) for x in sample_i]
-    end
-
-    # Correlation matrix 
-    correlator_zz = []
-
-    if args["compute-correlators"]
-        @info "Evaluating correlation function ..."
-        correlator_zz = 4 .* correlation_matrix(ψ, "Sz", "Sz") # renormalize to [-1, 1] range
-    end    
-
-    # Energies at t=T
-    energies = []
-    
-    if args["compute-energies"]
-        @info "Evaluating energies at t=T ..."
-
-        Δ_glob_ts = protocol[:global_detuning]
-        Δ_loc_ts = protocol[:local_detuning]
-        pattern = protocol[:pattern]
-    
-        energies = compute_energies(samples', Vij, Δ_glob_ts, Δ_loc_ts, pattern)
-    end    
-
-    results = Dict(
-        "samples" => samples,
-        "density" => density,
-        "summary" => summary_array
-    )
-
-    if args["compute-energies"]
-        results["energies"] = energies
-    end
-
-    if args["compute-correlators"]
-        results["correlator_zz"] = correlator_zz
-    end    
-
-    return results
-end
-
-args = parse_commandline()
-@info "Parsed command line arguments:"
-for (k,v) in args
-    @info "\t$k: $v"
-end
-
-experiment_path = args["experiment-path"]
-program_path    = args["program-path"]
-
-@info "JSON file to read: $program_path"
-ahs_json = JSON.parsefile(program_path)
-
-results = run(ahs_json, args)
-
-@info "Saving results"
-save_results(results, experiment_path)
-
-@info "Generating plots"
-if args["generate-plots"]
-    @info "Plotting results from $experiment_path"
-    plot_all(experiment_path)
-    @info "Plotting complete."
-end
\ No newline at end of file
diff --git a/plotter.jl b/plotter.jl
deleted file mode 100644
index 9f1496a..0000000
--- a/plotter.jl
+++ /dev/null
@@ -1,28 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-using ArgParse
-include("src/plotter.jl")
-
-function parse_commandline()
-    s = ArgParseSettings()
-    @add_arg_table! s begin
-        "--experiment-path"
-            help = "the directory in which to store all experiment data"
-            arg_type = String
-            default = joinpath("examples", "experiment_braket")
-    end
-    return parse_args(s)
-end
-
-
-parsed_args = parse_commandline()
-@info "Parsed command line arguments:"
-for (k,v) in parsed_args
-    @info "\t$k: $v"
-end
-experiment_path = parsed_args["experiment-path"]
-
-@info "Plotting results from $experiment_path"
-plot_all(experiment_path)
-@info "Plotting complete."
\ No newline at end of file
diff --git a/src/BraketAHS.jl b/src/BraketAHS.jl
new file mode 100644
index 0000000..c9e6988
--- /dev/null
+++ b/src/BraketAHS.jl
@@ -0,0 +1,16 @@
+
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+module BraketAHS
+
+export run
+export plot_all
+export plot_density, plot_correlator, plot_bitstrings
+export parse_ahs_program, save_results
+
+include("mps_utils.jl")
+include("plotter.jl")
+include("ahs_program.jl")
+
+end # module
diff --git a/src/ahs_program.jl b/src/ahs_program.jl
index 4737608..6ddf268 100644
--- a/src/ahs_program.jl
+++ b/src/ahs_program.jl
@@ -13,23 +13,20 @@ using Markdown
 using InteractiveUtils
 
 using Braket
-using Braket: AtomArrangement, AtomArrangementItem, TimeSeries, DrivingField, AwsDevice, AnalogHamiltonianSimulation, discretize, AnalogHamiltonianSimulationQuantumTaskResult
+using Braket: AtomArrangement, AtomArrangementItem, TimeSeries, DrivingField, Pattern
 using DataStructures, Statistics, Plots
 
 begin
 	a = 5.5e-6	
 	register = AtomArrangement()
-	push!(register, AtomArrangementItem((0.5, 0.5 + 1/√2) .* a))
-	push!(register, AtomArrangementItem((0.5 + 1/√2, 0.5) .* a))
-	push!(register, AtomArrangementItem((0.5 + 1/√2, -0.5) .* a))
-	push!(register, AtomArrangementItem((0.5, -0.5 - 1/√2) .* a))
-	push!(register, AtomArrangementItem((-0.5, -0.5 - 1/√2) .* a))
-	push!(register, AtomArrangementItem((-0.5 -1/√2, -0.5) .* a))
-	push!(register, AtomArrangementItem((-0.5 -1/√2, 0.5) .* a))
-	push!(register, AtomArrangementItem((-0.5, 0.5 + 1/√2) .* a))
+	push!(register, AtomArrangementItem((0., 0.) .* a))
+	push!(register, AtomArrangementItem((1., 0.) .* a))
+	push!(register, AtomArrangementItem((0., 1.) .* a))
+	push!(register, AtomArrangementItem((1., 1.) .* a))
 end
 
 
+
 begin
 	time_max                = 4e-6  # seconds
 	time_ramp               = 1e-7  # seconds
@@ -48,15 +45,23 @@ begin
 	Δ[time_ramp]            = Δ_start
 	Δ[time_max - time_ramp] = Δ_end
 	Δ[time_max]             = Δ_end
-	
+
+	Δ_loc                   = TimeSeries()
+	Δ_loc[0.0]              = 0.
+	Δ_loc[time_max]             	= 0.
+
 	ϕ           = TimeSeries()
 	ϕ[0.0]      = 0.0
 	ϕ[time_max] = 0.0
 end
 
 begin
-	drive                   = DrivingField(Ω, ϕ, Δ)
-	ahs_program             = AnalogHamiltonianSimulation(register, drive)
+	drive = DrivingField(Ω, ϕ, Δ)
+
+	pt = Pattern([0. for i in 1:length(register)])
+	shift = ShiftingField(Field(Δ_loc, pt))
+
+	ahs_program = AnalogHamiltonianSimulation(register, [drive, shift])
 end
 
 
@@ -64,7 +69,7 @@ json_str = JSON3.write(ir(ahs_program))
 json_obj = JSON.parse(json_str)
 
 # Define the file path
-file_path = "../examples/ahs_program.json"
+file_path = "examples/ahs_program.json"
 
 # Write the JSON object to a file
 open(file_path, "w") do file
diff --git a/src/mps_runner.jl b/src/mps_runner.jl
new file mode 100644
index 0000000..b27c4b2
--- /dev/null
+++ b/src/mps_runner.jl
@@ -0,0 +1,95 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+using ArgParse
+using ITensors
+using CSV, DataFrames
+using Dates
+using Missings
+using Random
+using Logging
+using JSON
+
+include("mps_utils.jl")
+include("plotter.jl")
+
+function parse_commandline()
+    s = ArgParseSettings()
+    @add_arg_table! s begin
+        "--program-path"
+            help = "the path to the AHS program JSON file"
+            arg_type = String
+            default = joinpath("examples", "ahs_program.json")
+        "--interaction-radius"
+            help = "the interaction radius in meters"
+            arg_type = Float64
+            default = 7e-6
+        "--experiment-path"
+            help = "the directory in which to store all experiment data"
+            arg_type = String
+            default = joinpath("examples", "experiment_braket")
+        "--cutoff"
+            help = "cutoff for SVD values in MPS evolution"
+            arg_type = Float64
+            default = 1e-7
+        "--shots"
+            help = "number of shots for sampling"
+            arg_type = Int
+            default = 1000
+        "--max-bond-dim"
+            help = "maximum bond dimension for MPS"
+            arg_type = Int
+            default = 16
+        "--compute-truncation-error"
+            help = "whether to compute the error induced by truncation at each step (computationally expensive)"
+            action = :store_true # default without this flag is false
+        "--tau"
+            help = "time evolution step size in seconds"
+            arg_type = Float64
+            default = 0.01e-6
+        "--n-tau-steps"
+            help = "number of time evolution steps to simulate"
+            arg_type = Int
+            default = 400
+        "--C6"
+            help = "C6 constant for van der Waals interaction between atoms in Rydberg state (Hz*m^6)"
+            arg_type = Float64
+            default = 5.42e-24
+        "--compute-correlators"
+            help = "Compute ZZ correlators at the end of the evolution (t=T)"
+            action = :store_true
+        "--compute-energies"
+            help = "Compute energies from samples at the end of the evolution (t=T)"
+            action = :store_true            
+        "--generate-plots"
+            help = "Generate plots after experiment is finished"
+            action = :store_true
+    end
+    return parse_args(s)
+end
+
+
+args = parse_commandline()
+        
+@info "Parsed command line arguments:"
+for (k,v) in args
+    @info "\t$k: $v"
+end
+
+experiment_path = args["experiment-path"]
+program_path    = args["program-path"]
+
+@info "JSON file to read: $program_path"
+ahs_json = JSON.parsefile(program_path)
+
+results = run(ahs_json, args)
+
+@info "Saving results"
+save_results(results, experiment_path)
+
+@info "Generating plots"
+if args["generate-plots"]
+    @info "Plotting results from $experiment_path"
+    plot_all(experiment_path)
+    @info "Plotting complete."
+end
diff --git a/src/mps_utils.jl b/src/mps_utils.jl
index 32a43dc..f7f288b 100644
--- a/src/mps_utils.jl
+++ b/src/mps_utils.jl
@@ -6,7 +6,6 @@ using CSV, DataFrames
 using Dates
 using Base.Filesystem
 using Missings
-using PythonCall
 using Random
 using JSON
 
@@ -238,10 +237,10 @@ function compute_MPS_evolution(ψ::MPS, circuit::Vector{Vector{ITensor}}, max_bo
         meas_array[:, i_τ] = 0.5 .+ Sz
         # Evolution step: psi update
         ψ_prev = ψ
-        ψ = apply(circuit[i_τ], ψ; cutoff=cutoff, maxdim=max_bond_dim)
+        ψ = ITensors.apply(circuit[i_τ], ψ; cutoff=cutoff, maxdim=max_bond_dim)
         normalize!(ψ)
         if compute_truncation_error
-            ψ_exact = apply(circuit[i_τ], ψ_prev; cutoff=1e-16)
+            ψ_exact = ITensors.apply(circuit[i_τ], ψ_prev; cutoff=1e-16)
             err_array[i_τ] = 1.0 - abs(inner(ψ_exact, ψ))
         end            
         @info "Step: $i_τ, current MPS bond dimension is $(maxlinkdim(ψ))"         
@@ -249,6 +248,94 @@ function compute_MPS_evolution(ψ::MPS, circuit::Vector{Vector{ITensor}}, max_bo
     return meas_array, err_array, ψ    
 end
 
+function run(ahs_json, args)
+
+    experiment_path = args["experiment-path"]
+    τ = args["tau"]
+    n_τ_steps = args["n-tau-steps"]
+    C6 = args["C6"]
+    interaction_R = args["interaction-radius"]
+    n_shots = args["shots"]    
+    Vij, protocol, N = parse_ahs_program(ahs_json, args)
+
+    @info "Preparing initial ψ MPS"
+    s = siteinds("S=1/2", N; conserve_qns=false)
+
+    # Initialize ψ to be a product state: Down state (Ground state)
+    ψ = MPS(s, n -> "Dn")
+    @info "Generating Trotterized circuit"
+    circuit = get_trotterized_circuit_2d(s, τ, n_τ_steps, N, Vij, protocol)
+
+    max_bond_dim = args["max-bond-dim"]
+    cutoff = args["cutoff"]
+    compute_truncation_error = args["compute-truncation-error"]
+
+    @info "Starting MPS evolution"
+    res = @timed begin
+        density, err_array, ψ = compute_MPS_evolution(ψ, circuit, max_bond_dim, cutoff, compute_truncation_error=compute_truncation_error)
+    end
+    summary_array = ["time: $(res.time)",
+                    "n_atoms: $N",
+                    "Trotter steps: $n_τ_steps",
+                    "interaction_R: $interaction_R",
+                    "MPS cutoff: $cutoff",
+                    "max_bond_dim: $max_bond_dim",
+                    "total truncation err: $(sum(err_array))"]
+    summary_string = join(summary_array, ", ")
+    write(joinpath(experiment_path, "summary.txt"), summary_string)
+
+    @info "Simulation complete. Elapsed time and memory used: $(res.time)."
+    @info "Number of atoms: $N, MPS cutoff: $cutoff, max_bond_dim: $max_bond_dim, Trotter steps: $n_τ_steps"
+
+    # Bitstring samples
+    @info "Sampling from final MPS state"
+    samples = Matrix{Int}(undef, N, n_shots)
+    for shot in 1:n_shots
+        sample_i = sample!(ψ) # Sampling bitstrings from a final psi(T)
+        # iTensor MPS sample outputs values [1, 2] for 2-level system
+        # Converting [1,2] -> [0,1]
+        @views samples[:, shot] = [(2 - x) for x in sample_i]
+    end
+
+    # Correlation matrix 
+    correlator_zz = []
+
+    if args["compute-correlators"]
+        @info "Evaluating correlation function ..."
+        correlator_zz = 4 .* correlation_matrix(ψ, "Sz", "Sz") # renormalize to [-1, 1] range
+    end    
+
+    # Energies at t=T
+    energies = []
+    
+    if args["compute-energies"]
+        @info "Evaluating energies at t=T ..."
+
+        Δ_glob_ts = protocol[:global_detuning]
+        Δ_loc_ts = protocol[:local_detuning]
+        pattern = protocol[:pattern]
+    
+        energies = compute_energies(samples', Vij, Δ_glob_ts, Δ_loc_ts, pattern)
+    end    
+
+    results = Dict(
+        "samples" => samples,
+        "density" => density,
+        "summary" => summary_array
+    )
+
+    if args["compute-energies"]
+        results["energies"] = energies
+    end
+
+    if args["compute-correlators"]
+        results["correlator_zz"] = correlator_zz
+    end    
+
+    return results
+end
+
+
 """
     save_results(experiment_path::String,
                  ψ::MPS,
diff --git a/test/runtests.jl b/test/runtests.jl
new file mode 100644
index 0000000..28feb69
--- /dev/null
+++ b/test/runtests.jl
@@ -0,0 +1,36 @@
+using BraketAHS: run, save_results
+using Test
+using JSON, JSON3
+
+@testset "BraketAHS.jl" begin
+    # Write your tests here.
+
+    json_str = """{"setup":{"ahs_register":{"sites":[["0.0","0.0"],["5.499999999999999856993733449161965e-6","0.0"],["0.0","5.499999999999999856993733449161965e-6"],["5.499999999999999856993733449161965e-6","5.499999999999999856993733449161965e-6"]],"filling":[1,1,1,1]}},"braketSchemaHeader":{"name":"braket.ir.ahs.program","version":"1"},"hamiltonian":{"shiftingFields":[{"magnitude":{"pattern":["0.0","0.0","0.0","0.0"],"time_series":{"values":["0.0","0.0"],"times":["0.0","3.999999999999999818992447303545035e-6"]}}}],"drivingFields":[{"phase":{"pattern":"uniform","time_series":{"values":["0.0","0.0"],"times":["0.0","3.999999999999999818992447303545035e-6"]}},"detuning":{"pattern":"uniform","time_series":{"values":["-3.15e7","-3.15e7","3.15e7","3.15e7"],"times":["0.0","9.999999999999999547481118258862587e-8","3.899999999999999929396754527743951e-6","3.999999999999999818992447303545035e-6"]}},"amplitude":{"pattern":"uniform","time_series":{"values":["0.0","6.3e6","6.3e6","0.0"],"times":["0.0","9.999999999999999547481118258862587e-8","3.899999999999999929396754527743951e-6","3.999999999999999818992447303545035e-6"]}}}]}}
+    """
+    ahs_json = JSON.parse(json_str)
+
+    args = Dict(
+        "experiment-path" => "test/",
+        "program-path" => "",
+        "interaction-radius" => 7e-6,
+        "cutoff" => 1e-7,
+        "shots" => 1000,
+        "max-bond-dim" => 16,
+        "compute-truncation-error" => false,
+        "tau" => 0.01e-6,        
+        "n-tau-steps" => 400,
+        "C6" => 5.42e-24,
+        "compute-correlators" => false,
+        "compute-energies" => false,
+        "generate-plots" => false
+    )
+
+    @show ahs_json
+
+    result = run(ahs_json, args)
+
+    mps_density = result["density"][:, end]
+    expected_density = [0.49558400861051327,0.4955839949350386,0.49558370707943783,0.49558323080865946]
+    @test isapprox(mps_density, expected_density; atol=1e-7)
+    # add apprx cmoparison isapprox(array1, array2; atol=tolerance)
+end