Skip to content

Commit

Permalink
Merge pull request #26 from jacanchaplais/patch/broadcast-unstruc-25
Browse files Browse the repository at this point in the history
Prevent broadcasting when passing unstructured data #25
  • Loading branch information
jacanchaplais authored Jun 12, 2024
2 parents 1f31fcb + b2c3569 commit d46a4b3
Showing 1 changed file with 12 additions and 19 deletions.
31 changes: 12 additions & 19 deletions heparchy/write/hdf.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
Provides the interface to write HEP data to the heparchy HDF5 format.
"""
import functools as fn
import itertools as it
import typing as ty
import warnings
from enum import Enum
Expand Down Expand Up @@ -290,12 +289,10 @@ def edges(self) -> ty.NoReturn:

@edges.setter
def edges(self, data: base.VoidVector) -> None:
self._mk_dset(
name="edges",
data=data,
shape=data.shape,
dtype=np.dtype([("src", "<i4"), ("dst", "<i4")]),
)
dtype = np.dtype([("src", "<i4"), ("dst", "<i4")])
if not data.dtype.names:
data = data.view(dtype).reshape(-1)
self._mk_dset(name="edges", data=data, shape=data.shape, dtype=dtype)
self._num_edges = len(data)

@property
Expand Down Expand Up @@ -328,13 +325,11 @@ def pmu(self) -> ty.NoReturn:

@pmu.setter
def pmu(self, data: base.VoidVector) -> None:
dtype = np.dtype([(name, "<f8") for name in "xyze"])
if not data.dtype.names:
data = data.view(dtype).reshape(-1)
self._set_num_pcls(data)
self._mk_dset(
name="pmu",
data=data,
shape=data.shape,
dtype=np.dtype(list(zip("xyze", it.repeat("<f8")))),
)
self._mk_dset(name="pmu", data=data, shape=data.shape, dtype=dtype)

@property
def color(self) -> ty.NoReturn:
Expand All @@ -343,13 +338,11 @@ def color(self) -> ty.NoReturn:

@color.setter
def color(self, data: base.AnyVector) -> None:
dtype = np.dtype([("color", "<i4"), ("anticolor", "<i4")])
if not data.dtype.names:
data = data.view(dtype).reshape(-1)
self._set_num_pcls(data)
self._mk_dset(
name="color",
data=data,
shape=data.shape,
dtype=np.dtype([("color", "<i4"), ("anticolor", "<i4")]),
)
self._mk_dset(name="color", data=data, shape=data.shape, dtype=dtype)

@property
def pdg(self) -> ty.NoReturn:
Expand Down

0 comments on commit d46a4b3

Please sign in to comment.