src.utils.ptransf.transf_octorf_augmentor

Classes

TransfOctoRFAugmentor(**kwargs)

class src.utils.ptransf.transf_octorf_augmentor.TransfOctoRFAugmentor(**kwargs)

Author:: Alberto M. Esmoris Pena

Data augmentor for the TransfOctoRF pipeline. Supports jittering, rotations, scaling on coordinate sub-tensors, and group-based feature disabling on feature sub-tensors.

The augmentor operates on a batch dict with keys:

'features': \(\mathbf{F} \in \mathbb{R}^{B \times K \times n_f}\) feature tensor.
'coordinates': \(\mathbf{X} \in \mathbb{R}^{B \times K \times 3}\) coordinate tensor.
'mask': \(\mathbf{m} \in \{0,1\}^{B \times K}\) boolean mask.

Coordinate augmentations (jitter, rotation, scaling) modify only \(\mathbf{X}\). Feature disable modifies only \(\mathbf{F}\). Transformations are applied sequentially in the order specified.

Jitter

Additive noise sampled independently per element:

\[\mathbf{x}_{bk}' = \mathbf{x}_{bk} + \boldsymbol{\epsilon}_{bk}, \quad \boldsymbol{\epsilon}_{bk} \sim \mathcal{D}\]

where \(\mathcal{D}\) is either \(\mathcal{N}(\mu, \sigma^2 \mathbf{I}_3)\) (normal) or \(\text{Uniform}(a, b)^3\) (uniform), configured via noise_distribution.

Rotation

A random rotation around a fixed unit axis \(\mathbf{k} \in \mathbb{R}^3\) by angle \(\theta \sim \mathcal{D}_\theta\). The rotation matrix is computed via Rodrigues’ formula:

\[\mathbf{R}(\mathbf{k}, \theta) = \mathbf{I}_3 + \sin\theta \, [\mathbf{k}]_\times + (1 - \cos\theta) \, [\mathbf{k}]_\times^2\]

where \([\mathbf{k}]_\times\) is the skew-symmetric cross-product matrix of \(\mathbf{k}\):

\[\begin{split}[\mathbf{k}]_\times = \begin{pmatrix} 0 & -k_z & k_y \\ k_z & 0 & -k_x \\ -k_y & k_x & 0 \end{pmatrix}\end{split}\]

Applied per batch element: \(\mathbf{x}_{bk}' = \mathbf{R}_b \, \mathbf{x}_{bk}\). This transformation is distance-preserving (isometric):

\[\lVert \mathbf{x}_{bi}' - \mathbf{x}_{bj}' \rVert = \lVert \mathbf{x}_{bi} - \mathbf{x}_{bj} \rVert \quad \forall\, i, j\]

Scaling

A random scalar factor \(s \sim \mathcal{D}_s\) applied uniformly to all coordinates within a batch element:

\[\mathbf{x}_{bk}' = s_b \, \mathbf{x}_{bk}\]

Distances are scaled proportionally: \(\lVert \mathbf{x}_{bi}' - \mathbf{x}_{bj}' \rVert = |s_b| \, \lVert \mathbf{x}_{bi} - \mathbf{x}_{bj} \rVert\).

Centroid Drop

Randomly removes a fraction of neighbor centroids from each receptive field, simulating missing data:

\[\begin{split}m_{bk}' = \begin{cases} 0 & \text{if } k \in \text{Drop}_b \\ m_{bk} & \text{otherwise} \end{cases}\end{split}\]

where \(\text{Drop}_b\) is a random subset of \(\lfloor r \cdot K \rfloor\) positions (excluding k=0 when protect_center is True). Dropped positions have their features and coordinates zeroed. Resampled per batch element for maximum diversity.

Feature Disable (Group-based)

Let \(\mathcal{G} = \{G_1, \ldots, G_M\}\) be a set of feature groups, where each \(G_i \subseteq \{0, \ldots, n_f - 1\}\) is a subset of feature indices with associated disable probability \(p_i \in [0, 1]\). At the start of each epoch, a Bernoulli trial determines which groups are disabled:

\[D_i \sim \text{Bernoulli}(p_i), \quad i = 1, \ldots, M\]

For all disabled groups (\(D_i = 1\)), the corresponding feature columns are zeroed for every sample in every batch throughout the epoch:

\[\begin{split}f_{bkj}' = \begin{cases} 0 & \text{if } j \in G_i \text{ and } D_i = 1 \text{ for some } i \\ f_{bkj} & \text{otherwise} \end{cases}\end{split}\]

Groups are evaluated independently. Features not belonging to any group are never zeroed. This acts as a structured form of dropout at the feature-group level, encouraging the model to not rely exclusively on any single group of correlated features.

Variables:

transformations (list) – List of transformation dicts.
feature_disable_groups (list) – List of group dicts.
disabled_groups (set) – Set of currently disabled group indices.

__init__(**kwargs)

Initialize a TransfOctoRFAugmentor.

Parameters:: kwargs – Configuration including transformations (list of transform dicts), feature_disable_groups (list of group dicts), centroid_drop (dict with drop_rate and protect_center), and fnames (list of feature names for resolving feature_names in disable groups).

on_epoch_start(): Resample which feature groups are disabled for this epoch. Called once per epoch by the training loop.

augment(batch_data, **kwargs)

Augment a batch. Applies in order:

Coordinate transformations (jitter, rotation, scale).
Centroid drop (randomly mask out neighbors).
Feature disable (zero disabled group columns).

Parameters:: batch_data – Dict with 'features', 'coordinates', 'mask' numpy arrays.
Returns:: Augmented batch_data (modified in-place).

apply_coord_transform(transf, coords)

Apply a single coordinate transformation.

Parameters:

transf – Transform dict with ‘type’ key.
coords – Coordinate array (B, K, 3).

Returns:

Transformed coordinates.

jitter(transf, coords)

Add noise to coordinates.

Parameters:

transf – Dict with noise_distribution.
coords – (B, K, 3) array.

Returns:

Jittered coordinates.

rotate(transf, coords)

Apply random rotation around a specified axis using Rodrigues’ rotation formula.

\[\mathbf{v}' = \mathbf{v} \cos\theta + (\mathbf{k} \times \mathbf{v}) \sin\theta + \mathbf{k} (\mathbf{k} \cdot \mathbf{v})(1 - \cos\theta)\]

Parameters:

transf – Dict with axis and angle_distribution.
coords – (B, K, 3) array.

Returns:

Rotated coordinates.

scale(transf, coords)

Apply random scaling to coordinates.

Parameters:

transf – Dict with scale_distribution.
coords – (B, K, 3) array.

Returns:

Scaled coordinates.

static rodrigues(axis, angle)

Compute rotation matrix via Rodrigues’ formula.

Parameters:

axis – Unit axis vector (3,).
angle – Rotation angle in radians.

Returns:

3x3 rotation matrix.

static sample_distribution(dist, shape)

Sample from a configured distribution.

Parameters:

dist – Dict with type and distribution parameters.
shape – Output shape.

Returns:

Numpy array of samples.

centroid_drop(coords, features, mask)

Randomly drop centroids from each receptive field by setting their mask to False and zeroing their features and coordinates.

A new random selection is made per batch element, maximizing data diversity across batches.

Parameters:

coords – (B, K, 3) coordinate array.
features – (B, K, n_f) feature array.
mask – (B, K) boolean mask.

Returns:

(coords, features, mask) with drops applied.

apply_feature_disable(features)

Zero out feature columns for disabled groups.

Parameters:: features – (B, K, n_f) array.
Returns:: Features with disabled group columns zeroed.