qlat field indexing and the NumPy bridge

Source: qlat/qlat/field_base.pyx, qlat/qlat/field_types.pyx

Note: Update this document when updating the source files.

This page covers the three indexing conventions that trip people up when moving between qlat fields and NumPy arrays:

  1. the flat local site index used by f[i],

  2. the global coordinate (_xg) API,

  3. the coordinate-ordered array layout for analysis code.

It also covers initialisation traps and the view/assign idioms.

1. Flat local site index

np.asarray(f) (equivalently f[:]) returns a writeable, zero-copy view with shape

(local_volume, multiplicity, *elem_shape)

C-contiguous. Axis 0 is the flat local site index: the first coordinate varies fastest, matching geo.coordinate_from_index. The site axis covers the local volume on each MPI rank, i.e. geo.node_site, not the global total_site. For local dimensions node_site = [Lx, Ly, Lz, Lt], local coordinate (x, y, z, t) maps to

i = x + Lx * (y + Ly * (z + Lz * t))

With a single rank, node_site == total_site; with several ranks each rank’s buffer has shape (prod(node_site), multiplicity, *elem_shape).

import numpy as np
import qlat as q

q.begin_with_mpi()
geo = q.Geometry(q.Coordinate([4, 4, 4, 4]))
f = q.Field(q.ElemTypeColorMatrix, geo, 1)

i = geo.index_from_coordinate(q.Coordinate([1, 2, 3, 0]))
xl = geo.coordinate_from_index(i)          # Coordinate([1, 2, 3, 0])

a = np.asarray(f)                          # (local_volume, 1, 3, 3)
a[i, 0]                                    # the element at local (1, 2, 3, 0)
f[i, 0] = q.ColorMatrix()                  # write through

f[i] and f[i, m] return views, so f[i][0, 0] = 1.0 writes to the field. A Coordinate is not an accepted index — passing one raises TypeError with a pointer to the APIs below.

2. Global coordinates (_xg)

The _xg methods take global coordinates (identical on every MPI rank) and are collective:

method

returns

f.get_elem_xg(xg_arr, m)

SelectedPoints with shape (N, 1, *elem_shape)

f.set_elem_xg(xg_arr, m, val)

writes element m at the coordinates

f.get_elems_xg(xg_arr)

SelectedPoints with all multiplicities

f.set_elems_xg(xg_arr, val)

writes all multiplicities

 
sp = f.get_elem_xg([[0, 0, 0, 0], [1, 2, 3, 0]], 0)
sp.shape        # (2, 1, 3, 3)
val = sp[:]     # (2, 1, 3, 3) NumPy array

Collective: xg_arr must be identical on all MPI processes. Passing different values on different ranks produces undefined results.

On a single rank with no decomposition, a global coordinate equals a local one; with multiple ranks, geo.coordinate_l_from_g(xg) converts.

3. Coordinate-ordered arrays

The buffer is flat over sites, so a.reshape(Lx, Ly, Lz, Lt, m, ...) does not give an [x, y, z, t]-indexed array: the flat site axis already has the first coordinate varying fastest, which is the opposite of NumPy’s C-contiguous convention for a (Lx, Ly, Lz, Lt, ...) array. The reshape must therefore use the reversed lattice dimensions, and then the four space-time axes are reversed back. Use the local dimensions from geo.node_site, not the global total_site:

loc = geo.node_site.to_list()           # [Lx, Ly, Lz, Lt] on this rank
rev = loc[::-1]                         # [Lt, Lz, Ly, Lx]
a = np.asarray(f)                       # (V, m, *elem_shape), V = prod(loc)
elem = a.shape[2:]                      # e.g. (3, 3)

# flat buffer -> arr[x, y, z, t, m, *elem]
arr = a.reshape(*rev, a.shape[1], *elem).transpose(3, 2, 1, 0, 4, 5, 6)

# sanity check: arr index is the geometric local coordinate
xl = geo.coordinate_from_index(0)
assert arr[tuple(xl.to_list()) + (0, 0, 0)] == f[geo.index_from_coordinate(xl), 0, 0, 0]

The result is a view with axes (Lx, Ly, Lz, Lt, multiplicity, *elem_shape) — so arr[x, y, z, t] is the element at that local coordinate. To cross-check against the library, compare with f.get_elem_xg([geo.coordinate_g_from_l(xl).to_list()], m), which takes global coordinates.

Writing back

The inverse operation needs no transpose after the reshape:

cfg = arr.copy()                        # edit at will, indexed [x, y, z, t]
idx = [d // 2 for d in loc]
cfg[tuple(idx) + (0, 0, 0)] = q.ColorMatrix()

a[:] = cfg.transpose(3, 2, 1, 0, 4, 5, 6).reshape(a.shape)

The transpose appears only on the coordinate-ordered side; it is its own inverse for the four space-time axes.

For a pure copy (rather than a view), use .copy() on arr.

4. Initialisation

qlat needs its global geometry node before a Geometry can be built from a total_site, so:

import qlat as q

q.begin_with_mpi()                       # or q.begin(id_node, size_node)
geo = q.Geometry(q.Coordinate([4, 4, 4, 4]))

Common traps:

  • q.Geometry(total_site) before begin* raises RuntimeError (previously it crashed with SIGFPE).

  • mpirun -n 1 python3 script.py does not help: mpirun only sets the environment, the process still has to call MPI_Init, and qlat additionally needs begin_comm to populate geon.

  • Geometry(id_node, size_node, node_site) works without begin* because it does not touch the global geon — useful for testing.

  • Calling begin_with_mpi() twice pushes a second communicator; call it once.

5. Summary of conventions

API

index meaning

layout

f[i], f[i, m], np.asarray(f)

flat local site index

(local_volume, multiplicity, *elem_shape), first coordinate fastest

f.get_elem_xg(xg, m)

global coordinate

returns SelectedPoints, shape (N, 1, *elem_shape)

coordinate-ordered arr (section 3)

[x, y, z, t]

(Lx, Ly, Lz, Lt, multiplicity, *elem_shape)

SelectedField / SelectedPoints buffer

selection index

(n_elems, multiplicity, *elem_shape) — axis 0 is a selection index, not a lattice coordinate