Aurora int16 safety

abacusnbody/data/_aurora_encodings.py

@@ -12,13 +12,40 @@
 - ``_unpack_ufloat8_35``: 8-bit unsigned float with 3-bit exponent, 5-bit
   mantissa. Used for velocity-dispersion-like fields.
 
+It also provides :func:`_check_int_range`, the overflow guard used when a
+math-able integer field is promoted to a narrower signed dtype on the reader
+side (see the per-module unpackers).
+
 This module is private. Callers should be other modules within
-:mod:`abacusnbody.data` (currently :mod:`output_particle` and :mod:`maplog`).
+:mod:`abacusnbody.data` (currently :mod:`output_particle`, :mod:`maplog`, and
+:mod:`healstruct`).
 """
 
 import numpy as np
 
 
+def _check_int_range(values, dtype, *, field):
+    """Raise :class:`OverflowError` if any element of ``values`` falls outside
+    the range representable by ``dtype``.
+
+    Comparison is done in ``values``' own (wider, unsigned) domain, so it is
+    well-defined and catches the silent wraparound that NEP 50 (NumPy >= 2.0)
+    would otherwise allow when assigning an out-of-range value into a
+    narrower/signed integer array. Used by the Aurora unpackers when promoting
+    a full-width ``uint32`` field (e.g. a multiplicity) to ``int32``.
+    """
+    if values.size == 0:
+        return
+    info = np.iinfo(dtype)
+    vmax = int(values.max())
+    vmin = int(values.min())
+    if vmax > info.max or vmin < info.min:
+        raise OverflowError(
+            f'{field}: values [{vmin}, {vmax}] do not fit '
+            f'{np.dtype(dtype).name} range [{info.min}, {info.max}]'
+        )
+
+
 def _unpack_vect32(packed):
     """
     Decode the vect32 velocity encoding: one uint32 → three signed floats.

abacusnbody/data/healstruct.py

@@ -69,10 +69,12 @@
 
 
 def _default_dtype(name):
+    # `dist_bin` and `count` are math-able (bin arithmetic, coaddition), so they
+    # are promoted to int32 for user safety.
     return {
         'pixel': np.uint64,
-        'dist_bin': np.uint16,
-        'count': np.uint32,
+        'dist_bin': np.int32,
+        'count': np.int32,
         'healstruct': np.uint64,
         'voxel_id': np.uint64,
     }[name]
@@ -107,8 +109,8 @@ def unpack_healstruct(data, out=None, fields=None):
         arrays have shape ``(N,)``. Default dtypes:
 
         - ``pixel``: ``uint64`` (36 bits, needs uint64)
-        - ``dist_bin``: ``uint16``
-        - ``count``: ``uint32`` (19 bits, does not fit in uint16)
+        - ``dist_bin``: ``int32`` (9-bit value, promoted for safe arithmetic)
+        - ``count``: ``int32`` (19-bit value, promoted for safe arithmetic)
         - ``healstruct``: ``uint64`` (the raw packed value)
         - ``voxel_id``: ``uint64`` (45 bits, the coaddition key)
     """

abacusnbody/data/maplog.py

@@ -25,7 +25,12 @@
 
 import numpy as np
 
-from ._aurora_encodings import _unpack_ufloat8_35, _unpack_ufloat8_44, _unpack_vect32
+from ._aurora_encodings import (
+    _check_int_range,
+    _unpack_ufloat8_35,
+    _unpack_ufloat8_44,
+    _unpack_vect32,
+)
 
 __all__ = [
     'unpack_maplog',
@@ -46,7 +51,7 @@
 
 # Position encoding: 21 bits per axis, centered at 2^20, 1/128 kpc/h quantum.
 _POS_OFFSET = 1 << 20  # 1048576
-_POS_SCALE = 128.0  # quanta per Mpc/h 
+_POS_SCALE = 128.0  # quanta per Mpc/h
 _POS_MASK = (1 << 21) - 1  # 0x1FFFFF
 
 
@@ -76,14 +81,17 @@ def _field_shape(name, N):
 def _default_dtype(name, float_dtype):
     if name in ('pos', 'vel', 'density', 'vel_disp', 'vel_rel'):
         return float_dtype
+    # Math-able fields are promoted to int32 for user safety.
+    # Bitflag/enum/identifier fields (`control`, `node_type`, `lc_label`, `pid`)
+    # keep their natural unsigned dtype.
     return {
-        'mult': np.uint32,
+        'mult': np.int32,
         'control': np.uint16,
         'node_type': np.uint8,
-        'timestep': np.uint16,
-        'mult_sec': np.uint32,
+        'timestep': np.int32,
+        'mult_sec': np.int32,
         'pid': np.uint16,
-        'length': np.uint16,
+        'length': np.int32,
         'lc_label': np.uint8,
     }[name]
 
@@ -213,6 +221,7 @@ def unpack_maplog(data, out=None, fields=None, float_dtype=np.float32):
         result['vel'][...] = _unpack_vect32(data[:, 2])
 
     if 'mult' in requested_set:
+        _check_int_range(data[:, 3], result['mult'].dtype, field='mult')
         result['mult'][...] = data[:, 3]
 
     if 'control' in requested_set:
@@ -227,7 +236,9 @@ def unpack_maplog(data, out=None, fields=None, float_dtype=np.float32):
     # Modality-dependent (zeroed where not applicable)
     if 'mult_sec' in requested_set:
         is_merger = node_type == NODE_MERGER
-        result['mult_sec'][...] = flex2 * is_merger
+        mult_sec_vals = flex2 * is_merger
+        _check_int_range(mult_sec_vals, result['mult_sec'].dtype, field='mult_sec')
+        result['mult_sec'][...] = mult_sec_vals
 
     if 'pid' in requested_set:
         is_formation = node_type == NODE_FORMATION

abacusnbody/data/output_particle.py

@@ -55,7 +55,8 @@ def _default_dtype(name, float_dtype):
     if name == 'is_map':
         return np.bool_
     if name == 'mult':
-        return np.uint32
+        # int32, not uint32: this is a math-able count, use signed type for user safety
+        return np.int32
     raise KeyError(name)
 
 
@@ -96,7 +97,7 @@ def unpack_output_particle(data, out=None, fields=None, float_dtype=np.float32):
         - ``density``: ``(N,)`` float
         - ``vel_disp``: ``(N,)`` float
         - ``is_map``: ``(N,)`` ``bool`` — True for MAP rows, False for DM
-        - ``mult``: ``(N,)`` ``uint32`` — multiplicity; 0 for DM
+        - ``mult``: ``(N,)`` ``int32`` — multiplicity; 0 for DM
         - ``rel_vel``: ``(N,)`` float — relative velocity; 0 for DM
     """
     if out is not None and fields is not None:

abacusnbody/data/read_abacus.py

@@ -168,7 +168,9 @@ def read_asdf(fn, load=None, colname=None, dtype=np.float32, verbose=True, **kwa
 
         handler = _HANDLERS[colname]
         handler_kwargs = {k: kwargs[k] for k in ('ppd',) if k in kwargs}
-        cols, nread = handler(data, header, load, dtype, colname=colname, **handler_kwargs)
+        cols, nread = handler(
+            data, header, load, dtype, colname=colname, **handler_kwargs
+        )
 
         table = Table(meta=meta)
         for name, col in cols.items():

pyproject.toml

@@ -17,7 +17,7 @@ authors = [{name = "Daniel Eisenstein"},
            ]
 description = "Python code to read halo catalogs and other Abacus N-body data products"
 dependencies = [
-    'numpy>=1.16',
+    'numpy>=2.0.0',
     'blosc>=1.9.2',
     'astropy>=4.0.0',
     'scipy>=1.5.0',