Fast computation of inertia tensors in Python

Question

I need to perform a very fast computation of inertia tensors in Python.

The problem is as follows:

Let xyz be a Numpy Array of shape: (samples, atoms, 3) that describes the 3 cartesian coordinates of each atoms in each samples.

The atoms in each sample must be splited into a set of molecules. I need the inertia tensor of each molecule in every sample.

Each molecule contains n atoms and they are ordered in such way that they are grouped together. That is:

   numpy.array( numpy.split(xyz, atoms/n, axis=1).shape )

returns an array of shape: (number of molecules, samples, atoms per molecule, 3), where the first index runs over the molecules.

Notes:

• It is very similar to the problem solved here, but they compute the inertia tensor of the whole system while I need the inertia tensors of its parts (inertia tensor of each molecule in the system).
• It would be preferable to obtain an array of inertia tensors.
• I think that numpy.einsum would be a good alternative (but I find it hard to implement)
• I can get an array of atomic masses fast and without problems.
• If it helps, samples is of size 100 to 500 approx., and atoms up to about 20 000 (but the operation runs in many chunks that made a total of about 1E5 samples).