Merge pull request #45 from albertomercurio:albertomercurio-patch-1

Introduced ProgressBar

src/QuantumToolbox.jl

@@ -36,6 +36,7 @@ include("spin_lattice.jl")
 include("arnoldi.jl")
 include("eigsolve.jl")
 include("negativity.jl")
+include("progress_bar.jl")
 
 export QuantumObject, Qobj, BraQuantumObject, KetQuantumObject, OperatorQuantumObject, SuperOperatorQuantumObject, TimeEvolutionSol
 export isket, isbra, isoper, issuper, ket2dm

src/progress_bar.jl

@@ -0,0 +1,53 @@
+export ProgressBar
+
+struct ProgressBar{CT,T1<:Integer, T2<:Real}
+    counter::CT
+    max_counts::T1
+    enable::Bool
+    bar_width::T1
+    start_time::T2
+    lock::ReentrantLock
+end
+
+function ProgressBar(max_counts; enable=true, bar_width=30)
+    return ProgressBar(Ref{Int64}(0), max_counts, enable, bar_width, time(), ReentrantLock())
+end
+
+function next!(p::ProgressBar)
+    p.counter[] >= p.max_counts && return
+
+    lock(p.lock)
+
+    p.counter[] += 1
+
+    !p.enable && return
+
+    counter = p.counter[]
+    max_counts = p.max_counts
+    bar_width = p.bar_width
+    start_time = p.start_time
+
+    percentage = counter / max_counts
+    precentage_100 = round(100 * percentage, digits=2)
+    progress = floor(Int, bar_width * percentage)
+
+    # Calculate the elapsed time is seconds
+    elapsed_time = time() - start_time
+    # Convert the elapsed time into a string in hours, minutes and seconds
+    elapsed_time_str = string(floor(Int, elapsed_time / 3600), "h ", floor(Int, (elapsed_time % 3600) / 60), "m ", floor(Int, elapsed_time % 60), "s")
+
+    # Calculate the estimated time of arrival
+    eta = elapsed_time / counter * (max_counts - counter)
+    # convert eta into a string in hours, minutes and seconds
+    eta_str = string(floor(Int, eta / 3600), "h ", floor(Int, (eta % 3600) / 60), "m ", floor(Int, eta % 60), "s")
+
+    # Construct the progress bar string
+    bar = "[" * repeat("=", progress) * repeat(" ", bar_width - progress) * "]"
+
+    print("\rProgress: $bar $precentage_100%  --- Elapsed Time: $elapsed_time_str  (ETA: $eta_str)     ")
+    flush(stdout)
+
+    unlock(p.lock)
+
+    return
+end

src/time_evolution/mcsolve.jl

@@ -63,7 +63,7 @@ function _mcsolve_prob_func(prob, i, repeat)
 
     prm = merge(internal_params, (expvals = similar(internal_params.expvals),
                 cache_mc = similar(internal_params.cache_mc), weights_mc = similar(internal_params.weights_mc), 
-                cumsum_weights_mc = similar(internal_params.weights_mc), random_n = Ref(rand()), progr_mc = ODEProgress(0), jump_times_which_idx = Ref(1),
+                cumsum_weights_mc = similar(internal_params.weights_mc), random_n = Ref(rand()), progr_mc = ProgressBar(size(internal_params.expvals, 2), enable=false), jump_times_which_idx = Ref(1),
                 jump_times = similar(internal_params.jump_times), jump_which = similar(internal_params.jump_which)))
 
     remake(prob, p=prm)
@@ -149,7 +149,7 @@ function mcsolveProblem(H::QuantumObject{MT1,OperatorQuantumObject},
 
     params2 = (expvals = expvals, e_ops_mc = e_ops2,
                 is_empty_e_ops_mc = isempty(e_ops),
-                progr_mc = ODEProgress(0), seeds = seeds,
+                progr_mc = ProgressBar(length(t_l), enable=false), seeds = seeds,
                 random_n = Ref(rand()), c_ops = get_data.(c_ops), cache_mc = cache_mc, 
                 weights_mc = weights_mc, cumsum_weights_mc = cumsum_weights_mc,
                 jump_times = jump_times, jump_which = jump_which,

src/time_evolution/mesolve.jl

@@ -34,6 +34,7 @@ end
         e_ops::AbstractVector=[],
         H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
         params::NamedTuple=NamedTuple(),
+        progress_bar::Bool=true,
         kwargs...)
 
 Generates the ODEProblem for the master equation time evolution of an open quantum system.
@@ -47,6 +48,7 @@ Generates the ODEProblem for the master equation time evolution of an open quant
 - `e_ops::AbstractVector=[]`: The list of the operators for which the expectation values are calculated.
 - `H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing`: The time-dependent Hamiltonian or Liouvillian.
 - `params::NamedTuple=NamedTuple()`: The parameters of the time evolution.
+- `progress_bar::Bool=true`: Whether to show the progress bar.
 - `kwargs...`: The keyword arguments for the ODEProblem.
 
 # Returns
@@ -60,6 +62,7 @@ function mesolveProblem(H::QuantumObject{MT1,HOpType},
     e_ops::Vector{QuantumObject{Te, OperatorQuantumObject}}=QuantumObject{MT1, OperatorQuantumObject}[],
     H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
     params::NamedTuple=NamedTuple(),
+    progress_bar::Bool=true,
     kwargs...) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix,Te<:AbstractMatrix,
     HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
     StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
@@ -72,7 +75,7 @@ function mesolveProblem(H::QuantumObject{MT1,HOpType},
     ρ0 = mat2vec(ket2dm(ψ0).data)
     L = liouvillian(H, c_ops).data
 
-    progr = ODEProgress(0)
+    progr = ProgressBar(length(t_l), enable=progress_bar)
     expvals = Array{ComplexF64}(undef, length(e_ops), length(t_l))
     e_ops2 = @. mat2vec(adjoint(get_data(e_ops)))
 
@@ -105,6 +108,7 @@ end
         e_ops::AbstractVector=[],
         H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
         params::NamedTuple=NamedTuple(),
+        progress_bar::Bool=true,
         kwargs...)
 
 Time evolution of an open quantum system using master equation.
@@ -118,6 +122,7 @@ Time evolution of an open quantum system using master equation.
 - `e_ops::AbstractVector`: List of operators for which to calculate expectation values.
 - `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: Time-dependent part of the Hamiltonian.
 - `params::NamedTuple`: Named Tuple of parameters to pass to the solver.
+- `progress_bar::Bool`: Whether to show the progress bar.
 - `kwargs...`: Additional keyword arguments to pass to the solver.
 
 # Returns
@@ -131,13 +136,14 @@ function mesolve(H::QuantumObject{MT1,HOpType},
     e_ops::Vector{QuantumObject{Te, OperatorQuantumObject}}=QuantumObject{MT1, OperatorQuantumObject}[],
     H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
     params::NamedTuple=NamedTuple(),
+    progress_bar::Bool=true,
     kwargs...) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix,Te<:AbstractMatrix,
     HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
     StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
     COpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}}
 
     prob = mesolveProblem(H, ψ0, t_l, c_ops; alg=alg, e_ops=e_ops,
-            H_t=H_t, params=params, kwargs...)
+            H_t=H_t, params=params, progress_bar=progress_bar, kwargs...)
 
     return mesolve(prob, alg; kwargs...)
 end

src/time_evolution/sesolve.jl

@@ -32,6 +32,7 @@ end
         e_ops::AbstractVector=[],
         H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
         params::NamedTuple=NamedTuple(),
+        progress_bar::Bool=true,
         kwargs...)
 
 Generates the ODEProblem for the Schrödinger time evolution of a quantum system.
@@ -44,6 +45,7 @@ Generates the ODEProblem for the Schrödinger time evolution of a quantum system
 - `e_ops::AbstractVector`: The list of operators to be evaluated during the evolution.
 - `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: The time-dependent Hamiltonian of the system. If `nothing`, the Hamiltonian is time-independent.
 - `params::NamedTuple`: The parameters of the system.
+- `progress_bar::Bool`: Whether to show the progress bar.
 - `kwargs...`: The keyword arguments passed to the `ODEProblem` constructor.
 
 # Returns
@@ -56,6 +58,7 @@ function sesolveProblem(H::QuantumObject{MT1,OperatorQuantumObject},
     e_ops::Vector{QuantumObject{MT2, OperatorQuantumObject}}=QuantumObject{MT1, OperatorQuantumObject}[],
     H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
     params::NamedTuple=NamedTuple(),
+    progress_bar::Bool=true,
     kwargs...) where {MT1<:AbstractMatrix,T2,MT2<:AbstractMatrix}
 
     H.dims != ψ0.dims && throw(ErrorException("The two operators don't have the same Hilbert dimension."))
@@ -65,7 +68,7 @@ function sesolveProblem(H::QuantumObject{MT1,OperatorQuantumObject},
     ϕ0 = get_data(ψ0)
     U = -1im * get_data(H)
 
-    progr = ODEProgress(0)
+    progr = ProgressBar(length(t_l), enable=progress_bar)
     expvals = Array{ComplexF64}(undef, length(e_ops), length(t_l))
     e_ops2 = get_data.(e_ops)
 
@@ -98,6 +101,7 @@ end
         e_ops::AbstractVector=[],
         H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
         params::NamedTuple=NamedTuple(),
+        progress_bar::Bool=true,
         kwargs...)
 
 Time evolution of a closed quantum system using the Schrödinger equation.
@@ -110,6 +114,7 @@ Time evolution of a closed quantum system using the Schrödinger equation.
 - `e_ops::AbstractVector`: List of operators for which to calculate expectation values.
 - `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: Time-dependent part of the Hamiltonian.
 - `params::NamedTuple`: Dictionary of parameters to pass to the solver.
+- `progress_bar::Bool`: Whether to show the progress bar.
 - `kwargs...`: Additional keyword arguments to pass to the solver.
 
 - Returns
@@ -122,10 +127,11 @@ function sesolve(H::QuantumObject{MT1,OperatorQuantumObject},
     e_ops::Vector{QuantumObject{MT2, OperatorQuantumObject}}=QuantumObject{MT1, OperatorQuantumObject}[],
     H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
     params::NamedTuple=NamedTuple(),
+    progress_bar::Bool=true,
     kwargs...) where {MT1<:AbstractMatrix,T2,MT2<:AbstractMatrix}
 
     prob = sesolveProblem(H, ψ0, t_l; alg=alg, e_ops=e_ops,
-            H_t=H_t, params=params, kwargs...)
+            H_t=H_t, params=params, progress_bar=progress_bar, kwargs...)
 
     return sesolve(prob, alg; kwargs...)
 end

src/time_evolution/time_evolution.jl

@@ -26,22 +26,6 @@ end
 
 LiouvillianDirectSolver(;tol=1e-16) = LiouvillianDirectSolver(tol)
 
-# It is needed to keep track of the index for saving the expectation values
-mutable struct ODEProgress{T<:Integer}
-    counter::T
-    # max_counts::T
-    # start_time::T2
-end
-
-function next!(p::ODEProgress)
-    p.counter += 1
-    # prog = round(100 * p.counter / p.max_counts, digits=2)
-    # eta = round((time() - p.start_time) / p.counter * (p.max_counts - p.counter), digits=2)
-    # # convert eta into a string in hours, minutes and seconds
-    # eta_str = string(round(Int, eta / 3600), "h ", round(Int, (eta % 3600) / 60), "m ", round(Int, eta % 60), "s")
-    # print("\rProgress: $prog%  ---  ETA: $eta_str")
-end
-
 abstract type LindbladJumpCallbackType end
 
 struct ContinuousLindbladJumpCallback <: LindbladJumpCallbackType

src/time_evolution/time_evolution_dynamical.jl

@@ -467,7 +467,7 @@ function _dsf_mcsolve_prob_func(prob, i, repeat)
     prm = merge(internal_params, (U = copy(internal_params.U), e_ops_mc = copy(internal_params.e_ops_mc),
                 c_ops = copy(internal_params.c_ops), expvals = similar(internal_params.expvals), 
                 cache_mc = similar(internal_params.cache_mc), weights_mc = similar(internal_params.weights_mc), 
-                cumsum_weights_mc = similar(internal_params.weights_mc), random_n = Ref(rand()), progr_mc = ODEProgress(0), jump_times_which_idx = Ref(1),
+                cumsum_weights_mc = similar(internal_params.weights_mc), random_n = Ref(rand()), progr_mc = ProgressBar(size(internal_params.expvals, 2), enable=false), jump_times_which_idx = Ref(1),
                 jump_times = similar(internal_params.jump_times), jump_which = similar(internal_params.jump_which),
                 αt_list = copy(internal_params.αt_list), dsf_cache1 = similar(internal_params.dsf_cache1),
                 dsf_cache2 = similar(internal_params.dsf_cache2), expv_cache = copy(internal_params.expv_cache),