osl_dynamics.models.dyneste

Dynamic Network States (DyNeStE).

See the model description for more details.

Classes

Config	Settings for DyNeStE.
Model	DyNeStE model class.

Module Contents

class osl_dynamics.models.dyneste.Config

Bases: osl_dynamics.models.mod_base.BaseModelConfig, osl_dynamics.models.inf_mod_base.VariationalInferenceModelConfig

Settings for DyNeStE.

Parameters:

• model_name (str) – Model name.

• n_states (int) – Number of states.

• n_channels (int) – Number of channels.

• sequence_length (int) – Length of sequence passed to the inference network and generative model.

• inference_rnn (str) – RNN to use, either 'gru' or 'lstm'.

• inference_n_layers (int) – Number of layers.

• inference_n_units (int) – Number of units.

• inference_normalization (str) – Type of normalization to use. Either None, 'batch' or 'layer'.

• inference_activation (str) – Type of activation to use after normalization and before dropout. E.g. 'relu', 'elu', etc.

• inference_dropout (float) – Dropout rate.

• inference_regularizer (str) – Regularizer.

• model_rnn (str) – RNN to use, either 'gru' or 'lstm'.

• model_n_layers (int) – Number of layers.

• model_n_units (int) – Number of units.

• model_normalization (str) – Type of normalization to use. Either None, 'batch' or 'layer'.

• model_activation (str) – Type of activation to use after normalization and before dropout. E.g. 'relu', 'elu', etc.

• model_dropout (float) – Dropout rate.

• model_regularizer (str) – Regularizer.

• learn_means (bool) – Should we make the mean vectors for each state trainable?

• learn_covariances (bool) – Should we make the covariance matrix for each state trainable?

• initial_means (np.ndarray) – Initialisation for mean vectors.

• initial_covariances (np.ndarray) – Initialisation for state covariances. If diagonal_covariances=True and full matrices are passed, the diagonal is extracted.

• covariances_epsilon (float) – Error added to state covariances for numerical stability.

• diagonal_covariances (bool) – Should we learn diagonal state covariances?

• means_regularizer (tf.keras.regularizers.Regularizer) – Regularizer for mean vectors.

• covariances_regularizer (tf.keras.regularizers.Regularizer) – Regularizer for covariance matrices.

• do_kl_annealing (bool) – Should we use KL annealing during training?

• kl_annealing_curve (str) – Type of KL annealing curve. Either 'linear' or 'tanh'.

• kl_annealing_sharpness (float) – Parameter to control the shape of the annealing curve if kl_annealing_curve='tanh'.

• n_kl_annealing_epochs (int) – Number of epochs to perform KL annealing.

• do_gs_annealing (bool) – Should we use temperature annealing for the Gumbel-Softmax distribution during training?

• gs_annealing_curve (str) – Type of Gumbel-Softmax temperature annealing curve. Either 'linear' or 'exp'.

• initial_gs_temperature (float) – Initial temperature for the Gumbel-Softmax distribution.

• final_gs_temperature (float) – Final temperature for the Gumbel-Softmax distribution after annealing.

• gs_annealing_slope (float) – Slope of the Gumbel-Softmax temperature annealing curve. Only used when gs_annealing_curve='exp'.

• n_gs_annealing_epochs (int) – Number of epochs to perform Gumbel-Softmax temperature annealing.

• init_method (str) – Initialization method. Defaults to ‘random_state_time_course’.

• n_init (int) – Number of initializations. Defaults to 3.

• n_init_epochs (int) – Number of epochs for each initialization. Defaults to 1.

• init_take (float) – Fraction of dataset to use in the initialization. Defaults to 1.0.

• batch_size (int) – Mini-batch size.

• learning_rate (float) – Learning rate.

• lr_decay (float) – Decay for learning rate. Default is 0.1. We use lr = learning_rate * exp(-lr_decay * epoch).

• gradient_clip (float) – Value to clip gradients by. This is the clipnorm argument passed to the Keras optimizer. Cannot be used if multi_gpu=True.

• n_epochs (int) – Number of training epochs.

• optimizer (str or tf.keras.optimizers.Optimizer) – Optimizer to use. 'adam' is recommended.

• loss_calc (str) – How should we collapse the time dimension in the loss? Either 'mean' or 'sum'.

• multi_gpu (bool) – Should be use multiple GPUs for training?

• strategy (str) – Strategy for distributed learning.

• best_of (int) – Number of full training runs to perform. A single run includes its own initialization and fitting from scratch.

model_name: str = 'DyNeStE'

inference_rnn: str = 'lstm'

inference_n_layers: int = 1

inference_n_units: int = None

inference_normalization: str = None

inference_activation: str = None

inference_dropout: float = 0.0

inference_regularizer: str = None

model_rnn: str = 'lstm'

model_n_layers: int = 1

model_n_units: int = None

model_normalization: str = None

model_activation: str = None

model_dropout: float = 0.0

model_regularizer: str = None

do_gs_annealing: bool = False

gs_annealing_curve: Literal['linear', 'exp'] = None

initial_gs_temperature: float = 1.0

final_gs_temperature: float = 0.01

gs_annealing_slope: float = None

n_gs_annealing_epochs: int = None

learn_means: bool = None

learn_covariances: bool = None

initial_means: numpy.ndarray = None

initial_covariances: numpy.ndarray = None

diagonal_covariances: bool = False

covariances_epsilon: float = None

means_regularizer: tensorflow.keras.regularizers.Regularizer = None

covariances_regularizer: tensorflow.keras.regularizers.Regularizer = None

init_method: str = 'random_state_time_course'

n_init: int = 3

n_init_epochs: int = 1

init_take: float = 1.0

validate_rnn_parameters()

Return type:

None

validate_observation_model_parameters()

Return type:

None

validate_gs_annealing_parameters()

Return type:

None

class osl_dynamics.models.dyneste.Model(config)

Bases: osl_dynamics.models.inf_mod_base.VariationalInferenceModelBase

DyNeStE model class.

Parameters:

config (osl_dynamics.models.dyneste.Config)

config_type

build_model()

Builds a keras model.

Return type:

None

fit(*args, gs_annealing_callback=None, **kwargs)

Wrapper for the standard keras fit method.

This function inherits fit() functions in ModelBase and VariationalInferenceModelBase.

Parameters:

• *args (arguments) – Arguments for ModelBase.fit() or VariationalInferenceModelBase.fit().

• gs_annealing_callback (bool, optional) – Should we anneal the Gumbel-Softmax temperature during training?

• **kwargs (keyword arguments, optional) – Keyword arguments for ModelBase.fit() or VariationalInferenceModelBase.fit().

Returns:

history – The training history.

Return type:

history

random_subset_initialization(training_data, n_epochs, n_init, take, n_kl_annealing_epochs=None, do_gs_annealing=None, **kwargs)

Random subset initialization.

This function inherits random_subset_initialization() in VariationalInferenceModelBase.

Parameters:

• training_data (tf.data.Dataset or osl_dynamics.data.Data) – Dataset to use for training.

• n_epochs (int) – Number of epochs to train the model.

• n_init (int) – Number of initializations.

• take (float) – Fraction of total batches to take.

• n_kl_annealing_epochs (int, optional) – Number of KL annealing epochs.

• do_gs_annealing (bool, optional) – Whether to anneal the Gumbel-Softmax temperature during initialization. Defaults to None, in which case the value set in the configuration will be used.

• kwargs (keyword arguments, optional) – Keyword arguments for the fit method.

Returns:

history – The training history of the best initialization.

Return type:

history

single_subject_initialization(training_data, n_epochs, n_init, n_kl_annealing_epochs=None, do_gs_annealing=None, **kwargs)

Initialization for the state means/covariances.

This function inherits single_subject_initialization() in VariationalInferenceModelBase.

Parameters:

• training_data (list of tf.data.Dataset or osl_dynamics.data.Data) – Datasets for each subject.

• n_epochs (int) – Number of epochs to train.

• n_init (int) – How many subjects should we train on?

• n_kl_annealing_epochs (int, optional) – Number of KL annealing epochs to use during initialization. If None then the KL annealing epochs in the config is used.

• kwargs (keyword arguments, optional) – Keyword arguments for the fit method.

• do_gs_annealing (Optional[bool])

Return type:

None

random_state_time_course_initialization(training_data, n_epochs, n_init, take=1, stay_prob=0.9, do_gs_annealing=None, **kwargs)

Random state time course initialization.

This function inherits random_state_time_course_initialization() in VariationalInferenceModelBase.

Parameters:

• training_data (tf.data.Dataset or osl_dynamics.data.Data) – Dataset to use for training.

• n_epochs (int) – Number of epochs to train the model.

• n_init (int) – Number of initializations.

• take (float, optional) – Fraction of total batches to take.

• stay_prob (float, optional) – Stay probability (diagonal for the transition probability matrix). Other states have uniform probability.

• do_gs_annealing (bool, optional) – Whether to anneal the Gumbel-Softmax temperature during initialization. Defaults to None, in which case the value set in the configuration will be used.

• kwargs (keyword arguments, optional) – Keyword arguments for the fit method.

Returns:

history – The training history of the best initialization.

Return type:

history

get_means()

Get the state means.

Returns:

means – State means.

Return type:

np.ndarary

get_covariances()

Get the state covariances.

Returns:

covariances – State covariances.

Return type:

np.ndarary

get_means_covariances()

Get the state means and covariances.

This is a wrapper for get_means and get_covariances.

Returns:

• means (np.ndarary) – State means.

• covariances (np.ndarray) – State covariances.

Return type:

Tuple[numpy.ndarray, numpy.ndarray]

get_observation_model_parameters()

Wrapper for get_means_covariances.

Return type:

Tuple[numpy.ndarray, numpy.ndarray]

set_means(means, update_initializer=True)

Set the state means.

Parameters:

• means (np.ndarray) – State means. Shape is (n_states, n_channels).

• update_initializer (bool) – Do we want to use the passed means when we re-initialize the model?

Return type:

None

set_covariances(covariances, update_initializer=True)

Set the state covariances.

Parameters:

• covariances (np.ndarray) – State covariances. Shape is (n_states, n_channels, n_channels).

• update_initializer (bool, optional) – Do we want to use the passed covariances when we re-initialize the model?

Return type:

None

set_means_covariances(means, covariances, update_initializer=True)

This is a wrapper for set_means and set_covariances.

Parameters:

• means (numpy.ndarray)

• covariances (numpy.ndarray)

• update_initializer (bool)

Return type:

None

set_observation_model_parameters(observation_model_parameters, update_initializer=True)

Wrapper for set_means_covariances.

Parameters:

• observation_model_parameters (tuple)

• update_initializer (bool)

Return type:

None

set_regularizers(training_dataset)

Set the means and covariances regularizer based on the training data.

A multivariate normal prior is applied to the mean vectors with mu=0, sigma=diag((range/2)**2). If config.diagonal_covariances=True, a log normal prior is applied to the diagonal of the covariances matrices with mu=0, sigma=sqrt(log(2*range)), otherwise an inverse Wishart prior is applied to the covariances matrices with nu=n_channels-1+0.1 and psi=diag(1/range).

Parameters:

training_dataset (tf.data.Dataset or osl_dynamics.data.Data) – Training dataset.

Return type:

None

sample_alpha(n_samples, states=None)

Uses the model RNN to sample a state probability time course, alpha.

Parameters:

• n_samples (int) – Number of samples to take.

• states (np.ndarray, optional) – One-hot state vectors to initialize the sampling with. Shape must be (sequence_length, n_states).

Returns:

alpha – Sampled alpha.

Return type:

np.ndarray