feat: make HarmonicEquation contain source system

Project.toml

@@ -1,7 +1,7 @@
 name = "QuestBase"
 uuid = "7e80f742-43d6-403d-a9ea-981410111d43"
 authors = ["Orjan Ameye <[email protected]>", "Jan Kosata <[email protected]>", "Javier del Pino <[email protected]>"]
-version = "0.3.4"
+version = "0.4.0"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

src/HarmonicEquation.jl

@@ -1,22 +1,25 @@
+# TODO replace link with HarmonicBalance link
 """
 $(TYPEDEF)
 
-Holds a set of algebraic equations governing the harmonics of a `DifferentialEquation`.
+Holds a set of algebraic equations governing the harmonics of a
+[`DifferentialEquation`](@ref). HarmonicEquation can also be constructed from a
+[`QuantumCumulants.MeanfieldEquations`](https://qojulia.github.io/QuantumCumulants.jl/stable/api/#QuantumCumulants.MeanfieldEquations)`.
 
 # Fields
 $(TYPEDFIELDS)
 """
-mutable struct HarmonicEquation
+mutable struct HarmonicEquation{T}
     """A set of equations governing the harmonics."""
     equations::Vector{Equation}
     """A set of variables describing the harmonics."""
     variables::Vector{HarmonicVariable}
     """The parameters of the equation set."""
     parameters::Vector{Num}
-    "The natural equation (before the harmonic ansatz was used)."
-    natural_equation::DifferentialEquation
     "The Jacobian of the natural equation."
     jacobian::Matrix{Num}
+    "The system where the HarmonicEquation is derived from."
+    source_equations::T
 
     # use a self-referential constructor with _parameters
     function HarmonicEquation(
@@ -25,12 +28,12 @@
         nat_eq::DifferentialEquation,
     )
         return (
-            x=new(
+            x=new{DifferentialEquation}(
                 equations,
                 variables,
                 Num[],
-                nat_eq,
                 dummy_symbolic_Jacobian(length(variables)),
+                nat_eq,
             );
             x.parameters=_parameters(x);
             x
@@ -42,24 +45,28 @@
         parameters::Vector{Num},
         natural_equation::DifferentialEquation,
     )
-        return new(
+        return new{DifferentialEquation}(
             equations,
             variables,
             parameters,
-            natural_equation,
             dummy_symbolic_Jacobian(length(variables)),
+            natural_equation,
         )
     end
     function HarmonicEquation(
         equations::Vector{Equation},
         variables::Vector{HarmonicVariable},
         parameters::Vector{Num},
         jacobian::Matrix{Num},
-    )
-        return new(equations, variables, parameters, DifferentialEquation(), jacobian)
+        source_equations::T,
+    ) where {T}
+        return new{T}(equations, variables, parameters, jacobian, source_equations)
     end
 end
 
+source(eom::HarmonicEquation) = eom.source_equations
+source_type(eom::HarmonicEquation{T}) where {T} = T
+
 "Get the parameters (not time nor variables) of a HarmonicEquation"
 function _parameters(eom::HarmonicEquation)
     all_symbols = flatten([

test/HarmonicEquation.jl

@@ -18,7 +18,8 @@ using QuestBase:
     dummy_symbolic_Jacobian,
     @eqtest,
     get_all_terms,
-    get_independent
+    get_independent,
+    source
 
 # Setup common test variables
 @variables t, T
@@ -42,13 +43,15 @@ hv2 = HarmonicVariable(v, "test", "v", Num(1.0), y)
         heq = heq1
         @test heq.equations == [eq1, eq2]
         @test heq.variables == [hv1, hv2]
-        @test heq.natural_equation == nat_eq
+        @test source(heq) == nat_eq
         @test heq.parameters == Num[]
         @test heq.jacobian isa Matrix{Num}
     end
 
-    heq3 = HarmonicEquation([eq1, eq2], [hv1, hv2], Num[], Num[1 1; 1 1])
-    @test isempty(heq3.natural_equation.harmonics)
+    heq3 = HarmonicEquation(
+        [eq1, eq2], [hv1, hv2], Num[], Num[1 1; 1 1], DifferentialEquation()
+    )
+    @test isempty(source(heq3).harmonics)
 end
 
 @testset "Parameter handling" begin