Поворот координат из структурированного массива NumPy без ошибки matmul

Как повернуть координаты в NumPy со структурированными массивами: извлеките поля через structured_to_unstructured и избегайте ошибки ufunc matmul при повороте.

Вращать координаты частиц прямо из структурированного массива NumPy выглядит заманчиво, особенно когда данные повторяют вывод LAMMPS. Но попытка применить матричное умножение непосредственно к наборам полей обычно заканчивается неприятным сюрпризом: математические ufunc в NumPy не работают по структурированным полям так, как можно ожидать. Ниже — аккуратный способ справиться с задачей, не нарушая расположение данных.

Как воспроизвести проблему

Здесь используется структурированный dtype для x, y, z, а поворот пытаются выполнить через матричное умножение над выбранными полями. Приведённый ниже код повторяет этот шаблон.

import numpy as np

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

cloud = 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=coord_sig,
)

cloud[["x", "y", "z"]] = cloud[["x", "y", "z"]] @ rot_mtx.T

Это вызывает ошибку типов в ufunc matmul.

ufunc 'matmul' did not contain a loop with signature matching types (dtype([('x', '<f8'), ('y', '<f8'), ('z', '<f8')]), dtype('float64')) -> None

В чём на самом деле проблема

NumPy не поддерживает выполнение операций по нескольким полям структурированного массива. Запись с полями удобна для организации разнородных данных в одном контейнере, но для многомерной линейной алгебры она подходит плохо. Поэтому матричное умножение по подмножеству полей и падает: операция ожидает обычный массив float формы (n, 3), а не структурированный dtype с именованными полями.

Есть и практический аспект производительности. Применение matmul к обычному двумерному массиву с типом float обычно самый быстрый путь: вызов уходит прямо в скомпилированные ядра без накладных расходов на обработку полей. В структурированных массивах это компромисс между удобством организации и скоростью вычислений.

И если в реальном наборе данных помимо координат есть дополнительные столбцы — целые числа, булевы значения, даже строки — сохранять структурированную схему оправданно. Просто не стоит запускать численные ядра непосредственно на таком контейнере.

Решение: преобразовать координаты в неструктурированный массив float, повернуть и записать обратно

Прямой и наглядный подход — получить представление координат как неструктурированного массива (n, 3) с типом float, выполнить поворот и скопировать результат обратно в поля. Ниже — код, который делает этот процесс явным.

import numpy as np
from numpy.lib.recfunctions import structured_to_unstructured as to_raw

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

cloud = 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=coord_sig,
)

pts = to_raw(cloud[["x", "y", "z"]], dtype=np.float64, copy=False)
pts = pts @ rot_mtx.T
cloud["x"] = pts[:, 0]
cloud["y"] = pts[:, 1]
cloud["z"] = pts[:, 2]

Так вы сохраняете данные в виде структурированного массива для всего, что хранится рядом с координатами, а сами вычисления запускаете на обычной матрице чисел, где NumPy особенно эффективен. Шаг преобразования прозрачен и легко объясним.

Почему это важно

При больших расчётах с десятками миллионов атомов важно не путать организацию данных и вычисления. Структурированные массивы остаются отличным выбором для разнородных записей, но численные ядра должны работать с однородными массивами. На практике матричное умножение быстрее всего выполняется на самостоятельном массиве float формы (n, 3), и преобразование координатных полей к такому виду — понятный и поддерживаемый способ этого добиться. Это соответствует устройству NumPy: записи — для хранения, ndarrays — для математики.

Итоги

Если при работе со структурированными массивами вы ловите ошибки ufunc 'matmul', не стоит спорить с системой типов. Извлеките блок координат через structured_to_unstructured, выполните преобразование на обычном массиве float и запишите результаты обратно в поля. Вы сохраните модель данных, код останется ясным, а вычисления — эффективными. И если в наборе появятся новые нечисловые поля, этот приём будет масштабироваться без изменений окружающей структуры.