aidsorb.modules.voxels#

torch.nn.Module’s for voxels processing.

References

[IntelliPore]

A. P. Sarikas, K. Gkagkas, and G. E. Froudakis, “IntelliPore: A Foundation Model for Gas Adsorption in Porous Materials” ADD IT IN IEEE FORMAT

[RetNeXt]

Sarikas, A. P., Gkagkas, K., & Froudakis, G. E. (2026). RetNeXt: A Pretrained Model for Transfer Learning Across the MOF Adsorption Space. Journal of Chemical Information and Modeling, 66(4), 2110–2116. https://doi.org/10.1021/acs.jcim.5c02698

class aidsorb.modules.voxels.IntelliPore(in_channels=1, n_outputs=1, *, pretrained=False)[source]#

Bases: Module

Architecture from the [IntelliPore] paper.

Note

pretrained=True is only compatible with in_channels=1, since the pretrained backbone was trained on single-channel images.

Parameters:

  • in_channels (int, default=1)

  • n_outputs (int or None, default=1) – Number of output units. If None, no linear head is added and the model acts as feature extractor.

  • pretrained (bool, default=False) – Whether to use pretrained weights for the backbone.

Examples

>>> model = IntelliPore(3, 100)
>>> x = torch.randn(4, 3, 32, 32, 32)
>>> model(x).shape  # Outputs
torch.Size([4, 100])
>>> model.backbone(x).shape  # Embeddings
torch.Size([4, 128])

>>> # Works with different grid sizes (adaptive pooling).
>>> x = torch.randn(8, 3, 24, 24, 24)
>>> model(x).shape
torch.Size([8, 100])

>>> # Acts as feature extractor.
>>> model = IntelliPore(in_channels=3, n_outputs=None)
>>> x = torch.randn(4, 3, 32, 32, 32)
>>> model(x).shape
torch.Size([4, 128])

>>> # Pretrained weights for the backbone.
>>> model = IntelliPore(pretrained=True)
forward(x)[source]#

Run the forward pass.

Parameters:

x (tensor of shape (B, C, H, W, D))

Returns:

out – If n_outputs=None return the embeddings of shape (B, 128), else the model outputs of shape (B, n_outputs).

Return type:

tensor

get_pretrained_weights()[source]#

Return the state dict of the pretrained backbone.

Return type:

OrderedDict

class aidsorb.modules.voxels.RetNeXt(in_channels=1, n_outputs=1, *, pretrained=False)[source]#

Bases: Module

Architecture from the [RetNeXt] paper.

Note

pretrained=True is only compatible with in_channels=1, since the pretrained backbone was trained on single-channel images.

Parameters:

  • in_channels (int, default=1)

  • n_outputs (int or None, default=1) – Number of output units. If None, no linear head is added and the model acts as feature extractor.

  • pretrained (bool, default=False) – Whether to use pretrained weights for the backbone.

Examples

>>> model = RetNeXt(3, 100)
>>> x = torch.randn(8, 3, 32, 32, 32)
>>> model(x).shape  # Outputs
torch.Size([8, 100])
>>> model.backbone(x).shape  # Embeddings
torch.Size([8, 128])

>>> # Works with different grid sizes (adaptive pooling).
>>> x = torch.randn(16, 3, 25, 25, 25)
>>> model(x).shape
torch.Size([16, 100])

>>> # Acts as feature extractor.
>>> model = IntelliPore(in_channels=3, n_outputs=None)
>>> x = torch.randn(4, 3, 32, 32, 32)
>>> model(x).shape
torch.Size([4, 128])

>>> # Pretrained weights for the backbone.
>>> model = RetNeXt(pretrained=True)
forward(x)[source]#

Run the forward pass.

Parameters:

x (tensor of shape (B, C, H, W, D))

Returns:

out – If n_outputs=None return the embeddings of shape (B, 128), else the model outputs of shape (B, n_outputs).

Return type:

tensor

get_pretrained_weights()[source]#

Return the state dict of the pretrained backbone.

Return type:

OrderedDict

aidsorb.modules.voxels.conv3d_block(in_channels, out_channels, config_activation=None, **kwargs)[source]#

Return a 3D convolutional block.

The block has the following form:

block = nn.Sequential(
    conv_layer,
    nn.BatchNorm3d(out_channels),
    activation_fn
    )
Parameters:

  • in_channels (int)

  • out_channels (int)

  • config_activation (dict, default=None) –

    Dictionary for configuring activation function. If None, the ReLU activation is used.

    • 'name' activation’s class name str

    • 'hparams' activation’s hyperparameters dict

  • **kwargs – Valid keyword arguments for Conv3d.

Returns:

block

Return type:

torch.nn.Sequential

Examples

>>> inp, out = 4, 8
>>> x = torch.randn(2, inp, 3, 3, 3)  # Shape (B, in_channels, H, W, D).
>>> config_afn = {'name': 'LeakyReLU', 'hparams': {'negative_slope': 0.9}}

>>> # Default activation function (ReLU).
>>> block = conv3d_block(inp, out, kernel_size=3)
>>> block(x).shape
torch.Size([2, 8, 1, 1, 1])
>>> block[2]
ReLU()

>>> # Custom activation function.
>>> block = conv3d_block(inp, out, config_afn, kernel_size=3)
>>> block(x).shape
torch.Size([2, 8, 1, 1, 1])
>>> block[2]
LeakyReLU(negative_slope=0.9)