feat: propagate source of Problem and result

Project.toml

@@ -1,7 +1,7 @@
 name = "HarmonicSteadyState"
 uuid = "1158f75c-a779-4b85-8bfb-8fcf6bf02ced"
 authors = ["Orjan Ameye <[email protected]>", "Jan Kosata <[email protected]>", "Javier del Pino <[email protected]>"]
-version = "0.2.6"
+version = "0.3.0"
 
 [deps]
 BijectiveHilbert = "91e7fc40-53cd-4118-bd19-d7fcd1de2a54"
@@ -46,7 +46,7 @@ DocStringExtensions = "0.9.4"
 Documenter = "1.4"
 ExplicitImports = "1.11"
 FunctionWrappers = "1.1.3"
-HarmonicBalance = "0.15.1"
+HarmonicBalance = "0.16.0"
 HomotopyContinuation = "2.12"
 JET = "0.9.18, 0.10"
 LinearAlgebra = "1.10"
@@ -62,7 +62,7 @@ PrecompileTools = "1.2"
 Printf = "1.10"
 ProgressMeter = "1.7.2"
 QuantumCumulants = "0.3.7"
-QuestBase = "0.3.4"
+QuestBase = "0.4.0"
 Random = "1.10"
 RuntimeGeneratedFunctions = "0.5.5"
 SciMLBase = "2.100"

ext/HarmonicBalanceExt.jl

@@ -53,9 +53,16 @@ Any substitution rules not specified in `res` can be supplied in `rules`."
 function HarmonicSteadyState.LinearResponse.ResponseMatrix(
     res::HarmonicSteadyState.Result; rules=Dict()
 )
+    eom = source(res.problem)
+    if isnothing(eom)
+        error("Cannot get the response matrix of the second order natural differential equations of a result with a problem with no source.")
+    elseif !isa(source(eom),QuestBase.DifferentialEquation)
+        error("Cannot get the response matrix the second order natural differential equations of a result where the source system is not second order differential equation.")
+    end
+
     # get the symbolic response matrix
     Symbolics.@variables Δ
-    M = get_response_matrix(res.problem.eom.natural_equation, Num(Δ))
+    M = get_response_matrix(source(eom), Num(Δ))
     M = QuestBase.substitute_all(M, merge(res.fixed_parameters, rules))
     symbols = HarmonicSteadyState._free_symbols(res)
 
@@ -65,7 +72,7 @@ function HarmonicSteadyState.LinearResponse.ResponseMatrix(
         f_im = Symbolics.build_function(el.im, args; expression=Val{false})
         (args...) -> f_re(args...) + im * f_im(args...)
     end
-    return ResponseMatrix(compiled_M, symbols, res.problem.eom.variables)
+    return ResponseMatrix(compiled_M, symbols, eom.variables)
 end
 
 """

ext/TimeEvolution/hysteresis_sweep.jl

@@ -61,7 +61,8 @@ function HarmonicSteadyState.follow_branch(
                 sol_dict[v] = var_values_noise[i]
             end
 
-            problem_t = ODEProblem(res.problem.eom, sol_dict; timespan=(0, tf))
+            harmonic_equation = source(res.problem)
+            problem_t = ODEProblem(harmonic_equation, sol_dict; timespan=(0, tf))
             res_t = solve(problem_t, OrdinaryDiffEqTsit5.Tsit5(); saveat=tf)
 
             # closest branch to final state

src/HarmonicSteadyState.jl

@@ -10,7 +10,9 @@ using QuestBase:
     substitute_all,
     get_independent_variables,
     d,
-    _remove_brackets
+    _remove_brackets,
+    source,
+    source_type
 
 # default global settings
 IM_TOL::Float64 = 1e-6

src/Jacobian.jl

@@ -77,7 +77,7 @@ end
 
 function get_implicit_Jacobian(p::SteadyStateProblem)
     return get_implicit_Jacobian(
-        p.eom; sym_order=_free_symbols(p), rules=p.fixed_parameters
+        source(p); sym_order=_free_symbols(p), rules=p.fixed_parameters
     )
 end
 

src/LimitCycles/LimitCycles.jl

@@ -8,7 +8,8 @@ using QuestBase:
     HarmonicVariable,
     get_independent_variables,
     declare_variable,
-    _remove_brackets
+    _remove_brackets,
+    source
 
 using DocStringExtensions: TYPEDSIGNATURES
 

src/LimitCycles/analysis.jl

@@ -5,7 +5,7 @@ function classify_unique!(res::Result, Δω; class_name="unique_cycle")
     c1 = classify_solutions(res, soln -> _is_physical(soln) && real(soln[i1]) >= 0)
 
     # 2nd degeneracy: ambiguity in the fixed phase, manifests as the sign of var
-    var = _remove_brackets(get_cycle_variables(res.problem.eom, Δω)[1])
+    var = _remove_brackets(get_cycle_variables(source(res.problem), Δω)[1])
     i2 = _symidx(var, res)
     c2 = classify_solutions(res, soln -> _is_physical(soln) && real(soln[i2]) >= 0)
 

src/LinearResponse/LinearResponse.jl

@@ -17,7 +17,7 @@ using HarmonicSteadyState:
     get_class,
     swept_parameters
 
-using QuestBase: HarmonicVariable, substitute_all
+using QuestBase: QuestBase, HarmonicVariable, substitute_all
 
 include("types.jl")
 include("utils.jl")

src/LinearResponse/Lorentzian_spectrum.jl

@@ -88,7 +88,8 @@ these are peaked far from the drive frequency.
 function JacobianSpectrum(
     res::Result{D,S,P}; index::Int, branch::Int, force=false
 ) where {D,S,P}
-    hvars = res.problem.eom.variables # fetch the vector of HarmonicVariable
+    eom = QuestBase.source(res.problem)
+    hvars = eom.variables # fetch the vector of HarmonicVariable
     # blank JacobianSpectrum for each variable
     all_spectra = Dict{Num,JacobianSpectrum{P}}([
         [nvar, JacobianSpectrum{P}()] for nvar in getfield.(hvars, :natural_variable)

src/Problem.jl

@@ -26,6 +26,7 @@ HomotopyContinuationProblem(
 mutable struct HomotopyContinuationProblem{
     ParType<:Number,
     Jac<:JacobianFunction(ComplexF64), # HC.jl only supports Float64
+    Source
 } <: SteadyStateProblem
     "The harmonic variables to be solved for."
     variables::Vector{Num}
@@ -43,19 +44,19 @@ mutable struct HomotopyContinuationProblem{
     """
     jacobian::Jac
     "The HarmonicEquation object used to generate this `HomotopyContinuationProblem`."
-    eom::HarmonicEquation
+    source::Source
 
     function HomotopyContinuationProblem(
         variables, parameters, swept::OrderedDict, fixed::OrderedDict{K,V}, system, jacobian
     ) where {K,V}
-        return new{V,typeof(jacobian)}(
+        return new{V,typeof(jacobian), Nothing}(
             variables, parameters, swept, fixed, system, jacobian
         )
     end # incomplete initialization for user-defined symbolic systems
     function HomotopyContinuationProblem(
         variables, parameters, swept::OrderedDict, fixed::OrderedDict{K,V}, system
     ) where {K,V}
-        return new{V,JacobianFunction(ComplexF64)}(
+        return new{V,JacobianFunction(ComplexF64), Nothing}(
             variables, parameters, swept, fixed, system
         )
     end # incomplete initialization for user-defined symbolic systems
@@ -66,9 +67,9 @@ mutable struct HomotopyContinuationProblem{
         fixed::OrderedDict{K,V},
         system,
         jacobian,
-        eom,
-    ) where {K,V}
-        return new{V,typeof(jacobian)}(
+        eom::T,
+    ) where {K,V, T}
+        return new{V,typeof(jacobian), T}(
             variables, parameters, swept, fixed, system, jacobian, eom
         )
     end
@@ -122,6 +123,9 @@ function Base.show(io::IO, p::HomotopyContinuationProblem)
     return nothing
 end
 
+QuestBase.source(p::HomotopyContinuationProblem) = p.source
+QuestBase.source_type(p::HomotopyContinuationProblem{P,J,Source}) where {P,J,Source} = Source
+
 # assume this order of variables in all compiled function (transform_solutions, Jacobians)
 function _free_symbols(p::HomotopyContinuationProblem)::Vector{Num}
     return cat(p.variables, collect(keys(p.swept_parameters)); dims=1)

src/Result.jl

@@ -7,15 +7,15 @@ Stores the steady states of a HarmonicEquation.
 $(TYPEDFIELDS)
 
 """
-mutable struct Result{D,SolType<:Number,ParType<:Number,F<:JacobianFunction(SolType)}
+mutable struct Result{D,SolType<:Number,ParType<:Number,F<:JacobianFunction(SolType), Source}
     "The variable values of steady-state solutions."
     solutions::Array{Vector{Vector{SolType}},D}
     "Values of all parameters for all solutions."
     swept_parameters::OrderedDict{Num,Vector{ParType}}
     "The parameters fixed throughout the solutions."
     fixed_parameters::OrderedDict{Num,ParType}
     "The `HomotopyContinuationProblem` used to generate this."
-    problem::HomotopyContinuationProblem{ParType,F}
+    problem::HomotopyContinuationProblem{ParType,F, Source}
     """
     Maps strings such as \"stable\", \"physical\" etc to arrays of values,
     classifying the solutions (see method `classify_solutions!`).
@@ -52,8 +52,9 @@ function Result(
     soltype = solution_type(solutions)
     partype = parameter_type(swept_parameters, fixed_parameters)
     dim = ndims(solutions)
+    Source = source_type(problem)
 
-    return Result{dim,soltype,partype,typeof(jacobian)}(
+    return Result{dim,soltype,partype,typeof(jacobian),Source}(
         solutions,
         swept_parameters,
         fixed_parameters,

src/transform_solutions.jl

@@ -253,7 +253,11 @@ function to_lab_frame(soln::OrderedDict, res::Result, nat_var::Num, times)
     else
         nat_var
     end
-    vars = filter(x -> isequal(x.natural_variable, var), res.problem.eom.variables)
+    eom = source(res.problem)
+    if !isa(source(eom), QuestBase.DifferentialEquation)
+        error("Cannot transform to lab frame for a equations of motion with a source not coming from the lab frame.")
+    end
+    vars = filter(x -> isequal(x.natural_variable, var), eom.variables)
 
     return if Symbolics.is_derivative(unwrap(nat_var))
         _to_lab_frame_velocity(soln, vars, times)

test/Problem.jl

@@ -16,7 +16,7 @@ using Test
         F,
         HarmonicSteadyState.JacobianFunction(ComplexF64)(x -> x),
     )
-    @test_throws UndefRefError prob.eom
+    @test isnothing(HarmonicSteadyState.source(prob))
 
     prob = HarmonicSteadyState.HomotopyContinuationProblem(
         [x, y], Num[], OrderedDict{Num,Vector{Float64}}(), OrderedDict{Num,Float64}(), F