mlspm.image#

mlspm.image.models#

class mlspm.image.models.ASDAFMNet(n_out: int = 3, activation: Module = LeakyReLU(negative_slope=0.1), padding_mode: str = 'replicate', last_relu: bool | Tuple[bool, ...] = [False, True, True], pretrained_weights: Literal['asdafm-light', 'asdafm-heavy'] | None = None)[source]#

Bases: Module

The model used in the paper “Automated structure discovery in atomic force microscopy”: https://doi.org/10.1126/sciadv.aay6913.

Two different sets of pretrained weights are available:

  • 'asdafm-light': trained on a set of molecules containing the elements H, C, N, O, and F.

  • 'asdafm-heavy': trained on a set of molecules additionally containing the elements Si, P, S, Cl, and Br.

Parameters:

  • n_out – number of output branches.

  • activation – Activation function used after each layer.

  • padding_mode – Type of padding in each convolution layer. 'zeros', 'reflect', 'replicate' or 'circular'.

  • last_relu – Whether to use a ReLU layer after the last convolution in each output branch. Either provide a single value that applies to all output branches or a list of values corresponding to each output branch.

  • pretrained_weights – Name of pretrained weights. If specified, load pretrained weights. Otherwise, weights are initialized randomly. If loading the weights, the other arguments should be left in their default values, so that the model hyperparameters match the training.

forward(x)[source]#

Define 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.

class mlspm.image.models.AttentionUNet(z_in: int, n_in: int = 1, n_out: int = 3, in_channels: int = 1, merge_block_channels: List[int] = [8], merge_block_depth: int = 2, conv3d_block_channels: List[int] = [8, 16, 32], conv3d_block_depth: int = 2, conv3d_dropouts: List[float] = [0.0, 0.0, 0.0], conv2d_block_channels: List[int] = [128], conv2d_block_depth: int = 3, conv2d_dropouts: List[float] = [0.1], attention_channels: List[int] = [32, 32, 32], upscale2d_block_channels: List[int] = [16, 16, 16], upscale2d_block_depth: int = 1, upscale2d_block_channels2: List[int] = [16, 16, 16], upscale2d_block_depth2: int = 2, split_conv_block_channels: List[int] = [16], split_conv_block_depth: int = 3, res_connections: bool = True, out_convs_channels: int | List[int] = 1, out_relus: bool | List[bool] = True, pool_type: Literal['avg', 'max'] = 'avg', pool_z_strides: List[int] = [2, 1, 2], padding_mode: Literal['zeros', 'reflect', 'replicate', 'circular'] = 'zeros', activation: Literal['relu', 'lrelu', 'elu'] | Module = 'lrelu', attention_activation: Literal['sigmoid', 'softmax'] = 'softmax', device: str = 'cuda')[source]#

Bases: Module

Pytorch 3D-to-2D U-net model with attention.

3D conv -> concatenate -> 3D conv/pool/dropout -> 2D conv/dropout -> 2D upsampling/conv with skip connections and attention. For multiple inputs, the inputs are first processed through separate 3D conv blocks before merging by concatenating along channel axis.

Parameters:

  • z_in – Size of input array in the z-dimension.

  • n_in – Number of input 3D images.

  • n_out – Number of output 2D maps.

  • in_channels – Number of channels in input array.

  • merge_block_channels – Number of channels in input merging 3D conv blocks.

  • merge_block_depth – Number of layers in each merge conv block.

  • conv3d_block_channels – Number channels in 3D conv blocks.

  • conv3d_block_depth – Number of layers in each 3D conv block.

  • conv3d_dropouts – Dropout rates after each conv3d block.

  • conv2d_block_channels – Number channels in 2D conv blocks.

  • conv2d_block_depth – Number of layers in each 2D conv block.

  • conv2d_dropouts – Dropout rates after each conv2d block.

  • attention_channels – Number of channels in conv layer within each attention block.

  • upscale2d_block_channels – Number of channels in each 2D conv block after upscale before skip connection.

  • upscale2d_block_depth – Number of layers in each 2D conv block after upscale before skip connection.

  • upscale2d_block_channels2 – Number of channels in each 2D conv block after skip connection.

  • upscale2d_block_depth2 – Number of layers in each 2D conv block after skip connection.

  • split_conv_block_channels – Number of channels in 2d conv blocks after splitting outputs.

  • split_conv_block_depth – Number of layers in each 2d conv block after splitting outputs.

  • res_connections – Whether to use residual connections in conv blocks.

  • out_convs_channels – Number of channels in splitted outputs.

  • out_relus – Whether to apply relu activation to the output 2D maps.

  • pool_type – Type of pooling to use.

  • pool_z_strides – Stride of pool layers in z direction.

  • padding_mode – Type of padding in each convolution layer.

  • activation – Activation to use after every layer except last one.

  • attention_activation – Type of activation to use for attention map.

  • device – Device to load model onto.

forward(x: List[Tensor]) → Tuple[List[Tensor], List[Tensor]][source]#

Do forward computation.

Parameters:

x – Input AFM images of shape (batch, channels, x, y, z).

Returns:

Tuple (outputs, attention_maps), where

  • outputs - Output arrays of shape (batch, out_convs_channels, x, y) or (batch, x, y) if out_convs_channels == 1.

  • attention_maps - Attention maps at each stage of skip-connections.

class mlspm.image.models.EDAFMNet(device: str = 'cuda', pretrained_weights: Literal['base', 'single-channel', 'CO-Cl', 'Xe-Cl', 'constant-noise', 'uniform-noise', 'no-gradient', 'matched-tips'] | None = None)[source]#

Bases: AttentionUNet

ED-AFM Attention U-net.

This is the model used in the ED-AFM paper for task of predicting electrostatics from AFM images. It is a subclass of the AttentionUNet class with specific hyperparameters.

The following pretrained weights are available:

  • 'base': The base model used for all predictions in the main ED-AFM paper and used for comparison in the various test in the supplementary information of the paper.

  • 'single-channel': Model trained on only a single CO-tip AFM input.

  • 'CO-Cl': Model trained on alternative tip combination of CO and Cl.

  • 'Xe-Cl': Model trained on alternative tip combination of Xe and Cl.

  • 'constant-noise': Model trained using constant noise amplitude instead of normally distributed amplitude.

  • 'uniform-noise': Model trained using uniform random noise amplitude instead of normally distributed amplitude.

  • 'no-gradient': Model trained without background-gradient augmentation.

  • 'matched-tips': Model trained on data with matched tip distance between CO and Xe, instead of independently randomized distances.

Parameters:

  • device – Device to load model onto.

  • trained_weights – If not None, load the specified pretrained weights to the model.