fix: manually implement controls for AbstractBasis

Author: AayushSabharwal

src/DataDrivenDiffEq.jl

@@ -19,7 +19,7 @@ using SymbolicUtils: operation, arguments, iscall, issym
 using Symbolics
 using Symbolics: scalarize, variable, value
 @reexport using ModelingToolkit: unknowns, parameters, independent_variable, observed,
-                                 controls, get_iv, get_observed
+                                 get_iv, get_observed
 
 using Random
 using QuadGK
@@ -98,7 +98,7 @@ include("./basis/utils.jl")
 include("./basis/type.jl")
 export Basis
 export jacobian, dynamics
-export implicit_variables, states
+export implicit_variables, states, controls
 export get_parameter_values, get_parameter_map
 
 include("./utils/basis_generators.jl")

src/basis/type.jl

@@ -263,6 +263,18 @@ function states(b::AbstractBasis)
     return getfield(b, :unknowns)
 end
 
+function controls(b::AbstractBasis)
+    ctrls = getfield(b, :ctrls)
+    systems = getfield(b, :systems)
+    isempty(systems) && return ctrls
+
+    ctrls = copy(ctrls)
+    for sys in systems
+        append!(ctrls, unknowns(sys, controls(sys)))
+    end
+    return ctrls
+end
+
 # For internal use
 is_implicit(b::Basis{X, <:Any}) where {X} = X
 is_controlled(b::Basis{<:Any, X}) where {X} = X