reflectorch 1.0.0__py3-none-any.whl

reflectorch/models/encoders/conv_encoder.py

@@ -0,0 +1,211 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from pathlib import Path
+
+import torch
+from torch import nn, load
+
+from reflectorch.models.activations import activation_by_name
+from reflectorch.paths import SAVED_MODELS_DIR
+
+__all__ = [
+    "ConvEncoder",
+    "ConvDecoder",
+    "ConvAutoencoder",
+    "ConvVAE",
+]
+
+logger = logging.getLogger(__name__)
+
+
+class ConvEncoder(nn.Module):
+    """A 1D CNN encoder / embedding network
+
+    Args:
+        in_channels (int, optional): the number of input channels. Defaults to 1.
+        hidden_channels (tuple, optional): the number of intermediate channels of each convolutional layer. Defaults to (32, 64, 128, 256, 512).
+        dim_latent (int, optional): the dimension of the output latent embedding. Defaults to 64.
+        dim_avpool (int, optional): the output size of the adaptive average pooling layer. Defaults to 1.
+        use_batch_norm (bool, optional): whether to use batch normalization. Defaults to True.
+        activation (str, optional): the type of activation function. Defaults to 'relu'.
+    """
+    def __init__(self,
+                 in_channels: int = 1,
+                 hidden_channels: tuple = (32, 64, 128, 256, 512),
+                 dim_latent: int = 64,
+                 dim_avpool: int = 1,
+                 use_batch_norm: bool = True,
+                 activation: str = 'relu',
+                 ):
+        super().__init__()
+
+        modules = []
+
+        activation = activation_by_name(activation)
+
+        for h in hidden_channels:
+            layers = [
+                nn.Conv1d(in_channels, out_channels=h, kernel_size=3, stride=2, padding=1),
+                activation(),
+            ]
+
+            if use_batch_norm:
+                layers.insert(1, nn.BatchNorm1d(h))
+
+            modules.append(nn.Sequential(*layers))
+            in_channels = h
+
+        self.core = nn.Sequential(*modules)
+        self.avpool = nn.AdaptiveAvgPool1d(dim_avpool)
+        self.fc = nn.Linear(hidden_channels[-1] * dim_avpool, dim_latent)
+
+    def forward(self, x):
+        """"""
+        if len(x.shape) < 3:
+            x = x.unsqueeze(1)
+        x = self.core(x)
+        x = self.avpool(x).view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+    def load_weights(self, path: str or Path = None, strict: bool = False):
+        if not path:
+            return
+
+        if isinstance(path, str):
+            if not path.endswith('.pt'):
+                path = path + '.pt'
+            path = SAVED_MODELS_DIR / path
+
+        if not path.is_file():
+            logger.error(f'File {str(path)} is not found.')
+            return
+        try:
+            state_dict = load(path)
+            self.load_state_dict(state_dict, strict=strict)
+        except Exception as err:
+            logger.exception(err)
+
+
+class ConvDecoder(nn.Module):
+    """A 1D CNN decoder
+
+    Args:
+        hidden_dims (tuple, optional): the number of intermediate channels of each convolutional layer. Defaults to (512, 256, 128, 64, 32).
+        latent_dim (int, optional): the dimension of the input latent embedding. Defaults to 64.
+        in_size (int, optional): the initial size for upscaling. Defaults to 8.
+        use_batch_norm (bool, optional): whether to use batch normalization. Defaults to True.
+        activation (str, optional): the type of activation function. Defaults to 'relu'.
+    """
+    def __init__(self, 
+                 hidden_channels: tuple = (512, 256, 128, 64, 32), 
+                 dim_latent: int = 64, 
+                 in_size: int = 8,
+                 use_batch_norm: bool = True,
+                 activation: str = 'relu',
+                 ):
+
+        super().__init__()
+
+        self.in_size = in_size
+        modules = []
+
+        self.decoder_input = nn.Linear(dim_latent, hidden_channels[0] * in_size)
+
+        activation = activation_by_name(activation)
+
+        for i in range(len(hidden_channels) - 1):
+            modules.append(
+                nn.Sequential(
+                    nn.ConvTranspose1d(
+                        hidden_channels[i],
+                        hidden_channels[i + 1],
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        output_padding=1,
+                    ),
+                    nn.BatchNorm1d(hidden_channels[i + 1]) if use_batch_norm else nn.Identity(),
+                    activation(),
+                )
+            )
+
+        self.decoder = nn.Sequential(*modules)
+
+        self.final_layer = nn.Sequential(
+            nn.ConvTranspose1d(hidden_channels[-1],
+                               hidden_channels[-1],
+                               kernel_size=3,
+                               stride=2,
+                               padding=1,
+                               output_padding=1),
+            nn.BatchNorm1d(hidden_channels[-1]) if use_batch_norm else nn.Identity(),
+            activation(),
+            nn.Conv1d(hidden_channels[-1], out_channels=1,
+                      kernel_size=3, padding=1)
+        )
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        x = self.decoder_input(x).view(batch_size, -1, self.in_size)
+        x = self.decoder(x)
+        x = self.final_layer(x).flatten(1)
+        return x
+
+
+class ConvAutoencoder(nn.Module):
+    """A 1D convolutional denoising autoencoder"""
+    def __init__(self,
+                 in_channels: int = 1,
+                 encoder_hidden_channels: tuple = (32, 64, 128, 256, 512),
+                 decoder_hidden_channels: tuple = (512, 256, 128, 64, 32),
+                 dim_latent: int = 64,
+                 dim_avpool: int = 1,
+                 use_batch_norm: bool = True,
+                 activation: str = 'relu',
+                 decoder_in_size: int = 8,
+                 **kwargs
+                 ):
+        super().__init__()
+        self.encoder = ConvEncoder(in_channels, encoder_hidden_channels, dim_latent, dim_avpool, use_batch_norm, activation, **kwargs)
+        self.decoder = ConvDecoder(decoder_hidden_channels, dim_latent, decoder_in_size, use_batch_norm, activation, **kwargs)
+
+    def forward(self, x):
+        return self.decoder(self.encoder(x))
+    
+class ConvVAE(nn.Module):
+    """A 1D convolutional variational autoencoder"""
+    def __init__(self,
+                 in_channels: int = 1,
+                 encoder_hidden_channels: tuple = (32, 64, 128, 256, 512),
+                 decoder_hidden_channels: tuple = (512, 256, 128, 64, 32),
+                 dim_latent: int = 64,
+                 dim_avpool: int = 1,
+                 use_batch_norm: bool = True,
+                 activation: str = 'relu',
+                 decoder_in_size: int = 8,
+                 **kwargs
+                 ):
+        super().__init__()
+        self.encoder = ConvEncoder(in_channels, encoder_hidden_channels, 2*dim_latent, dim_avpool, use_batch_norm, activation, **kwargs)
+        self.decoder = ConvDecoder(decoder_hidden_channels, dim_latent, decoder_in_size, use_batch_norm, activation, **kwargs)
+
+    def forward(self, x):
+        z_mu, z_logvar = self.encoder(x).chunk(2, dim=-1)
+        z = self.reparameterize(z_mu, z_logvar)
+
+        x_r_mu, x_r_logvar = self.decoder(z).chunk(2, dim=-1)
+        x = self.reparameterize(x_r_mu, x_r_logvar)
+
+        return x, (z_mu, z_logvar, x_r_mu, x_r_logvar)
+    
+    @staticmethod
+    def reparameterize(mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std).to(std)
+        return mu + eps * std

reflectorch/models/encoders/conv_res_net.py

@@ -0,0 +1,119 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch import nn
+from torch.nn import functional as F
+from torch.nn import init
+
+__all__ = [
+    'ConvResidualNet1D',
+]
+
+
+class ConvResidualBlock1D(nn.Module):
+    def __init__(
+            self,
+            channels,
+            activation=F.gelu,
+            dropout_probability=0.0,
+            use_batch_norm=False,
+            zero_initialization=True,
+            kernel_size: int = 3,
+            dilation: int = 1,
+            padding: int = 1,
+    ):
+        super().__init__()
+        self.activation = activation
+
+        self.use_batch_norm = use_batch_norm
+
+        if use_batch_norm:
+            self.batch_norm_layers = nn.ModuleList(
+                [nn.BatchNorm1d(channels, eps=1e-3) for _ in range(2)]
+            )
+        self.conv_layers = nn.ModuleList(
+            [nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding, dilation=dilation)
+             for _ in range(2)]
+        )
+        self.dropout = nn.Dropout(p=dropout_probability)
+
+        if zero_initialization:
+            init.uniform_(self.conv_layers[-1].weight, -1e-3, 1e-3)
+            init.uniform_(self.conv_layers[-1].bias, -1e-3, 1e-3)
+
+    def forward(self, inputs):
+        temps = inputs
+        if self.use_batch_norm:
+            temps = self.batch_norm_layers[0](temps)
+        temps = self.activation(temps)
+        temps = self.conv_layers[0](temps)
+        if self.use_batch_norm:
+            temps = self.batch_norm_layers[1](temps)
+
+        temps = self.activation(temps)
+        temps = self.dropout(temps)
+        temps = self.conv_layers[1](temps)
+
+        return inputs + temps
+
+
+class ConvResidualNet1D(nn.Module):
+    def __init__(
+            self,
+            in_channels: int = 1,
+            out_channels: int = 64,
+            hidden_channels: int = 128,
+            num_blocks=5,
+            activation=F.gelu,
+            dropout_probability=0.0,
+            use_batch_norm=True,
+            kernel_size: int = 3,
+            dilation: int = 1,
+            padding: int = 1,
+            avpool: int = 8,
+
+    ):
+        super().__init__()
+
+        self.hidden_channels = hidden_channels
+
+        self.initial_layer = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=hidden_channels,
+            kernel_size=1,
+            padding=0,
+        )
+        self.blocks = nn.ModuleList(
+            [
+                ConvResidualBlock1D(
+                    channels=hidden_channels,
+                    activation=activation,
+                    dropout_probability=dropout_probability,
+                    use_batch_norm=use_batch_norm,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    padding=padding,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+        self.avpool = nn.AdaptiveAvgPool1d(avpool)
+
+        self.final_layer = nn.Linear(
+            hidden_channels * avpool, out_channels
+        )
+
+    def forward(self, x):
+        temps = self.initial_layer(x.unsqueeze(1))
+
+        for block in self.blocks:
+            temps = block(temps)
+
+        temps = self.avpool(temps).view(temps.size(0), -1)
+        outputs = self.final_layer(temps)
+
+        return outputs

reflectorch/models/encoders/fno.py

@@ -0,0 +1,127 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from reflectorch.models.activations import activation_by_name
+
+class SpectralConv1d(nn.Module):
+    def __init__(self, in_channels, out_channels, modes):
+        super().__init__()
+
+        """
+        1D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
+        """
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.modes = modes  #Number of Fourier modes to multiply, at most floor(N/2) + 1
+
+        self.scale = (1 / (in_channels*out_channels))
+        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, modes, dtype=torch.cfloat))
+
+    # Complex multiplication
+    def compl_mul1d(self, input, weights):
+        # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
+        return torch.einsum("bix,iox->box", input, weights)
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        #Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfft(x)
+
+        # Multiply relevant Fourier modes
+        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1,  device=x.device, dtype=torch.cfloat)
+        out_ft[:, :, :self.modes] = self.compl_mul1d(x_ft[:, :, :self.modes], self.weights1)
+
+        #Return to physical space
+        x = torch.fft.irfft(out_ft, n=x.size(-1))
+        return x
+    
+
+class FnoEncoder(nn.Module):
+    """An embedding network based on the Fourier Neural Operator (FNO) architecture
+
+    .. image:: ../documentation/fig_reflectometry_embedding_networks.png
+        :width: 400px
+        :align: center
+
+    Args:
+        ch_in (int): number of input channels
+        dim_embedding (int): dimension of the output embedding
+        modes (int): number of Fourier modes
+        width_fno (int): number of channels of the intermediate representations
+        n_fno_blocks (int): number of FNO blocks
+        activation (str): the activation function
+        fusion_self_attention (bool): if ``True`` a fusion layer is used after the FNO blocks to produce the final embedding
+    """
+    def __init__(
+            self, 
+            ch_in: int = 2, 
+            dim_embedding: int = 128, 
+            modes: int = 32, 
+            width_fno: int = 64, 
+            n_fno_blocks: int = 6, 
+            activation: str = 'gelu',
+            fusion_self_attention: bool = False,
+            ):
+        super().__init__()
+
+
+        self.ch_in = ch_in
+        self.dim_embedding = dim_embedding
+        
+        self.modes = modes
+        self.width_fno = width_fno
+        self.n_fno_blocks = n_fno_blocks
+        self.activation = activation_by_name(activation)()
+        self.fusion_self_attention = fusion_self_attention
+        
+
+        self.fc0 = nn.Linear(ch_in, width_fno) #(r(q), q)
+        self.spectral_convs = nn.ModuleList([SpectralConv1d(in_channels=width_fno, out_channels=width_fno, modes=modes) for _ in range(n_fno_blocks)])
+        self.w_convs = nn.ModuleList([nn.Conv1d(in_channels=width_fno, out_channels=width_fno, kernel_size=1) for _ in range(n_fno_blocks)])
+        self.fc_out = nn.Linear(width_fno, dim_embedding)
+
+        if fusion_self_attention:
+            self.fusion = FusionSelfAttention(width_fno, 2*width_fno)
+        
+    def forward(self, x):
+        """"""
+
+        x = x.permute(0, 2, 1) #(B, D, S) -> (B, S, D)
+        x = self.fc0(x)
+        x = x.permute(0, 2, 1) #(B, S, D) -> (B, D, S) 
+
+        for i in range(self.n_fno_blocks):
+            x1 = self.spectral_convs[i](x)
+            x2 = self.w_convs[i](x)
+            
+            x = x1 + x2
+            x = self.activation(x)
+
+        if self.fusion_self_attention:
+            x = x.permute(0, 2, 1)
+            x = self.fusion(x)  
+        else:
+            x = x.mean(dim=-1)
+
+        x = self.fc_out(x)
+        
+        return x
+    
+class FusionSelfAttention(nn.Module):
+    def __init__(self, 
+                 embed_dim: int = 64, 
+                 hidden_dim: int = 64,
+                 activation=nn.Tanh,
+                 ):
+        super().__init__()
+        self.fuser = nn.Sequential(nn.Linear(embed_dim, hidden_dim), 
+                                   activation(),
+                                   nn.Linear(hidden_dim, 1, bias=False))
+        
+    def forward(self, c):  # (batch_size x seq_len x embed_dim)
+        a = self.fuser(c)
+        alpha = torch.exp(a)
+        alpha = alpha/alpha.sum(dim=1, keepdim=True)
+        return (alpha*c).sum(dim=1)  # (batch_size x embed_dim)

reflectorch/models/encoders/transformers.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+from torch import nn
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(
+            self,
+            dim: int = 64,
+            nhead: int = 8,
+            num_encoder_layers: int = 4,
+            num_decoder_layers: int = 2,
+            dim_feedforward: int = 512,
+            dropout: float = 0.01,
+            activation: str = 'gelu',
+            in_dim: int = 2,
+            out_dim: int = None,
+    ):
+
+        super().__init__()
+
+        self.in_projector = nn.Linear(in_dim, dim)
+
+        self.dim = dim
+
+        self.transformer = nn.Transformer(
+            dim, nhead=nhead,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation=activation
+        )
+
+        if out_dim:
+            self.out_projector = nn.Linear(dim, out_dim)
+        else:
+            self.out_projector = None
+
+    def forward(self, src, tgt, src_key_padding_mask=None, tgt_key_padding_mask=None, **kwargs):
+        src = self.in_projector(src.transpose(1, 2)).transpose(0, 1)
+
+        res = self.transformer(
+            src, tgt, src_key_padding_mask=src_key_padding_mask, tgt_key_padding_mask=tgt_key_padding_mask, **kwargs
+        )
+
+        if self.out_projector:
+            res = self.out_projector(res).squeeze(-1)
+
+        return res.squeeze(0)

reflectorch/models/networks/__init__.py

@@ -0,0 +1,18 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+from reflectorch.models.networks.mlp_networks import (
+    NetworkWithPriorsConvEmb,
+    NetworkWithPriorsFnoEmb,
+)
+from reflectorch.models.networks.residual_net import ResidualMLP
+
+
+__all__ = [
+    "ResidualMLP",
+    "NetworkWithPriorsConvEmb",
+    "NetworkWithPriorsFnoEmb",
+]