ReorderData

SmoothParticleNets

Description

The ReorderData layer is fairly simple. The layer reorders a given tensor based on a tensor containing the indices for the data in the first tensor. More formally, assume that DATA is a BxNxD tensor containing N D-dimensional data points (e.g., XYZ particle locations) over B batches. Let IDXS be a BxN tensor, where each IDXS[i, :] contains the numbers 0 to N-1 in some arbitrary order. This layer then returns DATA where the second dimension has been rearranged according to IDXS. This is equivalent to

DATA[i, :, :] = DATA[i, IDXS[i, :], :]

in PyTorch syntax, however this layer is specialized for this specific kind of indexing resulting in a faster implementation. This layer is designed as a helper layer for the ParticleCollision layer.

ReorderData is implemented as a subclass of torch.nn.Module. This allows it to be used in the same manner as any other PyTorch layer (e.g., conv2d). Additionally, this layer computes graidents, so it can be used in a backward pass.

Example

Assume locs is a BxNxD tensor containing the locations of N D-dimensional particles across B batches and vel is a same size tensor containing the particles’ velocities.

# ReorderData is most commonly used in conjunction with ParticleCollision.
coll = ParticleCollision(ndim, radius)
# Set reverse=True. ParticleCollision calls ReorderData internally, so we want to undo that reordering when we're done.
reorder = ReorderData(reverse=True)
# PartileCollision reorders locs and vel.
locs, vel, idxs, neighbors = coll(locs, vel)
# Perform desired operations with locs, vel, neighbors...
# When we're done, return locs and vel to their original order using ReorderData.
locs, vel = reorder(idxs, locs, vel)

Documentation

ReorderData provides two functions: a constructor and forward. Forward is called by calling the layer object itself (in the same manner as any standard PyTorch layer).