@@ -5,24 +5,20 @@ using KernelAbstractions, Adapt
55using StaticArrays
66using PoissonRandom, Random
77
8- function SciMLBase. __solve (
9- ensembleprob:: SciMLBase.AbstractEnsembleProblem ,
8+ function SciMLBase. __solve (ensembleprob:: SciMLBase.AbstractEnsembleProblem ,
109 alg:: SimpleTauLeaping ,
1110 ensemblealg:: EnsembleGPUKernel ;
1211 trajectories,
1312 seed = nothing ,
1413 dt = error (" dt is required for SimpleTauLeaping." ),
15- kwargs...
16- )
14+ kwargs... )
1715 if trajectories == 1
18- return SciMLBase. __solve (
19- ensembleprob, alg, EnsembleSerial (); trajectories = 1 ,
20- seed, dt, kwargs...
21- )
16+ return SciMLBase. __solve (ensembleprob, alg, EnsembleSerial (); trajectories = 1 ,
17+ seed, dt, kwargs... )
2218 end
2319
2420 ensemblealg. backend === nothing ? backend = CPU () :
25- backend = ensemblealg. backend
21+ backend = ensemblealg. backend
2622
2723 jump_prob = ensembleprob. prob
2824
@@ -35,42 +31,36 @@ function SciMLBase.__solve(
3531
3632 # Run vectorized solve
3733 ts,
38- us = vectorized_solve (
39- probs, jump_prob, SimpleTauLeaping (); backend, trajectories, seed, dt
40- )
34+ us = vectorized_solve (
35+ probs, jump_prob, SimpleTauLeaping (); backend, trajectories, seed, dt)
4136
4237 # Convert to CPU for inspection
4338 _ts = Array (ts)
4439 _us = Array (us)
4540
46- time = @elapsed sol = [
47- begin
48- ts = @view _ts[:, i]
49- us = @view _us[:, :, i]
50- sol_idx = findlast (x -> x != probs[i]. prob. tspan[1 ], ts)
51- if sol_idx === nothing
52- @error " No solution found" tspan = probs[i]. tspan[1 ] ts
53- error (" Batch solve failed" )
54- end
55- @views ensembleprob. output_func (
56- SciMLBase. build_solution (
57- probs[i]. prob,
58- alg,
59- ts[1 : sol_idx],
60- [us[j, :] for j in 1 : sol_idx],
61- k = nothing ,
62- stats = nothing ,
63- calculate_error = false ,
64- retcode = sol_idx !=
65- length (ts) ?
66- ReturnCode. Terminated :
67- ReturnCode. Success
68- ),
69- i
70- )[1 ]
71- end
72- for i in eachindex (probs)
73- ]
41+ time = @elapsed sol = [begin
42+ ts = @view _ts[:, i]
43+ us = @view _us[:, :, i]
44+ sol_idx = findlast (x -> x != probs[i]. prob. tspan[1 ], ts)
45+ if sol_idx === nothing
46+ @error " No solution found" tspan= probs[i]. tspan[1 ] ts
47+ error (" Batch solve failed" )
48+ end
49+ @views ensembleprob. output_func (
50+ SciMLBase. build_solution (probs[i]. prob,
51+ alg,
52+ ts[1 : sol_idx],
53+ [us[j, :] for j in 1 : sol_idx],
54+ k = nothing ,
55+ stats = nothing ,
56+ calculate_error = false ,
57+ retcode = sol_idx !=
58+ length (ts) ?
59+ ReturnCode. Terminated :
60+ ReturnCode. Success),
61+ i)[1 ]
62+ end
63+ for i in eachindex (probs)]
7464 return SciMLBase. EnsembleSolution (sol, time, true )
7565end
7666
9282@kernel function simple_tau_leaping_kernel (
9383 @Const (probs_data), _us, _ts, dt, @Const (rj_data),
9484 current_u_buf, rate_cache_buf, counts_buf, local_dc_buf,
95- seed:: UInt64
96- )
85+ seed:: UInt64 )
9786 i = @index (Global, Linear)
9887
9988 # Get thread-local buffers
125114
126115 # Get input/output arrays
127116 ts_view = @inbounds view (_ts, :, i)
128- us_view = @inbounds view (_us, :, :, i)
117+ us_view = @inbounds view (_us,:,:, i)
129118
130119 # Initialize first time step and state
131120 @inbounds ts_view[1 ] = tspan[1 ]
135124
136125 # Main loop
137126 for j in 2 : n
138- tprev = tspan[1 ] + (j - 2 ) * dt
127+ tprev = tspan[1 ] + (j- 2 ) * dt
139128
140129 # Compute rates and scale by dt
141130 rate (rate_cache, current_u, p, tprev)
@@ -154,24 +143,20 @@ end
154143 @inbounds for k in 1 : state_dim
155144 us_view[j, k] = current_u[k]
156145 end
157- @inbounds ts_view[j] = tspan[1 ] + (j - 1 ) * dt
146+ @inbounds ts_view[j] = tspan[1 ] + (j- 1 ) * dt
158147 end
159148end
160149
161150# Vectorized solve function
162- function vectorized_solve (
163- probs, prob:: JumpProblem , alg:: SimpleTauLeaping ;
164- backend, trajectories, seed, dt, kwargs...
165- )
151+ function vectorized_solve (probs, prob:: JumpProblem , alg:: SimpleTauLeaping ;
152+ backend, trajectories, seed, dt, kwargs... )
166153 # Extract common jump data
167154 rj = prob. regular_jump
168155 rj_data = JumpData (rj. rate, rj. c, rj. numjumps)
169156
170157 # Extract trajectory-specific data without static typing
171- probs_data = [
172- TrajectoryData (SA{eltype (p. prob. u0)}[p. prob. u0... ], p. prob. p, p. prob. tspan)
173- for p in probs
174- ]
158+ probs_data = [TrajectoryData (SA{eltype (p. prob. u0)}[p. prob. u0... ], p. prob. p, p. prob. tspan)
159+ for p in probs]
175160
176161 # Adapt to GPU
177162 probs_data_gpu = adapt (backend, probs_data)
@@ -212,15 +197,13 @@ function vectorized_solve(
212197 KernelAbstractions. synchronize (backend)
213198
214199 # Seed for Poisson sampling
215- seed = seed === nothing ? UInt64 (12345 ) : UInt64 (seed)
200+ seed = seed === nothing ? UInt64 (12345 ) : UInt64 (seed);
216201
217202 # Launch main kernel
218203 kernel = simple_tau_leaping_kernel (backend)
219- main_event = kernel (
220- probs_data_gpu, us, ts, dt, rj_data_gpu,
204+ main_event = kernel (probs_data_gpu, us, ts, dt, rj_data_gpu,
221205 current_u_buf, rate_cache_buf, counts_buf, local_dc_buf, seed;
222- ndrange = n_trajectories
223- )
206+ ndrange = n_trajectories)
224207 KernelAbstractions. synchronize (backend)
225208
226209 return ts, us
0 commit comments