Context

class braintools.cogtask.Context(key=None)

Mutable trial-level state container shared across phases.

The Context object is the backbone of inter-phase communication. It holds: - Trial-level state that persists across phases (e.g., ground_truth, stimulus_value) - Time tracking (current timestep, phase boundaries) - Input/output buffers for trial data - RNG for reproducible randomness

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> ctx = Context(key=jax.random.PRNGKey(42))
>>> ctx['ground_truth'] = 1
>>> ctx['stimulus_direction'] = 0.5 * jnp.pi
>>> print(ctx['ground_truth'])
1

>>> # In a phase, write into the buffer with a functional update
>>> # (JAX arrays are immutable; ``.at[...].set(...)`` returns a copy):
>>> # ctx.inputs = ctx.inputs.at[ctx.phase_start:ctx.phase_end, :].set(stimulus_encoding)

Parameters:

key (Array | None) – JAX PRNGKey for random number generation. If None, draws fresh randomness from brainstate’s default RNG.

Note

The time step (dt) is obtained from brainstate.environ.get_dt().

clear()

Reset state for a new trial.

property dt: float | Quantity

Time step in milliseconds (from brainstate.environ.get_dt()).

get(key, default=None)

Get a value from the trial state with a default.

Return type:

Any

property phase_duration: int

Duration of current phase in timesteps.

phase_slice()

Get a slice object for the current phase’s time range.

Return type:

slice

property phase_time: int

Current timestep relative to phase start.

property state: Dict[str, Any]

Get a copy of the trial state dictionary.

property total_steps: int

Total number of timesteps in the trial (based on input buffer).

update(**kwargs)

Update multiple values in the trial state.