diff --git a/Project.toml b/Project.toml
index 4f28b90..9aeee50 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,8 +5,12 @@ version = "1.0.0-DEV"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DiffEqGPU = "071ae1c0-96b5-11e9-1965-c90190d839ea"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+NeuralPDE = "315f7962-48a3-4962-8226-d0f33b1235f0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
diff --git a/src/PSOGPU.jl b/src/PSOGPU.jl
index 650420b..f1c7a69 100644
--- a/src/PSOGPU.jl
+++ b/src/PSOGPU.jl
@@ -1,6 +1,6 @@
 module PSOGPU
 
-using SciMLBase, StaticArrays, Setfield, CUDA
+using SciMLBase, StaticArrays, Setfield, CUDA, ComponentArrays, Lux, Functors
 
 import DiffEqGPU: GPUTsit5, vectorized_asolve, make_prob_compatible
 
@@ -117,5 +117,16 @@ function Base.typemax(::Type{PSOGPU.PSOParticle{T1, T2}}) where {T1, T2}
         typemax(T2))
 end
 
+function vector_to_parameters(ps_new::AbstractVector, ps::ComponentArray)
+    @assert length(ps_new) == Lux.parameterlength(ps)
+    i = 1
+    function get_ps(x)
+        z = reshape(view(ps_new, i:(i + length(x) - 1)), size(x))
+        i += length(x)
+        return z
+    end
+    return fmap(get_ps, ps)
+end
+
 export ParallelPSOKernel, ParallelSyncPSO, OptimizationProblem, solve
 end
diff --git a/src/pso_cpu.jl b/src/pso_cpu.jl
index 84ff81b..9afda60 100644
--- a/src/pso_cpu.jl
+++ b/src/pso_cpu.jl
@@ -1,14 +1,14 @@
 # Based on: https://stackoverflow.com/questions/65342388/why-my-code-in-julia-is-getting-slower-for-higher-iteration
 
 mutable struct Particle{T}
-    position::Array{T, 1}
-    velocity::Array{T, 1}
+    position::AbstractVector{T}
+    velocity::AbstractVector{T}
     cost::T
-    best_position::Array{T, 1}
+    best_position::AbstractVector{T}
     best_cost::T
 end
 mutable struct Gbest{T}
-    position::Array{T, 1}
+    position::AbstractVector{T}
     cost::T
 end
 
@@ -19,7 +19,11 @@ function _init_particles(prob, population)
     cost_func = prob.f
 
     if lb === nothing || (all(isinf, lb) && all(isinf, ub))
-        gbest_position = Array{eltype(prob.u0), 1}(undef, dim)
+        if prob.u0 isa ComponentVector
+            gbest_position = ComponentArray(undef, getaxes(prob.u0))
+        else
+            gbest_position = Array{eltype(prob.u0), 1}(undef, dim)
+        end
         for i in 1:dim
             if abs(prob.u0[i]) > 0
                 gbest_position[i] = prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i])
@@ -30,12 +34,17 @@ function _init_particles(prob, population)
     else
         gbest_position = uniform(dim, lb, ub)
     end
+
     gbest = Gbest(gbest_position, cost_func(gbest_position, prob.p))
 
     particles = Particle[]
     for i in 1:population
         if lb === nothing || (all(isinf, lb) && all(isinf, ub))
-            position = Array{eltype(prob.u0), 1}(undef, dim)
+            if prob.u0 isa ComponentVector
+                position = ComponentArray(undef, getaxes(prob.u0))
+            else
+                position = Array{eltype(prob.u0), 1}(undef, dim)
+            end
             for i in 1:dim
                 if abs(prob.u0[i]) > 0
                     position[i] = prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i])
diff --git a/src/utils.jl b/src/utils.jl
index c7a0d12..1410943 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -14,7 +14,11 @@ function init_particles(prob, n_particles)
     p = prob.p
 
     if lb === nothing || (all(isinf, lb) && all(isinf, ub))
-        gbest_position = Array{eltype(prob.u0), 1}(undef, dim)
+        if prob.u0 isa ComponentVector
+            gbest_position = ComponentArray(undef, getaxes(prob.u0))
+        else
+            gbest_position = Array{eltype(prob.u0), 1}(undef, dim)
+        end
         for i in 1:dim
             if abs(prob.u0[i]) > 0
                 gbest_position[i] = prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i])
@@ -31,7 +35,11 @@ function init_particles(prob, n_particles)
     particles = PSOParticle[]
     for i in 1:n_particles
         if lb === nothing || (all(isinf, lb) && all(isinf, ub))
-            position = Array{eltype(prob.u0), 1}(undef, dim)
+            if prob.u0 isa ComponentVector
+                position = ComponentArray(undef, getaxes(prob.u0))
+            else
+                position = Array{eltype(prob.u0), 1}(undef, dim)
+            end
             for i in 1:dim
                 if abs(prob.u0[i]) > 0
                     position[i] = prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i])
@@ -65,7 +73,11 @@ function init_particles(prob, population, ::CPU)
     cost_func = prob.f
 
     if lb === nothing || (all(isinf, lb) && all(isinf, ub))
-        gbest_position = Array{eltype(prob.u0), 1}(undef, dim)
+        if prob.u0 isa ComponentVector
+            gbest_position = ComponentArray(undef, getaxes(prob.u0))
+        else
+            gbest_position = Array{eltype(prob.u0), 1}(undef, dim)
+        end
         for i in 1:dim
             if abs(prob.u0[i]) > 0
                 gbest_position[i] = prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i])
@@ -82,7 +94,11 @@ function init_particles(prob, population, ::CPU)
     particles = Particle[]
     for i in 1:population
         if lb === nothing || (all(isinf, lb) && all(isinf, ub))
-            position = Array{eltype(prob.u0), 1}(undef, dim)
+            if prob.u0 isa ComponentVector
+                position = ComponentArray(undef, getaxes(prob.u0))
+            else
+                position = Array{eltype(prob.u0), 1}(undef, dim)
+            end
             for i in 1:dim
                 if abs(prob.u0[i]) > 0
                     position[i] = prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i])
diff --git a/test/neuralode.jl b/test/neuralode.jl
new file mode 100644
index 0000000..e69de29
diff --git a/test/ode.jl b/test/ode.jl
new file mode 100644
index 0000000..fb19223
--- /dev/null
+++ b/test/ode.jl
@@ -0,0 +1,63 @@
+using StaticArrays, SciMLBase
+
+function f(u, p, t)
+    dx = p[1] * u[1] - u[1] * u[2]
+    dy = -3 * u[2] + u[1] * u[2]
+    return SVector{2}(dx,dy)
+end
+
+u0 = @SArray [1.0; 1.0]
+tspan = (0.0, 10.0)
+p = @SArray [1.5]
+prob = ODEProblem(f, u0, tspan, p)
+
+using OrdinaryDiffEq
+sol = solve(prob, Tsit5())
+t = collect(range(0f0, stop = 10.0, length = 200))
+using RecursiveArrayTools # for VectorOfArray
+randomized = VectorOfArray([(sol(t[i]) + 0.01randn(2)) for i in 1:length(t)])
+
+data = convert(Array, randomized)
+
+n_particles = 1000
+
+maxiters = 1000
+
+n = 1
+
+lb = @SArray ones(n)
+
+lb = -10*lb
+
+ub = @SArray ones(n)
+
+ub = 10*ub
+
+
+function loss(u,p)
+    odeprob, t = p
+    prob = remake(odeprob; p = u)
+    pred = Array(solve(prob, Tsit5(), saveat = t))
+    sum(abs2, data .- pred)
+end
+
+optprob = OptimizationProblem(loss, [1.0], (prob,t);lb, ub)
+
+using PSOGPU
+using CUDA
+
+
+gbest, particles = PSOGPU.init_particles(optprob, n_particles)
+
+gpu_data = cu([SVector{length(prob.u0), eltype(prob.u0)}(@view data[:,i]) for i in 1:length(t)])
+
+
+gpu_particles = cu(particles)
+
+
+@time gsol = PSOGPU.parameter_estim_ode!(prob,
+    gpu_particles,
+    gbest,
+    gpu_data,
+    lb,
+    ub; saveat = t, dt = 0.1)
diff --git a/test/pinn.jl b/test/pinn.jl
new file mode 100644
index 0000000..cc3fb18
--- /dev/null
+++ b/test/pinn.jl
@@ -0,0 +1,45 @@
+using NeuralPDE, Lux, Optimization, OptimizationOptimJL, PSOGPU
+import ModelingToolkit: Interval
+
+@parameters t, x
+@variables u(..)
+Dxx = Differential(x)^2
+Dtt = Differential(t)^2
+Dt = Differential(t)
+
+#2D PDE
+C = 1
+eq = Dtt(u(t, x)) ~ C^2 * Dxx(u(t, x))
+
+# Initial and boundary conditions
+bcs = [u(t, 0) ~ 0.0,# for all t > 0
+    u(t, 1) ~ 0.0,# for all t > 0
+    u(0, x) ~ x * (1.0 - x), #for all 0 < x < 1
+    Dt(u(0, x)) ~ 0.0] #for all  0 < x < 1]
+
+# Space and time domains
+domains = [t ∈ Interval(0.0, 1.0),
+    x ∈ Interval(0.0, 1.0)]
+# Discretization
+dx = 0.1
+
+# Neural network
+chain = Lux.Chain(Dense(2, 10, Lux.σ), Dense(10, 5, Lux.σ), Dense(5, 1))
+discretization = PhysicsInformedNN(chain, GridTraining(dx))
+
+@named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])
+prob = discretize(pde_system, discretization)
+
+callback = function (p, l)
+    println("Current loss is: $l")
+    return false
+end
+
+# optimizer
+opt = OptimizationOptimJL.BFGS()
+@time res = Optimization.solve(prob, opt; callback = callback, maxiters = 1200)
+phi = discretization.phi
+
+opt = ParallelPSOKernel(1000; gpu = true, threaded = true)
+@time res = Optimization.solve(prob, opt; maxiters = 1200)
+phi = discretization.phi
\ No newline at end of file