States - Structuring task steps#

States are how you can structure which happens when in a task.

States are typically written after the configuration of your setup. The following example shows 5 optional methods that can be added to a state to run at a specific context.

from neurokraken import Neurokraken, State
from neurokraken.configurators import devices, Display
from neurokraken.controls import get

serial_out = {'led': devices.direct_on(pin=3, start_value=False)}

nk = Neurokraken(serial_in={}, serial_out=serial_out, log_dir='./', display=Display())

get.white_background = False # we will make use of this created variable later

# States and task code starts here

class My_state(State):
    def pre_task(self):
        pass

    def on_start(self):
        pass

    def loop_main(self):
        return False, 0
    
    def loop_visual(self, sketch):
        pass

    def on_end(self):
        pass

task = {'state_0': My_state(next_state='state_0')}

nk.load_task(task)

nk.run()

The line class My_state(State): does 2 things:

  • We decide on the class name “My_state”. This allows use to later use it as part of the task with My_state(next_state=<str>)

  • (State) ensures that neurokraken can the class as a state, by automatically adding internal properties.

Afterwards we can add one or more of the shown methods to add context-specific code to our task. The (self) argument enables access to a shared self. variable space across the class.

At the end we create a task dictionary where we map names (like here 'state_0') to the states we want to use (My_state()) Even in single state tasks, States require an argument next_state=<str> that names the successor state or list of possible successor states, with the progression typically decided within loop_main() or by a maximum time provided to the state My_state(..., max_time_s=3).

Each method has access to self, allowing you to share variables and their usage across methods of the same state, and access to get. allowing you to access and interact with global variables like the get.white_background = False defined in the example above. Additionally, if you are familiar with python’s global keyword or specified file imports and namespacing, these approaches can provide further ways of accessing individual variables across different functions.

pre_task(self)#

Code in here is executed once before the task begins.

This is useful for setting up variables used later later in the state as self.<variable_name>, as well as loading files like state-related images to be used in the task. A good pattern is to move compute-intensive code that only needs to be run once (like loading files) outside of repeatedly run methods like loop_main() or on_start().

def pre_task(self):
    self.number = 5

pre_task can be run with a second argument sketch, which will provide your code access to py5 sketch functions. This is useful for loading images, or 3D models to later be used in the loop_visual().

def pre_task(self, sketch):
    self.image = sketch.load_image('C:/Path/To/My/Image.png')

on_start(self)#

This method is run every time the state is entered - typically following the previous state’s end.

This is a good spot to set or reset a variable for the subsequent loop_main() or loop_visual() and to change a controlled device like a servo motor or valve to a different value for this state or a followup state.

class My_state(State):
    def on_start(self):
        get.send_out('led', True)
        self.choice = 'none' # if a choice is made we will update it to 'left' or 'right' in loop_main()

loop_main(self)#

This method is where your main state logic takes place - reading sensors, deciding outcomes, logging custom data, and then either continueing to the next iteration of loop_main() or to to on_end() and the next state.

loop_main runs as fast as your system allows, expectably thousands of times per second, allowing for tight control of task progression.

Notably loop_main() has to return a bool and a number, i.e. return False, 0

The bool True/False signals whether the state is completed and you want to move on to the next state i.e. when you get.read_in() a touch choice from the subject, or find its position entered a target radius.

The number decides which follow-up state option to continue to when a list of options is provided. Thus if there is only a single follow-up option we are always returning index 0.

class Choice(State):
    def loop_main(self):
        if get.read_in(...):
            # bad choice made, continue to the timeout penalty state 0
            return True, 0
        elif get.read_in(...):
            # good choice made, continue to the reward state 1
            return True, 1
        else:
            # no choice made, rerun loop_main()
            return False, 0

class Reward(State):
    ...

class Timeout(State):
    ...

task = {'choosing': Choice(next_state=['timeout', 'reward']),
        'timeout': Timeout(next_state='choosing'),
        'reward':   Reward(next_state='choosing')}

loop_main() can be filled with your own python code to make complex decisions and changes based on a subject’s current get.read_in actions, the history get.log, get.time_ms, or state or global variables.

A common pattern is for the main loop to change state self. or global get. variables like the starting example’s get.white_background=False that are then utilized by other code like loop_visual() with their updating values.

Returning False means the state was not yet completed and loop_main() will loop run again. If you return True the task will move on to the next state, decided by the returned number.

Notably a state can also be provided with a maximum duration, and if this many seconds have passed the state will progress regardles of whether loop_main() returned True or False.

'choosing': Choice(next_state=['timeout', 'reward'], max_time_s=10)

You can manually read the state’s self.start_time variable for the task millisecond when the state started. The duration the state has already been running thus is get.time_ms - self.start_time

loop_visual(self, sketch)#

When using a display in your task, this function allows you to draw custom visual stimuli using py5’s drawing function through the sketch argument.

def loop_visual(self, sketch):
    if get.white_background:
        sketch.background(255, 255, 255) # R, G, B
    else:
        sketch.background(0, 0, 0)

You can find the py5 reference here. Py5 makes it easy to draw interactive custom 2D or 3D stimuli, animations, and environments using simple shapes, text, or loaded images and 3d models.

on_end(self)#

This function runs after conclusion of a state, due to loop_main() returning True to signal the state being finished, or due to a state’s duration exceeding the provided seconds max_time_s=.

def on_end(self):
    print('the state has concluded')

Further methods#

As the code of the state class is under your control you can add further methods to call from scheduled methods like loop_main() or on_start() to organize your code.

class My_state(State):
    def sum(self, a, b):
        return a + b

    def on_start(self):
         number = self.sum(5, 3)
         print(number) # 8

If you are familiar with python classes and the __init__() method, you can also add one to your class to create different variations of the same State with specific differences. I.e. a task might contain 2 versions of a state, one with a rewarding self.probability=0.9 and one with a less positive self.probability=0.1. An example of this can be found in the frequently used patterns.

Organizing state flow - The task dictionary#

The task dictionary contains named states of our task, with each state also naming its successor or list of possible succesors. We can use this to create an example task that blinks an LED every 5 seconds the following way.

class Light_on(State):
    def on_start(self):
        get.send_out('led', True)
    
class Light_off(State):
    def on_start(self):
        get.send_out('led', False)

task = {'on' :  Light_on(next_state='off', max_time_s=5),
        'off': Light_off(next_state='on',  max_time_s=5)}

nk.load_task(task)

Often a subject will make a choice and based on the choice we want to progress to splitting outcome states. For this we can provide a list of next_state options.

For example we can reorganize our slow blinking LED example into a trial starting with a 2 second “Start” state that decides whether the next state will be Light_on or Light_off, with either of these states returning to another decidiging Start state after their 5 seconds.

import random

class Light_on(State):
    def on_start(self):
        get.send_out('led', True)
    
class Light_off(State):
    def on_start(self):
        get.send_out('led', False)

class Start(State):
    def on_start(self):
        self.light_on = True if random.random() > 0.5 else False
    
    def loop_main(self):
        next_state_idx = 1 if self.light_on else 0
        return False, next_state_idx

task = {'start': Start(next_state=['off', 'on'], max_time_s=2),
        'off': Light_off(next_state='start', max_time_s=5),
        'on' :  Light_on(next_state='start', max_time_s=5)}
        
nk.load_task(task)

The state flow can be as extensive, branching and complex as you want it to be, i.e. a task may progress through a sequence of 7 states before returning to the first one with the outcome of a middle state determining a normal progression, a reset to the first state or skipping ahead of the next states.

You can also add a single forever state or a group of states which are unreachable from the next_state= options of the core task. These state’s can still manually be switched to or from with get.progress_state('<state_name>') from a UI button to create a manual standby state or special use states within the same task, to enable use cases similar to blocks.

State arguments#

States always have to receive their next_state argument. Further they can receive 4 optional arguments to extend their functionality.

next_state#

str or list

The next_state or list of next state options to progress to upon the main loop returning True or a state’s provided max_time_s running out.

task = {'start': Start(next_state=['off', 'on'], max_time_s=2),
        'off': Light_off(next_state='start', max_time_s=5),
        'on' :  Light_on(next_state='start', max_time_s=5)}

max_time_s#

float or tuple[float,float] or Callable->float

The maximum duration in seconds for which a state will run before progressing regardless of wether the main loop returns the trial to be finished True. Defaults to 1_000_000.

You can provide either a number (float/int), a 2 entry interval (tuple/list) for a random time within the interval every time the state runs, or a function returning a float for a customized dynamic timing.

Light_on(next_state='start', max_time_s=5)

Light_on(next_state='start', max_time_s=[2.5, 7.5])

def five_seconds():
    return 5
Light_on(next_state='start', max_time_s=five_seconds)

run_at_start and run_at_end#

Callable or list

These 2 arguments allow you to provide a function or list of functions to run after on_start() or after on_end() respectively.

This can be useful to organize complex task code, i.e. variables and logic might be defined in on_start() and on_end() but all logging and complex analysis is defined in run_at_start= and run_at_end= allowing you to separate the task itself from calculations for live plots or post-experiment analysis.

class Start(State):
    def on_start(self):
        get.light_on = True if random.random() > 0.5 else False
        self.light_on = get.light_on # extra copy to the state self for later in this example

# add an empty list to log light decisions of this example task
get.log['light_decisions'] = []

def log_light():
    # append an entry with the time and decision to this list
    get.log['light_decisions'].append([get.time_ms, get.light_on])

task = {'pause': Start(next_state=['off', 'on'], max_time_s=3, run_at_start=log_light),
        ... }

A function provided to run_at_start= or run_at_end= can optionally contain an argument to access the respective state/variables added to its self.

def log_light(state):
    # add the time and decision to this list
    get.log['light_decisions'].append([get.time_ms, state.light_on])

run_at_end functions can also be provided with a 2nd argument, a boolean to read True/False whether a state was ended by the main loop or timed out.

Trials#

States can provide basic trial management.

Frequently tasks are organized in a repeating pattern of states and it makes sense to organize and log data accordingly. get.log contains a list of dictionaries under get.log['trials'] that can be used for this purpose.

You can provide one or more states at the end of your task pattern with the argument trial_complete=True and every time one of these states culminates another entry will be added to the log trials list.

You can access the current trial number as len(get.log['trials']), the current (thus most recent and last) trial with get.log['trials'][-1], and the most recent 5 trials with get.log['trials'][-5:].

Each trial by default contains a ['start'] time entry, i.e. get.log['trials'][-1]['start'] for the start of the most recent trial. As a flexible python dictionary the trial entry can manually be extended with every information relevant to the task or successive analysis.

Finally, neurokraken.load_task() can be provided with a function as run_post_trial=, that will run at the end of a trial. This time point is well suited for updating trial specific values and lists relevant for the future of the task or live plots and visualizations. For example one might maintain a list of time-value-pairs of a subject’s reaction times, or actions during the last 30 trials to then regularly plot this list in a live-UI.

class Light_on(State):
    ...

class Light_off(State):
    ...

class Start(State):
    def on_start(self):
        print(f'entering trial number {len(get.log['trials'])}')
        get.light_on = True if random.random() > 0.5 else False
        # save the information to the current trial dictionary
        get.log['trials'][-1]['light'] = get.light_on

task = {'start': Start(next_state=['off', 'on'], max_time_s=2),
        'on' :  Light_on(next_state='start', max_time_s=5, trial_complete=True),
        'off': Light_off(next_state='start', max_time_s=5, trial_complete=True)}

get.average_light = 0.5

def recalculate_average_light():
    was_light_on = []
    for trial in get.log['trials']:
        was_light_on.append(trial['light'])
    true_count = sum(was_light_on) # True counts as 1, False as 0
    get.average_light = true_count / len(get.log['trials'])

nk.load_task(task, run_post_trial=recalculate_average_light)

Alternative blocks of states#

typically a task is loaded as a dictionary of states. However you can also provide a nested dictionary allowing you to organize groups of states and their progressions into individual state blocks.

As a next_state= can only be within the same block the only way to switch between blocks is manually through get.switch_block('<block_name>').

from neurokraken import State
class State_A(State):
    ...
class State_B(State):
    ...
class State_C(State):
    ...

task = {
'consistent_block' :  {'single_state': State_A(next_state='single_state')}
'alternating_block':  {'B': State_B(next_state='C'),
                       'C': State_C(next_state='B')}
}
nk.load_task(task)
...

# within a UI-triggered element switch the active block
get.switch_block('alternating_block')

The first block will be the one executed when you run your task, unless a different one was provided to neurokraken.load_task(..., start_block='<block_name>'.

An interesting use case of blocks is to use python’s deepcopy to create an independent copy of a block and to manually change variables within selected states of the copy in order to create distinct and switchable versions of a given task.

Autologging#

Progressions in the current state, trial, and block are automatically saved in get.log[‘states’], get.log[‘trials’] and get.log[‘blocks’].

State and block changes are maintained as a list of (time, state/block-name) tuples.

Trials are maintained as a list of extendable dictionaries with a preexisting start entry.

load_task arguments#

Neurokraken.load_task() can receive optional arguments besides the task dictionary for extended control of task elements, as shown in the usage of trials above.

start_block :str#

The name of the block to start the run with. If not provided the task will begin with the first block

run_at_start :Callable#

The provided function will be called at the very beginning of the task before the states.

The run_at_start function will be followed by an additional communication loop before the task begins, allowing you to start the task with controlled devices in a desired state.

def move_back_servo():
    get.send_out('my_servo', 160)

nk.load_task(..., run_at_start=move_back_servo)

run_at_quit :Callable#

THe provided function will be run once the task is ended, i.e. through get.quit() or Ctrl+Alt+Q.

After the function an additional communication loop will be performed, allowing to set the final position of teensy-side controlled devices at shutdown regardless of the state during which the task was quit.

def close_valve():
    get.send_out('my_valve', False)

nk.load_task(..., run_at_quit=close_valve)

run_post_trial :Callable#

This function will be called upon trial completion as shown in the trials subchapter.

specialized arguments#

main_as_sketch :bool#

Whether to run the main loop as a py5 sketch or a python while-True: - loop. Defaults to True for a py5 sketch.

It was found that using a py5 sketch allows for more stable main loop timing during high workload tasks (cameras, visual, live-AI). The ability to run as a while-True: loop is still available with main_as_sketch=False.

permanent_states :list#

A list of classes containing a .run() function. The respective .run() function will be executed at the end of each main loop iteration.

run_at_visual_start :Callable#

When using py5’s shader functions it was found that load_shader() needs to be run by the same sketch that will display it. A method provided to run_at_visual_start allows its code access to the task’s loop_visual() sketch during setup for this niche situation.

def method_on_visual_setup(sketch):
    pass

nk.load_task(..., run_at_visual_start=method_on_visual_setup)

experiment :Container#

Used by the runner mode to provide get.experiment access to the chosen experiment .py file and variables therein.