Finding Centre Of Mass Of Tensor (tensorflow)

Is there an efficient way to find the centre of mass of a tensor? I'm working with N stacked volumes (Nx64x64x64) and would like to obtain an Nx3 tensor with the x,y,z position of

Solution 1:

Following the formula, you should just need to multiply each coordinate by the corresponding mass, sum everything and divide by the total mass:

import tensorflow as tf

# Input volumes
volumes = tf.placeholder(tf.float32, [None, 64, 64, 64])
# Make array of coordinates (each row contains three coordinates)
ii, jj, kk = tf.meshgrid(tf.range(64), tf.range(64), tf.range(64), indexing='ij')
coords = tf.stack([tf.reshape(ii, (-1,)), tf.reshape(jj, (-1,)), tf.reshape(kk, (-1,))], axis=-1)
coords = tf.cast(coords, tf.float32)
# Rearrange input into one vector per volume
volumes_flat = tf.reshape(volumes, [-1, 64 * 64 * 64, 1])
# Compute total mass for each volume
total_mass = tf.reduce_sum(volumes_flat, axis=1)
# Compute centre of mass
centre_of_mass = tf.reduce_sum(volumes_flat * coords, axis=1) / total_mass

Solution 2:

This is a perfect use case for numpy.apply_over_axes:

my_tensor = np.apply_over_axes( np.mean, my_tensor, (1,2,3))

This will return an array of shape (N,1,1,1) with the average along each axis.