Merge pull request #57 from legend-exp/dev

Support setting a fill value when "exploding" `VectorOfVectors` into NumPy arrays in `.view_as("np")`
legend-exp · Jan 15, 2024 · 0cf3eaa · 0cf3eaa
2 parents 579cea6 + 24c2dce
commit 0cf3eaa
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/src/lgdo/types/vectorofvectors.py b/src/lgdo/types/vectorofvectors.py
@@ -434,7 +434,7 @@ def __repr__(self) -> str:
     def to_aoesa(
         self,
         max_len: int = None,
-        fill_val: int | float = np.nan,
+        fill_val: bool | int | float = np.nan,
         preserve_dtype: bool = False,
     ) -> aoesa.ArrayOfEqualSizedArrays:
         """Convert to :class:`ArrayOfEqualSizedArrays`.
@@ -477,7 +477,11 @@ def to_aoesa(
         return aoesa.ArrayOfEqualSizedArrays(nda=nda, attrs=self.getattrs())
 
     def view_as(
-        self, library: str, with_units: bool = False, preserve_dtype: bool = False
+        self,
+        library: str,
+        with_units: bool = False,
+        fill_val: bool | int | float = np.nan,
+        preserve_dtype: bool = False,
     ) -> pd.DataFrame | np.NDArray | ak.Array:
         r"""View the vector data as a third-party format data structure.
 
@@ -488,21 +492,24 @@ def view_as(
         - ``pd``: returns a :class:`pandas.Series` (supported through the
           ``awkward-pandas`` package)
         - ``np``: returns a :class:`numpy.ndarray`, padded with zeros to make
-          it rectangular.
+          it rectangular. This implies memory re-allocation.
         - ``ak``: returns an :class:`ak.Array`. ``self.cumulative_length`` is
           currently re-allocated for technical reasons.
 
         Notes
         -----
         Awkward array views partially involve memory re-allocation (the
-        `cumulative_length`\ s).
+        `cumulative_length`\ s), while NumPy "exploded" views clearly imply a
+        full copy.
 
         Parameters
         ----------
         library
             format of the returned data view.
         with_units
             forward physical units to the output data.
+        fill_val
+            forwarded to :meth:`.to_aoesa`, if `library` is ``np``.
         preserve_dtype
             forwarded to :meth:`.to_aoesa`, if `library` is ``np``.
 
@@ -535,7 +542,7 @@ def view_as(
 
         if library == "np":
             if preserve_dtype:
-                return self.to_aoesa(fill_val=0, preserve_dtype=True).view_as(
+                return self.to_aoesa(fill_val=fill_val, preserve_dtype=True).view_as(
                     "np", with_units=with_units
                 )
             else:

diff --git a/tests/types/test_vectorofvectors.py b/tests/types/test_vectorofvectors.py
@@ -324,7 +324,7 @@ def test_view(lgdo_vov):
     assert isinstance(np_arr, np.ndarray)
     assert np.issubdtype(np_arr.dtype, np.floating)
 
-    np_arr = lgdo_vov.view_as("np", with_units=False, preserve_dtype=True)
+    np_arr = lgdo_vov.view_as("np", with_units=False, fill_val=0, preserve_dtype=True)
     assert np.issubdtype(np_arr.dtype, np.integer)
 
     np_arr = lgdo_vov.view_as("pd", with_units=False)