Neural Network Methods

Neural network-based dimensionality reduction methods including autoencoders.

Standard Architectures

class driada.dim_reduction.neural.AE(*args, **kw)[source]

Standard Autoencoder for non-linear dimensionality reduction.

Combines an encoder and decoder to learn a compressed representation of high-dimensional data through reconstruction. The model is trained to minimize reconstruction error.

Parameters:

  • orig_dim (int) – Original input dimension (number of features).

  • inter_dim (int) – Intermediate hidden layer dimension for both encoder and decoder.

  • code_dim (int) – Dimension of the latent representation (bottleneck).

  • enc_kwargs (dict) – Additional parameters for the encoder (e.g., dropout rate).

  • dec_kwargs (dict) – Additional parameters for the decoder (e.g., dropout rate).

  • device (torch.device) – Device to run the model on (CPU or CUDA).

encoder

The encoder network.

Type:

Encoder

decoder

The decoder network.

Type:

Decoder

orig_dim

Original data dimension.

Type:

int

inter_dim

Hidden layer dimension.

Type:

int

code_dim

Latent space dimension.

Type:

int

Examples

>>> import torch
>>> device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
>>> ae = AE(orig_dim=100, inter_dim=50, code_dim=10,
...         enc_kwargs={'dropout': 0.2}, dec_kwargs={'dropout': 0.2},
...         device=device)
>>> data = torch.randn(32, 100).to(device)
>>> reconstructed = ae(data)
>>> print(reconstructed.shape)
torch.Size([32, 100])
>>> latent = ae.get_code_embedding(data)
>>> print(latent.shape)  # Note: transposed output
(10, 32)

Notes

The encoder produces unbounded latent codes, while the decoder outputs unbounded reconstructions. This design is suitable for general-purpose dimensionality reduction of unbounded data.

See also

VAE

Variational autoencoder for probabilistic encoding.

Encoder

The encoder component.

Decoder

The decoder component.

__init__(orig_dim, inter_dim, code_dim, enc_kwargs, dec_kwargs, device)[source]

Initialize the autoencoder.

Creates encoder and decoder networks with the specified architecture.

Parameters:

  • orig_dim (int) – Original input dimension.

  • inter_dim (int) – Hidden layer dimension for both networks.

  • code_dim (int) – Latent representation dimension.

  • enc_kwargs (dict) – Encoder parameters (e.g., {‘dropout’: 0.2}).

  • dec_kwargs (dict) – Decoder parameters (e.g., {‘dropout’: 0.2}).

  • device (torch.device) – Device for computations.

forward(features)[source]

Forward pass through the autoencoder.

Encodes input data to latent representation and then decodes it back to reconstruct the original data.

Parameters:

features (torch.Tensor) – Input data of shape (batch_size, orig_dim).

Returns:

Reconstructed data of shape (batch_size, orig_dim).

Return type:

torch.Tensor

Notes

The forward pass performs: input → encoder → latent → decoder → reconstruction. Both latent codes and reconstructions are unbounded.

get_code_embedding(input_)[source]

Extract latent representation from input data.

Parameters:

input (torch.Tensor) – Input data of shape (batch_size, orig_dim).

Returns:

Latent representation of shape (code_dim, batch_size). Note: Output is transposed for compatibility with DRIADA conventions.

Return type:

numpy.ndarray

Notes

This method only runs the encoder portion and returns the latent codes as a numpy array. The transpose operation converts from PyTorch’s (batch, features) to DRIADA’s (features, samples) format.

class driada.dim_reduction.neural.VAE(*args, **kw)[source]

Variational Autoencoder for probabilistic dimensionality reduction.

Implements a VAE that learns a probabilistic mapping to a latent space. Unlike standard autoencoders, VAEs learn a distribution over the latent space, enabling generation of new samples and providing uncertainty estimates.

The encoder outputs parameters of a Gaussian distribution (mean and log variance), and the latent code is sampled from this distribution using the reparameterization trick.

Parameters:

  • orig_dim (int) – Original input dimension (number of features).

  • inter_dim (int) – Intermediate hidden layer dimension for both encoder and decoder.

  • code_dim (int) – Dimension of the latent representation (bottleneck).

  • enc_kwargs (dict, optional) – Additional parameters for the encoder (e.g., dropout rate).

  • dec_kwargs (dict, optional) – Additional parameters for the decoder (e.g., dropout rate).

  • device (torch.device, optional) – Device to run the model on. Defaults to CUDA if available, else CPU.

encoder

The encoder network that outputs mean and log variance.

Type:

VAEEncoder

decoder

The decoder network.

Type:

Decoder

orig_dim

Original data dimension.

Type:

int

inter_dim

Hidden layer dimension.

Type:

int

code_dim

Latent space dimension.

Type:

int

Examples

>>> import torch
>>> vae = VAE(orig_dim=100, inter_dim=50, code_dim=10,
...           enc_kwargs={'dropout': 0.2}, dec_kwargs={'dropout': 0.2})
>>> data = torch.randn(32, 100)
>>> reconstructed, mean, log_var = vae(data)
>>> # Compute VAE loss
>>> recon_loss = torch.nn.functional.mse_loss(reconstructed, data)
>>> kl_loss = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())
>>> vae_loss = recon_loss + kl_loss
>>> # Generate new samples
>>> z = torch.randn(32, 10)
>>> generated = vae.decoder(z)

Notes

The VAE loss consists of two terms: 1. Reconstruction loss (e.g., MSE or BCE) 2. KL divergence between the learned distribution and a standard Gaussian

The encoder internally outputs 2*code_dim features which are split into mean and log variance parameters for the latent Gaussian distribution.

See also

AE

Standard deterministic autoencoder.

VAEEncoder

The probabilistic encoder component.

__init__(orig_dim, inter_dim, code_dim, enc_kwargs=None, dec_kwargs=None, device=None)[source]

Initialize the Variational Autoencoder.

Creates a VAE with encoder outputting distribution parameters (mean and log variance) and decoder for reconstruction. The encoder output dimension is doubled to accommodate both parameters.

Parameters:

  • orig_dim (int) – Original input dimension (number of features).

  • inter_dim (int) – Hidden layer dimension for both encoder and decoder.

  • code_dim (int) – Dimension of the latent representation. The encoder will output 2 * code_dim features (code_dim for mean, code_dim for log variance).

  • enc_kwargs (dict, optional) – Additional encoder parameters (e.g., {‘dropout’: 0.2}). Defaults to empty dict if None.

  • dec_kwargs (dict, optional) – Additional decoder parameters (e.g., {‘dropout’: 0.2}). Defaults to empty dict if None.

  • device (torch.device, optional) – Device for computations. If None, encoder/decoder handle device selection.

Notes

The encoder output dimension is set to 2 * code_dim to enable the VAE to learn both mean and log variance parameters for the latent Gaussian distribution. These parameters are later split in the get_code method.

reparameterization(mu, log_var)[source]

Reparameterization trick for VAE.

Samples from the latent distribution N(mu, sigma^2) in a way that allows backpropagation through the sampling operation by expressing the sample as a deterministic function of the parameters and a separate noise variable.

Parameters:

  • mu (torch.Tensor) – Mean of the latent Gaussian distribution, shape (batch_size, code_dim).

  • log_var (torch.Tensor) – Log variance of the latent Gaussian distribution, shape (batch_size, code_dim).

Returns:

Sampled latent vector, shape (batch_size, code_dim).

Return type:

torch.Tensor

Examples

>>> import torch
>>> # Create a VAE instance to access the method
>>> vae = VAE(orig_dim=100, inter_dim=50, code_dim=20,
...           enc_kwargs={}, dec_kwargs={}, device=torch.device('cpu'))
>>> mu = torch.zeros(32, 10)
>>> log_var = torch.ones(32, 10) * -2  # Small variance
>>> z = vae.reparameterization(mu, log_var)
>>> print(z.shape)
torch.Size([32, 10])

Notes

The reparameterization trick transforms sampling from N(mu, sigma^2) into: z = mu + sigma * epsilon, where epsilon ~ N(0, I)

This allows gradients to flow through mu and log_var during backpropagation while maintaining the stochasticity through the random epsilon.

get_code(features)[source]

Extract latent code from input features using VAE encoding.

Encodes input features through the VAE encoder which outputs concatenated mean and log variance parameters. These are reshaped and separated, then used to sample from the latent distribution via reparameterization.

Parameters:

features (torch.Tensor) – Input data of shape (batch_size, orig_dim).

Returns:

  • code : Sampled latent representation, shape (batch_size, code_dim)

  • mu : Mean of latent distribution, shape (batch_size, code_dim)

  • log_var : Log variance of latent distribution, shape (batch_size, code_dim)

Return type:

tuple of torch.Tensor

Examples

>>> import torch
>>> # Create a VAE with latent dimension 10 (code_dim=20 for mean+logvar)
>>> vae = VAE(orig_dim=100, inter_dim=50, code_dim=20,
...           enc_kwargs={}, dec_kwargs={},
...           device=torch.device('cpu'))
>>> features = torch.randn(32, 100)
>>> code, mu, log_var = vae.get_code(features)
>>> print(code.shape, mu.shape, log_var.shape)
torch.Size([32, 20]) torch.Size([32, 20]) torch.Size([32, 20])

Notes

The encoder outputs a tensor of shape (batch_size, 2 * code_dim) which is reshaped to (batch_size, 2, code_dim) where: - [:, 0, :] contains the mean parameters - [:, 1, :] contains the log variance parameters

forward(features)[source]

Forward pass through the VAE.

Performs a complete forward pass: encoding input to latent distribution parameters, sampling from the distribution, and decoding back to reconstruction space. Returns both reconstruction and distribution parameters needed for VAE loss computation.

Parameters:

features (torch.Tensor) – Input data of shape (batch_size, orig_dim).

Returns:

  • reconstructed : Reconstructed data, shape (batch_size, orig_dim)

  • mu : Mean of latent distribution, shape (batch_size, code_dim)

  • log_var : Log variance of latent distribution, shape (batch_size, code_dim)

Return type:

tuple of torch.Tensor

Examples

>>> import torch
>>> import torch.nn.functional as F
>>> # Create a simple VAE instance for demonstration
>>> vae = VAE(orig_dim=100, inter_dim=50, code_dim=20,
...           enc_kwargs={}, dec_kwargs={},
...           device=torch.device('cpu'))
>>> data = torch.randn(32, 100)
>>> recon, mu, log_var = vae(data)
>>> # Compute VAE loss
>>> recon_loss = F.mse_loss(recon, data)
>>> kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
>>> vae_loss = recon_loss + kl_loss
>>> print(f"Reconstruction shape: {recon.shape}")
Reconstruction shape: torch.Size([32, 100])
>>> print(f"Latent mean shape: {mu.shape}")
Latent mean shape: torch.Size([32, 20])
>>> print(f"Latent log variance shape: {log_var.shape}")
Latent log variance shape: torch.Size([32, 20])

Notes

The mu and log_var are needed to compute the KL divergence loss: KL = -0.5 * sum(1 + log_var - mu^2 - exp(log_var))

The total VAE loss is: L = reconstruction_loss + beta * KL_loss where beta is a hyperparameter controlling the regularization strength.

See also

get_code

For encoding only without reconstruction.

reparameterization

The sampling mechanism.

get_code_embedding(input_, use_mean=True)[source]

Extract latent representation from input data.

Returns either the mean of the latent distribution (deterministic) or a sample from it (stochastic), transposed to match DRIADA conventions.

Parameters:

  • input (torch.Tensor) – Input data of shape (batch_size, orig_dim).

  • use_mean (bool, default=True) – If True, returns the mean of the latent distribution (deterministic). If False, returns a sample from the distribution (stochastic).

Returns:

Latent representation of shape (code_dim, batch_size).

Return type:

numpy.ndarray

Examples

>>> import torch
>>> import numpy as np
>>> from driada.dim_reduction.neural import VAE
>>> # Create VAE
>>> vae = VAE(orig_dim=100, inter_dim=50, code_dim=20,
...           enc_kwargs={}, dec_kwargs={},
...           device=torch.device('cpu'))
>>> data = torch.randn(32, 100)
>>> # Get deterministic embedding (mean)
>>> embedding = vae.get_code_embedding(data, use_mean=True)
>>> embedding2 = vae.get_code_embedding(data, use_mean=True)
>>> print(np.allclose(embedding, embedding2))  # True - deterministic
True
>>> # Get stochastic embedding (sampled)
>>> _ = torch.manual_seed(42)  # For reproducibility in doctest
>>> embedding3 = vae.get_code_embedding(data, use_mean=False)
>>> _ = torch.manual_seed(43)  # Different seed
>>> embedding4 = vae.get_code_embedding(data, use_mean=False)
>>> print(np.allclose(embedding3, embedding4))  # False - different samples
False
>>> print(embedding.shape)  # Note: transposed output
(20, 32)

Notes

  • Output is transposed: (batch, features) → (features, samples)

  • use_mean=True is recommended for visualization, downstream tasks, and when you need consistent embeddings

  • use_mean=False captures the uncertainty in the latent representation

See also

get_code

Returns code, mean, and log variance as tensors.

AE.get_code_embedding

Always deterministic (standard autoencoder).

Building Blocks

class driada.dim_reduction.neural.Encoder(*args, **kw)[source]

Neural network encoder for dimensionality reduction.

Implements a two-layer neural network that encodes high-dimensional data into a lower-dimensional latent representation. Used as the encoder component in autoencoders.

Parameters:

  • orig_dim (int) – Original input dimension (number of features).

  • inter_dim (int) – Intermediate hidden layer dimension.

  • code_dim (int) – Output dimension of the encoded representation (latent space).

  • kwargs (dict) – Additional parameters. Supports key ‘dropout’ (float, optional): dropout rate from 0 to 1. Default is no dropout.

  • device (torch.device, optional) – Device to run the model on. Defaults to CUDA if available, else CPU.

encoder_hidden_layer

First linear transformation layer.

Type:

nn.Linear

encoder_output_layer

Second linear transformation to latent space.

Type:

nn.Linear

dropout

Dropout layer for regularization.

Type:

nn.Dropout

Raises:

ValueError – If dropout rate is not in the range [0, 1).

Examples

>>> import torch
>>> encoder = Encoder(orig_dim=100, inter_dim=50, code_dim=10,
...                   kwargs={'dropout': 0.2})
>>> data = torch.randn(32, 100)  # batch of 32 samples
>>> latent = encoder(data)
>>> print(latent.shape)
torch.Size([32, 10])

Notes

The encoder uses LeakyReLU activation for the hidden layer. The output layer has no activation function, producing unbounded latent codes to maximize the representational capacity of the latent space.

See also

Decoder

The corresponding decoder network.

VAEEncoder

Variational encoder for probabilistic latent representations.

__init__(orig_dim, inter_dim, code_dim, kwargs, device=None)[source]

Initialize the encoder network.

Sets up the two-layer neural network architecture with optional dropout regularization and moves the model to the specified device.

Parameters:

  • orig_dim (int) – Original input dimension (number of features).

  • inter_dim (int) – Intermediate hidden layer dimension.

  • code_dim (int) – Output dimension of the encoded representation.

  • kwargs (dict) – Additional parameters, supports ‘dropout’ key.

  • device (torch.device, optional) – Target device for computations.

Raises:

ValueError – If dropout rate is not in the range [0, 1).

forward(features)[source]

Forward pass through the encoder.

Applies the two-layer neural network transformation to encode input features into a lower-dimensional latent representation.

Parameters:

features (torch.Tensor) – Input tensor of shape (batch_size, orig_dim).

Returns:

Encoded representation of shape (batch_size, code_dim). Values are unbounded to maximize representational capacity.

Return type:

torch.Tensor

Notes

The forward pass applies the following transformations: 1. Linear transformation to hidden dimension 2. Dropout regularization (if enabled) 3. LeakyReLU activation 4. Linear transformation to latent dimension

The output is unbounded to allow full representational capacity in the latent space.

class driada.dim_reduction.neural.VAEEncoder(*args, **kw)[source]

Variational encoder that outputs parameters for latent Gaussian distribution.

Unlike standard autoencoders, VAE encoders output parameters (mean and log variance) for a Gaussian distribution in the latent space, enabling probabilistic sampling and regularization via KL divergence.

Parameters:

  • orig_dim (int) – Original input dimension (number of features).

  • inter_dim (int) – Intermediate hidden layer dimension.

  • code_dim (int) – Latent space dimension. Note: this should be 2*latent_dim if you want latent_dim means and latent_dim log variances. The encoder outputs code_dim values total.

  • kwargs (dict) – Additional parameters. Supports key ‘dropout’ (float, optional): dropout rate from 0 to 1. Default is no dropout.

  • device (torch.device, optional) – Device to run the model on. Defaults to CUDA if available, else CPU.

encoder_hidden_layer

First linear transformation layer.

Type:

nn.Linear

encoder_output_layer

Second linear transformation to latent parameters (outputs code_dim values).

Type:

nn.Linear

dropout

Dropout layer for regularization.

Type:

nn.Dropout

Raises:

ValueError – If dropout rate is not in the range [0, 1).

Examples

>>> import torch
>>> # For 10-dim latent space, need code_dim=20 (10 means + 10 log variances)
>>> vae_encoder = VAEEncoder(orig_dim=100, inter_dim=50, code_dim=20,
...                          kwargs={'dropout': 0.2})
>>> data = torch.randn(32, 100)  # batch of 32 samples
>>> params = vae_encoder(data)
>>> print(params.shape)
torch.Size([32, 20])

Notes

The output layer does not use sigmoid activation (unlike standard AE) because it needs to output unconstrained means and log variances for the Gaussian distribution.

See also

VAE

Complete variational autoencoder that uses this encoder.

Encoder

Standard encoder with bounded outputs.

__init__(orig_dim, inter_dim, code_dim, kwargs, device=None)[source]

Initialize the variational encoder network.

Sets up the two-layer neural network architecture that outputs parameters for a Gaussian distribution in latent space.

Parameters:

  • orig_dim (int) – Original input dimension (number of features).

  • inter_dim (int) – Intermediate hidden layer dimension.

  • code_dim (int) – Total output dimension (should be 2*latent_dim for mean and log variance).

  • kwargs (dict) – Additional parameters, supports ‘dropout’ key.

  • device (torch.device, optional) – Target device for computations.

Raises:

ValueError – If dropout rate is not in the range [0, 1).

forward(features)[source]

Forward pass through the VAE encoder.

Applies the two-layer neural network transformation to encode input features into parameters for a Gaussian distribution.

Parameters:

features (torch.Tensor) – Input tensor of shape (batch_size, orig_dim).

Returns:

Encoded representation of shape (batch_size, code_dim). Contains concatenated parameters for the latent Gaussian: typically first half for means, second half for log variances.

Return type:

torch.Tensor

Notes

Unlike standard encoders, VAE encoders output unconstrained values (no sigmoid activation) since they represent distribution parameters. The output should be reshaped to extract means and log variances.

class driada.dim_reduction.neural.Decoder(*args, **kw)[source]

Neural network decoder for dimensionality reduction.

Implements a two-layer neural network that decodes low-dimensional latent representations back to the original high-dimensional space. Used as the decoder component in autoencoders.

Parameters:

  • code_dim (int) – Input dimension of the latent representation.

  • inter_dim (int) – Intermediate hidden layer dimension.

  • orig_dim (int) – Output dimension (same as original data dimension).

  • kwargs (dict) – Additional parameters. Supports key ‘dropout’ (float, optional): dropout rate from 0 to 1. Default is no dropout.

  • device (torch.device, optional) – Device to run the model on. Defaults to CUDA if available, else CPU.

decoder_hidden_layer

First linear transformation from latent space.

Type:

nn.Linear

decoder_output_layer

Second linear transformation to original space.

Type:

nn.Linear

dropout

Dropout layer for regularization.

Type:

nn.Dropout

Raises:

ValueError – If dropout rate is not in the range [0, 1).

Examples

>>> import torch
>>> decoder = Decoder(code_dim=10, inter_dim=50, orig_dim=100,
...                   kwargs={'dropout': 0.2})
>>> latent = torch.randn(32, 10)  # batch of 32 latent codes
>>> reconstructed = decoder(latent)
>>> print(reconstructed.shape)
torch.Size([32, 100])

Notes

The decoder uses LeakyReLU activation for the hidden layer and no activation function on the output layer, allowing it to output unbounded values for reconstruction.

See also

Encoder

The corresponding encoder network.

AE

Complete autoencoder using this decoder.

__init__(code_dim, inter_dim, orig_dim, kwargs, device=None)[source]

Initialize the decoder network.

Sets up the two-layer neural network architecture with optional dropout regularization and moves the model to the specified device.

Parameters:

  • code_dim (int) – Input dimension of the latent representation.

  • inter_dim (int) – Intermediate hidden layer dimension.

  • orig_dim (int) – Output dimension (same as original data).

  • kwargs (dict) – Additional parameters, supports ‘dropout’ key.

  • device (torch.device, optional) – Target device for computations.

Raises:

ValueError – If dropout rate is not in the range [0, 1).

forward(features)[source]

Forward pass through the decoder.

Applies the two-layer neural network transformation to decode latent representations back to the original data space.

Parameters:

features (torch.Tensor) – Latent representation tensor of shape (batch_size, code_dim).

Returns:

Reconstructed data of shape (batch_size, orig_dim).

Return type:

torch.Tensor

Notes

The forward pass applies the following transformations: 1. Linear transformation to hidden dimension 2. Dropout regularization (if enabled) 3. LeakyReLU activation 4. Linear transformation to original dimension 5. No output activation (unbounded reconstruction)

Flexible Architectures

class driada.dim_reduction.flexible_ae.FlexibleAutoencoderBase(loss_components=None, device=None, logger=None, _latent_dim=None)[source]

Abstract base class for autoencoders with flexible loss composition.

Provides common infrastructure for loss management, device handling, and logging. Subclasses must implement encode(), forward(), and compute_loss() methods with their specific signatures.

Parameters:

  • loss_components (list of dict, optional) – List of loss component configurations. Each dict should have: - “name”: str, name of the loss type - “weight”: float, weight for this loss - Additional parameters specific to each loss type

  • device (torch.device, optional) – Device to run the model on. Defaults to CUDA if available, else CPU.

  • logger (logging.Logger, optional) – Logger instance for tracking training progress.

  • _latent_dim (int | None)

loss_registry

Registry for creating loss components.

Type:

LossRegistry

losses

List of loss components (dynamically modifiable).

Type:

list

device

Device for computations.

Type:

torch.device

logger

Logger instance.

Type:

logging.Logger

Notes

This base class handles: - Loss system initialization and management - Device selection and management - Logger setup - Common utilities for loss computation

Subclasses define: - Encoder/decoder architecture - encode() method signature - forward() method signature - Specific loss computation logic

__init__(loss_components=None, device=None, logger=None, _latent_dim=None)[source]

Initialize the base autoencoder infrastructure.

Parameters:

  • loss_components (list of dict, optional) – Loss component configurations. If None, subclasses should provide defaults.

  • device (torch.device, optional) – Device for computations. If None, auto-selects CUDA if available.

  • logger (logging.Logger, optional) – Logger instance. If None, creates default logger.

  • _latent_dim (int, optional) – Latent dimension, auto-injected into loss configs that need it.

abstract encode(x)[source]

Encode input to latent representation.

Parameters:

x (torch.Tensor) – Input data.

Return type:

Latent representation (type depends on subclass).

abstract forward(x)[source]

Forward pass through the autoencoder.

Parameters:

x (torch.Tensor) – Input data.

Return type:

Output (type depends on subclass).

abstract compute_loss(inputs, indices=None, **extra_args)[source]

Compute total loss from all components.

Parameters:

  • inputs (torch.Tensor) – Input data.

  • indices (torch.Tensor, optional) – Batch indices for data-dependent losses.

  • **extra_args – Additional arguments for specific losses.

Return type:

Tuple[MagicMock, Dict[str, float]]

Returns:

  • total_loss (torch.Tensor) – Weighted sum of all loss components.

  • loss_dict (dict) – Individual loss values for logging.

class driada.dim_reduction.flexible_ae.ModularAutoencoder(input_dim, latent_dim, hidden_dim=100, encoder_config=None, decoder_config=None, loss_components=None, device=None, logger=None)[source]

Flexible autoencoder with modular loss composition.

This autoencoder supports dynamic addition of loss components for various training objectives. It maintains backward compatibility while enabling advanced techniques like disentanglement and regularization.

Parameters:

  • input_dim (int) – Dimension of input data.

  • latent_dim (int) – Dimension of latent representation.

  • hidden_dim (int, default=100) – Dimension of hidden layers.

  • encoder_config (dict, optional) – Configuration for encoder (e.g., dropout rate).

  • decoder_config (dict, optional) – Configuration for decoder.

  • loss_components (list of dict, optional) – List of loss component configurations. Each dict should have: - “name”: str, name of the loss type - “weight”: float, weight for this loss - Additional parameters specific to each loss type If not provided, defaults to [{“name”: “reconstruction”, “weight”: 1.0}].

  • device (torch.device, optional) – Device to run the model on. Defaults to CUDA if available, else CPU.

  • logger (logging.Logger, optional) – Logger instance for tracking training progress.

encoder

Neural network encoder module.

Type:

Encoder

decoder

Neural network decoder module.

Type:

Decoder

loss_registry

Registry for creating loss components.

Type:

LossRegistry

losses

List of loss components (dynamically modifiable).

Type:

nn.ModuleList

input_dim

Input data dimension.

Type:

int

latent_dim

Latent representation dimension.

Type:

int

hidden_dim

Hidden layer dimension.

Type:

int

device

Device for computations.

Type:

torch.device

logger

Logger instance.

Type:

logging.Logger

Examples

>>> # Standard autoencoder with correlation loss
>>> ae = ModularAutoencoder(
...     input_dim=100, latent_dim=10,
...     loss_components=[
...         {"name": "reconstruction", "weight": 1.0},
...         {"name": "correlation", "weight": 0.1}
...     ]
... )

>>> # Autoencoder with sparsity and orthogonality constraints
>>> import numpy as np
>>> data = np.random.randn(100, 1000)  # 100 features, 1000 samples
>>> ae = ModularAutoencoder(
...     input_dim=100, latent_dim=10,
...     loss_components=[
...         {"name": "reconstruction", "weight": 1.0},
...         {"name": "sparse", "weight": 0.1, "sparsity_target": 0.05},
...         {"name": "orthogonality", "weight": 0.05, "external_data": data}
...     ]
... )

Notes

  • Loss components are stored in an nn.ModuleList for proper PyTorch integration

  • Default reconstruction loss is automatically added if no components specified

  • The get_latent_representation method transposes output for DRIADA compatibility

  • Loss components can be dynamically added/removed during training

See also

FlexibleVAE

Variational version with probabilistic encoding.

LossRegistry

Available loss components.

__init__(input_dim, latent_dim, hidden_dim=100, encoder_config=None, decoder_config=None, loss_components=None, device=None, logger=None)[source]

Initialize the modular autoencoder.

Parameters:

  • input_dim (int) – Dimension of input data. Must be positive.

  • latent_dim (int) – Dimension of latent representation. Must be positive.

  • hidden_dim (int, default=100) – Dimension of hidden layers. Must be positive.

  • encoder_config (dict, optional) – Configuration for encoder (e.g., {“dropout”: 0.2}).

  • decoder_config (dict, optional) – Configuration for decoder.

  • loss_components (list of dict, optional) – Loss component configurations. If None, uses default reconstruction loss.

  • device (torch.device, optional) – Device for computations. If None, auto-selects CUDA if available.

  • logger (logging.Logger, optional) – Logger instance. If None, creates default logger.

Raises:

ValueError – If any dimension is not positive. If loss component name is not registered. If loss component dict is malformed.

forward(x)[source]

Forward pass through the autoencoder.

Parameters:

x (torch.Tensor) – Input data, shape (batch_size, input_dim).

Returns:

Reconstructed data, shape (batch_size, input_dim).

Return type:

torch.Tensor

Notes

Applies encoder then decoder. Behavior affected by training/eval mode (dropout). No input validation performed.

encode(x)[source]

Encode input to latent representation.

Parameters:

x (torch.Tensor) – Input data, shape (batch_size, input_dim).

Returns:

Latent code, shape (batch_size, latent_dim).

Return type:

torch.Tensor

Notes

Direct passthrough to encoder network. Applies Linear->ReLU->Dropout->Linear. Behavior affected by training/eval mode.

decode(z)[source]

Decode latent representation to output.

Parameters:

z (torch.Tensor) – Latent code, shape (batch_size, latent_dim).

Returns:

Reconstructed data, shape (batch_size, input_dim). Output constrained to [0, 1] by sigmoid activation.

Return type:

torch.Tensor

Notes

Direct passthrough to decoder network. Final sigmoid activation constrains output to [0, 1] range.

compute_loss(inputs, indices=None, **extra_args)[source]

Compute total loss from all components.

Parameters:

  • inputs (torch.Tensor) – Input data, shape (batch_size, input_dim).

  • indices (torch.Tensor, optional) – Batch indices for data-dependent losses.

  • **extra_args – Additional arguments passed to all loss components.

Return type:

Tuple[MagicMock, Dict[str, float]]

Returns:

  • total_loss (torch.Tensor) – Weighted sum of all loss components. Scalar tensor on model device.

  • loss_dict (dict) – Individual loss values with keys “{ClassName}_{index}” and “{ClassName}_{index}_weighted”. Includes “total_loss”.

Notes

Performs encode->decode forward pass. Uses base class helper for loss aggregation. All extra_args passed to every loss component.

get_latent_representation(x)[source]

Get latent representation for data.

Parameters:

x (torch.Tensor) – Input data, shape (batch_size, input_dim).

Returns:

Latent representation, shape (latent_dim, batch_size). Note: Transposed for DRIADA compatibility.

Return type:

np.ndarray

Notes

Runs in no_grad mode. Returns detached numpy array on CPU.

class driada.dim_reduction.flexible_ae.FlexibleVAE(input_dim, latent_dim, hidden_dim=100, encoder_config=None, decoder_config=None, loss_components=None, device=None, logger=None)[source]

Flexible Variational Autoencoder with modular loss composition.

Supports variational inference with flexible loss components for advanced disentanglement techniques. Uses VAEEncoder that outputs mean and log variance parameters.

Parameters:

  • input_dim (int) – Dimension of input data.

  • latent_dim (int) – Dimension of latent representation.

  • hidden_dim (int, default=100) – Dimension of hidden layers.

  • encoder_config (dict, optional) – Configuration for encoder.

  • decoder_config (dict, optional) – Configuration for decoder.

  • loss_components (list of dict, optional) – List of loss component configurations. Defaults to reconstruction + beta_vae losses (each with weight 1.0 and beta 1.0).

  • device (torch.device, optional) – Device to run the model on.

  • logger (logging.Logger, optional) – Logger instance.

Notes

  • Uses VAEEncoder with output dimension 2*latent_dim for mean and log_var

  • Supports deterministic mode via use_mean parameter in get_latent_representation

  • forward() returns (recon, mu, log_var) unlike parent’s single tensor

  • Default includes standard VAE losses if none specified

Examples

>>> # β-VAE for disentanglement
>>> vae = FlexibleVAE(
...     input_dim=100, latent_dim=10,
...     loss_components=[
...         {"name": "reconstruction", "weight": 1.0},
...         {"name": "beta_vae", "weight": 1.0, "beta": 4.0}
...     ]
... )

>>> # TC-VAE with decomposed ELBO
>>> vae = FlexibleVAE(
...     input_dim=100, latent_dim=10,
...     loss_components=[
...         {"name": "reconstruction", "weight": 1.0},
...         {"name": "tc_vae", "weight": 1.0, "alpha": 1.0, "beta": 5.0, "gamma": 1.0}
...     ]
... )
__init__(input_dim, latent_dim, hidden_dim=100, encoder_config=None, decoder_config=None, loss_components=None, device=None, logger=None)[source]

Initialize FlexibleVAE.

Parameters:

  • input_dim (int) – Dimension of input data. Must be positive.

  • latent_dim (int) – Dimension of latent representation. Must be positive.

  • hidden_dim (int, default=100) – Dimension of hidden layers. Must be positive.

  • encoder_config (dict, optional) – Configuration for VAE encoder.

  • decoder_config (dict, optional) – Configuration for decoder.

  • loss_components (list of dict, optional) – Loss configurations. If None, uses reconstruction + beta_vae.

  • device (torch.device, optional) – Device for computations. Auto-selects CUDA if available.

  • logger (logging.Logger, optional) – Logger instance.

Notes

Uses VAEEncoder with 2*latent_dim output for mean and log variance. Default loss configuration includes both reconstruction and KL divergence.

reparameterize(mu, log_var)[source]

Reparameterization trick for VAE.

Parameters:

  • mu (torch.Tensor) – Mean of latent distribution, shape (batch_size, latent_dim).

  • log_var (torch.Tensor) – Log variance of latent distribution, shape (batch_size, latent_dim).

Returns:

Sampled latent code, shape (batch_size, latent_dim).

Return type:

torch.Tensor

Notes

Always samples stochastically. No numerical stability checks for extreme log_var values. Use get_latent_representation(use_mean=True) for deterministic behavior.

decode(z)[source]

Decode latent representation to output.

Parameters:

z (torch.Tensor) – Latent code, shape (batch_size, latent_dim).

Returns:

Reconstructed data, shape (batch_size, input_dim). Output constrained to [0, 1] by sigmoid activation.

Return type:

torch.Tensor

Notes

Direct passthrough to decoder network. Final sigmoid activation constrains output to [0, 1] range.

encode(x)[source]

Encode input to latent distribution parameters.

Parameters:

x (torch.Tensor) – Input data, shape (batch_size, input_dim).

Return type:

Tuple[MagicMock, MagicMock, MagicMock]

Returns:

  • z (torch.Tensor) – Sampled latent code, shape (batch_size, latent_dim).

  • mu (torch.Tensor) – Mean of latent distribution, shape (batch_size, latent_dim).

  • log_var (torch.Tensor) – Log variance of latent distribution, shape (batch_size, latent_dim).

Notes

Encoder must output exactly 2*latent_dim features. First half is interpreted as mean, second half as log variance. Always samples via reparameterization.

forward(x)[source]

Forward pass through the VAE.

Parameters:

x (torch.Tensor) – Input data, shape (batch_size, input_dim).

Return type:

Tuple[MagicMock, MagicMock, MagicMock]

Returns:

  • recon (torch.Tensor) – Reconstructed data, shape (batch_size, input_dim).

  • mu (torch.Tensor) – Mean of latent distribution, shape (batch_size, latent_dim).

  • log_var (torch.Tensor) – Log variance of latent distribution, shape (batch_size, latent_dim).

Notes

Returns tuple unlike parent’s single tensor. Always uses sampled z for reconstruction, not the mean.

compute_loss(inputs, indices=None, **extra_args)[source]

Compute total loss from all components.

Parameters:

  • inputs (torch.Tensor) – Input data, shape (batch_size, input_dim).

  • indices (torch.Tensor, optional) – Batch indices for data-dependent losses.

  • **extra_args – Additional arguments for specific losses.

Return type:

Tuple[MagicMock, Dict[str, float]]

Returns:

  • total_loss (torch.Tensor) – Weighted sum of all loss components.

  • loss_dict (dict) – Individual loss values for logging.

Notes

Performs single forward pass. Passes mu and log_var to all loss components. Uses base class aggregation helper.

get_latent_representation(x, use_mean=True)[source]

Get latent representation for data.

Parameters:

  • x (torch.Tensor) – Input data, shape (batch_size, input_dim).

  • use_mean (bool, default=True) – If True, return mean of latent distribution (deterministic). If False, return sampled latent code (stochastic).

Returns:

Latent representation, shape (latent_dim, batch_size). Transposed for DRIADA compatibility.

Return type:

np.ndarray

Notes

Default behavior is deterministic (use_mean=True) for reproducible embeddings. Set use_mean=False to capture uncertainty via sampling.

Loss Functions

class driada.dim_reduction.losses.AELoss(weight=1.0, **kwargs)[source]

Base class for all autoencoder loss components.

Each loss component computes a specific objective (e.g., reconstruction, disentanglement, sparsity) and has an associated weight for balancing multiple objectives.

Parameters:

weight (float)

__init__(weight=1.0, **kwargs)[source]

Initialize loss component.

Parameters:

  • weight (float, default=1.0) – Weight for this loss component when combining multiple losses.

  • **kwargs – Additional parameters specific to each loss type.

abstract compute(code, recon, inputs, **kwargs)[source]

Compute the loss value.

Parameters:

  • code (torch.Tensor) – Latent representation, shape (batch_size, code_dim).

  • recon (torch.Tensor) – Reconstructed outputs, shape (batch_size, input_dim).

  • inputs (torch.Tensor) – Original inputs, shape (batch_size, input_dim).

  • **kwargs – Additional tensors/parameters needed by specific losses (e.g., mu and log_var for VAE losses).

Returns:

Scalar loss value.

Return type:

torch.Tensor

property weight: float

Get the loss weight for balancing multiple objectives.

The weight determines the relative importance of this loss component when combined with other losses in a multi-objective optimization. Higher weights increase the influence of this loss on the total loss.

Returns:

Current weight value for this loss component.

Return type:

float

class driada.dim_reduction.losses.ReconstructionLoss(loss_type='mse', weight=1.0)[source]

Standard reconstruction loss for autoencoders (MSE or BCE).

The reconstruction loss measures how well the autoencoder can reconstruct the input data from its latent representation. This is the primary loss component for training autoencoders.

Supports two loss types: - MSE (Mean Squared Error): For continuous data reconstruction - BCE (Binary Cross-Entropy): For binary data or data in [0,1] range

The loss encourages the decoder to accurately reconstruct inputs from the encoded representations, ensuring information preservation.

Parameters:

  • loss_type ({'mse', 'bce'}, default='mse') – Type of reconstruction loss to use.

  • weight (float, default=1.0) – Weight for this loss component in multi-objective optimization.

Examples

>>> # For continuous data
>>> loss = ReconstructionLoss(loss_type='mse', weight=1.0)
>>>
>>> # For binary or probabilistic data
>>> loss = ReconstructionLoss(loss_type='bce', weight=2.0)
__init__(loss_type='mse', weight=1.0)[source]

Initialize reconstruction loss.

Parameters:

  • loss_type (str, default="mse") – Type of reconstruction loss (“mse” or “bce”).

  • weight (float, default=1.0) – Loss weight.

Raises:

ValueError – If loss_type is not ‘mse’ or ‘bce’.

compute(code, recon, inputs, **kwargs)[source]

Compute reconstruction loss between input and reconstruction.

Parameters:

  • code (torch.Tensor) – Latent representation (unused for reconstruction loss).

  • recon (torch.Tensor) – Reconstructed data, shape (batch_size, n_features).

  • inputs (torch.Tensor) – Original input data, shape (batch_size, n_features).

  • **kwargs – Additional arguments (unused).

Returns:

Scalar loss value. MSE for continuous data or BCE for binary data, depending on loss_type specified in __init__.

Return type:

torch.Tensor

class driada.dim_reduction.losses.BetaVAELoss(beta=4.0, weight=1.0)[source]

β-VAE loss for disentanglement via increased KL penalty.

Implements the β-VAE objective which modifies the standard VAE loss by scaling the KL divergence term with a factor β > 1. This encourages the model to learn disentangled representations where each latent dimension captures at most one factor of variation.

The full β-VAE loss (when combined with reconstruction) is: L = Reconstruction + β * KL(q(z|x)||p(z))

Notes

  • β = 1 recovers the standard VAE

  • β > 1 encourages disentanglement

  • Too high β can hurt reconstruction quality

  • Typical values: β ∈ [4, 10] for disentanglement tasks

References

Higgins, I., et al. (2017). β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR 2017.

Parameters:
__init__(beta=4.0, weight=1.0)[source]

Initialize β-VAE loss.

Parameters:

  • beta (float, default=4.0) – Beta parameter controlling KL penalty strength. Must be positive for valid disentanglement.

  • weight (float, default=1.0) – Loss weight.

Raises:

ValueError – If beta <= 0.

compute(code, recon, inputs, mu=None, log_var=None, **kwargs)[source]

Compute β-weighted KL divergence loss.

Calculates the KL divergence between the learned posterior q(z|x) and the standard normal prior p(z) = N(0,I), scaled by the β parameter. Higher β values encourage greater disentanglement at the cost of reconstruction quality.

Parameters:

  • code (torch.Tensor) – Latent representation (unused, included for interface consistency).

  • recon (torch.Tensor) – Reconstructed outputs (unused).

  • inputs (torch.Tensor) – Original inputs (unused).

  • mu (torch.Tensor) – Mean of the approximate posterior, shape (batch_size, latent_dim).

  • log_var (torch.Tensor) – Log variance of the approximate posterior, shape (batch_size, latent_dim).

  • **kwargs – Additional arguments (unused).

Returns:

β-weighted KL divergence loss, averaged over the batch.

Return type:

torch.Tensor

Raises:

ValueError – If mu or log_var are not provided.

Notes

The KL divergence for a Gaussian posterior is: KL(q(z|x)||p(z)) = -0.5 * Σ(1 + log(σ²) - μ² - σ²) where the sum is over latent dimensions.

References

Higgins, I., et al. (2017). β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR 2017.

class driada.dim_reduction.losses.CorrelationLoss(weight=1.0)[source]

Correlation loss to encourage decorrelated latent features.

This loss minimizes correlations between different dimensions of the latent code, encouraging the autoencoder to learn a disentangled representation where each latent dimension captures independent factors of variation.

The loss is computed as the sum of squared off-diagonal elements of the correlation matrix of latent codes. A fully decorrelated representation would have a diagonal correlation matrix (identity matrix after normalization).

This regularization is particularly useful for:

  • Learning interpretable latent representations

  • Preventing redundancy in latent dimensions

  • Improving generalization by reducing overfitting

Mathematical formulation: L_corr = (1/P) * sum_{i≠j} abs(corr(z_i, z_j)) where z_i is the i-th dimension of the latent code across the batch, and P is the number of off-diagonal pairs.

Parameters:

weight (float, default=1.0) – Weight for this loss component. Higher values enforce stronger decorrelation but may harm reconstruction quality.

Notes

The correlation is computed across the batch dimension, so larger batch sizes provide more accurate correlation estimates. Requires batch_size >= 2.

__init__(weight=1.0)[source]

Initialize correlation loss.

Parameters:

weight (float, default=1.0) – Loss weight.

compute(code, recon, inputs, **kwargs)[source]

Compute average pairwise correlation amplitude in latent code.

Encourages decorrelated latent features by penalizing correlations between different dimensions of the latent representation.

Parameters:

  • code (torch.Tensor) – Latent representation, shape (batch_size, code_dim).

  • recon (torch.Tensor) – Reconstructed data (unused).

  • inputs (torch.Tensor) – Original input data (unused).

  • **kwargs – Additional arguments (unused).

Returns:

Average absolute correlation across all pairs of latent dimensions. Returns 0 if code_dim = 1 or batch_size < 2.

Return type:

torch.Tensor

Notes

Requires batch_size >= 2 for correlation computation. Returns 0 for single-sample batches.

class driada.dim_reduction.losses.OrthogonalityLoss(external_data=None, weight=1.0)[source]

Orthogonality loss to minimize correlation with external data (MI proxy).

FUTURE: Replace correlation-based approach with proper mutual information estimation (e.g., using GCMI or KSG estimators from information module). Current implementation uses correlation as a crude proxy for MI.

Parameters:
__init__(external_data=None, weight=1.0)[source]

Initialize orthogonality loss.

Parameters:

  • external_data (np.ndarray, optional) – External data to minimize correlation with, shape (n_features, n_samples).

  • weight (float, default=1.0) – Loss weight.

compute(code, recon, inputs, indices=None, **kwargs)[source]

Compute correlation between latent code and external data.

Used as a proxy for mutual information minimization. Encourages the latent representation to be orthogonal (uncorrelated) with provided external variables.

FUTURE: Implement proper MI computation using driada.information module estimators instead of correlation-based approximation.

Parameters:

  • code (torch.Tensor) – Latent representation, shape (batch_size, code_dim).

  • recon (torch.Tensor) – Reconstructed data (unused).

  • inputs (torch.Tensor) – Original input data (unused).

  • indices (torch.Tensor, optional) – Batch indices to select corresponding external data columns. If None, assumes first batch_size columns of external_data.

  • **kwargs – Additional arguments (unused).

Returns:

Average absolute correlation between latent code and external data. Returns 0 if no external data provided.

Return type:

torch.Tensor

Notes

External data should have shape (n_features, n_samples) where n_samples should be >= batch_size if indices is None.

class driada.dim_reduction.losses.SparsityLoss(sparsity_target=0.05, weight=1.0)[source]

Sparsity loss to encourage sparse latent representations.

Implements a sparsity constraint on the latent activations by penalizing the KL divergence between the average activation and a target sparsity level. This encourages the model to use only a subset of latent dimensions for each input, leading to more interpretable representations.

The loss is based on the KL divergence between: - ρ̂: average activation of each latent unit (across the batch) - ρ: target sparsity level (e.g., 0.05 for 5% average activation)

References

Ng, A. (2011). Sparse autoencoder. CS294A Lecture notes, Stanford University.

Parameters:
__init__(sparsity_target=0.05, weight=1.0)[source]

Initialize sparsity loss.

Parameters:

  • sparsity_target (float, default=0.05) – Target average activation level.

  • weight (float, default=1.0) – Loss weight.

compute(code, recon, inputs, **kwargs)[source]

Compute KL divergence between actual and target sparsity.

Calculates the sparsity penalty based on the average activation of each latent unit compared to the target sparsity level.

Parameters:

  • code (torch.Tensor) – Latent representation after activation (e.g., sigmoid), shape (batch_size, latent_dim). Values should be in [0, 1].

  • recon (torch.Tensor) – Reconstructed outputs (unused).

  • inputs (torch.Tensor) – Original inputs (unused).

  • **kwargs – Additional arguments (unused).

Returns:

KL divergence between actual and target sparsity, summed over all latent dimensions.

Return type:

torch.Tensor

Notes

For each latent unit j: KL(ρ||ρ̂_j) = ρ log(ρ/ρ̂_j) + (1-ρ) log((1-ρ)/(1-ρ̂_j)) where ρ̂_j is the average activation of unit j over the batch.

Uses clamping for numerical stability to avoid log(0).

class driada.dim_reduction.losses.LossRegistry[source]

Registry for dynamically managing loss components.

Provides a centralized system for registering and creating loss functions for autoencoders. Supports dynamic registration of custom losses and maintains a catalog of available loss types.

The registry pattern allows for easy extension with new loss types without modifying existing code. All registered losses must inherit from AELoss.

losses

Mapping from loss names to their class types.

Type:

Dict[str, Type[AELoss]]

Examples

>>> registry = LossRegistry()
>>> # Create a standard reconstruction loss
>>> recon_loss = registry.create('reconstruction', loss_type='mse')
>>>
>>> # Register a custom loss
>>> class MyCustomLoss(AELoss):
...     def compute(self, code, recon, inputs, **kwargs):
...         return torch.tensor(0.0)
>>> registry.register('custom', MyCustomLoss)
>>> custom_loss = registry.create('custom', weight=2.0)
__init__()[source]

Initialize the registry with default loss types.

Creates an empty loss registry and populates it with standard loss functions commonly used in autoencoders: - ‘reconstruction’: Standard reconstruction loss (MSE/BCE) - ‘correlation’: Decorrelation loss for latent features - ‘kl’: KL divergence for variational autoencoders - ‘activity’: L1/L2 activity regularization - ‘jacobian’: Jacobian regularization for contractive autoencoders

The registry can be extended with custom losses after initialization.

register(name, loss_class)[source]

Register a new loss type.

Parameters:

  • name (str) – Name identifier for the loss.

  • loss_class (Type[AELoss]) – Loss class (must inherit from AELoss).

Raises:

ValueError – If loss_class does not inherit from AELoss.

create(name, **kwargs)[source]

Create a loss instance by name.

Parameters:

  • name (str) – Registered name of the loss.

  • **kwargs – Parameters to pass to the loss constructor.

Returns:

Instantiated loss component.

Return type:

AELoss

Raises:

ValueError – If the loss name is not registered.

Data Handling

class driada.dim_reduction.neural.NeuroDataset(*args, **kw)[source]

PyTorch Dataset wrapper for neural activity data.

Wraps neural data matrices for use with PyTorch DataLoader, enabling efficient batching and sampling during neural network training.

Parameters:

  • data (ndarray) – Input data matrix of shape (n_features, n_samples). Will be transposed internally to (n_samples, n_features) for PyTorch compatibility.

  • transform (callable, optional) – Optional transform to be applied on each sample.

data

Transposed data matrix of shape (n_samples, n_features).

Type:

ndarray

transform

Transform function to apply to samples.

Type:

callable or None

Examples

>>> import numpy as np
>>> from torch.utils.data import DataLoader
>>> # Create dataset with 100 neurons, 1000 time points
>>> data = np.random.randn(100, 1000)
>>> dataset = NeuroDataset(data)
>>> # Create DataLoader for batching
>>> loader = DataLoader(dataset, batch_size=32, shuffle=True)
>>> for batch_data, batch_idx in loader:
...     print(batch_data.shape)  # (32, 100) - batch_size x n_features
...     break
torch.Size([32, 100])

Notes

The dataset returns tuples of (sample, index) where the index can be used for tracking which samples were selected during training.

__init__(data, transform=None)[source]

Initialize the neural dataset.

Transposes the input data from (n_features, n_samples) to (n_samples, n_features) for PyTorch compatibility.

Parameters:

  • data (ndarray) – Input data matrix of shape (n_features, n_samples).

  • transform (callable, optional) – Optional transform function to apply to each sample.

__len__()[source]

Return the number of samples in the dataset.

Returns:

Number of samples (n_samples).

Return type:

int

__getitem__(idx)[source]

Retrieve a sample and its index from the dataset.

Parameters:

idx (int or torch.Tensor) – Index of the sample to retrieve. If tensor, will be converted to Python list/int for numpy indexing.

Returns:

  • samplendarray

    Data sample of shape (n_features,), optionally transformed.

  • idxint

    The index of the retrieved sample.

Return type:

tuple

Notes

Returns both the sample and its index to allow tracking of which samples were used during training. This can be useful for debugging or sample weighting schemes.