mdlearn.nn.models.aae.point_3d_aae

Adversarial Autoencoder for 3D point cloud data (3dAAE)

Classes

AAE3d(*args, **kwargs)
AAE3dTrainer(num_points[, num_features, ...])	Trainer class to fit a 3D Adversarial Autoencoder (3dAAE) model.

class mdlearn.nn.models.aae.point_3d_aae.AAE3d(*args: Any, **kwargs: Any)

__init__(num_points: int, num_features: int = 0, latent_dim: int = 20, encoder_bias: bool = True, encoder_relu_slope: float = 0.0, encoder_filters: list[int] = [64, 128, 256, 256, 512], encoder_kernels: list[int] = [5, 5, 3, 1, 1], decoder_bias: bool = True, decoder_relu_slope: float = 0.0, decoder_affine_widths: list[int] = [64, 128, 512, 1024], discriminator_bias: bool = True, discriminator_relu_slope: float = 0.0, discriminator_affine_widths: list[int] = [512, 128, 64])

Adversarial Autoencoder module for point cloud data from the “Adversarial Autoencoders for Compact Representations of 3D Point Clouds” paper and adapted to work on atomic coordinate data in the “AI-Driven Multiscale Simulations Illuminate Mechanisms of SARS-CoV-2 Spike Dynamics” paper. Inherits from mdlearn.nn.models.aae.AAE.

Parameters:

• num_points (int) – Number of input points in point cloud.

• num_features (int, optional) – Number of scalar features per point in addition to 3D coordinates, by default 0

• latent_dim (int, optional) – Latent dimension of the encoder, by default 20

• encoder_bias (bool, optional) – Use a bias term in the encoder Conv1d layers, by default True.

• encoder_relu_slope (float, optional) – If greater than 0.0, will use LeakyReLU activiation in the encoder with negative_slope set to relu_slope, by default 0.0

• encoder_filters (List[int], optional) – Encoder Conv1d filter sizes, by default [64, 128, 256, 256, 512]

• encoder_kernels (List[int], optional) – Encoder Conv1d kernel sizes, by default [5, 5, 3, 1, 1]

• decoder_bias (bool, optional) – Use a bias term in the decoder Linear layers, by default True

• decoder_relu_slope (float, optional) – If greater than 0.0, will use LeakyReLU activiation in the decoder with negative_slope set to relu_slope, by default 0.0

• decoder_affine_widths (List[int], optional) – Decoder Linear layers in_features, by default [64, 128, 512, 1024]

• discriminator_bias (bool, optional) – Use a bias term in the discriminator Linear layers, by default True.

• discriminator_relu_slope (float, optional) – If greater than 0.0, will use LeakyReLU activiation in the discriminator with negative_slope set to relu_slope, by default 0.0

• discriminator_affine_widths (List[int], optional) – Discriminator Linear layers in_features, by default [512, 128, 64]

critic_loss(real_logits: torch.Tensor, fake_logits: torch.Tensor) → torch.Tensor

Classification loss (critic) function.

Parameters:

• real_logits (torch.Tensor) – Discriminator output logits from prior distribution.

• fake_logits (torch.Tensor) – Discriminator output logits from encoded latent vectors.

Returns:

torch.Tensor – Classification loss i.e. the difference between logit means.

decoder_loss(fake_logit: torch.Tensor) → torch.Tensor

Decoder/Generator loss.

Parameters:

fake_logit (torch.Tensor) – Output of discriminator.

Returns:

torch.Tensor – Negative mean of the fake logits.

forward(x: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

Forward pass of encoder and decoder.

Parameters:

x (torch.Tensor) – Input point cloud data.

Returns:

tuple[torch.Tensor, torch.Tensor] – The \(z\)-latent vector, and the recon_x reconstruction.

gp_loss(noise: torch.Tensor, z: torch.Tensor) → torch.Tensor

Gradient penalty loss function.

Parameters:

• noise ([type]) – Random noise sampled from prior distribution.

• z ([type]) – Encoded latent vectors.

Returns:

torch.Tensor – The gradient penalty loss.

recon_loss(x: torch.Tensor, recon_x: torch.Tensor) → torch.Tensor

Reconstruction loss using ChamferLoss.

Parameters:

• x (torch.Tensor) – The original input tensor.

• recon_x (torch.Tensor) – The reconstructed output tensor.

Returns:

torch.Tensor – Reconstruction loss measured by Chamfer distance.

class mdlearn.nn.models.aae.point_3d_aae.AAE3dTrainer(num_points: int, num_features: int = 0, latent_dim: int = 20, encoder_bias: bool = True, encoder_relu_slope: float = 0.0, encoder_filters: list[int] = [64, 128, 256, 256, 512], encoder_kernels: list[int] = [5, 5, 3, 1, 1], decoder_bias: bool = True, decoder_relu_slope: float = 0.0, decoder_affine_widths: list[int] = [64, 128, 512, 1024], discriminator_bias: bool = True, discriminator_relu_slope: float = 0.0, discriminator_affine_widths: list[int] = [512, 128, 64], disc_optimizer_name: str = 'Adam', disc_optimizer_hparams: dict[str, Any] = {'lr': 0.0001}, ae_optimizer_name: str = 'Adam', ae_optimizer_hparams: dict[str, Any] = {'lr': 0.0001}, cms_transform: bool = False, noise_mu: float = 0.0, noise_std: float = 0.2, lambda_gp: float = 10.0, lambda_rec: float = 1.0, init_weights: str | Path | None = None, seed: int = numpy.random.default_rng.integers, num_data_workers: int = 0, prefetch_factor: int = 2, split_pct: float = 0.8, split_method: str = 'random', batch_size: int = 64, inference_batch_size: int = 64, shuffle: bool = True, device: str = 'cpu', epochs: int = 100, verbose: bool = False, clip_grad_max_norm: float = 10.0, checkpoint_log_every: int = 10, plot_log_every: int = 10, plot_n_samples: int = 10000, plot_method: str | None = None, train_subsample_pct: float = 1.0, valid_subsample_pct: float = 1.0, use_wandb: bool = False)

Trainer class to fit a 3D Adversarial Autoencoder (3dAAE) model.

__init__(num_points: int, num_features: int = 0, latent_dim: int = 20, encoder_bias: bool = True, encoder_relu_slope: float = 0.0, encoder_filters: list[int] = [64, 128, 256, 256, 512], encoder_kernels: list[int] = [5, 5, 3, 1, 1], decoder_bias: bool = True, decoder_relu_slope: float = 0.0, decoder_affine_widths: list[int] = [64, 128, 512, 1024], discriminator_bias: bool = True, discriminator_relu_slope: float = 0.0, discriminator_affine_widths: list[int] = [512, 128, 64], disc_optimizer_name: str = 'Adam', disc_optimizer_hparams: dict[str, Any] = {'lr': 0.0001}, ae_optimizer_name: str = 'Adam', ae_optimizer_hparams: dict[str, Any] = {'lr': 0.0001}, cms_transform: bool = False, noise_mu: float = 0.0, noise_std: float = 0.2, lambda_gp: float = 10.0, lambda_rec: float = 1.0, init_weights: str | Path | None = None, seed: int = numpy.random.default_rng.integers, num_data_workers: int = 0, prefetch_factor: int = 2, split_pct: float = 0.8, split_method: str = 'random', batch_size: int = 64, inference_batch_size: int = 64, shuffle: bool = True, device: str = 'cpu', epochs: int = 100, verbose: bool = False, clip_grad_max_norm: float = 10.0, checkpoint_log_every: int = 10, plot_log_every: int = 10, plot_n_samples: int = 10000, plot_method: str | None = None, train_subsample_pct: float = 1.0, valid_subsample_pct: float = 1.0, use_wandb: bool = False)

Parameters:

• num_points (int) – Number of input points in point cloud.

• num_features (int, optional) – Number of scalar features per point in addition to 3D coordinates, by default 0

• latent_dim (int, optional) – Latent dimension of the encoder, by default 20

• encoder_bias (bool, optional) – Use a bias term in the encoder Conv1d layers, by default True.

• encoder_relu_slope (float, optional) – If greater than 0.0, will use LeakyReLU activiation in the encoder with negative_slope set to relu_slope, by default 0.0

• encoder_filters (list[int], optional) – Encoder Conv1d filter sizes, by default [64, 128, 256, 256, 512]

• encoder_kernels (list[int], optional) – Encoder Conv1d kernel sizes, by default [5, 5, 3, 1, 1]

• decoder_bias (bool, optional) – Use a bias term in the decoder Linear layers, by default True

• decoder_relu_slope (float, optional) – If greater than 0.0, will use LeakyReLU activiation in the decoder with negative_slope set to relu_slope, by default 0.0

• decoder_affine_widths (list[int], optional) – Decoder Linear layers in_features, by default [64, 128, 512, 1024]

• discriminator_bias (bool, optional) – Use a bias term in the discriminator Linear layers, by default True.

• discriminator_relu_slope (float, optional) – If greater than 0.0, will use LeakyReLU activiation in the discriminator with negative_slope set to relu_slope, by default 0.0

• discriminator_affine_widths (list[int], optional) – Discriminator Linear layers in_features, by default [512, 128, 64]

• disc_optimizer_name (str, default=”Adam”) – Name of the PyTorch optimizer to use for the discriminator.

• disc_optimizer_hparams (dict[str, Any], default={“lr”: 0.0001}) – Dictionary of hyperparameters to pass to the chosen PyTorch optimizer.

• ae_optimizer_name (str, default=”Adam”) – Name of the PyTorch optimizer to use for the autoencoder.

• ae_optimizer_hparams (dict[str, Any], default={“lr”: 0.0001}) – Dictionary of hyperparameters to pass to the chosen PyTorch optimizer.

• cms_transform (bool, default=False) – If True, will subtract center of mass from batch and shift and scale batch by the full dataset statistics.

• noise_mu (float, default=0.0) – Mean of the prior distribution.

• noise_std (float, default=0.2) – Standard deviation of the prior distribution.

• lambda_gp (float, default=10.0) – Factor to scale gradient penalty loss by during training such that loss = critic_loss + lambda_gp * gp_loss.

• lambda_rec (float, default=1.0) – Factor to scale reconstruction loss by during training such that loss = decoder_loss + lambda_rec * recon_loss.

• init_weights (PathLike | None, default=None) – Path to a specific model checkpoint file to load model weights for initialization (does not load optimizer states).

• seed (int, default=np.random.default_rng().integers(2**31 - 1, dtype=int)) – Random seed for torch, numpy, and random module.

• num_data_workers (int, default=0) – How many subprocesses to use for data loading. 0 means that the data will be loaded in the main process.

• prefetch_factor (int, by default=2) – Number of samples loaded in advance by each worker. 2 means there will be a total of 2 * num_workers samples prefetched across all workers.

• split_pct (float, default=0.8) – Proportion of data set to use for training. The rest goes to validation.

• split_method (str, default=”random”) – Method to split the data. For random split use “random”, for a simple partition, use “partition”.

• batch_size (int, default=64) – Mini-batch size for training.

• inference_batch_size (int, default=64) – Mini-batch size for inference.

• shuffle (bool, default=True) – Whether to shuffle training data or not.

• device (str, default=”cpu”) – Specify training hardware either cpu or cuda for GPU devices.

• epochs (int, default=100) – Number of epochs to train for.

• verbose (bool, default False) – If True, will print training and validation loss at each epoch.

• clip_grad_max_norm (float, default=10.0) – Max norm of the gradients for gradient clipping for more information see: torch.nn.utils.clip_grad_norm_ documentation. Note that this parameter is ignored in the AAE3d model.

• checkpoint_log_every (int, default=10) – Epoch interval to log a checkpoint file containing the model weights, optimizer, and scheduler parameters.

• plot_log_every (int, default=10) – Epoch interval to log a visualization plot of the latent space.

• plot_n_samples (int, default=10000) – Number of validation samples to use for plotting.

• plot_method (str | None, default=None) – The method for visualizing the latent space or if visualization should not be run, set plot_method=None. If using "TSNE", it will attempt to use the RAPIDS.ai GPU implementation and will fallback to the sklearn CPU implementation if RAPIDS.ai is unavailable. A fast alternative is to plot the raw embeddings (or up to the first 3 dimensions if D > 3) using "raw".

• train_subsample_pct (float, default=1.0) – Percentage of training data to use during hyperparameter sweeps.

• valid_subsample_pct (float, default=1.0) – Percentage of validation data to use during hyperparameter sweeps.

• use_wandb (bool, default=False) – If True, will log results to wandb. Metric keys include train_disc_loss, train_ae_loss, and valid_recon_loss.

Raises:

• ValueError – split_pct should be between 0 and 1.

• ValueError – train_subsample_pct should be between 0 and 1.

• ValueError – valid_subsample_pct should be between 0 and 1.

• ValueError – Specified device as cuda, but it is unavailable.

fit(X: numpy.typing.ArrayLike, scalars: dict[str, numpy.typing.ArrayLike] = {}, output_path: str | Path = './', checkpoint: str | Path | None = None) → None

Trains the autoencoder on the input data X.

Parameters:

• X (ArrayLike) – Input point cloud data of shape (N, 3, num_points) where N is the number of data examples, 3 is the x, y, z coordinates of each point, and num_points is the number of points in the point cloud (e.g. number of residues in a protein structure).

• scalars (Dict[str, ArrayLike], default={}) – Dictionary of scalar arrays. For instance, the root mean squared deviation (RMSD) for each feature vector can be passed via {"rmsd": np.array(...)}. The dimension of each scalar array should match the number of input feature vectors N.

• output_path (PathLike, default=”./”) – Path to write training results to. Makes an output_path/checkpoints folder to save model checkpoint files, and output_path/plots folder to store latent space visualizations.

• checkpoint (Optional[PathLike], default=None) – Path to a specific model checkpoint file to restore training.

Raises:

• TypeError – If scalars is not type dict. A common error is to pass output_path as the second argument.

• NotImplementedError – If using a learning rate scheduler other than ReduceLROnPlateau, a step function will need to be implemented.

predict(X: numpy.typing.ArrayLike, inference_batch_size: int | None = None, checkpoint: str | Path | None = None) → tuple[numpy.typing.ArrayLike, float]

Predict using the LinearAE

Parameters:

• X (ArrayLike) – Input point cloud data of shape (N, 3, num_points) where N is the number of data examples, 3 is the x, y, z coordinates of each point, and num_points is the number of points in the point cloud (e.g. number of residues in a protein structure).

• inference_batch_size (int, default=None) – The batch size for inference (if None uses the value specified during Trainer construction).

• checkpoint (PathLike | None, default=None) – Path to a specific model checkpoint file.

Returns:

tuple[ArrayLike, float] – The z latent vectors corresponding to the input data X and the average reconstruction loss.