neural_lam.models.latent.base_encoder

Abstract base class for latent encoders.

Classes

BaseLatentEncoder

Abstract encoder mapping an input grid representation to a Gaussian

Module Contents

class neural_lam.models.latent.base_encoder.BaseLatentEncoder(latent_dim: int, output_dist: str = 'isotropic')

Bases: torch.nn.Module

Abstract encoder mapping an input grid representation to a Gaussian distribution over a latent variable defined on mesh nodes.

Subclasses implement compute_dist_params(), which returns the raw parameters used to build the output distribution. With output_dist="isotropic" only the mean is produced (unit variance); with output_dist="diagonal" both mean and a positive std are output.

Set up output dimensionality for the chosen distribution type.

Parameters:

• latent_dim (int) – Dimensionality of the latent variable at each mesh node.

• output_dist (str) – Type of output distribution: "isotropic" (mean only, unit variance) or "diagonal" (mean and per-dimension std).

abstract compute_dist_params(grid_rep: torch.Tensor, *args: Any, **kwargs: Any) → torch.Tensor

Compute raw distribution parameters from the grid representation.

Parameters:

• grid_rep (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Grid input representation.

• **kwargs – Additional inputs used by concrete encoders (e.g. graph embeddings).

Returns:

Shape (B, num_mesh_nodes, output_dim). Raw parameters of the latent distribution.

Return type:

torch.Tensor

forward(grid_rep: torch.Tensor, **kwargs: Any) → torch.distributions.Normal

Compute the Gaussian distribution over the latent variable.

Parameters:

• grid_rep (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Grid input representation.

• **kwargs – Additional inputs forwarded to compute_dist_params().

Returns:

Distribution over the latent variable, with batch shape (B, num_mesh_nodes, latent_dim).

Return type:

torch.distributions.Normal

output_dist = 'isotropic'