HiddenTreeState

class brainstate.HiddenTreeState(value, **kwargs)

A pytree of multiple hidden states for eligibility trace-based learning.

Note

The value in this state class behaves likes a dictionary/sequence of hidden states. However, the state is actually stored as a single dimensionless array.

There are two ways to define the hidden states.

1. The first is to define a sequence of hidden states.

import brainunit as u
value = [np.random.randn(10, 10) * u.mV,
         np.random.randn(10, 10) * u.mA,
         np.random.randn(10, 10) * u.mS]
state = HiddenTreeState(value)

Then, you can retrieve the hidden state value with the following method.

state.get_value(0)  # get the first hidden state
# or
state.get_value('0')  # get the hidden state with the name '0'

You can write the hidden state value with the following method.

state.set_value({0: np.random.randn(10, 10) * u.mV})  # set the first hidden state
# or
state.set_value({'1': np.random.randn(10, 10) * u.mA})  # set the hidden state with the name '1'
# or
state.set_value([np.random.randn(10, 10) * u.mV,
                 np.random.randn(10, 10) * u.mA,
                 np.random.randn(10, 10) * u.mS])  # set all hidden state value
# or
state.set_value({
    0: np.random.randn(10, 10) * u.mV,
    1: np.random.randn(10, 10) * u.mA,
    2: np.random.randn(10, 10) * u.mS
})  # set all hidden state value

2. The second is to define a dictionary of hidden states.

import brainunit as u
value = {'v': np.random.randn(10, 10) * u.mV,
         'i': np.random.randn(10, 10) * u.mA,
         'g': np.random.randn(10, 10) * u.mS}
state = HiddenTreeState(value)

Then, you can retrieve the hidden state value with the following method.

state.get_value('v')  # get the hidden state with the name 'v'
# or
state.get_value('i')  # get the hidden state with the name 'i'

You can write the hidden state value with the following method.

state.set_value({'v': np.random.randn(10, 10) * u.mV})  # set the hidden state with the name 'v'
# or
state.set_value({'i': np.random.randn(10, 10) * u.mA})  # set the hidden state with the name 'i'
# or
state.set_value([np.random.randn(10, 10) * u.mV,
                 np.random.randn(10, 10) * u.mA,
                 np.random.randn(10, 10) * u.mS])  # set all hidden state value
# or
state.set_value({
    'v': np.random.randn(10, 10) * u.mV,
    'g': np.random.randn(10, 10) * u.mA,
    'i': np.random.randn(10, 10) * u.mS
})  # set all hidden state value

Note

Avoid using HiddenTreeState.value to get the state value, or HiddenTreeState.value = to assign the state value.

Instead, use HiddenTreeState.get_value() and HiddenTreeState.set_value(). This is because .value loss hidden state units and other information, and it is only dimensionless data.

This design aims to ensure that any etrace hidden state has only one array.

Parameters:

value (Dict[str, Array | ndarray | bool | number | bool | int | float | complex | Quantity] | Sequence[Array | ndarray | bool | number | bool | int | float | complex | Quantity]) – The values of the hidden states.

get_value(item)

Get the value of the hidden state with the key.

Parameters:

item (str | int) – The key of the hidden state. - If int, the index of the hidden state. - If str, the name of the hidden state.

Return type:

Array | ndarray | bool | number | bool | int | float | complex | Quantity

property num_state: int

The number of hidden states.

set_value(val)

Set the value of the hidden state with the specified item.

This method updates the hidden state values based on the provided dictionary or sequence. The values are set according to the indices or names specified in the input.

Parameters:

• str (val (Dict[int |) – A dictionary or sequence containing the new values for the hidden states. - If a dictionary, keys can be integers (indices) or strings (names) of the hidden states. - If a sequence, it is converted to a dictionary with indices as keys.

• Sequence[ArrayLike]) (ArrayLike] |) – A dictionary or sequence containing the new values for the hidden states. - If a dictionary, keys can be integers (indices) or strings (names) of the hidden states. - If a sequence, it is converted to a dictionary with indices as keys.

Returns:

None

Return type:

None

property varshape: Tuple[int, ...]

The shape of each hidden state variable.