legend-pydataobj 1.11.7__py3-none-any.whl → 1.12.0a2__py3-none-any.whl

lgdo/lh5/tools.py

@@ -1,16 +1,10 @@
 from __future__ import annotations
 
 import fnmatch
-import glob
 import logging
-import os
 from copy import copy
-from warnings import warn
 
 import h5py
-import numpy as np
-import pandas as pd
-from numpy.typing import NDArray
 
 from . import utils
 from .store import LH5Store
@@ -223,108 +217,3 @@ def show(
             break
 
         key = k_new
-
-
-def load_nda(
-    f_list: str | list[str],
-    par_list: list[str],
-    lh5_group: str = "",
-    idx_list: list[NDArray | list | tuple] | None = None,
-) -> dict[str, NDArray]:
-    r"""Build a dictionary of :class:`numpy.ndarray`\ s from LH5 data.
-
-    Given a list of files, a list of LH5 table parameters, and an optional
-    group path, return a NumPy array with all values for each parameter.
-
-    Parameters
-    ----------
-    f_list
-        A list of files. Can contain wildcards.
-    par_list
-        A list of parameters to read from each file.
-    lh5_group
-        group path within which to find the specified parameters.
-    idx_list
-        for fancy-indexed reads. Must be one index array for each file in
-        `f_list`.
-
-    Returns
-    -------
-    par_data
-        A dictionary of the parameter data keyed by the elements of `par_list`.
-        Each entry contains the data for the specified parameter concatenated
-        over all files in `f_list`.
-    """
-    warn(
-        "load_nda() is deprecated. "
-        "Please replace it with LH5Store.read(...).view_as('np'), "
-        "or just read_as(..., 'np'). "
-        "load_nda() will be removed in a future release.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-
-    if isinstance(f_list, str):
-        f_list = [f_list]
-        if idx_list is not None:
-            idx_list = [idx_list]
-    if idx_list is not None and len(f_list) != len(idx_list):
-        msg = f"f_list length ({len(f_list)}) != idx_list length ({len(idx_list)})!"
-        raise ValueError(msg)
-
-    # Expand wildcards
-    f_list = [f for f_wc in f_list for f in sorted(glob.glob(os.path.expandvars(f_wc)))]
-
-    sto = LH5Store()
-    par_data = {par: [] for par in par_list}
-    for ii, ff in enumerate(f_list):
-        f = sto.gimme_file(ff, "r")
-        for par in par_list:
-            if f"{lh5_group}/{par}" not in f:
-                msg = f"'{lh5_group}/{par}' not in file {ff}"
-                raise RuntimeError(msg)
-
-            if idx_list is None:
-                data, _ = sto.read(f"{lh5_group}/{par}", f)
-            else:
-                data, _ = sto.read(f"{lh5_group}/{par}", f, idx=idx_list[ii])
-            if not data:
-                continue
-            par_data[par].append(data.nda)
-    return {par: np.concatenate(par_data[par]) for par in par_list}
-
-
-def load_dfs(
-    f_list: str | list[str],
-    par_list: list[str],
-    lh5_group: str = "",
-    idx_list: list[NDArray | list | tuple] | None = None,
-) -> pd.DataFrame:
-    """Build a :class:`pandas.DataFrame` from LH5 data.
-
-    Given a list of files (can use wildcards), a list of LH5 columns, and
-    optionally the group path, return a :class:`pandas.DataFrame` with all
-    values for each parameter.
-
-    See Also
-    --------
-    :func:`load_nda`
-
-    Returns
-    -------
-    dataframe
-        contains columns for each parameter in `par_list`, and rows containing
-        all data for the associated parameters concatenated over all files in
-        `f_list`.
-    """
-    warn(
-        "load_dfs() is deprecated. "
-        "Please replace it with LH5Store.read(...).view_as('pd'), "
-        "or just read_as(..., 'pd'). "
-        "load_dfs() will be removed in a future release.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-    return pd.DataFrame(
-        load_nda(f_list, par_list, lh5_group=lh5_group, idx_list=idx_list)
-    )

lgdo/types/array.py

@@ -6,7 +6,7 @@ corresponding utilities.
 from __future__ import annotations
 
 import logging
-from collections.abc import Iterator
+from collections.abc import Collection, Iterator
 from typing import Any
 
 import awkward as ak
@@ -17,12 +17,12 @@ import pint_pandas  # noqa: F401
 
 from .. import utils
 from ..units import default_units_registry as u
-from .lgdo import LGDO
+from .lgdo import LGDOCollection
 
 log = logging.getLogger(__name__)
 
 
-class Array(LGDO):
+class Array(LGDOCollection):
     r"""Holds an :class:`numpy.ndarray` and attributes.
 
     :class:`Array` (and the other various array types) holds an `nda` instead
@@ -78,11 +78,7 @@ class Array(LGDO):
         elif isinstance(nda, Array):
             nda = nda.nda
 
-        elif not isinstance(nda, np.ndarray):
-            nda = np.array(nda)
-
         self.nda = nda
-        self.dtype = self.nda.dtype
 
         super().__init__(attrs)
 
@@ -96,18 +92,91 @@ class Array(LGDO):
         return dt + "<" + nd + ">{" + et + "}"
 
     def __len__(self) -> int:
-        return len(self.nda)
-
-    def resize(self, new_size: int) -> None:
-        new_shape = (new_size,) + self.nda.shape[1:]
-        return self.nda.resize(new_shape, refcheck=True)
+        return self._size
+
+    @property
+    def nda(self):
+        return self._nda[: self._size, ...] if self._nda.shape != () else self._nda
+
+    @nda.setter
+    def nda(self, value):
+        self._nda = value if isinstance(value, np.ndarray) else np.array(value)
+        self._size = len(self._nda) if self._nda.shape != () else 0
+
+    @property
+    def dtype(self):
+        return self._nda.dtype
+
+    @property
+    def shape(self):
+        return (len(self),) + self._nda.shape[1:]
+
+    def reserve_capacity(self, capacity: int) -> None:
+        "Set size (number of rows) of internal memory buffer"
+        if capacity < len(self):
+            msg = "Cannot reduce capacity below Array length"
+            raise ValueError(msg)
+        self._nda.resize((capacity,) + self._nda.shape[1:], refcheck=False)
+
+    def get_capacity(self) -> int:
+        "Get capacity (i.e. max size before memory must be re-allocated)"
+        return len(self._nda)
+
+    def trim_capacity(self) -> None:
+        "Set capacity to be minimum needed to support Array size"
+        self.reserve_capacity(np.prod(self.shape))
+
+    def resize(self, new_size: int | Collection[int], trim=False) -> None:
+        """Set size of Array in rows. Only change capacity if it must be
+        increased to accommodate new rows; in this case double capacity.
+        If trim is True, capacity will be set to match size. If new_size
+        is an int, do not change size of inner dimensions.
+
+        If new_size is a collection, internal memory will be re-allocated, so
+        this should be done only rarely!"""
+
+        if isinstance(new_size, Collection):
+            self._size = new_size[0]
+            self._nda.resize(new_size)
+        else:
+            self._size = new_size
+
+            if trim and new_size != self.get_capacity:
+                self.reserve_capacity(new_size)
+
+            # If capacity is not big enough, set to next power of 2 big enough
+            if new_size > self.get_capacity():
+                self.reserve_capacity(int(2 ** (np.ceil(np.log2(new_size)))))
 
     def append(self, value: np.ndarray) -> None:
-        self.resize(len(self) + 1)
-        self.nda[-1] = value
+        "Append value to end of array (with copy)"
+        self.insert(len(self), value)
 
     def insert(self, i: int, value: int | float) -> None:
-        self.nda = np.insert(self.nda, i, value)
+        "Insert value into row i (with copy)"
+        if i > len(self):
+            msg = f"index {i} is out of bounds for array with size {len(self)}"
+            raise IndexError(msg)
+
+        value = np.array(value)
+        if value.shape == self.shape[1:]:
+            self.resize(len(self) + 1)
+            self[i + 1 :] = self[i:-1]
+            self[i] = value
+        elif value.shape[1:] == self.shape[1:]:
+            self.resize(len(self) + len(value))
+            self[i + len(value) :] = self[i : -len(value)]
+            self[i : i + len(value)] = value
+        else:
+            msg = f"Could not insert value with shape {value.shape} into Array with shape {self.shape}"
+            raise ValueError(msg)
+
+    def replace(self, i: int, value: int | float) -> None:
+        "Replace value at row i"
+        if i >= len(self):
+            msg = f"index {i} is out of bounds for array with size {len(self)}"
+            raise IndexError(msg)
+        self[i] = value
 
     def __getitem__(self, key):
         return self.nda[key]

lgdo/types/encoded.py

@@ -11,12 +11,12 @@ from numpy.typing import NDArray
 
 from .. import utils
 from .array import Array
-from .lgdo import LGDO
+from .lgdo import LGDOCollection
 from .scalar import Scalar
 from .vectorofvectors import VectorOfVectors
 
 
-class VectorOfEncodedVectors(LGDO):
+class VectorOfEncodedVectors(LGDOCollection):
     """An array of variable-length encoded arrays.
 
     Used to represent an encoded :class:`.VectorOfVectors`. In addition to an
@@ -92,6 +92,17 @@ class VectorOfEncodedVectors(LGDO):
 
         return False
 
+    def reserve_capacity(self, *capacity: int) -> None:
+        self.encoded_data.reserve_capacity(*capacity)
+        self.decoded_size.reserve_capacity(capacity[0])
+
+    def get_capacity(self) -> tuple:
+        return (self.decoded_size.get_capacity, *self.encoded_data.get_capacity())
+
+    def trim_capacity(self) -> None:
+        self.encoded_data.trim_capacity()
+        self.decoded_size.trim_capacity()
+
     def resize(self, new_size: int) -> None:
         """Resize vector along the first axis.
 
@@ -102,21 +113,6 @@ class VectorOfEncodedVectors(LGDO):
         self.encoded_data.resize(new_size)
         self.decoded_size.resize(new_size)
 
-    def append(self, value: tuple[NDArray, int]) -> None:
-        """Append a 1D encoded vector at the end.
-
-        Parameters
-        ----------
-        value
-            a tuple holding the encoded array and its decoded size.
-
-        See Also
-        --------
-        .VectorOfVectors.append
-        """
-        self.encoded_data.append(value[0])
-        self.decoded_size.append(value[1])
-
     def insert(self, i: int, value: tuple[NDArray, int]) -> None:
         """Insert an encoded vector at index `i`.
 
@@ -282,7 +278,7 @@ class VectorOfEncodedVectors(LGDO):
         raise ValueError(msg)
 
 
-class ArrayOfEncodedEqualSizedArrays(LGDO):
+class ArrayOfEncodedEqualSizedArrays(LGDOCollection):
     """An array of encoded arrays with equal decoded size.
 
     Used to represent an encoded :class:`.ArrayOfEqualSizedArrays`. In addition
@@ -349,14 +345,23 @@ class ArrayOfEncodedEqualSizedArrays(LGDO):
 
         return False
 
-    def resize(self, new_size: int) -> None:
+    def reserve_capacity(self, *capacity: int) -> None:
+        self.encoded_data.reserve_capacity(capacity)
+
+    def get_capacity(self) -> tuple:
+        return self.encoded_data.get_capacity()
+
+    def trim_capacity(self) -> None:
+        self.encoded_data.trim_capacity()
+
+    def resize(self, new_size: int, trim: bool = False) -> None:
         """Resize array along the first axis.
 
         See Also
         --------
         .VectorOfVectors.resize
         """
-        self.encoded_data.resize(new_size)
+        self.encoded_data.resize(new_size, trim)
 
     def append(self, value: NDArray) -> None:
         """Append a 1D encoded array at the end.

lgdo/types/histogram.py

@@ -424,7 +424,7 @@ class Histogram(Struct):
             dict.__setitem__(self, name, obj)
         else:
             msg = "histogram fields cannot be mutated "
-            raise TypeError(msg)
+            raise AttributeError(msg)
 
     def __getattr__(self, name: str) -> None:
         # do not allow for new attributes on this

lgdo/types/lgdo.py

@@ -92,3 +92,53 @@ class LGDO(ABC):
 
     def __repr__(self) -> str:
         return self.__class__.__name__ + f"(attrs={self.attrs!r})"
+
+
+class LGDOCollection(LGDO):
+    """Abstract base class representing a LEGEND Collection Object (LGDO).
+    This defines the interface for classes used as table columns.
+    """
+
+    @abstractmethod
+    def __init__(self, attrs: dict[str, Any] | None = None) -> None:
+        super().__init__(attrs)
+
+    @abstractmethod
+    def __len__(self) -> int:
+        """Provides ``__len__`` for this array-like class."""
+
+    @abstractmethod
+    def reserve_capacity(self, capacity: int) -> None:
+        """Reserve capacity (in rows) for later use. Internal memory buffers
+        will have enough entries to store this many rows.
+        """
+
+    @abstractmethod
+    def get_capacity(self) -> int:
+        "get reserved capacity of internal memory buffers in rows"
+
+    @abstractmethod
+    def trim_capacity(self) -> None:
+        """set capacity to only what is required to store current contents
+        of LGDOCollection
+        """
+
+    @abstractmethod
+    def resize(self, new_size: int, trim: bool = False) -> None:
+        """Return this LGDO's datatype attribute string."""
+
+    def append(self, val) -> None:
+        "append val to end of LGDOCollection"
+        self.insert(len(self), val)
+
+    @abstractmethod
+    def insert(self, i: int, val) -> None:
+        "insert val into LGDOCollection at position i"
+
+    @abstractmethod
+    def replace(self, i: int, val) -> None:
+        "replace item at position i with val in LGDOCollection"
+
+    def clear(self, trim: bool = False) -> None:
+        "set size of LGDOCollection to zero"
+        self.resize(0, trim=trim)

lgdo/types/table.py

@@ -19,7 +19,7 @@ from pandas.io.formats import format as fmt
 
 from .array import Array
 from .arrayofequalsizedarrays import ArrayOfEqualSizedArrays
-from .lgdo import LGDO
+from .lgdo import LGDO, LGDOCollection
 from .scalar import Scalar
 from .struct import Struct
 from .vectorofvectors import VectorOfVectors
@@ -27,13 +27,9 @@ from .vectorofvectors import VectorOfVectors
 log = logging.getLogger(__name__)
 
 
-class Table(Struct):
+class Table(Struct, LGDOCollection):
     """A special struct of arrays or subtable columns of equal length.
 
-    Holds onto an internal read/write location ``loc`` that is useful in
-    managing table I/O using functions like :meth:`push_row`, :meth:`is_full`,
-    and :meth:`clear`.
-
     Note
     ----
     If you write to a table and don't fill it up to its total size, be sure to
@@ -49,7 +45,7 @@ class Table(Struct):
 
     def __init__(
         self,
-        col_dict: Mapping[str, LGDO] | pd.DataFrame | ak.Array | None = None,
+        col_dict: Mapping[str, LGDOCollection] | pd.DataFrame | ak.Array | None = None,
         size: int | None = None,
         attrs: Mapping[str, Any] | None = None,
     ) -> None:
@@ -65,7 +61,7 @@ class Table(Struct):
         col_dict
             instantiate this table using the supplied mapping of column names
             and array-like objects. Supported input types are: mapping of
-            strings to LGDOs, :class:`pd.DataFrame` and :class:`ak.Array`.
+            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
@@ -85,7 +81,8 @@ class Table(Struct):
             col_dict = _ak_to_lgdo_or_col_dict(col_dict)
 
         # call Struct constructor
-        super().__init__(obj_dict=col_dict, attrs=attrs)
+        Struct.__init__(self, obj_dict=col_dict)
+        LGDOCollection.__init__(self, attrs=attrs)
 
         # if col_dict is not empty, set size according to it
         # if size is also supplied, resize all fields to match it
@@ -93,13 +90,10 @@ class Table(Struct):
         if col_dict is not None and len(col_dict) > 0:
             self.resize(new_size=size, do_warn=(size is None))
 
-        # if no col_dict, just set the size (default to 1024)
+        # if no col_dict, just set the size
         else:
             self.size = size if size is not None else None
 
-        # always start at loc=0
-        self.loc = 0
-
     def datatype_name(self) -> str:
         return "table"
 
@@ -107,7 +101,31 @@ class Table(Struct):
         """Provides ``__len__`` for this array-like class."""
         return self.size
 
-    def resize(self, new_size: int | None = None, do_warn: bool = False) -> None:
+    def reserve_capacity(self, capacity: int | list) -> None:
+        "Set size (number of rows) of internal memory buffer"
+        if isinstance(capacity, int):
+            for obj in self.values():
+                obj.reserve_capacity(capacity)
+        else:
+            if len(capacity) != len(self.keys()):
+                msg = "List of capacities must have same length as number of keys"
+                raise ValueError(msg)
+
+            for obj, cap in zip(self.values(), capacity):
+                obj.reserve_capacity(cap)
+
+    def get_capacity(self) -> int:
+        "Get list of capacities for each key"
+        return [v.get_capacity() for v in self.values()]
+
+    def trim_capacity(self) -> int:
+        "Set capacity to be minimum needed to support Array size"
+        for v in self.values():
+            v.trim_capacity()
+
+    def resize(
+        self, new_size: int | None = None, do_warn: bool = False, trim: bool = False
+    ) -> None:
         # if new_size = None, use the size from the first field
         for field, obj in self.items():
             if new_size is None:
@@ -119,21 +137,20 @@ class Table(Struct):
                         f"with size {len(obj)} != {new_size}"
                     )
                 if isinstance(obj, Table):
-                    obj.resize(new_size)
+                    obj.resize(new_size, trim)
                 else:
-                    obj.resize(new_size)
+                    obj.resize(new_size, trim)
         self.size = new_size
 
-    def push_row(self) -> None:
-        self.loc += 1
-
-    def is_full(self) -> bool:
-        return self.loc >= self.size
-
-    def clear(self) -> None:
-        self.loc = 0
+    def insert(self, i: int, vals: dict) -> None:
+        "Insert vals into table at row i. Vals is a mapping from table key to val"
+        for k, ar in self.items():
+            ar.insert(i, vals[k])
+        self.size += 1
 
-    def add_field(self, name: str, obj: LGDO, use_obj_size: bool = False) -> None:
+    def add_field(
+        self, name: str, obj: LGDOCollection, use_obj_size: bool = False
+    ) -> None:
         """Add a field (column) to the table.
 
         Use the name "field" here to match the terminology used in
@@ -170,7 +187,9 @@ class Table(Struct):
             new_size = len(obj) if use_obj_size else self.size
             self.resize(new_size=new_size)
 
-    def add_column(self, name: str, obj: LGDO, use_obj_size: bool = False) -> None:
+    def add_column(
+        self, name: str, obj: LGDOCollection, use_obj_size: bool = False
+    ) -> None:
         """Alias for :meth:`.add_field` using table terminology 'column'."""
         self.add_field(name, obj, use_obj_size=use_obj_size)
 
@@ -201,8 +220,10 @@ class Table(Struct):
             set to ``False`` to turn off warnings associated with mismatched
             `loc` parameter or :meth:`add_column` warnings.
         """
-        if other_table.loc != self.loc and do_warn:
-            log.warning(f"other_table.loc ({other_table.loc}) != self.loc({self.loc})")
+        if len(other_table) != len(self) and do_warn:
+            log.warning(
+                f"len(other_table) ({len(other_table)}) != len(self) ({len(self)})"
+            )
         if cols is None:
             cols = other_table.keys()
         for name in cols: