trphysx.utils

trphysx.utils.metrics

class trphysx.utils.metrics.Metrics(file_path: str = '.', file_name: str = 'log_metrics.h5')

Bases: object

Data class for storing training errors

Parameters:
  • file_path (str, optional) – Path to write logging files
  • file_name (str, optional) – Log file name
file_path = '.'
file_name = 'log_metrics.h5'
push(**kwargs) → None

Pushes elements in kwargs into the attributes of this class

Parameters:**kwargs – Attributes to save
writeToHDF5(file_name: str = None) → None

Write the classes attributes to HDF5 file

Parameters:file_name (str, optional) – File name to write to
appendToHDF5(file_name: str = None) → None

Appends the classes attributes to HDF5 file

Parameters:file_name (str, optional) – File name to write to
delHDF5(file_name: str = None) → None

Deletes hdf5 file if it exists

Parameters:file_name (str, optional) – File name

trphysx.utils.trainer

trphysx.utils.trainer.set_seed(seed: int) → None

Set random seed

Parameters:seed (int) – random seed
class trphysx.utils.trainer.Trainer(model: trphysx.transformer.phys_transformer_helpers.PhysformerTrain, args: trphysx.config.args.TrainingArguments, optimizers: Tuple[torch.optim.optimizer.Optimizer, torch.optim.lr_scheduler._LRScheduler], train_dataset: torch.utils.data.dataset.Dataset = None, eval_dataset: torch.utils.data.dataset.Dataset = None, embedding_model: trphysx.embedding.embedding_model.EmbeddingModel = None, viz: trphysx.viz.viz_model.Viz = None)

Bases: object

Generalized trainer for physics transformer models

Parameters:
  • model (PhysformerTrain) – Transformer with training head
  • args (TrainingArguments) – Training arguments
  • optimizers (Tuple[Optimizer, Scheduler], optional) – Tuple of Pytorch optimizer and lr scheduler.
  • train_dataset (Dataset, optional) – Training dataset. Defaults to None.
  • eval_dataset (Dataset, optional) – Eval/Validation dataset. Defaults to None.
  • embedding_model (EmbeddingModel, optional) – Embedding model. Used for recovering states during state evaluation of the model. Defaults to None.
  • viz (Viz, optional) – Visualization class. Defaults to None.
get_train_dataloader(train_dataset: torch.utils.data.dataset.Dataset = None) → torch.utils.data.dataloader.DataLoader

Creates a training dataloader. Overload for unusual training cases.

Parameters:train_dataset (Dataset, optional) – Optional training dataset. If none is provided, the class training data will be used. Defaults to None.
Raises:ValueError – If both the dataset parameter and class dataset have not been provided
Returns:Training dataloader
Return type:DataLoader
get_eval_dataloader(eval_dataset: torch.utils.data.dataset.Dataset = None) → torch.utils.data.dataloader.DataLoader

Creates a evaluation dataloader used for validation or testing of model.

Parameters:eval_dataset (Dataset, optional) – Optional eval dataset. If none is provided, the class eval data will be used. Defaults to None.
Raises:ValueError – If both the dataset parameter and class dataset have not been provided
Returns:Evaluation dataloader
Return type:DataLoader
train() → None

Trains the transformer model

training_step(model: trphysx.transformer.phys_transformer_helpers.PhysformerTrain, inputs: Dict[str, Any]) → Tuple[float, torch.Tensor, torch.Tensor]

Calls a forward pass of the training model and backprops for a single time-step

Parameters:
  • model (PhysformerTrain) – Transformer model with training head, could be
  • inputs (Dict[str, Any]) – Dictionary of model inputs for forward pass
Returns:

Tuple containing: loss value, hidden states

of transformer, attention states of the transformer.
Return type:

Tuple[float, Tensor, Tensor]
evaluate(epoch: int = None) → Dict[str, float]

Run evaluation and return metrics.

Parameters:epoch (int, optional) – Current epoch, used for naming figures. Defaults to None.
Returns:Dictionary of prediction metrics
Return type:Dict[str, float]
eval_step(model: trphysx.transformer.phys_transformer_helpers.PhysformerTrain, inputs: Dict[str, Any]) → Tuple[float, torch.Tensor, torch.Tensor]

Calls a eval pass of the training model.

Parameters:
  • model (PhysformerTrain) – Transformer model with training head
  • inputs (Dict[str, Any]) – Dictionary of model inputs for forward pass
Returns:

Tuple containing: prediction error value,

time-step error, predicted embeddings.
Return type:

Tuple[float, Tensor, Tensor]
eval_states(pred_embeds: torch.Tensor, states: Any, epoch: int = None, plot_id: int = 0, plot: bool = True) → float

Evaluates the predicted states by recovering the state space from the predicted embedding vectors. Can be overloaded for cases with special methods for recovering the state field.

Parameters:
  • pred_embeds (Tensor) – [B, T, n_embed] Predicted embedded vectors
  • states (Any) – Target states / data for recovery
  • epoch (int, optional) – Current epoch, used for naming figures. Defaults to None.
  • plot_id (int, optional) – Secondary plotting id to distinguish between numerical cases. Defaults to 0.
  • plot (bool, optional) – Plot models states. Defaults to True.
Returns:

Predicted state MSE error
Return type:

float