pythonarraysnumpy

How to apply rotations to structured numpy arrays?


I'm using structured arrays to store atoms data produced by LAMMPS (I'm using a structured array that follows its format). I need to rotate the positions:

import numpy as np

transform = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
dtype = np.dtype([("x", np.float64), ("y", np.float64), ("z", np.float64)])

atoms = np.array(
    [
        (0.0, 0.0, 0.0),
        (1.0, 0.0, 0.0),
        (0.0, 1.0, 0.0),
        (1.0, 1.0, 1.0),
    ],
    dtype=dtype,
)

atoms[["x", "y", "z"]] = atoms[["x", "y", "z"]] @ transform.T

But this produces:

Traceback (most recent call last):
  File "c:\Users\acgc99\Desktop\rotation.py", line 16, in <module>
    atoms[["x", "y", "z"]] = atoms[["x", "y", "z"]] @ transform.T
                             ~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~
numpy._core._exceptions._UFuncNoLoopError: ufunc 'matmul' did not contain a loop with signature matching types (dtype([('x', '<f8'), ('y', '<f8'), ('z', '<f8')]), dtype('float64')) -> None

I can convert to unstructured arrays, but I guess that doing that change multiple times is not efficient when working with tens of millions of atoms.


Solution

  • Thanks for all answers, but I decided to use `from numpy.lib.recfunctions import structured_to_unstructured as str2unstr` since it might be a more direct and clear way of getting the same result.

    pos = str2unstr(atoms [["x", "y", "z"]], dtype=np.float64, copy=False)
    pos = transform.apply(pos)  # atoms[["x", "y", "z"]] @ transform.T
    atoms ["x"] = pos[:, 0]
    atoms ["y"] = pos[:, 1]
    atoms ["z"] = pos[:, 2]