MolecularDiffusion.modules.models.diffsmol.shape_ae¶

DiffSMol point-cloud shape autoencoder (VN-DGCNN encoder + implicit signed-distance decoder).

Ported from DiffSMol source/models/shape_pointcloud_modelAE.py. Only the encoder is used by the integration: it maps 512 points sampled from a molecular surface mesh to a (128, 3) SO(3)-equivariant latent that rotates with the molecule. That latent is DiffSMol’s conditioning signal.

Two deliberate departures from upstream, both required for correctness:

``blocks`` is an ``nn.ModuleList``, not a plain Python list. Upstream stores the 6 encoder blocks in a bare list, so PyTorch never registers them: they are absent from state_dict(), were never saved into the released se.pt, and are therefore re-randomised on every process start. That makes upstream’s shape latent irreproducible across runs – fatal here, where the precompute, the training run and the generation run are three separate processes that must agree on the same embedding function. Registering them properly plus vendoring a checkpoint that contains them (see checkpoints/shape_ae_pointcloud.pt) fixes it.
No ``easydict``. The vendored checkpoint stores a plain dict config.

The AE is frozen and eval-only here; nothing in this file is trained.

Attributes¶

`DEFAULT_SHAPE_AE_CHECKPOINT`
`EPS`

Classes¶

`DecoderInner`	Implicit decoder: (query points, VN latent) -> signed distance.
`PointCloudAE`	VN point-cloud autoencoder. Only `encode` is used downstream.
`VNDGCNNEncoder`	VN-DGCNN encoder: centered point cloud -> `(latent_dim, 3)` latent.
`VNResnetEncoder`	Alternative VN-ResNet encoder (`encoder: VN_Resnet`). Not used by the

Functions¶

load_shape_ae(→ PointCloudAE)

Load the frozen shape autoencoder from a vendored checkpoint.

Module Contents¶

class MolecularDiffusion.modules.models.diffsmol.shape_ae.DecoderInner(dim: int = 3, z_dim: int = 128, hidden_size: int = 128, layer_num: int = 4, loss_type: str = 'occupancy', leaky: bool = False)¶

Bases: torch.nn.Module

Implicit decoder: (query points, VN latent) -> signed distance.

Unused by training/generation – kept so the vendored checkpoint loads strictly and so gradient-guidance can be added later without a re-port.

forward(p: torch.Tensor, z: torch.Tensor) → torch.Tensor¶

actvn¶

blocks¶

fc_in¶

fc_out¶

layer_num = 4¶

loss_type = 'occupancy'¶

z_dim = 128¶

class MolecularDiffusion.modules.models.diffsmol.shape_ae.PointCloudAE(config: Dict[str, Any])¶

Bases: torch.nn.Module

VN point-cloud autoencoder. Only encode is used downstream.

encode(point_clouds: torch.Tensor) → torch.Tensor¶

point_clouds: [B, N, 3], surface-centered.

Returns the equivariant latent [B, latent_dim, 3].

forward(point_clouds: torch.Tensor, query_points: torch.Tensor) → torch.Tensor¶

encoder¶

generator¶

loss_type¶

class MolecularDiffusion.modules.models.diffsmol.shape_ae.VNDGCNNEncoder(hidden_dim: int, latent_dim: int, layer_num: int, num_k: int)¶

Bases: torch.nn.Module

VN-DGCNN encoder: centered point cloud -> (latent_dim, 3) latent.

forward(x: torch.Tensor) → torch.Tensor¶

x: [B, 1, N, 3] (already surface-centered).

Returns [B, latent_dim, 3].

blocks¶

conv_c¶

conv_pos¶

layer_num¶

num_k¶

pool¶

class MolecularDiffusion.modules.models.diffsmol.shape_ae.VNResnetEncoder(hidden_dim: int, latent_dim: int, layer_num: int, num_k: int)¶

Bases: torch.nn.Module

Alternative VN-ResNet encoder (encoder: VN_Resnet). Not used by the shipped checkpoint; kept for config parity.

forward(x: torch.Tensor) → torch.Tensor¶

actvn_c¶

blocks¶

conv_pos¶

fc_c¶

fc_pos¶

layer_num¶

num_k¶

pool¶

MolecularDiffusion.modules.models.diffsmol.shape_ae.load_shape_ae(checkpoint: str | None = None, device: str | torch.device = 'cpu') → PointCloudAE¶

Load the frozen shape autoencoder from a vendored checkpoint.

The checkpoint is a plain {"config": {...}, "model": state_dict} dict (no easydict). Parameters are detached and the module is put in eval mode – the AE is a fixed featuriser, never trained here.

MolecularDiffusion.modules.models.diffsmol.shape_ae.DEFAULT_SHAPE_AE_CHECKPOINT = ''¶

MolecularDiffusion.modules.models.diffsmol.shape_ae.EPS = 1e-06¶