onyxengine.modeling

The models and training configs for the Engine.

Models

class onyxengine.modeling.MLPConfig(*, type: Literal['mlp'] = 'mlp', outputs: List[Output], inputs: List[Input | State], dt: float = 0.0, sequence_length: int = 1, hidden_layers: int = 2, hidden_size: int = 32, activation: Literal['relu', 'gelu', 'tanh', 'sigmoid'] = 'relu', dropout: float = 0.0, bias: bool = True)[source]

Configuration class for the MLP model.

Parameters:

  • type (str) – Model type = ‘mlp’, immutable.

  • outputs (List[Output]) – List of output variables.

  • inputs (List[Input | State]) – List of input variables.

  • dt (float) – Time step for the model.

  • sequence_length (int) – Length of the input sequence (default is 1).

  • hidden_layers (int) – Number of hidden layers (default is 2).

  • hidden_size (int) – Size of each hidden layer (default is 32).

  • activation (Literal['relu', 'gelu', 'tanh', 'sigmoid']) – Activation function (default is ‘relu’).

  • dropout (float) – Dropout rate for layers (default is 0.0).

  • bias (bool) – Whether to use bias in layers (default is True).

class onyxengine.modeling.RNNConfig(*, type: Literal['rnn'] = 'rnn', outputs: List[Output], inputs: List[Input | State], dt: float = 0.0, sequence_length: int = 1, rnn_type: Literal['RNN', 'LSTM', 'GRU'] = 'LSTM', hidden_layers: int = 2, hidden_size: int = 32, dropout: float = 0.0, bias: bool = True)[source]

Configuration class for the RNN model.

Parameters:

  • type (str) – Model type = ‘rnn’, immutable.

  • outputs (List[Output]) – List of output variables.

  • inputs (List[Input | State]) – List of input variables.

  • dt (float) – Time step for the model.

  • sequence_length (int) – Length of input sequences (default is 1).

  • rnn_type (Literal['RNN', 'LSTM', 'GRU']) – Type of RNN to use (default is ‘LSTM’).

  • hidden_layers (int) – Number of hidden layers in the RNN (default is 2).

  • hidden_size (int) – Number of hidden units in each layer (default is 32).

  • dropout (float) – Dropout rate (default is 0.0).

  • bias (bool) – Whether or not to include bias in the RNN (default is True).

class onyxengine.modeling.TransformerConfig(*, type: Literal['transformer'] = 'transformer', outputs: List[Output], inputs: List[Input | State], dt: float = 0.0, sequence_length: int = 1, n_layer: int = 1, n_head: int = 4, n_embd: int = 32, dropout: float = 0.0, bias: bool = True)[source]

Configuration class for the Transformer model.

Parameters:

  • type (str) – Model type = ‘transformer’, immutable.

  • outputs (List[Output]) – List of output variables.

  • inputs (List[Input | State]) – List of input variables.

  • dt (float) – Time step for the model.

  • sequence_length (int) – Length of the input sequence (default is 1).

  • n_layer (int) – Number of transformer layers (default is 1).

  • n_head (int) – Number of attention heads (default is 4).

  • n_embd (int) – Size of the embedding dimension (default is 32).

  • dropout (float) – Dropout rate for layers (default is 0.0).

  • bias (bool) – Whether to use bias in layers (default is True).

Model Simulation

class onyxengine.modeling.Output(*, type: Literal['output'] = 'output', name: str, scale: None | Literal['mean'] | List[float] = 'mean', train_mean: float | None = None, train_std: float | None = None, train_min: float | None = None, train_max: float | None = None)[source]

A standard output feature to be used by the model.

Parameters:

  • name (str) – Name of the output feature.

  • scale (Union[None, Literal['mean'], List[float]]) –

    Scale for the output feature:

    • None: Feature is not scaled.

    • ’mean’: Feature is scaled to have a mean of 0 and std of 1. (Default).

    • List[float]: Feature is scaled from its real-world [min, max] to a range of [-1, 1].

class onyxengine.modeling.Input(*, type: Literal['input'] = 'input', name: str, scale: None | Literal['mean'] | List[float] = 'mean', train_mean: float | None = None, train_std: float | None = None, train_min: float | None = None, train_max: float | None = None)[source]

A standard input feature to be used by the model.

Parameters:

  • name (str) – Name of the input feature.

  • scale (Union[None, Literal['mean'], List[float]]) –

    Scale for the output feature:

    • None: Feature is not scaled.

    • ’mean’: Feature is scaled to have a mean of 0 and std of 1. (Default).

    • List[float]: Feature is scaled from its real-world [min, max] to a range of [-1, 1].

class onyxengine.modeling.State(*, type: Literal['state'] = 'state', name: str, scale: None | Literal['mean'] | List[float] = 'mean', train_mean: float | None = None, train_std: float | None = None, train_min: float | None = None, train_max: float | None = None, relation: Literal['output', 'delta', 'derivative'], parent: str)[source]

A state feature that can be derived from a parent feature through different relationships (output, delta, or derivative).

Parameters:

  • name (str) – Name of the state feature.

  • relation (Literal['output', 'delta', 'derivative']) –

    Method to solve for the feature:

    • ’output’: Feature is the direct output of the model

    • ’delta’: Feature is the change/delta of the parent value

    • ’derivative’: Feature is the derivative of the parent value

  • parent (str) – Name of the parent feature from which this state is derived

  • scale (Union[None, Literal['mean'], List[float]]) –

    Scale for the output feature:

    • None: Feature is not scaled.

    • ’mean’: Feature is scaled to have a mean of 0 and std of 1. (Default).

    • List[float]: Feature is scaled from its real-world [min, max] to a range of [-1, 1].

Model Training

class onyxengine.modeling.TrainingConfig(*, type: Literal['training_config'] = 'training_config', training_iters: int = 3000, train_batch_size: int = 32, train_val_split_ratio: float = 0.9, test_dataset_size: int = 500, checkpoint_type: Literal['single_step', 'multi_step'] = 'single_step', optimizer: AdamWConfig | SGDConfig = AdamWConfig(type='adamw', lr=0.0003, weight_decay=0.01), lr_scheduler: None | CosineDecayWithWarmupConfig | CosineAnnealingWarmRestartsConfig = None)[source]

Configuration for the training of a model.

Parameters:

  • training_iters (int) – Number of training iterations (default is 3000).

  • train_batch_size (int) – Batch size for training (default is 32).

  • train_val_split_ratio (float) – Ratio of training data to validation data (default is 0.9).

  • test_dataset_size (int) – Number of samples in the test dataset (default is 500).

  • checkpoint_type (Literal['single_step', 'multi_step']) – Type of checkpointing (default is ‘single_step’).

  • optimizer (Union[AdamWConfig, SGDConfig]) – Optimizer configuration (default is AdamWConfig()).

  • lr_scheduler (Union[None, CosineDecayWithWarmupConfig, CosineAnnealingWarmRestartsConfig]) – Learning rate scheduler configuration (default is None).

class onyxengine.modeling.OptimizationConfig(*, type: Literal['optimization_config'] = 'optimization_config', training_iters: int = 3000, train_batch_size: int = 32, train_val_split_ratio: float = 0.9, test_dataset_size: int = 500, checkpoint_type: Literal['single_step', 'multi_step'] = 'single_step', opt_models: List[MLPOptConfig | RNNOptConfig | TransformerOptConfig] = [], opt_optimizers: List[AdamWOptConfig | SGDOptConfig] = [], opt_lr_schedulers: List[None | CosineDecayWithWarmupOptConfig | CosineAnnealingWarmRestartsOptConfig] = [None], num_trials: int = 10)[source]

Configuration for the optimization of models.

Parameters:

  • training_iters (int) – Number of training iterations (default is 3000).

  • train_batch_size (int) – Batch size for training (default is 32).

  • train_val_split_ratio (float) – Ratio of training data to validation data (default is 0.9).

  • test_dataset_size (int) – Number of samples in the test dataset (default is 500).

  • checkpoint_type (Literal['single_step', 'multi_step']) – Type of checkpointing (default is ‘single_step’).

  • opt_models (List[Union[MLPOptConfig, RNNOptConfig, TransformerOptConfig]]) – List of model optimization configurations.

  • opt_optimizers (List[Union[AdamWOptConfig, SGDOptConfig]]) – List of optimizer optimization configurations.

  • opt_lr_schedulers (List[Union[None, CosineDecayWithWarmupOptConfig, CosineAnnealingWarmRestartsOptConfig]]) – List of learning rate scheduler optimization configurations.

  • num_trials (int) – Number of optimization trials (default is 10).

class onyxengine.modeling.AdamWConfig(*, type: Literal['adamw'] = 'adamw', lr: float = 0.0003, weight_decay: float = 0.01)[source]

Configuration for the AdamW optimizer.

Parameters:

  • lr (float) – Learning rate (default is 3e-4).

  • weight_decay (float) – Weight decay (default is 1e-2).

class onyxengine.modeling.SGDConfig(*, type: Literal['sgd'] = 'sgd', lr: float = 0.0003, weight_decay: float = 0.01, momentum: float = 0.9)[source]

Configuration for the SGD optimizer.

Parameters:

  • lr (float) – Learning rate (default is 3e-4).

  • weight_decay (float) – Weight decay (default is 1e-2).

  • momentum (float) – Momentum (default is 0.9).

class onyxengine.modeling.CosineDecayWithWarmupConfig(*, type: Literal['cosine_decay_with_warmup'] = 'cosine_decay_with_warmup', max_lr: float = 0.0003, min_lr: float = 3e-05, warmup_iters: int = 200, decay_iters: int = 1000)[source]

Configuration for learning rate scheduler with cosine decay and linear warmup.

Parameters:

  • max_lr (float) – Maximum learning rate (default is 3e-4).

  • min_lr (float) – Minimum learning rate (default is 3e-5).

  • warmup_iters (int) – Number of warmup iterations (default is 200).

  • decay_iters (int) – Number of decay iterations (default is 1000).

class onyxengine.modeling.CosineAnnealingWarmRestartsConfig(*, type: Literal['cosine_annealing_warm_restarts'] = 'cosine_annealing_warm_restarts', T_0: int = 2000, T_mult: int = 1, eta_min: float = 3e-05)[source]

Configuration for learning rate scheduler with cosine annealing and warm restarts.

Parameters:

  • T_0 (int) – Initial period of learning rate decay (default is 2000).

  • T_mult (int) – Multiplicative factor for the period of learning rate decay (default is 1).

  • eta_min (float) – Minimum learning rate (default is 3e-5).