Merge pull request #207 from legend-exp/units

Support forwarding `units` ak.Array parameters to LGDO attributes

Author: gipert

src/lgdo/types/array.py

@@ -40,20 +40,19 @@ class Array(LGDOCollection):
 
     def __init__(
         self,
-        nda: np.ndarray = None,
+        nda: np.ndarray | ak.Array | None = None,
         shape: tuple[int, ...] = (),
-        dtype: np.dtype = None,
+        dtype: np.dtype | None = None,
         fill_val: float | int | None = None,
         attrs: dict[str, Any] | None = None,
     ) -> None:
         """
         Parameters
         ----------
         nda
-            An :class:`numpy.ndarray` to be used for this object's internal
-            array. Note: the array is used directly, not copied. If not
-            supplied, internal memory is newly allocated based on the shape and
-            dtype arguments.
+            An :class:`numpy.ndarray` or :class:`ak.Array` to be used for this
+            object's internal array. If the Awkward array carries a ``units``
+            parameter, it will forwarded as LGDO attribute.
         shape
             A numpy-format shape specification for shape of the internal
             ndarray. Required if `nda` is ``None``, otherwise unused.
@@ -65,7 +64,21 @@ def __init__(
             the array is allocated with all elements set to the corresponding
             fill value. If `nda` is not ``None``, this parameter is ignored.
         attrs
-            A set of user attributes to be carried along with this LGDO.
+            A set of user attributes to be carried along with this LGDO. These
+            attributes have always precedence over all the others (e.g. those
+            carried by `nda`).
+
+        Warning
+        -------
+        This constructor has partial units support. It supports fishing `units`
+        from Awkward Array parameters but not (yet) from e.g. NumPy+Pint
+        arrays. In any case, the user can always attach units later by
+        modifying the dictionary held by :attr:`attrs`.
+
+        Note
+        ----
+        The array is used directly, not copied. If not supplied, internal
+        memory is newly allocated based on the shape and dtype arguments.
         """
         if nda is None:
             if fill_val is None:
@@ -75,6 +88,26 @@ def __init__(
             else:
                 nda = np.full(shape, fill_val, dtype=dtype)
 
+        elif isinstance(nda, ak.Array):
+            # units: we don't just forward all ak parameters, there might be
+            # some weird thing in there
+            units = ak.parameters(nda).get("units", None)
+            if units is not None:
+                if attrs is None:
+                    attrs = {}
+
+                # give precedence to the user units
+                attrs = {"units": units} | attrs
+
+            if nda.type.content.parameters.get("__array__") == "string":
+                # Variable length strings aren't quite up to snuff yet, so pad the
+                # fixed-width string length in case we want to update the array
+                # TODO: numpy 2.2.5 fixes this; but it required python v3.10 or higher
+                s_len = np.max(ak.num(nda))
+                nda = np.array(nda, dtype=f"<U{s_len * 2}")
+            else:
+                nda = ak.to_numpy(nda)  # this is zero-copy
+
         elif isinstance(nda, Array):
             nda = nda.nda
 

src/lgdo/types/table.py

@@ -64,17 +64,24 @@ def __init__(
             instantiate this table using the supplied mapping of column names
             and array-like objects. Supported input types are: mapping of
             strings to LGDOCollections, :class:`pd.DataFrame` and :class:`ak.Array`.
-            Note 1: no copy is performed, the objects are used directly (unless
-            :class:`ak.Array` is provided).  Note 2: if `size` is not ``None``,
-            all arrays will be resized to match it.  Note 3: if the arrays have
-            different lengths, all will be resized to match the length of the
-            first array.
         attrs
             A set of user attributes to be carried along with this LGDO.
 
         Notes
         -----
-        the :attr:`loc` attribute is initialized to 0.
+        - The :attr:`loc` attribute is initialized to 0.
+        - No copy is performed, the objects are used directly (unless
+          :class:`ak.Array` is provided).
+        - If `size` is not ``None``, all arrays will be resized to match it.
+        - If the arrays have different lengths, all will be resized to match the
+          length of the first array.
+
+        Warning
+        -------
+        This constructor has partial units support. It supports fishing `units`
+        from Awkward Array parameters but not (yet) from e.g. NumPy+Pint
+        arrays. In any case, the user can always attach units later by
+        modifying the dictionary held by :attr:`attrs`.
         """
         if isinstance(col_dict, pd.DataFrame):
             col_dict = {k: Array(v) for k, v in col_dict.items()}
@@ -624,12 +631,13 @@ def view_as(
 def _ak_to_lgdo_or_col_dict(array: ak.Array):
     if isinstance(array.type.content, ak.types.RecordType):
         return {field: _ak_to_lgdo_or_col_dict(array[field]) for field in array.fields}
-    if array.type.content.parameters.get("__array__") == "string":
-        # Variable length strings aren't quite up to snuff yet, so pad the
-        # fixed-width string length in case we want to update the array
-        # TODO: numpy 2.2.5 fixes this; but it required python v3.10 or higher
-        s_len = np.max(ak.num(array))
-        return Array(np.array(array, dtype=f"<U{s_len * 2}"))
-    if isinstance(array.type.content, ak.types.NumpyType):
-        return Array(ak.to_numpy(array))
+
+    # be smart and just use Array when it makes sense
+    if (
+        isinstance(array.type.content, ak.types.NumpyType)
+        or array.type.content.parameters.get("__array__") == "string"
+    ):
+        return Array(array)
+
+    # otherwise fallback to VoV
     return VectorOfVectors(array)

src/lgdo/types/vectorofvectors.py

@@ -111,6 +111,16 @@ def __init__(
             if not isinstance(data, ak.Array):
                 data = ak.Array(data)
 
+            # units: we don't just forward all ak parameters, there might be
+            # some weird thing in there
+            units = ak.parameters(data).get("units", None)
+            if units is not None:
+                if attrs is None:
+                    attrs = {}
+
+                # give precedence to the user units
+                attrs = {"units": units} | attrs
+
             if data.ndim < 2:
                 # treat as a single-row VoV
                 data = ak.Array([data])

tests/types/test_array.py

@@ -28,6 +28,20 @@ def test_init():
     assert array.attrs == attrs | {"datatype": "array<1>{real}"}
 
 
+def test_ak_init():
+    orig = ak.with_parameter([1, 2, 3, 4], "units", "mm")
+    assert ak.parameters(orig)["units"] == "mm"
+
+    array = Array(orig)
+    assert isinstance(array.nda, np.ndarray)
+    assert array.attrs["units"] == "mm"
+    assert (array.nda == orig.to_numpy()).all()
+
+    # check that this is a view
+    array.nda[1] = -1
+    assert orig[1] == -1
+
+
 def test_resize_and_capacity():
     array = Array(nda=np.array([1, 2, 3, 4]))
     assert array.get_capacity() == 4
@@ -103,3 +117,8 @@ def test_pickle():
     assert ex.attrs["attr1"] == 1
     assert ex.attrs["datatype"] == obj.attrs["datatype"]
     assert np.all(ex.nda == np.array([1, 2, 3, 4]))
+
+
+def test_string_array():
+    array = Array(ak.Array(["e", "sticazzi", "non", "ce", "li", "metti?"]))
+    assert array.dtype == "<U16"

tests/types/test_table.py

@@ -63,6 +63,7 @@ def test_ak_array_init():
             "a": [1, 2, 3, 4],
             "b": [[1, 2], [3], [4], [5, 6, 7]],
             "c": {"f1": [[], [5], [3, 7, 6], []], "f2": [5, 6, 7, 8]},
+            "d": ["boh", "hello", "there", "!"],
         }
     )
     tbl = Table(array)
@@ -73,6 +74,28 @@ def test_ak_array_init():
     assert isinstance(tbl.c.f2, lgdo.Array)
 
 
+def test_ak_array_init_attrs():
+    array = ak.Array(
+        {
+            "a": ak.with_parameter([1, 2, 3, 4], "units", "mm"),
+            "b": ak.with_parameter([[1, 2], [3], [4], [5, 6, 7]], "units", "keV"),
+            "c": {
+                "f1": [[], [5], [3, 7, 6], []],
+                "f2": ak.with_parameter([5, 6, 7, 8], "units", "C"),
+            },
+        }
+    )
+    tbl = Table(array)
+    assert isinstance(tbl.a, lgdo.Array)
+    assert tbl.a.attrs["units"] == "mm"
+    assert isinstance(tbl.b, lgdo.VectorOfVectors)
+    assert tbl.b.attrs["units"] == "keV"
+    assert isinstance(tbl.c, Table)
+    assert isinstance(tbl.c.f1, lgdo.VectorOfVectors)
+    assert isinstance(tbl.c.f2, lgdo.Array)
+    assert tbl.c.f2.attrs["units"] == "C"
+
+
 def test_datatype_name():
     tbl = Table()
     assert tbl.datatype_name() == "table"

tests/types/test_vectorofvectors.py

@@ -100,6 +100,12 @@ def test_init(testvov):
     )
 
 
+def test_init_with_units():
+    v = ak.with_parameter([[1], [2, 3]], "units", "mm")
+    vov = VectorOfVectors(v)
+    assert vov.attrs["units"] == "mm"
+
+
 def test_eq(testvov):
     assert testvov.v2d == VectorOfVectors(
         [[1, 2], [3, 4, 5], [2], [4, 8, 9, 7], [5, 3, 1]]