trphysx.embedding

trphysx.embedding.embedding_auto

class trphysx.embedding.embedding_auto.AutoEmbeddingModel

Bases: object

Helper class for initializing of loading various embedding models.

Raises:EnvironmentError – If direct initialization of this class is attempted.
classmethod init_model(model_name: str, config: trphysx.config.configuration_phys.PhysConfig) → trphysx.embedding.embedding_model.EmbeddingModel

Initialize embedding model. Currently supports: “lorenz”, “cylinder”, “grayscott”

Parameters:
  • model_name (str) – Keyword/name of embedding model
  • config (PhysConfig) – Transformer configuration class
Raises:

ValueError – If model_name is not a supported model type

Returns:

Initialized embedding model

Return type:

(EmbeddingModel)

classmethod init_trainer(model_name: str, config: trphysx.config.configuration_phys.PhysConfig) → trphysx.embedding.embedding_model.EmbeddingTrainingHead

Initialize embedding model with a training head. Currently supports: “lorenz”, “cylinder”, “grayscott”

Parameters:
  • model_name (str) – Keyword/name of embedding model
  • config (PhysConfig) – Transformer configuration class
Raises:

KeyError – If model_name is not a supported trainer model types

Returns:

Initialized embedding model trainer

Return type:

(EmbeddingTrainer)

classmethod load_model(model_name: str, config: trphysx.config.configuration_phys.PhysConfig, file_or_path_directory: Optional[str] = None, epoch: int = 0) → trphysx.embedding.embedding_model.EmbeddingModel

Initialize and load embedding model from memory. Currently supports: “lorenz”, “cylinder”, “grayscott”

Parameters:
  • model_name (str) – Keyword/name of embedding model
  • config (PhysConfig) – Transformer configuration class
  • file_or_path_directory (str, optional) – embedding model file or directory path
  • epoch (int, optional) – Epoch to load model from, only used if function is provided a directory
Raises:

ValueError – If model_name is not a supported model type

Returns:

Initialized embedding model with loaded weights

Return type:

(EmbeddingModel)

trphysx.embedding.embedding_cylinder

class trphysx.embedding.embedding_cylinder.CylinderEmbedding(config: trphysx.config.configuration_phys.PhysConfig)

Bases: trphysx.embedding.embedding_model.EmbeddingModel

Embedding Koopman model for the 2D flow around a cylinder system

Parameters:config (PhysConfig) – Configuration class with transformer/embedding parameters
model_name = 'embedding_cylinder'
forward(x: torch.Tensor, visc: torch.Tensor) → Tuple[torch.Tensor]

Forward pass

Parameters:
  • x (Tensor) – [B, 3, H, W] Input feature tensor
  • visc (Tensor) – [B] Viscosities of the fluid in the mini-batch
Returns:

Tuple containing:

(Tensor): [B, config.n_embd] Koopman observables
(Tensor): [B, 3, H, W] Recovered feature tensor

Return type:

(TensorTuple)

embed(x: torch.Tensor, visc: torch.Tensor) → torch.Tensor

Embeds tensor of state variables to Koopman observables

Parameters:
  • x (Tensor) – [B, 3, H, W] Input feature tensor
  • visc (Tensor) – [B] Viscosities of the fluid in the mini-batch
Returns:

[B, config.n_embd] Koopman observables

Return type:

(Tensor)

recover(g: torch.Tensor) → torch.Tensor

Recovers feature tensor from Koopman observables

Parameters:g (Tensor) – [B, config.n_embd] Koopman observables
Returns:[B, 3, H, W] Physical feature tensor
Return type:(Tensor)
koopmanOperation(g: torch.Tensor, visc: torch.Tensor) → torch.Tensor

Applies the learned Koopman operator on the given observables

Parameters:
  • g (Tensor) – [B, config.n_embd] Koopman observables
  • visc (Tensor) – [B] Viscosities of the fluid in the mini-batch
Returns:

[B, config.n_embd] Koopman observables at the next time-step

Return type:

Tensor

koopmanOperator

Current Koopman operator

Parameters:requires_grad (bool, optional) – If to return with gradient storage. Defaults to True
Returns:Full Koopman operator tensor
Return type:Tensor
koopmanDiag
class trphysx.embedding.embedding_cylinder.CylinderEmbeddingTrainer(config: trphysx.config.configuration_phys.PhysConfig)

Bases: trphysx.embedding.embedding_model.EmbeddingTrainingHead

Training head for the Lorenz embedding model

Parameters:config (PhysConfig) – Configuration class with transformer/embedding parameters
forward(states: torch.Tensor, viscosity: torch.Tensor) → Tuple[float]

Trains model for a single epoch

Parameters:
  • states (Tensor) – [B, T, 3, H, W] Time-series feature tensor
  • viscosity (Tensor) – [B] Viscosities of the fluid in the mini-batch
Returns:

Tuple containing:

(float): Koopman based loss of current epoch
(float): Reconstruction loss

Return type:

FloatTuple

evaluate(states: torch.Tensor, viscosity: torch.Tensor) → Tuple[float, torch.Tensor, torch.Tensor]

Evaluates the embedding models reconstruction error and returns its predictions.

Parameters:
  • states (Tensor) – [B, T, 3, H, W] Time-series feature tensor
  • viscosity (Tensor) – [B] Viscosities of the fluid in the mini-batch
Returns:

Test error, Predicted states, Target states

Return type:

Tuple[Float, Tensor, Tensor]

trphysx.embedding.embedding_grayscott

class trphysx.embedding.embedding_grayscott.GrayScottEmbedding(config: trphysx.config.configuration_phys.PhysConfig)

Bases: trphysx.embedding.embedding_model.EmbeddingModel

Embedding Koopman model for the 3D Gray-Scott system

Parameters:config (PhysConfig) – Configuration class with transformer/embedding parameters

Note

For more information on the Gray-Scott model see “Complex Patterns in a Simple System” by John E. Pearson; https://doi.org/10.1126/science.261.5118.189

model_name = 'embedding_grayscott'
forward(x: torch.Tensor) → Tuple[torch.Tensor]

Forward pass

Parameters:x (Tensor) – [B, 2, H, W, D] Input feature tensor
Returns:Tuple containing:
(Tensor): [B, config.n_embd] Koopman observables
(Tensor): [B, 2, H, W, D] Recovered feature tensor
Return type:TensorTuple
embed(x: torch.Tensor) → torch.Tensor

Embeds tensor of state variables to Koopman observables

Parameters:x (Tensor) – [B, 2, H, W, D] Input feature tensor
Returns:[B, config.n_embd] Koopman observables
Return type:Tensor
recover(g: torch.Tensor) → torch.Tensor

Recovers feature tensor from Koopman observables

Parameters:g (Tensor) – [B, config.n_embd] Koopman observables
Returns:[B, 2, H, W, D] Physical feature tensor
Return type:(Tensor)
koopmanOperation(g: torch.Tensor) → torch.Tensor

Applies the learned Koopman operator on the given observables

Parameters:g (Tensor) – [B, config.n_embd] Koopman observables
Returns:[B, config.n_embd] Koopman observables at the next time-step
Return type:(Tensor)
koopmanOperator

Current Koopman operator

Parameters:requires_grad (bool, optional) – If to return with gradient storage. Defaults to True
Returns:Full Koopman operator tensor
Return type:Tensor
koopmanDiag
class trphysx.embedding.embedding_grayscott.GrayScottEmbeddingTrainer(config: trphysx.config.configuration_phys.PhysConfig)

Bases: trphysx.embedding.embedding_model.EmbeddingTrainingHead

Training head for the Gray-Scott embedding model

Parameters:config (PhysConfig) – Configuration class with transformer/embedding parameters
forward(states: torch.Tensor) → Tuple[float]

Trains model for a single epoch

Parameters:states (Tensor) – [B, T, 3, H, W] Time-series feature tensor
Returns:Tuple containing:
(float): Koopman based loss of current epoch
(float): Reconstruction loss
Return type:FloatTuple

trphysx.embedding.embedding_lorenz

class trphysx.embedding.embedding_lorenz.LorenzEmbedding(config: trphysx.config.configuration_phys.PhysConfig)

Bases: trphysx.embedding.embedding_model.EmbeddingModel

Embedding Koopman model for the Lorenz ODE system

Parameters:config (PhysConfig) – Configuration class with transformer/embedding parameters
model_name = 'embedding_lorenz'
forward(x: torch.Tensor) → Tuple[torch.Tensor]

Forward pass

Parameters:x (Tensor) – [B, 3] Input feature tensor
Returns:Tuple containing:
(Tensor): [B, config.n_embd] Koopman observables
(Tensor): [B, 3] Recovered feature tensor
Return type:TensorTuple
embed(x: torch.Tensor) → torch.Tensor

Embeds tensor of state variables to Koopman observables

Parameters:x (Tensor) – [B, 3] Input feature tensor
Returns:[B, config.n_embd] Koopman observables
Return type:Tensor
recover(g: torch.Tensor) → torch.Tensor

Recovers feature tensor from Koopman observables

Parameters:g (Tensor) – [B, config.n_embd] Koopman observables
Returns:[B, 3] Physical feature tensor
Return type:Tensor
koopmanOperation(g: torch.Tensor) → torch.Tensor

Applies the learned Koopman operator on the given observables

Parameters:g (Tensor) – [B, config.n_embd] Koopman observables
Returns:[B, config.n_embd] Koopman observables at the next time-step
Return type:(Tensor)
koopmanOperator

Current Koopman operator

Parameters:requires_grad (bool, optional) – If to return with gradient storage. Defaults to True
Returns:Full Koopman operator tensor
Return type:(Tensor)
koopmanDiag
class trphysx.embedding.embedding_lorenz.LorenzEmbeddingTrainer(config: trphysx.config.configuration_phys.PhysConfig)

Bases: trphysx.embedding.embedding_model.EmbeddingTrainingHead

Training head for the Lorenz embedding model

Parameters:config (PhysConfig) – Configuration class with transformer/embedding parameters
forward(states: torch.Tensor) → Tuple[float]

Trains model for a single epoch

Parameters:states (Tensor) – [B, T, 3] Time-series feature tensor
Returns:Tuple containing:
(float): Koopman based loss of current epoch
(float): Reconstruction loss
Return type:FloatTuple
evaluate(states: torch.Tensor) → Tuple[float, torch.Tensor, torch.Tensor]

Evaluates the embedding models reconstruction error and returns its predictions.

Parameters:states (Tensor) – [B, T, 3] Time-series feature tensor
Returns:Test error, Predicted states, Target states
Return type:Tuple[Float, Tensor, Tensor]

trphysx.embedding.embedding_model

class trphysx.embedding.embedding_model.EmbeddingModel(config: trphysx.config.configuration_phys.PhysConfig)

Bases: torch.nn.modules.module.Module

Parent class for embedding models that handle the projection of the physical systems states into a vector representation

Parameters:config (PhysConfig) – Configuration class with transformer/embedding parameters
model_name = 'embedding_model'
embed(x)
recover(x)
koopmanOperator
koopmanDiag
input_dims
embed_dims
num_parameters

Get number of learnable parameters in model

devices

Get list of unique device(s) model exists on

save_model(save_directory: str, epoch: int = 0) → None

Saves embedding model to the specified directory.

Parameters:
  • save_directory (str) – Folder directory to save state dictionary to.
  • epoch (int, optional) – Epoch of current model for file name. Defaults to 0.
Raises:

FileNotFoundError – If provided path is a file

load_model(file_or_path_directory: str, epoch: int = 0) → None

Load a embedding model from the specified file or path

Parameters:
  • file_or_path_directory (str) – File or folder path to load state dictionary from.
  • epoch (int, optional) – Epoch of current model for file name, used if folder path is provided. Defaults to 0.
Raises:

FileNotFoundError – If provided file or directory could not be found.

class trphysx.embedding.embedding_model.EmbeddingTrainingHead

Bases: torch.nn.modules.module.Module

Parent class for training head for embedding models

forward(*args, **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

evaluate(*args, **kwargs)
save_model(*args, **kwargs)

Saves the embedding model

load_model(*args, **kwargs)

Load the embedding model