Parsing of time-dependent data, plot overload (#31)

jkosata · web-flow · commit 75528ae02613 · 2022-05-26T10:51:35.000+02:00
* transform_solutions for time-dependent results

* plot overloaded for Result and ODESolution

* examples and docs changes

* Jacobian bugfix
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "HarmonicBalance"
 uuid = "e13b9ff6-59c3-11ec-14b1-f3d2cc6c135e"
 authors = ["Jan Kosata <kosataj@phys.ethz.ch>", "Javier del Pino <jdelpino@phys.ethz.ch>"]
-version = "0.5.1"
+version = "0.5.2"
 
 [deps]
 BijectiveHilbert = "91e7fc40-53cd-4118-bd19-d7fcd1de2a54"
diff --git a/README.md b/README.md
@@ -48,7 +48,7 @@ Classes: stable, physical, Hopf, binary_labels
 ```
 
 ```julia
-plot_1D_solutions(solutions, x="ω", y="sqrt(u1^2 + v1^2)");
+plot(solutions, x="ω", y="sqrt(u1^2 + v1^2)");
 ```
 
 <img src="/docs/images/DuffingPlot.png" width="500">
diff --git a/docs/src/examples/linear_response.md b/docs/src/examples/linear_response.md
@@ -29,7 +29,7 @@ fixed = (α => 1, ω0 => 1.0, γ => 1E-2, F => 1E-6)   # fixed parameters
 swept = ω => LinRange(0.9, 1.1, 100)           # range of parameter values
 solutions = get_steady_states(harmonic_eq, swept, fixed)
 
-plot_1D_solutions(solutions, x="ω", y="sqrt(u1^2 + v1^2)")
+plot(solutions, x="ω", y="sqrt(u1^2 + v1^2)")
 ```
 
 ```@raw html
@@ -54,7 +54,7 @@ fixed = (α => 1, ω0 => 1.0, γ => 1E-2, F => 1E-2)   # fixed parameters
 swept = ω => LinRange(0.9, 1.1, 100)           # range of parameter values
 solutions = get_steady_states(harmonic_eq, swept, fixed)
 
-plot_1D_solutions(solutions, x="ω", y="sqrt(u1^2 + v1^2)");
+plot(solutions, x="ω", y="sqrt(u1^2 + v1^2)");
 ```
 ```@raw html
 <img style="display: block; margin: 0 auto;" src="../../assets/linear_response/Duffing_nonlin_amp.png" width="450" alignment="left" \>
@@ -81,7 +81,7 @@ fixed = (α => 1., ω0 => 1.0, γ => 1E-2, ω => 1)   # fixed parameters
 swept = F => 10 .^ LinRange(-6, -1, 200)           # range of parameter values
 solutions = get_steady_states(harmonic_eq, swept, fixed)
 
-plot_1D_solutions(solutions, x="F", y="sqrt(u1^2 + v1^2)");
+plot(solutions, x="F", y="sqrt(u1^2 + v1^2)");
 HarmonicBalance.xscale("log") # use log scale on x
 
 LinearResponse.plot_jacobian_spectrum(solutions, x, 
diff --git a/docs/src/examples/single_Duffing.md b/docs/src/examples/single_Duffing.md
@@ -104,7 +104,7 @@ The "Classes" are boolean labels classifying each solution point, which may be u
 
 We now want to visualize the results. Here we plot the solution amplitude, ``\sqrt{U^2 + V^2}`` against the drive frequency ``\omega``: 
 ```julia
-plot_1D_solutions(solutions, x="ω", y="sqrt(u1^2 + v1^2)")
+plot(solutions, x="ω", y="sqrt(u1^2 + v1^2)")
 ```
 ```@raw html
 <img style="display: block; margin: 0 auto;" src="../../assets/duffing_single.png" width="450" alignment="center" \>
@@ -170,8 +170,8 @@ Classes: stable, physical, Hopf, binary_labels
 
 Although 9 branches were found in total, only 3 remain physical (real-valued). Let us visualise the amplitudes corresponding to the two harmonics, ``\sqrt{U_1^2 + V_1^2}`` and ``\sqrt{U_2^2 + V_2^2}`` :
 ```julia
-plot_1D_solutions(solutions, x="ω", y="sqrt(u1^2 + v1^2)")
-plot_1D_solutions(solutions, x="ω", y="sqrt(u2^2 + v2^2)")
+plot(solutions, x="ω", y="sqrt(u1^2 + v1^2)")
+plot(solutions, x="ω", y="sqrt(u2^2 + v2^2)")
 ```
 ![fig3](./../assets/duff_nonpert_w_3w.png)
 
diff --git a/docs/src/examples/single_parametron_1D.md b/docs/src/examples/single_parametron_1D.md
@@ -55,15 +55,15 @@ HarmonicBalance.save("parametron_result.jld2", soln);
 
 During the execution of `get_steady_states`, different solution branches are classified by their proximity in complex space, with subsequent filtering of real (physically accceptable solutions). In addition, the stability properties of each steady state is assesed from the eigenvalues of the Jacobian matrix. All this information can be succintly represented in a 1D plot via
 ```julia
-save_dict = HarmonicBalance.plot_1D_solutions(soln, x="ω", y="sqrt(u1^2 + v1^2)", plot_only=["physical"]);
+save_dict = HarmonicBalance.plot(soln, x="ω", y="sqrt(u1^2 + v1^2)", plot_only=["physical"]);
 ```
-where `save_dict` is a dictionary that contains the plotted data and can be also exported if desired by setting a filename through the argument `filename` in `plot_1D_solutions`. A call to the above function produces the following figure
+where `save_dict` is a dictionary that contains the plotted data and can be also exported if desired by setting a filename through the argument `filename` in `plot`. A call to the above function produces the following figure
 
 ![fig1](./../assets/single_parametron_1D.png)
 
 The user can also can also introduce custom clases based on parameter conditions. Here we show some arbitrary example
 ```julia
-plt = HarmonicBalance.plot_1D_solutions(soln, x="ω", y="sqrt(u1^2 + v1^2)", marker_classification="ω^15 * sqrt(u1^2 + v1^2) < 0.1")
+plt = HarmonicBalance.plot(soln, x="ω", y="sqrt(u1^2 + v1^2)", marker_classification="ω^15 * sqrt(u1^2 + v1^2) < 0.1")
 ```
 producing 
 
diff --git a/docs/src/examples/single_parametron_time_dep.md b/docs/src/examples/single_parametron_time_dep.md
@@ -30,7 +30,7 @@ fixed_parameters = ParameterList(Ω => 1.0,γ => 1E-2, λ => 5E-2, F => 1E-3,  
 Finally, we solve the harmonic equations and represent the solutions  by
 
 ```julia
-time_dep = HarmonicBalance.TimeEvolution.ODEProblem(averagedEOM, fixed_parameters, sweep=HarmonicBalance.TimeEvolution.ParameterSweep(), x0 = x0, timespan = times);
+time_dep = HarmonicBalance.TimeEvolution.(averagedEOM, fixed_parameters, sweep=HarmonicBalance.TimeEvolution.ParameterSweep(), x0 = x0, timespan = times);
 time_soln = HarmonicBalance.TimeEvolution.solve(time_dep, saveat=dt);
 HarmonicBalance.plot(getindex.(time_soln.u, 1), getindex.(time_soln.u,2))
 HarmonicBalance.xlabel("u",fontsize=20)
diff --git a/docs/src/examples/time_dependent.md b/docs/src/examples/time_dependent.md
@@ -67,7 +67,7 @@ DifferentialEquations.jl takes it from here - we only need to use `solve`.
 
 ```julia
 time_evo = solve(ode_problem, saveat=1.);
-plot(getindex.(time_evo.u, 1), getindex.(time_evo.u,2))
+plot(time_evo, ["u1", "v1"], harmonic_eqs)
 ```
 
 Running the above code with `x0 = [0., 0.]` and `x0 = [0.2, 0.2]` gives the plots
@@ -79,7 +79,7 @@ Let us compare this to the steady state diagram.
 ```julia
 range = ω => LinRange(0.9, 1.1, 100)           # range of parameter values
 solutions = get_steady_states(harmonic_eqs, range, fixed)
-plot_1D_solutions(solutions, x="ω", y="sqrt(u1^2 + v1^2)*sign(u1)");
+plot(solutions, x="ω", y="sqrt(u1^2 + v1^2)*sign(u1)");
 ```
 ```@raw html
 <img style="display: block; margin: 0 auto;" src="../../assets/time_dependent/parametron_response.png" width="600" alignment="center" \>
@@ -101,7 +101,7 @@ Let us now define a new `ODEProblem` which incorporates `sweep` and again use `s
 ```julia
 ode_problem = ODEProblem(harmonic_eqs, fixed, sweep=sweep, x0=[0.1;0.0], timespan=(0, 2E4))
 time_evo = solve(prob, saveat=100)
-plot(time_evo.t, sqrt.(getindex.(time_evo.u,1).^2 .+ getindex.(time_evo,2).^2))
+plot(time_evo, "sqrt(u1^2 + v1^2)", harmonic_eqs)
 ```
 ```@raw html
 <img style="display: block; margin: 0 auto;" src="../../assets/time_dependent/sweep_omega.png" width="450" alignment="center" \>
@@ -190,7 +190,7 @@ Solution branches:   9
 
 Let us first see the steady states.
 ```julia
-plt = plot_1D_solutions(soln, x="F0", y="u1^2 + v1^2", yscale=:log);
+plt = plot(soln, x="F0", y="u1^2 + v1^2", yscale=:log);
 ```
 ```@raw html
 <img style="display: block; margin: 0 auto; padding-bottom: 20px" src="../../assets/time_dependent/pra_ss.png" width="1000" alignment="center"\>
@@ -212,7 +212,7 @@ TDsoln = solve(TDproblem, saveat=1); # space datapoints by 1
 ```
 Inspecting for example $u_1(T)$ as a function of time,
 ```julia
-plot(TDsoln.t, getindex.(TDsoln.u, 1))
+plot(time_evo, "u1", harmonic_eqs)
 ```
 ```@raw html
 <img style="display: block; margin: 0 auto; padding-bottom: 20px" src="../../assets/time_dependent/lc_sweep.png" width="450" alignment="center" \>
diff --git a/docs/src/manual/plotting.md b/docs/src/manual/plotting.md
@@ -11,12 +11,17 @@ HarmonicBalance.transform_solutions
 Any function of the steady state solutions may be plotted. 
 In 1D, the solutions are colour-coded according to the branches obtained by `sort_solutions`. 
 
+Note: from v0.5.2, `plot(r::Result, x::String, y::String)` can be used to call `plot_1D_solutions` and `plot_2D_solutions` as needed.
+To plot a function `y` of a time-dependent result `r`, the syntax is `plot(r::OrdinaryDiffEq.ODECompositeSolution, y, he::HarmonicEquation)`. For `y::String`, `y` is parsed into a function and plotted as a function of time.
+
 ```@docs
 HarmonicBalance.plot_1D_solutions
 HarmonicBalance.plot_1D_jacobian_eigenvalues
 HarmonicBalance.plot_2D_solutions
 ```
 
+
+
 ## Plotting phase diagrams (2D)
 
 In many problems, rather than in any property of the solutions themselves, we are interested in the phase diagrams, encoding the number of (stable) solutions in different regions of the parameter space. We provide functions to tackle solutions calculated over 2D parameter grids.
diff --git a/src/HarmonicBalance.jl b/src/HarmonicBalance.jl
@@ -44,6 +44,7 @@ module HarmonicBalance
     include("sorting.jl")
     include("classification.jl")
     include("saving.jl")
+    include("transform_solutions.jl")
     include("plotting_static.jl")
     include("plotting_interactive.jl")
     include("hysteresis_sweep.jl")
diff --git a/src/modules/HC_wrapper/homotopy_interface.jl b/src/modules/HC_wrapper/homotopy_interface.jl
@@ -47,7 +47,7 @@ function Problem(eom::HarmonicEquation; Jacobian=true)
     elseif Jacobian == "implicit"
         # compute the Jacobian implicitly
         J = HarmonicBalance.LinearResponse.get_implicit_Jacobian(eom)
-    elseif Jacobian == "false"
+    elseif Jacobian == "false" || Jacobian == false
         dummy_J(arg) = I(1)
         J = dummy_J
     else
diff --git a/src/modules/TimeEvolution/ODEProblem.jl b/src/modules/TimeEvolution/ODEProblem.jl
@@ -1,5 +1,6 @@
-using LinearAlgebra
-
+using LinearAlgebra, PyPlot
+import HarmonicBalance: transform_solutions, plot
+export transform_solutions, plot
 
 """ 
 
@@ -15,7 +16,7 @@ Creates an ODEProblem object used by DifferentialEquations.jl from the equations
 `fixed_parameters` must be a dictionary mapping parameters+variables to numbers (possible to use a solution index, e.g. solutions[x][y] for branch y of solution x).
 If `x0` is specified, it is used as an initial condition; otherwise the values from `fixed_parameters` are used.
 """
-function ODEProblem(eom::HarmonicEquation, fixed_parameters; sweep::ParameterSweep=ParameterSweep(), x0::Vector=[], timespan::Tuple, perturb_initial=0.0)
+function ODEProblem(eom::HarmonicEquation, fixed_parameters; sweep::ParameterSweep=ParameterSweep(), x0::Vector=[], timespan::Tuple, perturb_initial=0.0, kwargs...)
 
     if !is_rearranged(eom) # check if time-derivatives of the variable are on the right hand side
         eom = HarmonicBalance.rearrange_standard(eom)
@@ -43,7 +44,7 @@ function ODEProblem(eom::HarmonicEquation, fixed_parameters; sweep::ParameterSwe
     # the initial condition is x0 if specified, taken from fixed_parameters otherwise
     initial = isempty(x0) ? real.(collect(values(fixed_parameters))[1:length(vars)]) * (1-perturb_initial) : x0
 
-    return DifferentialEquations.ODEProblem(f!, initial, timespan)
+    return DifferentialEquations.ODEProblem(f!, initial, timespan; kwargs...)
 end
 
 
@@ -61,4 +62,19 @@ function is_stable(soln::StateDict, eom::HarmonicEquation; timespan, tol=1E-1, p
     solution = solve(problem)
     dist = norm(solution[end] - solution[1]) / (norm(solution[end]) + norm(solution[1]))
     !is_real(solution[end]) || !is_real(solution[1]) ? error("the solution is complex!") : dist < tol
+end
+
+
+transform_solutions(soln::OrdinaryDiffEq.ODECompositeSolution, f::String, harm_eq::HarmonicEquation) = transform_solutions(soln.u, f, harm_eq)
+transform_solutions(s::OrdinaryDiffEq.ODECompositeSolution, funcs::Vector{String}, he::HarmonicEquation) = [transform_solutions(s, f, he) for f in funcs]
+
+
+function plot(soln::OrdinaryDiffEq.ODECompositeSolution, funcs, harm_eq::HarmonicEquation)
+    if funcs isa String || length(funcs) == 1
+        plot(soln.t, transform_solutions(soln, funcs, harm_eq))
+    elseif length(funcs) == 2
+        plot(transform_solutions(soln, funcs, harm_eq)...)
+    else
+        error("Invalid plotting argument: ", funcs)
+    end
 end
diff --git a/src/plotting_static.jl b/src/plotting_static.jl
@@ -3,9 +3,12 @@ using PyCall
 using Latexify
 using JLD2
 using Base
-export plot_1D_solutions, plot_2D_phase_diagram, transform_solutions
+export plot_1D_solutions, plot_2D_phase_diagram
 export _set_plotting_settings, _prepare_colorbar, _prettify_label
 
+import PyPlot: plot
+export plot
+
 
 "Set global plotting settings"
 function _set_plotting_settings()
@@ -65,46 +68,6 @@ function filter_solutions(solution::Vector,  booleans)
 end
 
 
-"""
-$(TYPEDSIGNATURES)
-
-Takes a `Result` object and a string `f` representing a Symbolics.jl expression.
-Returns an array with the values of `f` evaluated for the respective solutions.
-Additional substitution rules can be specified in `rules` in the format `("a" => val)` or `(a => val)`
-"""
-function transform_solutions(res::Result, f::String; rules=Dict())
-    # a string is used as input - a macro would not "see" the user's namespace while the user's namespace does not "see" the variables
-    transformed = [Vector{ComplexF64}(undef, length(res.solutions[1])) for k in res.solutions] # preallocate
-        
-    # define variables in rules in this namespace
-    new_keys = declare_variable.(string.(keys(Dict(rules)))) 
-    expr = f isa String ? Num(eval(Meta.parse(f))) : f
-
-    fixed_subs = merge(res.fixed_parameters, Dict(zip(new_keys, values(Dict(rules)))))
-    expr = substitute_all(expr, Dict(fixed_subs))
-
-    vars = res.problem.variables
-    all_symbols = cat(vars, collect(keys(res.swept_parameters)), dims=1)
-    comp_func = build_function(expr, all_symbols) 
-    f = eval(comp_func)
-
-    # preallocate an array for the numerical values, rewrite parts of it
-    # when looping through the solutions
-    vals = Vector{ComplexF64}(undef, length(all_symbols))
-    n_vars = length(vars)
-    n_pars = length(all_symbols) - n_vars
-
-    for idx in CartesianIndices(res.solutions)
-        params_values = res.swept_parameters[Tuple(idx)...]
-        vals[end-n_pars+1:end] .= params_values # param values are common to all branches
-        for (branch,soln) in enumerate(res.solutions[idx])
-            vals[1:n_vars] .= soln
-            transformed[idx][branch] = Base.invokelatest(f, vals)
-        end
-    end
-    return transformed
-end
-
 """
 $(TYPEDSIGNATURES)
 
@@ -589,4 +552,22 @@ function plot_2D_phase_diagram(res::Result; stable=false,observable="nsols",ax=n
     !isnothing(filename) ? JLD2.save(_jld2_name(filename),save_dict) : nothing
     return save_dict,im,Nmax
 
-end
+end
+
+
+##########
+# Plot multiple dispatch
+##########
+
+function plot(res::Result; x::String, y::String, kwargs...)
+    if length(size(res.solutions)) == 1
+        plot_1D_solutions(res; x=x, y=y, kwargs...)
+    elseif length(size(res.solutions)) == 2
+        plot_2D_solutions(res; x=x, y=y, kwargs...)
+    else
+        error("error")
+    end
+end
+
+
+plot(res::Result, x::String, y::String; kwargs...) = plot(res; x=x, y=y, kwargs...)
diff --git a/src/transform_solutions.jl b/src/transform_solutions.jl
@@ -0,0 +1,62 @@
+export transform_solutions
+
+
+_parse_expression(exp) = exp isa String ? Num(eval(Meta.parse(exp))) : exp
+
+
+"""
+$(TYPEDSIGNATURES)
+
+Takes a `Result` object and a string `f` representing a Symbolics.jl expression.
+Returns an array with the values of `f` evaluated for the respective solutions.
+Additional substitution rules can be specified in `rules` in the format `("a" => val)` or `(a => val)`
+"""
+function transform_solutions(res::Result, f::String; rules=Dict())
+    # a string is used as input - a macro would not "see" the user's namespace while the user's namespace does not "see" the variables
+    transformed = [Vector{ComplexF64}(undef, length(res.solutions[1])) for k in res.solutions] # preallocate
+        
+    # define variables in rules in this namespace
+    new_keys = declare_variable.(string.(keys(Dict(rules)))) 
+    expr = f isa String ? _parse_expression(f) : f
+
+    fixed_subs = merge(res.fixed_parameters, Dict(zip(new_keys, values(Dict(rules)))))
+    expr = substitute_all(expr, Dict(fixed_subs))
+
+    vars = res.problem.variables
+    all_symbols = cat(vars, collect(keys(res.swept_parameters)), dims=1)
+    comp_func = build_function(expr, all_symbols) 
+    f = eval(comp_func)
+
+    # preallocate an array for the numerical values, rewrite parts of it
+    # when looping through the solutions
+    vals = Vector{ComplexF64}(undef, length(all_symbols))
+    n_vars = length(vars)
+    n_pars = length(all_symbols) - n_vars
+
+    for idx in CartesianIndices(res.solutions)
+        params_values = res.swept_parameters[Tuple(idx)...]
+        vals[end-n_pars+1:end] .= params_values # param values are common to all branches
+        for (branch,soln) in enumerate(res.solutions[idx])
+            vals[1:n_vars] .= soln
+            transformed[idx][branch] = Base.invokelatest(f, vals)
+        end
+    end
+    return transformed
+end
+
+
+# a simplified version meant to work with arrays of solutions
+# cannot parse parameter values  mainly meant for time-dependent results
+function transform_solutions(soln::Vector, f::String, harm_eq::HarmonicEquation)
+
+    vars = _remove_brackets(get_variables(harm_eq))
+    transformed = Vector{ComplexF64}(undef, length(soln))
+
+    # parse the input with Symbolics
+    expr = HarmonicBalance._parse_expression(f)
+
+    rule(u) = Dict(zip(vars, u))
+
+    transformed = map( x -> substitute_all(expr, rule(x)), soln)
+    return convert(typeof(soln[1]), transformed)
+end
