add SciMLStructures and relax params type

Author: albertomercurio

Project.toml

@@ -21,6 +21,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
@@ -53,6 +54,7 @@ Random = "1"
 Reexport = "1"
 SciMLBase = "2"
 SciMLOperators = "0.3"
+SciMLStructures = "1.5.0"
 SparseArrays = "1"
 SpecialFunctions = "2"
 StaticArraysCore = "1"

src/QuantumToolbox.jl

@@ -45,6 +45,7 @@ import SciMLOperators:
     IdentityOperator,
     update_coefficients!,
     concretize
+import SciMLStructures: isscimlstructure, ismutablescimlstructure, hasportion, canonicalize, replace, replace!, Tunable
 import LinearSolve: LinearProblem, SciMLLinearSolveAlgorithm, KrylovJL_MINRES, KrylovJL_GMRES
 import DiffEqBase: get_tstops
 import DiffEqCallbacks: PeriodicCallback, PresetTimeCallback, TerminateSteadyState
@@ -87,6 +88,7 @@ include("qobj/superoperators.jl")
 include("qobj/synonyms.jl")
 
 # time evolution
+include("time_evolution/time_evo_parameters.jl")
 include("time_evolution/time_evolution.jl")
 include("time_evolution/mesolve.jl")
 include("time_evolution/lr_mesolve.jl")

src/time_evolution/sesolve.jl

@@ -62,7 +62,7 @@ _sesolve_make_U_QobjEvo(H) = QobjEvo(H, -1im)
         ψ0::QuantumObject{DT2,KetQuantumObject},
         tlist::AbstractVector;
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::NamedTuple = NamedTuple(),
+        params::Union{NamedTuple, AbstractVector} = NamedTuple(),
         progress_bar::Union{Val,Bool} = Val(true),
         inplace::Union{Val,Bool} = Val(true),
         kwargs...,
@@ -80,7 +80,7 @@ Generate the ODEProblem for the Schrödinger time evolution of a quantum system:
 - `ψ0`: Initial state of the system ``|\psi(0)\rangle``.
 - `tlist`: List of times at which to save either the state or the expectation values of the system.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
-- `params`: `NamedTuple` of parameters to pass to the solver.
+- `params`: `NamedTuple` or `AbstractVector` of parameters to pass to the solver.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
 - `inplace`: Whether to use the inplace version of the ODEProblem. The default is `Val(true)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
@@ -101,7 +101,7 @@ function sesolveProblem(
     ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::NamedTuple = NamedTuple(),
+    params::Union{NamedTuple,AbstractVector} = NamedTuple(),
     progress_bar::Union{Val,Bool} = Val(true),
     inplace::Union{Val,Bool} = Val(true),
     kwargs...,
@@ -148,7 +148,7 @@ end
         tlist::AbstractVector;
         alg::OrdinaryDiffEqAlgorithm = Tsit5(),
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::NamedTuple = NamedTuple(),
+        params::Union{NamedTuple, AbstractVector} = NamedTuple(),
         progress_bar::Union{Val,Bool} = Val(true),
         inplace::Union{Val,Bool} = Val(true),
         kwargs...,
@@ -167,7 +167,7 @@ Time evolution of a closed quantum system using the Schrödinger equation:
 - `tlist`: List of times at which to save either the state or the expectation values of the system.
 - `alg`: The algorithm for the ODE solver. The default is `Tsit5()`.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
-- `params`: `NamedTuple` of parameters to pass to the solver.
+- `params`: `NamedTuple` or `AbstractVector` of parameters to pass to the solver.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
 - `inplace`: Whether to use the inplace version of the ODEProblem. The default is `Val(true)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
@@ -190,12 +190,21 @@ function sesolve(
     tlist::AbstractVector;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::NamedTuple = NamedTuple(),
+    params::Union{NamedTuple,AbstractVector} = NamedTuple(),
     progress_bar::Union{Val,Bool} = Val(true),
     inplace::Union{Val,Bool} = Val(true),
     kwargs...,
 ) where {DT1,DT2}
-    prob = sesolveProblem(H, ψ0, tlist; e_ops = e_ops, params = params, progress_bar = progress_bar, inplace = inplace, kwargs...)
+    prob = sesolveProblem(
+        H,
+        ψ0,
+        tlist;
+        e_ops = e_ops,
+        params = params,
+        progress_bar = progress_bar,
+        inplace = inplace,
+        kwargs...,
+    )
 
     return sesolve(prob, alg)
 end

src/time_evolution/time_evo_parameters.jl

@@ -0,0 +1,54 @@
+# This function should be implemented after Julia v1.12
+Base.@constprop :aggressive function _delete_field(a::NamedTuple{an}, field::Symbol) where {an}
+    names = Base.diff_names(an, (field,))
+    return NamedTuple{names}(a)
+end
+
+struct QuantumTimeEvoParameters{TE<:AbstractMatrix,PT<:ProgressBar,ParT}
+    expvals::TE
+    progr::PT
+    params::ParT
+
+    function QuantumTimeEvoParameters(expvals, progr, params)
+        _expvals = expvals
+        _progr = progr
+        _params = params
+
+        # We replace the fields if they are aleady in the `params` struct
+        # Then, we remove them from the `params` struct
+        if :expvals ∈ fieldnames(typeof(_params))
+            _expvals = _params.expvals
+            _params = _delete_field(_params, :expvals)
+        end
+        if :progr ∈ fieldnames(typeof(_params))
+            _progr = _params.progr
+            _params = _delete_field(_params, :progr)
+        end
+
+        return new{typeof(_expvals),typeof(_progr),typeof(_params)}(_expvals, _progr, _params)
+    end
+end
+
+#=
+By defining a custom `getproperty` method for the `QuantumTimeEvoParameters` struct, we can access the fields of `params` directly.
+=#
+function Base.getproperty(obj::QuantumTimeEvoParameters, field::Symbol)
+    if field ∈ fieldnames(typeof(obj))
+        getfield(obj, field)
+    elseif field ∈ fieldnames(typeof(obj.params))
+        getfield(obj.params, field)
+    else
+        throw(KeyError("Field $field not found in QuantumTimeEvoParameters or params."))
+    end
+end
+
+#=
+It also supports `params` as a `Vector`, so we implement the `getindex` method for the `QuantumTimeEvoParameters` struct.
+=#
+Base.getindex(obj::QuantumTimeEvoParameters, i::Int) = getindex(obj.params, i)
+
+Base.length(obj::QuantumTimeEvoParameters) = length(obj.params)
+
+function Base.merge(a::QuantumTimeEvoParameters, b::NamedTuple)
+    return QuantumTimeEvoParameters(a.expvals, a.progr, merge(a.params, b))
+end

src/time_evolution/time_evolution.jl

@@ -5,54 +5,6 @@ export liouvillian_floquet, liouvillian_generalized
 const DEFAULT_ODE_SOLVER_OPTIONS = (abstol = 1e-8, reltol = 1e-6, save_everystep = false, save_end = true)
 const DEFAULT_SDE_SOLVER_OPTIONS = (abstol = 1e-2, reltol = 1e-2, save_everystep = false, save_end = true)
 
-# This function should be implemented after Julia v1.12
-Base.@constprop :aggressive function _delete_field(a::NamedTuple{an}, field::Symbol) where {an}
-    names = Base.diff_names(an, (field,))
-    return NamedTuple{names}(a)
-end
-
-struct QuantumTimeEvoParameters{TE<:AbstractMatrix,PT<:ProgressBar,ParT}
-    expvals::TE
-    progr::PT
-    params::ParT
-
-    function QuantumTimeEvoParameters(expvals, progr, params)
-        _expvals = expvals
-        _progr = progr
-        _params = params
-
-        # We replace the fields if they are aleady in the `params` struct
-        # Then, we remove them from the `params` struct
-        if :expvals ∈ fieldnames(typeof(_params))
-            _expvals = _params.expvals
-            _params = _delete_field(_params, :expvals)
-        end
-        if :progr ∈ fieldnames(typeof(_params))
-            _progr = _params.progr
-            _params = _delete_field(_params, :progr)
-        end
-
-        return new{typeof(_expvals),typeof(_progr),typeof(_params)}(_expvals, _progr, _params)
-    end
-end
-
-#=
-By defining a custom `getproperty` method for the `QuantumTimeEvoParameters` struct, we can access the fields of `params` directly.
-=#
-function Base.getproperty(obj::QuantumTimeEvoParameters, field::Symbol)
-    if field ∈ fieldnames(typeof(obj))
-        getfield(obj, field)
-    elseif field ∈ fieldnames(typeof(obj.params))
-        getfield(obj.params, field)
-    else
-        throw(KeyError("Field $field not found in QuantumTimeEvoParameters or params."))
-    end
-end
-
-function Base.merge(a::QuantumTimeEvoParameters, b::NamedTuple)
-    return QuantumTimeEvoParameters(a.expvals, a.progr, merge(a.params, b))
-end
-
 struct QuantumTimeEvoProblem{PT,TT<:AbstractVector,DT<:AbstractVector}
     prob::PT
     times::TT

test/core-test/time_evolution.jl

@@ -16,7 +16,7 @@
         sol2 = sesolve(H, psi0, t_l, progress_bar = Val(false))
         sol3 = sesolve(H, psi0, t_l, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
         sol_string = sprint((t, s) -> show(t, "text/plain", s), sol)
-        @test prob.f.f isa MatrixOperator
+        @test prob.prob.f.f isa MatrixOperator
         @test sum(abs.(sol.expect[1, :] .- sin.(η * t_l) .^ 2)) / length(t_l) < 0.1
         @test length(sol.times) == length(t_l)
         @test length(sol.states) == 1