trphysx.embedding¶
Subpackages¶
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:
-
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:
-
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:
-
classmethod
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 tensorReturn 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 lossReturn 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 tensorReturn 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 lossReturn 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 tensorReturn 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 lossReturn 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
-