diff --git a/examples/dopa-sweep.py b/examples/dopa-sweep.py index 5961907..bb46af2 100644 --- a/examples/dopa-sweep.py +++ b/examples/dopa-sweep.py @@ -13,17 +13,27 @@ ) # initial conditions -y0 = jp.array([0., -2.0, 0.0, 0.0, 0.0, 0.0]) +y0 = jp.array([0.1, -60.0, 0.0, 0.0, 0.0, 0.0]) # run sweep end_time = 256.0 pkeys, ys = sweep_node(y0, params, T=end_time, cores=4) +# for large sweeps, ys may not yet be filled with outputs +ys.block_until_ready() + # pkeys provides the names for the extra dims of ys result print(pkeys, ys.shape) # now similar for network sweep n_nodes = 8 -Ci, Ce, Cd = jp.zeros((3, n_nodes)) +Ci, Ce, Cd = vb.randn(3, n_nodes, n_nodes) pkeys, ys = sweep_network(y0, params, Ci, Ce, Cd, T=end_time, cores=4) +ys.block_until_ready() print(pkeys, ys.shape) + +# plot some of it +import matplotlib.pyplot as pl + +pl.plot(ys[0,0,:,0], 'k') +pl.show() diff --git a/vbjax/app/dopa/sweep.py b/vbjax/app/dopa/sweep.py index bb7780a..0f70555 100644 --- a/vbjax/app/dopa/sweep.py +++ b/vbjax/app/dopa/sweep.py @@ -5,6 +5,19 @@ import vbjax as vb +# TODO can't continue simulations b/c init doesn't vmap w/ params +# TODO need to run long time tavg and bold +# TODO allow predefined set of params for e.g. random sample +# TODO regional parameters? +# TODO abstract over simulation configuration e.g. noise +# TODO ISPC impl for benchmark comparison +# TODO use delays if available +# TODO implement memory limits & test them +# TODO memmap raw/tavg/bold direct to file +# TODO automate inner vs outer loop over parameters + + + def sweep_node(init, params, T=10.0, dt=0.01, sigma=1e-3, seed=42, cores=4): "Run sweep for single dopa node on params matrix"