Model discovery using Sparse Regression on a UDE model

Hi everyone,

I am working on modeling the unknown physics of a non-autonomous system. The system has 5 states and 2 control inputs. The dynamics of states 1 to 3 are known precisely but those of states 4 and 5 are unknown. I collected some simulation data and modeled the unknown dynamics with a feedforward neural network. Upon training, the neural network approximates the unknown physics quite well as shown below.

Now my next step is to use Sparse Regression for model disvovery. I am using the following code

@variables u[1:5] c[1:2]
u = collect(u)
c = collect(c)

h = Num[polynomial_basis(u, 2); polynomial_basis(c, 2)]
basis = Basis(h, u, controls = c)
## Ŷ is the neural network output for the sequence of states X̂
nn_problem = DirectDataDrivenProblem(X̂[:,1:end-1], Ŷ, U=U)

sampler = DataProcessing(split = 0.8, shuffle = true, batchsize = 30, rng = rng)
λs = exp10.(-10:0.1:0)
opt = STLSQ(λs)
res = solve(nn_problem, basis, opt,
            options = DataDrivenCommonOptions(data_processing = sampler, digits = 1))


nn_eqs = get_basis(res)
# Evaluate for some state and control input 
nn_eqs(X, get_parameter_map(nn_eqs), U[:,1])

but I get this error

MethodError: no method matching (::Basis{false, true})(::Vector{Float64}, ::Vector{Pair{Sym{Real, Base.ImmutableDict{DataType, Any}}, Float64}}, ::Vector{Float64})

I would really appreciate guidance to mitigate this error.

Start with the callable Basis expression shown in the example, then inspect how get_basis(res) and get_parameter_map(nn_eqs) are intended to be used. Reproduce the MethodError using the supplied DirectDataDrivenProblem, STLSQ solve, and evaluation call; done means the evaluation path accepts the state, parameter map, and control input without the reported dispatch error.