pybes3.besio

Rare Help for Users

There is rare help for users in this page. It is recommended to see BESIII Data Reading.

besio

concatenate_raw(files, n_block_per_batch=10000, sub_detectors=None, max_workers=None, verbose=False)

Concatenate multiple raw binary files into ak.Array

Parameters:

Name	Type	Description	Default
files	Union[Union[str, Path], list[Union[str, Path]]]	files to be read.	required
n_block_per_batch	int	The number of blocks to read per batch. Defaults to 1000.	10000
sub_detectors	Optional[list[str]]	List of sub-detectors to read. Defaults to None, which means read all sub-detectors.	None
max_workers	Optional[int]	The maximum number of worker threads to use for reading the data. Defaults to None, which means use the default number of worker threads.	None
verbose	bool	Show reading process.	False

Returns:

Type	Description
Array	Concatenated raw data array.

Source code in src/pybes3/besio/raw_io.py

def concatenate(
    files: Union[Union[str, Path], list[Union[str, Path]]],
    n_block_per_batch: int = 10000,
    sub_detectors: Optional[list[str]] = None,
    max_workers: Optional[int] = None,
    verbose: bool = False,
) -> ak.Array:
    """
    Concatenate multiple raw binary files into `ak.Array`

    Parameters:
        files (Union[Union[str, Path], list[Union[str, Path]]]): files to be read.
        n_block_per_batch (int, optional): The number of blocks to read per batch. Defaults to 1000.
        sub_detectors (Optional[list[str]]): List of sub-detectors to read. Defaults to `None`, which means read all sub-detectors.
        max_workers (Optional[int]): The maximum number of worker threads to use for reading the data. Defaults to `None`, which means use the default number of worker threads.
        verbose (bool): Show reading process.

    Returns:
        Concatenated raw data array.
    """

    if not isinstance(files, list):
        files = glob.glob(files)

    files = [str(Path(file).resolve()) for file in files if _is_raw(file)]

    if len(files) == 0:
        raise ValueError("No valid raw files found")

    res = []
    for i, f in enumerate(files):
        if verbose:
            print(f"\rreading file {i+1}/{len(files)} ...", end="")

        res.append(
            RawBinaryReader(f).arrays(-1, n_block_per_batch, sub_detectors, max_workers)
        )

    if verbose:
        print()

    return ak.concatenate(res)

wrap_uproot()

Wraps the uproot functions to use the BES interpretation.

Source code in src/pybes3/besio/root_io.py

def wrap_uproot():
    """
    Wraps the uproot functions to use the BES interpretation.
    """
    wrap_uproot_interpretation()
    wrap_uproot_TBranchElement_branches()

open(file, **kwargs)

A wrapper around uproot.open that automatically calls wrap_uproot before opening the file.

Parameters:

Name	Type	Description	Default
file	str | Path | IO | dict[str | Path | IO, str]	The file to open.	required
**kwargs	dict	Additional arguments to pass to uproot.open.	{}

Returns:

Type	Description
Any	The uproot file object.

Source code in src/pybes3/besio/__init__.py

def open(file, **kwargs) -> Any:
    """
    A wrapper around `uproot.open` that automatically calls `wrap_uproot` before opening the file.

    Parameters:
        file (str | Path | IO | dict[str | Path | IO, str]): The file to open.
        **kwargs (dict): Additional arguments to pass to `uproot.open`.

    Returns:
        The uproot file object.
    """
    wrap_uproot()
    return uproot.open(file, **kwargs)

concatenate(files, branch, **kwargs)

A wrapper around uproot.concatenate that automatically calls wrap_uproot before concatenating the files.

Parameters:

Name	Type	Description	Default
files	list[str | Path | IO, str]	The files to concatenate.	required
branch	str	The branch to concatenate.	required
**kwargs	dict	Additional arguments to pass to uproot.concatenate.	{}

Returns:

Type	Description
Any	The concatenated array.

Source code in src/pybes3/besio/__init__.py

def concatenate(files, branch: str, **kwargs) -> Any:
    """
    A wrapper around `uproot.concatenate` that automatically calls `wrap_uproot` before concatenating the files.

    Parameters:
        files (list[str | Path | IO, str]): The files to concatenate.
        branch (str): The branch to concatenate.
        **kwargs (dict): Additional arguments to pass to `uproot.concatenate`.

    Returns:
        The concatenated array.
    """
    wrap_uproot()
    return uproot.concatenate({str(f): branch for f in files}, **kwargs)

open_raw(file)

Open a raw binary file.

Parameters:

Name	Type	Description	Default
file	str	The file to open.	required

Returns:

Type	Description
RawBinaryReader	The raw binary reader.

Source code in src/pybes3/besio/__init__.py

def open_raw(file: str) -> RawBinaryReader:
    """
    Open a raw binary file.

    Parameters:
        file (str): The file to open.

    Returns:
        (RawBinaryReader): The raw binary reader.
    """
    return RawBinaryReader(file)

root_io

BaseReader

Base class for all readers.

Source code in src/pybes3/besio/root_io.py

class BaseReader:
    """
    Base class for all readers.
    """

    @classmethod
    def priority(cls) -> int:
        """
        The priority of the reader. Higher priority means the reader will be
        used first.
        """
        return 20

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name: str,
        cls_streamer_info: dict,
        all_streamer_info: dict,
        item_path: str = "",
    ) -> dict:
        raise NotImplementedError("This method should be implemented in subclasses.")

    @classmethod
    def get_reader_instance(cls, tree_config: dict) -> bcpp.BaseReader:
        """
        Args:
            tree_config (dict): The configuration dictionary for the reader.

        Returns:
            An instance of the appropriate reader class.
        """
        raise NotImplementedError("This method should be implemented in subclasses.")

    @classmethod
    def reconstruct_array(
        cls,
        raw_data: Union[np.ndarray, tuple, list, None],
        tree_config: dict,
    ) -> Union[ak.Array, None]:
        """
        Args:
            raw_data (Union[np.ndarray, tuple, list, None]): The raw data to be
                recovered.
            tree_config (dict): The configuration dictionary for the reader.

        Returns:
            awkward.Array: The recovered data as an awkward array.
        """
        raise NotImplementedError("This method should be implemented in subclasses.")

priority() classmethod

The priority of the reader. Higher priority means the reader will be used first.

Source code in src/pybes3/besio/root_io.py

@classmethod
def priority(cls) -> int:
    """
    The priority of the reader. Higher priority means the reader will be
    used first.
    """
    return 20

get_reader_instance(tree_config) classmethod

Parameters:

Name	Type	Description	Default
tree_config	dict	The configuration dictionary for the reader.	required

Returns:

Type	Description
BaseReader	An instance of the appropriate reader class.

Source code in src/pybes3/besio/root_io.py

@classmethod
def get_reader_instance(cls, tree_config: dict) -> bcpp.BaseReader:
    """
    Args:
        tree_config (dict): The configuration dictionary for the reader.

    Returns:
        An instance of the appropriate reader class.
    """
    raise NotImplementedError("This method should be implemented in subclasses.")

reconstruct_array(raw_data, tree_config) classmethod

Parameters:

Name	Type	Description	Default
raw_data	Union[ndarray, tuple, list, None]	The raw data to be recovered.	required
tree_config	dict	The configuration dictionary for the reader.	required

Returns:

Type	Description
Union[Array, None]	awkward.Array: The recovered data as an awkward array.

Source code in src/pybes3/besio/root_io.py

@classmethod
def reconstruct_array(
    cls,
    raw_data: Union[np.ndarray, tuple, list, None],
    tree_config: dict,
) -> Union[ak.Array, None]:
    """
    Args:
        raw_data (Union[np.ndarray, tuple, list, None]): The raw data to be
            recovered.
        tree_config (dict): The configuration dictionary for the reader.

    Returns:
        awkward.Array: The recovered data as an awkward array.
    """
    raise NotImplementedError("This method should be implemented in subclasses.")

CTypeReader

Bases: BaseReader

This class reads C++ primitive types from a binary parser.

Source code in src/pybes3/besio/root_io.py

class CTypeReader(BaseReader):
    """
    This class reads C++ primitive types from a binary parser.
    """

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if top_type_name in num_typenames:
            ctype = num_typenames[top_type_name]
            return {
                "reader": ReaderType.CType,
                "name": cls_streamer_info["fName"],
                "ctype": ctype,
            }
        else:
            return None

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.CType:
            return None

        ctype = tree_config["ctype"]
        return {
            "i1": bcpp.Int8Reader,
            "i2": bcpp.Int16Reader,
            "i4": bcpp.Int32Reader,
            "i8": bcpp.Int64Reader,
            "u1": bcpp.UInt8Reader,
            "u2": bcpp.UInt16Reader,
            "u4": bcpp.UInt32Reader,
            "u8": bcpp.UInt64Reader,
            "f": bcpp.FloatReader,
            "d": bcpp.DoubleReader,
        }[ctype](tree_config["name"])

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.CType:
            return None

        return ak.Array(raw_data)

STLSequenceReader

Bases: BaseReader

This class reads STL sequence (vector, array) from a binary parser.

Source code in src/pybes3/besio/root_io.py

class STLSequenceReader(BaseReader):
    """
    This class reads STL sequence (vector, array) from a binary parser.
    """

    @staticmethod
    def get_sequence_element_typename(type_name: str) -> str:
        """
        Get the element type name of a vector type.

        e.g. vector<vector<int>> -> vector<int>
        """
        type_name = (
            type_name.replace("std::", "").replace("< ", "<").replace(" >", ">").strip()
        )
        return re.match(r"^(vector|array)<(.*)>$", type_name).group(2)

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if top_type_name not in ["vector", "array"]:
            return None

        fName = cls_streamer_info["fName"]
        fTypeName = cls_streamer_info["fTypeName"]
        element_type = cls.get_sequence_element_typename(fTypeName)
        element_info = {
            "fName": fName,
            "fTypeName": element_type,
        }

        element_tree_config = gen_tree_config(
            element_info,
            all_streamer_info,
            item_path,
        )

        top_element_type = get_top_type_name(element_type)
        if top_element_type in stl_typenames:
            element_tree_config["is_top"] = False

        return {
            "reader": ReaderType.STLSequence,
            "name": fName,
            "element_reader": element_tree_config,
        }

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.STLSequence:
            return None

        element_reader = get_reader_instance(tree_config["element_reader"])
        is_top = tree_config.get("is_top", True)
        return bcpp.STLSeqReader(tree_config["name"], is_top, element_reader)

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.STLSequence:
            return None

        counts, element_raw_data = raw_data
        element_data = reconstruct_array(
            element_raw_data,
            tree_config["element_reader"],
        )
        return ak.unflatten(element_data, counts)

get_sequence_element_typename(type_name) staticmethod

Get the element type name of a vector type.

e.g. vector> -> vector

Source code in src/pybes3/besio/root_io.py

@staticmethod
def get_sequence_element_typename(type_name: str) -> str:
    """
    Get the element type name of a vector type.

    e.g. vector<vector<int>> -> vector<int>
    """
    type_name = (
        type_name.replace("std::", "").replace("< ", "<").replace(" >", ">").strip()
    )
    return re.match(r"^(vector|array)<(.*)>$", type_name).group(2)

STLMapReader

Bases: BaseReader

This class reads std::map from a binary parser.

Source code in src/pybes3/besio/root_io.py

class STLMapReader(BaseReader):
    """
    This class reads std::map from a binary parser.
    """

    @staticmethod
    def get_map_key_val_typenames(type_name: str) -> tuple[str, str]:
        """
        Get the key and value type names of a map type.

        e.g. map<int, vector<int>> -> (int, vector<int>)
        """
        type_name = (
            type_name.replace("std::", "").replace("< ", "<").replace(" >", ">").strip()
        )
        return re.match(r"^(map|unordered_map|multimap)<(.*),(.*)>$", type_name).groups()[1:3]

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if top_type_name not in ["map", "unordered_map", "multimap"]:
            return None

        fTypeName = cls_streamer_info["fTypeName"]
        key_type_name, val_type_name = cls.get_map_key_val_typenames(fTypeName)

        fName = cls_streamer_info["fName"]
        key_info = {
            "fName": "key",
            "fTypeName": key_type_name,
        }

        val_info = {
            "fName": "val",
            "fTypeName": val_type_name,
        }

        key_tree_config = gen_tree_config(key_info, all_streamer_info, item_path)
        if get_top_type_name(key_type_name) in stl_typenames:
            key_tree_config["is_top"] = False

        val_tree_config = gen_tree_config(val_info, all_streamer_info, item_path)
        if get_top_type_name(val_type_name) in stl_typenames:
            val_tree_config["is_top"] = False

        return {
            "reader": ReaderType.STLMap,
            "name": fName,
            "key_reader": key_tree_config,
            "val_reader": val_tree_config,
        }

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.STLMap:
            return None

        key_cpp_reader = get_reader_instance(tree_config["key_reader"])
        val_cpp_reader = get_reader_instance(tree_config["val_reader"])
        is_top = tree_config.get("is_top", True)
        return bcpp.STLMapReader(
            tree_config["name"],
            is_top,
            key_cpp_reader,
            val_cpp_reader,
        )

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.STLMap:
            return None

        key_tree_config = tree_config["key_reader"]
        val_tree_config = tree_config["val_reader"]
        counts, key_raw_data, val_raw_data = raw_data
        key_data = reconstruct_array(key_raw_data, key_tree_config)
        val_data = reconstruct_array(val_raw_data, val_tree_config)

        return ak.unflatten(
            ak.zip(
                {
                    key_tree_config["name"]: key_data,
                    val_tree_config["name"]: val_data,
                },
                with_name="pair",
            ),
            counts,
        )

get_map_key_val_typenames(type_name) staticmethod

Get the key and value type names of a map type.

e.g. map> -> (int, vector)

Source code in src/pybes3/besio/root_io.py

@staticmethod
def get_map_key_val_typenames(type_name: str) -> tuple[str, str]:
    """
    Get the key and value type names of a map type.

    e.g. map<int, vector<int>> -> (int, vector<int>)
    """
    type_name = (
        type_name.replace("std::", "").replace("< ", "<").replace(" >", ">").strip()
    )
    return re.match(r"^(map|unordered_map|multimap)<(.*),(.*)>$", type_name).groups()[1:3]

STLStringReader

Bases: BaseReader

This class reads std::string from a binary parser.

Source code in src/pybes3/besio/root_io.py

class STLStringReader(BaseReader):
    """
    This class reads std::string from a binary parser.
    """

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if top_type_name != "string":
            return None

        return {
            "reader": ReaderType.STLString,
            "name": cls_streamer_info["fName"],
        }

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.STLString:
            return None

        return bcpp.STLStringReader(
            tree_config["name"],
            tree_config.get("is_top", True),
        )

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.STLString:
            return None

        counts, data = raw_data
        return ak.enforce_type(ak.unflatten(data, counts), "string")

TArrayReader

Bases: BaseReader

This class reads TArray from a binary paerser.

TArray includes TArrayC, TArrayS, TArrayI, TArrayL, TArrayF, and TArrayD. Corresponding ctype is u1, u2, i4, i8, f, and d.

Source code in src/pybes3/besio/root_io.py

class TArrayReader(BaseReader):
    """
    This class reads TArray from a binary paerser.

    TArray includes TArrayC, TArrayS, TArrayI, TArrayL, TArrayF, and TArrayD.
    Corresponding ctype is u1, u2, i4, i8, f, and d.
    """

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if top_type_name not in tarray_typenames:
            return None

        ctype = tarray_typenames[top_type_name]
        return {
            "reader": ReaderType.TArray,
            "name": cls_streamer_info["fName"],
            "ctype": ctype,
        }

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.TArray:
            return None

        ctype = tree_config["ctype"]

        return {
            "i1": bcpp.TArrayCReader,
            "i2": bcpp.TArraySReader,
            "i4": bcpp.TArrayIReader,
            "i8": bcpp.TArrayLReader,
            "f": bcpp.TArrayFReader,
            "d": bcpp.TArrayDReader,
        }[ctype](tree_config["name"])

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.TArray:
            return None

        counts, data = raw_data
        return ak.unflatten(data, counts)

TStringReader

Bases: BaseReader

This class reads TString from a binary parser.

Source code in src/pybes3/besio/root_io.py

class TStringReader(BaseReader):
    """
    This class reads TString from a binary parser.
    """

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if top_type_name != "TString":
            return None

        return {
            "reader": ReaderType.TString,
            "name": cls_streamer_info["fName"],
        }

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.TString:
            return None

        return bcpp.TStringReader(tree_config["name"])

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.TString:
            return None

        counts, data = raw_data
        offsets = np.zeros(len(counts) + 1, dtype=counts.dtype)
        np.cumsum(counts, out=offsets[1:])

        return ak.Array(
            awkward.contents.ListOffsetArray(
                awkward.index.Index(offsets),
                awkward.contents.NumpyArray(data, parameters={"__array__": "char"}),
                parameters={"__array__": "string"},
            )
        )

TObjectReader

Bases: BaseReader

This class reads TObject from a binary parser.

It will not record any data.

Source code in src/pybes3/besio/root_io.py

class TObjectReader(BaseReader):
    """
    This class reads TObject from a binary parser.

    It will not record any data.
    """

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if top_type_name != "BASE":
            return None

        fType = cls_streamer_info["fType"]
        if fType != 66:
            return None

        return {
            "reader": ReaderType.TObject,
            "name": cls_streamer_info["fName"],
        }

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.TObject:
            return None

        return bcpp.TObjectReader(tree_config["name"])

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        return None

CArrayReader

Bases: BaseReader

This class reads a C-array from a binary parser.

Source code in src/pybes3/besio/root_io.py

class CArrayReader(BaseReader):
    """
    This class reads a C-array from a binary parser.
    """

    @classmethod
    def priority(cls):
        return 100  # This reader should be called first

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if cls_streamer_info.get("fArrayDim", 0) == 0:
            return None

        fName = cls_streamer_info["fName"]
        fTypeName = cls_streamer_info["fTypeName"]
        fArrayDim = cls_streamer_info["fArrayDim"]
        fMaxIndex = cls_streamer_info["fMaxIndex"]

        element_streamer_info = cls_streamer_info.copy()
        element_streamer_info["fArrayDim"] = 0

        element_tree_config = gen_tree_config(
            element_streamer_info,
            all_streamer_info,
        )

        flat_size = np.prod(fMaxIndex[:fArrayDim])
        assert flat_size > 0, f"flatten_size should be greater than 0, but got {flat_size}"

        # c-type number or TArray
        if top_type_name in num_typenames or top_type_name in tarray_typenames:
            return {
                "reader": ReaderType.CArray,
                "name": fName,
                "is_obj": False,
                "element_reader": element_tree_config,
                "flat_size": flat_size,
                "fMaxIndex": fMaxIndex,
                "fArrayDim": fArrayDim,
            }

        # TSTring
        elif top_type_name == "TString":
            return {
                "reader": ReaderType.CArray,
                "name": fName,
                "is_obj": True,
                "element_reader": element_tree_config,
                "flat_size": flat_size,
                "fMaxIndex": fMaxIndex,
                "fArrayDim": fArrayDim,
            }

        # STL
        elif top_type_name in stl_typenames:
            element_tree_config["is_top"] = False
            return {
                "reader": ReaderType.CArray,
                "name": fName,
                "is_obj": True,
                "flat_size": flat_size,
                "element_reader": element_tree_config,
                "fMaxIndex": fMaxIndex,
                "fArrayDim": fArrayDim,
            }

        else:
            raise ValueError(f"Unknown type: {top_type_name} for C-array: {fTypeName}")

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        reader_type = tree_config["reader"]
        if reader_type != ReaderType.CArray:
            return None

        element_reader = get_reader_instance(tree_config["element_reader"])

        return bcpp.CArrayReader(
            tree_config["name"],
            tree_config["is_obj"],
            tree_config["flat_size"],
            element_reader,
        )

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.CArray:
            return None

        element_tree_config = tree_config["element_reader"]
        fMaxIndex = tree_config["fMaxIndex"]
        fArrayDim = tree_config["fArrayDim"]
        shape = [fMaxIndex[i] for i in range(fArrayDim)]

        element_data = reconstruct_array(
            raw_data,
            element_tree_config,
        )

        for s in shape[::-1]:
            element_data = ak.unflatten(element_data, s)

        return element_data

BaseObjectReader

Bases: BaseReader

Base class is what a custom class inherits from. It has fNBytes(uint32), fVersion(uint16) at the beginning.

Source code in src/pybes3/besio/root_io.py

class BaseObjectReader(BaseReader):
    """
    Base class is what a custom class inherits from.
    It has fNBytes(uint32), fVersion(uint16) at the beginning.
    """

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        if top_type_name != "BASE":
            return None

        fType = cls_streamer_info["fType"]
        if fType != 0:
            return None

        fName = cls_streamer_info["fName"]
        sub_streamers: list = all_streamer_info[fName]

        sub_tree_configs = [
            gen_tree_config(s, all_streamer_info, item_path) for s in sub_streamers
        ]

        return {
            "reader": ReaderType.BaseObject,
            "name": fName,
            "sub_readers": sub_tree_configs,
        }

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.BaseObject:
            return None

        sub_readers = [get_reader_instance(s) for s in tree_config["sub_readers"]]
        return bcpp.BaseObjectReader(tree_config["name"], sub_readers)

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.BaseObject:
            return None

        sub_tree_configs = tree_config["sub_readers"]

        arr_dict = {}
        for s_cfg, s_data in zip(sub_tree_configs, raw_data):
            s_name = s_cfg["name"]
            s_reader_type = s_cfg["reader"]

            if s_reader_type == ReaderType.TObject:
                continue

            arr_dict[s_name] = reconstruct_array(s_data, s_cfg)

        return ak.Array(arr_dict)

ObjectHeaderReader

Bases: BaseReader

This class read an object starting with an object header.

Source code in src/pybes3/besio/root_io.py

class ObjectHeaderReader(BaseReader):
    """
    This class read an object starting with an object header.
    """

    @classmethod
    def priority(cls):
        return 0  # should be called last

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        sub_streamers: list = all_streamer_info[top_type_name]
        sub_tree_configs = [
            gen_tree_config(s, all_streamer_info, item_path) for s in sub_streamers
        ]
        return {
            "reader": ReaderType.ObjectHeader,
            "name": top_type_name,
            "sub_readers": sub_tree_configs,
        }

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.ObjectHeader:
            return None

        sub_readers = [get_reader_instance(s) for s in tree_config["sub_readers"]]
        return bcpp.ObjectHeaderReader(tree_config["name"], sub_readers)

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.ObjectHeader:
            return None

        sub_tree_configs = tree_config["sub_readers"]

        arr_dict = {}
        for s_cfg, s_data in zip(sub_tree_configs, raw_data):
            s_name = s_cfg["name"]
            s_reader_type = s_cfg["reader"]

            if s_reader_type == ReaderType.TObject:
                continue

            arr_dict[s_name] = reconstruct_array(s_data, s_cfg)

        return ak.Array(arr_dict)

EmptyReader

Bases: BaseReader

This class does nothing.

Source code in src/pybes3/besio/root_io.py

class EmptyReader(BaseReader):
    """
    This class does nothing.
    """

    @classmethod
    def gen_tree_config(
        cls,
        top_type_name,
        cls_streamer_info,
        all_streamer_info,
        item_path,
    ):
        return None

    @classmethod
    def get_reader_instance(cls, tree_config: dict):
        if tree_config["reader"] != ReaderType.Empty:
            return None

        return bcpp.EmptyReader(tree_config["name"])

    @classmethod
    def reconstruct_array(cls, raw_data, tree_config):
        if tree_config["reader"] != ReaderType.Empty:
            return None

        return np.empty(shape=(0,))

Bes3Interpretation

Bases: Interpretation

Custom interpretation for Bes3 data.

Source code in src/pybes3/besio/root_io.py

class Bes3Interpretation(uproot.interpretation.Interpretation):
    """
    Custom interpretation for Bes3 data.
    """

    target_branches: set[str] = set(Bes3TObjArrayReader.bes3_branch2types.keys())

    def __init__(
        self,
        branch: uproot.behaviors.TBranch.TBranch,
        context: dict,
        simplify: bool,
    ):
        """
        Args:
            branch (:doc:`uproot.behaviors.TBranch.TBranch`): The ``TBranch`` to
                interpret as an array.
            context (dict): Auxiliary data used in deserialization.
            simplify (bool): If True, call
                :ref:`uproot.interpretation.objects.AsObjects.simplify` on any
                :doc:`uproot.interpretation.objects.AsObjects` to try to get a
                more efficient interpretation.

        Accept arguments from `uproot.interpretation.identify.interpretation_of`.
        """
        self._branch = branch
        self._context = context
        self._simplify = simplify

        # simplify streamer information
        self.all_streamer_info: dict[str, list[dict]] = {}
        for k, v in branch.file.streamers.items():
            cur_infos = [i.all_members for i in next(iter(v.values())).member("fElements")]
            self.all_streamer_info[k] = cur_infos

    @classmethod
    def match_branch(
        cls,
        branch: uproot.behaviors.TBranch.TBranch,
        context: dict,
        simplify: bool,
    ) -> bool:
        """
        Args:
            branch (:doc:`uproot.behaviors.TBranch.TBranch`): The ``TBranch`` to
                interpret as an array.
            context (dict): Auxiliary data used in deserialization.
            simplify (bool): If True, call
                :ref:`uproot.interpretation.objects.AsObjects.simplify` on any
                :doc:`uproot.interpretation.objects.AsObjects` to try to get a
                more efficient interpretation.

        Accept arguments from `uproot.interpretation.identify.interpretation_of`,
        determine whether this interpretation can be applied to the given branch.
        """
        full_path = regularize_object_path(branch.object_path)
        return full_path in cls.target_branches

    @property
    def typename(self) -> str:
        """
        The name of the type of the interpretation.
        """
        return self._branch.streamer.typename

    @property
    def cache_key(self) -> str:
        """
        The cache key of the interpretation.
        """
        return id(self)

    def __repr__(self) -> str:
        """
        The string representation of the interpretation.
        """
        return f"AsBes3Custom({self.typename})"

    def final_array(
        self,
        basket_arrays,
        entry_start,
        entry_stop,
        entry_offsets,
        library,
        branch,
        options,
    ):
        """
        Concatenate the arrays from the baskets and return the final array.
        """

        awkward = uproot.extras.awkward()

        basket_entry_starts = np.array(entry_offsets[:-1])
        basket_entry_stops = np.array(entry_offsets[1:])

        basket_start_idx = np.where(basket_entry_starts <= entry_start)[0].max()
        basket_end_idx = np.where(basket_entry_stops >= entry_stop)[0].min()

        arr_to_concat = [basket_arrays[i] for i in range(basket_start_idx, basket_end_idx + 1)]
        tot_array = awkward.concatenate(arr_to_concat)

        relative_entry_start = entry_start - basket_entry_starts[basket_start_idx]
        relative_entry_stop = entry_stop - basket_entry_starts[basket_start_idx]

        return tot_array[relative_entry_start:relative_entry_stop]

    def basket_array(
        self,
        data,
        byte_offsets,
        basket,
        branch,
        context,
        cursor_offset,
        library,
        interp_options,
    ):
        assert library.name == "ak", "Only awkward arrays are supported"

        full_branch_path = regularize_object_path(branch.object_path)

        # generate reader config
        tree_config = gen_tree_config_from_type_name(
            branch.streamer.typename, self.all_streamer_info, full_branch_path
        )

        # get reader
        reader = get_reader_instance(tree_config)

        # do read
        raw_data = bcpp.read_data(data, byte_offsets, reader)

        # recover raw data
        raw_ak_arr = reconstruct_array(raw_data, tree_config)

        # preprocess awkward array and return
        return preprocess_subbranch(full_branch_path, raw_ak_arr)

typename property

The name of the type of the interpretation.

cache_key property

The cache key of the interpretation.

__init__(branch, context, simplify)

Parameters:

Name	Type	Description	Default
branch		doc:uproot.behaviors.TBranch.TBranch): The TBranch to interpret as an array.	required
context	dict	Auxiliary data used in deserialization.	required
simplify	bool	If True, call :ref:uproot.interpretation.objects.AsObjects.simplify on any :doc:uproot.interpretation.objects.AsObjects to try to get a more efficient interpretation.	required

Accept arguments from uproot.interpretation.identify.interpretation_of.

Source code in src/pybes3/besio/root_io.py

def __init__(
    self,
    branch: uproot.behaviors.TBranch.TBranch,
    context: dict,
    simplify: bool,
):
    """
    Args:
        branch (:doc:`uproot.behaviors.TBranch.TBranch`): The ``TBranch`` to
            interpret as an array.
        context (dict): Auxiliary data used in deserialization.
        simplify (bool): If True, call
            :ref:`uproot.interpretation.objects.AsObjects.simplify` on any
            :doc:`uproot.interpretation.objects.AsObjects` to try to get a
            more efficient interpretation.

    Accept arguments from `uproot.interpretation.identify.interpretation_of`.
    """
    self._branch = branch
    self._context = context
    self._simplify = simplify

    # simplify streamer information
    self.all_streamer_info: dict[str, list[dict]] = {}
    for k, v in branch.file.streamers.items():
        cur_infos = [i.all_members for i in next(iter(v.values())).member("fElements")]
        self.all_streamer_info[k] = cur_infos

match_branch(branch, context, simplify) classmethod

Parameters:

Name	Type	Description	Default
branch		doc:uproot.behaviors.TBranch.TBranch): The TBranch to interpret as an array.	required
context	dict	Auxiliary data used in deserialization.	required
simplify	bool	If True, call :ref:uproot.interpretation.objects.AsObjects.simplify on any :doc:uproot.interpretation.objects.AsObjects to try to get a more efficient interpretation.	required

Accept arguments from uproot.interpretation.identify.interpretation_of, determine whether this interpretation can be applied to the given branch.

Source code in src/pybes3/besio/root_io.py

@classmethod
def match_branch(
    cls,
    branch: uproot.behaviors.TBranch.TBranch,
    context: dict,
    simplify: bool,
) -> bool:
    """
    Args:
        branch (:doc:`uproot.behaviors.TBranch.TBranch`): The ``TBranch`` to
            interpret as an array.
        context (dict): Auxiliary data used in deserialization.
        simplify (bool): If True, call
            :ref:`uproot.interpretation.objects.AsObjects.simplify` on any
            :doc:`uproot.interpretation.objects.AsObjects` to try to get a
            more efficient interpretation.

    Accept arguments from `uproot.interpretation.identify.interpretation_of`,
    determine whether this interpretation can be applied to the given branch.
    """
    full_path = regularize_object_path(branch.object_path)
    return full_path in cls.target_branches

__repr__()

The string representation of the interpretation.

Source code in src/pybes3/besio/root_io.py

def __repr__(self) -> str:
    """
    The string representation of the interpretation.
    """
    return f"AsBes3Custom({self.typename})"

final_array(basket_arrays, entry_start, entry_stop, entry_offsets, library, branch, options)

Concatenate the arrays from the baskets and return the final array.

Source code in src/pybes3/besio/root_io.py

def final_array(
    self,
    basket_arrays,
    entry_start,
    entry_stop,
    entry_offsets,
    library,
    branch,
    options,
):
    """
    Concatenate the arrays from the baskets and return the final array.
    """

    awkward = uproot.extras.awkward()

    basket_entry_starts = np.array(entry_offsets[:-1])
    basket_entry_stops = np.array(entry_offsets[1:])

    basket_start_idx = np.where(basket_entry_starts <= entry_start)[0].max()
    basket_end_idx = np.where(basket_entry_stops >= entry_stop)[0].min()

    arr_to_concat = [basket_arrays[i] for i in range(basket_start_idx, basket_end_idx + 1)]
    tot_array = awkward.concatenate(arr_to_concat)

    relative_entry_start = entry_start - basket_entry_starts[basket_start_idx]
    relative_entry_stop = entry_stop - basket_entry_starts[basket_start_idx]

    return tot_array[relative_entry_start:relative_entry_stop]

gen_tree_config(cls_streamer_info, all_streamer_info, item_path='')

Generate reader configuration for a class streamer information.

The content it returns should be:

{
    "reader": ReaderType,
    "name": str,
    "ctype": str, # for CTypeReader, TArrayReader
    "element_reader": dict, # reader config of the element, for STLVectorReader, SimpleCArrayReader, TObjectCArrayReader
    "flat_size": int, # for SimpleCArrayReader, TObjectCArrayReader
    "fMaxIndex": list[int], # for SimpleCArrayReader, TObjectCArrayReader
    "fArrayDim": int, # for SimpleCArrayReader, TObjectCArrayReader
    "key_reader": dict, # reader config of the key, for STLMapReader
    "val_reader": dict, # reader config of the value, for STLMapReader
    "sub_readers": list[dict], # for BaseObjectReader, ObjectHeaderReader
    "is_top_level": bool, # for STLVectorReader, STLMapReader, STLStringReader
}

Parameters:

Name	Type	Description	Default
cls_streamer_info	dict	Class streamer information.	required
all_streamer_info	dict	All streamer information.	required
item_path	str	Path to the item.	''

Returns:

Name	Type	Description
dict	dict	Reader configuration.

Source code in src/pybes3/besio/root_io.py

def gen_tree_config(
    cls_streamer_info: dict,
    all_streamer_info: dict,
    item_path: str = "",
) -> dict:
    """
    Generate reader configuration for a class streamer information.

    The content it returns should be:

    ```python
    {
        "reader": ReaderType,
        "name": str,
        "ctype": str, # for CTypeReader, TArrayReader
        "element_reader": dict, # reader config of the element, for STLVectorReader, SimpleCArrayReader, TObjectCArrayReader
        "flat_size": int, # for SimpleCArrayReader, TObjectCArrayReader
        "fMaxIndex": list[int], # for SimpleCArrayReader, TObjectCArrayReader
        "fArrayDim": int, # for SimpleCArrayReader, TObjectCArrayReader
        "key_reader": dict, # reader config of the key, for STLMapReader
        "val_reader": dict, # reader config of the value, for STLMapReader
        "sub_readers": list[dict], # for BaseObjectReader, ObjectHeaderReader
        "is_top_level": bool, # for STLVectorReader, STLMapReader, STLStringReader
    }
    ```

    Args:
        cls_streamer_info (dict): Class streamer information.
        all_streamer_info (dict): All streamer information.
        item_path (str): Path to the item.

    Returns:
        dict: Reader configuration.
    """
    fName = cls_streamer_info["fName"]
    item_path = fName if item_path == "" else f"{item_path}.{fName}"

    for reader in sorted(readers, key=lambda x: x.priority(), reverse=True):
        top_type_name = get_top_type_name(cls_streamer_info["fTypeName"])
        tree_config = reader.gen_tree_config(
            top_type_name,
            cls_streamer_info,
            all_streamer_info,
            item_path,
        )
        if tree_config is not None:
            return tree_config

    raise ValueError(f"Unknown type: {cls_streamer_info['fTypeName']} for {item_path}")

get_symetric_matrix_idx(i, j, ndim)

Returns the index of the similarity matrix given the row and column indices.

The matrix is assumed to be symmetric-like. (i, j) -> index relationship is:

	i=0	i=1	i=2
j=0	0
j=1	1	2
j=2	3	4	5

Parameters:

Name	Type	Description	Default
i	Union[int, Array, ndarray]	The row index or array of row indices.	required
j	Union[int, Array, ndarray]	The column index or array of column indices.	required
ndim	int	The dimension of the similarity matrix.	required

Returns:

Type	Description
int	The index or array of indices corresponding to the given row and column indices.

Raises:

Type	Description
ValueError	If the row and column indices are not of the same type, or if one of them is not an integer.
ValueError	If the row or column indices are greater than or equal to the dimension of the similarity matrix.
ValueError	If the row or column indices are negative.

Source code in src/pybes3/besio/root_io.py

def get_symetric_matrix_idx(
    i: Union[int, ak.Array, np.ndarray], j: Union[int, ak.Array, np.ndarray], ndim: int
) -> int:
    """
    Returns the index of the similarity matrix given the row and column indices.

    The matrix is assumed to be symmetric-like. (i, j) -> index relationship is:

    |     | i=0 | i=1 | i=2 |
    | :-: | :-: | :-: | :-: |
    | j=0 |  0  |     |     |
    | j=1 |  1  |  2  |     |
    | j=2 |  3  |  4  |  5  |

    Parameters:
        i (Union[int, ak.Array, np.ndarray]): The row index or array of row indices.
        j (Union[int, ak.Array, np.ndarray]): The column index or array of column indices.
        ndim (int): The dimension of the similarity matrix.

    Returns:
        The index or array of indices corresponding to the given row and column indices.

    Raises:
        ValueError: If the row and column indices are not of the same type, or if one of them is not an integer.
        ValueError: If the row or column indices are greater than or equal to the dimension of the similarity matrix.
        ValueError: If the row or column indices are negative.
    """
    # Check type
    return_type: Literal["ak", "np"] = "ak"
    if type(i) != type(j):
        if isinstance(i, int):
            return_type = "np" if isinstance(j, np.ndarray) else "ak"
            i = ak.ones_like(j) * i
        elif isinstance(j, int):
            return_type = "np" if isinstance(i, np.ndarray) else "ak"
            j = ak.ones_like(i) * j
        else:
            raise ValueError(
                "i and j should be the same type, or one of them should be an integer."
            )
    else:
        return_type = "np" if isinstance(i, np.ndarray) else "ak"

    i, j = ak.sort([i, j], axis=0)
    res = j * (j + 1) // 2 + i

    # Check dimension
    if ak.any([i >= ndim, j >= ndim]):
        raise ValueError(
            "Indices i and j should be less than the dimension of the similarity matrix."
        )
    if ak.any([i < 0, j < 0]):
        raise ValueError("Indices i and j should be non-negative.")

    if return_type == "np" and isinstance(res, ak.Array):
        res = res.to_numpy()

    return res

expand_zipped_symetric_matrix(arr)

Recover a flattened simplified symmetric matrix represented as a 1D array back to a 2D matrix. This function assumes the last dimension of the input array is the flattened symmetric matrix, and will transform array

[[a11, a12, a22, a13, a23, a33],
 [b11, b12, b22, b13, b23, b33]]

to

[[[a11, a12, a13],
  [a12, a22, a23],
  [a13, a23, a33]],

 [[b11, b12, b13],
  [b12, b22, b23],
  [b13, b23, b33]]]

Parameters:

Name	Type	Description	Default
arr	Union[Array, ndarray]	The input array representing the flattened simplified symmetric matrix.	required

Returns:

Type	Description
Union[Array, ndarray]	The reshaped symmetric matrix as a 2D array.

Raises:

Type	Description
ValueError	If the input array does not have a symmetric shape.

Source code in src/pybes3/besio/root_io.py

def expand_zipped_symetric_matrix(
    arr: Union[ak.Array, np.ndarray],
) -> Union[ak.Array, np.ndarray]:
    """
    Recover a flattened simplified symmetric matrix represented as a 1D array back to a 2D matrix.
    This function assumes the last dimension of the input array is the flattened symmetric matrix,
    and will transform array

    ```
    [[a11, a12, a22, a13, a23, a33],
     [b11, b12, b22, b13, b23, b33]]
    ```

    to

    ```
    [[[a11, a12, a13],
      [a12, a22, a23],
      [a13, a23, a33]],

     [[b11, b12, b13],
      [b12, b22, b23],
      [b13, b23, b33]]]
    ```

    Parameters:
        arr (Union[ak.Array, np.ndarray]): The input array representing the flattened simplified symmetric matrix.

    Returns:
        The reshaped symmetric matrix as a 2D array.

    Raises:
        ValueError: If the input array does not have a symmetric shape.
    """

    # Get the number of elements in the symmetric matrix
    if isinstance(arr, ak.Array):
        type_strs = [i.strip() for i in arr.typestr.split("*")[:-1]]
        n_err_elements = int(type_strs[-1])
        raw_shape = _extract_index(arr.layout)[:-1]
        flat_arr = _flat_to_numpy(arr).flatten().reshape(-1, n_err_elements)
    else:
        n_err_elements = arr.shape[-1]
        raw_shape = arr.shape[:-1]
        flat_arr = arr.reshape(-1, n_err_elements)

    ndim_err = (np.sqrt(1 + 8 * n_err_elements) - 1) / 2
    if not ndim_err.is_integer():
        raise ValueError("The array does not have a symmetric shape.")
    ndim_err = int(ndim_err)

    # Preapre output array
    n_raw_len = len(flat_arr.flatten())
    n_out_len = n_raw_len // n_err_elements * (ndim_err**2)
    raw_out = np.zeros(n_out_len, dtype=flat_arr.dtype).reshape(-1, ndim_err, ndim_err)

    # Fill error matrix
    for i in range(ndim_err):
        for j in range(ndim_err):
            idx = get_symetric_matrix_idx(i, j, ndim_err)
            raw_out[:, i, j] = flat_arr[:, idx]

    # Reshape the output array to match the original shape
    if isinstance(arr, ak.Array):
        res = _recover_shape(ak.Array(raw_out), raw_shape)
    else:
        res = raw_out.reshape(*raw_shape, ndim_err, ndim_err)

    return res

expand_subbranch_symetric_matrix(sub_br_arr, matrix_fields)

Recover simplified symmetric matrix back to 2D matrix from specified fields of a branch array.

Parameters:

Name	Type	Description	Default
sub_br_arr	Array	Subbranch array that need to be recovered.	required
matrix_fields	Union[str, set[str]]	Name of list of names of fields to be recovered.	required

Returns:

Type	Description
Array	An array with recovered fields.

Source code in src/pybes3/besio/root_io.py

def expand_subbranch_symetric_matrix(
    sub_br_arr: ak.Array, matrix_fields: Union[str, set[str]]
) -> ak.Array:
    """
    Recover simplified symmetric matrix back to 2D matrix from specified fields of a branch array.

    Parameters:
        sub_br_arr: Subbranch array that need to be recovered.
        matrix_fields: Name of list of names of fields to be recovered.

    Returns:
        An array with recovered fields.
    """
    if isinstance(matrix_fields, str):
        matrix_fields = {matrix_fields}
    matrix_fields = set(matrix_fields)

    raw_shape = _extract_index(sub_br_arr.layout)

    res_dict = {}
    for field_name in sub_br_arr.fields:
        flat_sub_arr = sub_br_arr[field_name]
        for _ in range(len(raw_shape)):
            flat_sub_arr = ak.flatten(flat_sub_arr)

        if field_name in matrix_fields:
            res_dict[field_name] = expand_zipped_symetric_matrix(flat_sub_arr)
        else:
            res_dict[field_name] = flat_sub_arr

    res_arr = _recover_shape(ak.Array(res_dict), raw_shape)
    return res_arr

process_digi_subbranch(org_arr)

Processes the TRawData subbranch of the input awkward array and returns a new array with the subbranch fields merged into the top level.

Parameters:

Name	Type	Description	Default
org_arr	Array	The input awkward array containing the TRawData subbranch.	required

Returns:

Type	Description
Array	A new awkward array with the fields of TRawData merged into the top level.

Raises:

Type	Description
AssertionError	If TRawData is not found in the input array fields.

Source code in src/pybes3/besio/root_io.py

def process_digi_subbranch(org_arr: ak.Array) -> ak.Array:
    """
    Processes the `TRawData` subbranch of the input awkward array and returns a new array with the subbranch fields
    merged into the top level.

    Parameters:
        org_arr (ak.Array): The input awkward array containing the `TRawData` subbranch.

    Returns:
        A new awkward array with the fields of `TRawData` merged into the top level.

    Raises:
        AssertionError: If `TRawData` is not found in the input array fields.
    """
    assert "TRawData" in org_arr.fields, "TRawData not found in the input array"

    fields = {}
    for field_name in org_arr.fields:
        if field_name == "TRawData":
            for raw_field_name in org_arr[field_name].fields:
                fields[raw_field_name] = org_arr[field_name][raw_field_name]
        else:
            fields[field_name] = org_arr[field_name]

    return ak.Array(fields)

wrap_uproot()

Wraps the uproot functions to use the BES interpretation.

Source code in src/pybes3/besio/root_io.py

def wrap_uproot():
    """
    Wraps the uproot functions to use the BES interpretation.
    """
    wrap_uproot_interpretation()
    wrap_uproot_TBranchElement_branches()

raw_io

RawBinaryReader

Source code in src/pybes3/besio/raw_io.py

class RawBinaryReader:
    def __init__(
        self,
        file: str,
    ):
        self.file = str(Path(file).resolve())
        self._file = open(file, "rb")

        self.file_version: int = -1
        self.file_number: int = -1
        self.file_date: int = -1
        self.file_time: int = -1

        self.app_name: str = "None"
        self.app_tag: str = "None"

        self.run_number: int = -1
        self.max_events: int = -1
        self.rec_enable: int = -1
        self.trigger_type: int = -1
        self.detector_mask: int = -1
        self.beam_type: int = -1
        self.beam_energy: int = -1

        self.entries: int = -1

        self.data_start: int = 0  # in char
        self.data_end: int = 0  # in char
        self.file_size: int = 0  # in char
        self.data_size: int = 0  # in char

        self.event_starts: np.ndarray = np.empty(0, dtype=np.uint32)  # in char
        self.event_stops: np.ndarray = np.empty(0, dtype=np.uint32)  # in char
        self.max_event_offset: int = 0
        self.current_entry: int = -1

        self._preprocess_file()

    def arrays(
        self,
        n_blocks: int = -1,
        n_block_per_batch: int = 1000,
        sub_detectors: Optional[list[str]] = None,
        max_workers: Optional[int] = None,
    ) -> ak.Array:
        """
        Read and return arrays of data from the BES raw file.

        Parameters:
            n_blocks (int, optional): The number of blocks to read. Defaults to -1, which means read all blocks.
            n_block_per_batch (int, optional): The number of blocks to read per batch. Defaults to 1000.
            sub_detectors (Optional[list[str]]): List of sub-detectors to read. Defaults to `None`, which means read all sub-detectors.
            max_workers (Optional[int]): The maximum number of worker threads to use for reading the data. Defaults to `None`, which means use the default number of worker threads.

        Returns:
            An Awkward Array containing the read data.
        """

        self._reset_cursor()

        if sub_detectors is None:
            sub_detectors = []

        executor = ThreadPoolExecutor(max_workers=max_workers)

        n_total_blocks_read = 0

        futures: list[Future] = []
        while n_total_blocks_read < n_blocks or (
            n_blocks == -1 and self._file.tell() < self.data_end
        ):
            n_tlock_to_read = (
                min(n_blocks - n_total_blocks_read, n_block_per_batch)
                if n_blocks != -1
                else n_block_per_batch
            )

            batch_data, n_read = self._read_batch(n_tlock_to_read)
            futures.append(executor.submit(read_bes_raw, batch_data, sub_detectors))
            n_total_blocks_read += n_read

        res = []
        for future in futures:
            org_dict = future.result()
            res.append(_raw_dict_to_ak(org_dict))

        return ak.concatenate(res)

    def _read(self) -> int:
        return int.from_bytes(self._file.read(4), "little")

    def _skip(self, n: int = 1) -> None:
        self._file.seek(4 * n, 1)

    def _preprocess_file(self):
        # file header
        assert self._read() == BesFlag.FILE_START, "Invalid start flag"
        self._skip()

        self.file_version = self._read()
        self.file_number = self._read()
        self.file_date = self._read()
        self.file_time = self._read()
        self._skip(2)

        # file name
        assert self._read() == BesFlag.FILE_NAME, "Invalid file name flag"

        nchar_name = self._read()
        nbytes_name = np.ceil(nchar_name / 4).astype(int)
        self.file_name = self._file.read(nbytes_name * 4).decode("utf-8").strip()

        nchar_tag = self._read()
        nbytes_tag = np.ceil(nchar_tag / 4).astype(int)
        self.file_tag = self._file.read(nbytes_tag * 4).decode("utf-8").strip()

        # run parameters
        assert self._read() == BesFlag.RUN_PARAMS, "Invalid run params flag"
        self._skip()

        self.run_number = self._read()
        self.max_events = self._read()
        self.rec_enable = self._read()
        self.trigger_type = self._read()
        self.detector_mask = self._read()
        self.beam_type = self._read()
        self.beam_energy = self._read()

        # other information
        self.data_start = self._file.tell()
        self._file.seek(0, 2)
        self.file_size = self._file.tell()
        self.data_end = self.file_size - 10 * 4
        self.data_size = self.data_end - self.data_start

        # read file tail
        self._file.seek(-10 * 4, 2)
        assert self._read() == BesFlag.FILE_TAIL_START, "Invalid file tail start flag"
        self._skip(3)
        self.entries = self._read()
        self._skip(4)
        assert self._read() == BesFlag.FILE_END, "Invalid file end flag"

        self._reset_cursor()

    def _reset_cursor(self):
        self._file.seek(self.data_start)
        self.current_entry = 0

    def _skip_event(self):
        flag = self._read()
        if flag == BesFlag.DATA_SEPERATOR:
            self._skip(3)
            flag = self._read()

        assert flag == BesFlag.FULL_EVENT_FRAGMENT, "Invalid event fragment flag"

        total_size = self._read()

        if self.current_entry > self.max_event_offset:
            self.event_starts[self.current_entry] = self._file.tell() - 4 * 2
            self.event_stops[self.current_entry] = (
                self.event_starts[self.current_entry] + total_size
            )

        self._skip(total_size - 2)
        self.current_entry += 1

    def _read_batch(self, n_blocks: int):
        pos_start = self._file.tell()
        block_counter = 0
        for _ in range(n_blocks):
            if self._file.tell() >= self.data_end:
                assert self._file.tell() == self.data_end, "Invalid data end"
                break

            assert self._read() == BesFlag.DATA_SEPERATOR, "Invalid data seperator flag"
            self._skip(2)
            block_size = self._read()
            self._skip(block_size // 4)
            block_counter += 1

        pos_end = self._file.tell()

        self._file.seek(pos_start, 0)
        batch_data = np.frombuffer(self._file.read(pos_end - pos_start), dtype=np.uint32)

        return batch_data, block_counter

    def __repr__(self) -> str:
        return (
            f"BesRawReader\n"
            f"- File: {self.file}\n"
            f"- Run Number: {self.run_number}\n"
            f"- Entries: {self.entries}\n"
            f"- File Size: {self.file_size//1024//1024} MB\n"
        )

arrays(n_blocks=-1, n_block_per_batch=1000, sub_detectors=None, max_workers=None)

Read and return arrays of data from the BES raw file.

Parameters:

Name	Type	Description	Default
n_blocks	int	The number of blocks to read. Defaults to -1, which means read all blocks.	-1
n_block_per_batch	int	The number of blocks to read per batch. Defaults to 1000.	1000
sub_detectors	Optional[list[str]]	List of sub-detectors to read. Defaults to None, which means read all sub-detectors.	None
max_workers	Optional[int]	The maximum number of worker threads to use for reading the data. Defaults to None, which means use the default number of worker threads.	None

Returns:

Type	Description
Array	An Awkward Array containing the read data.

Source code in src/pybes3/besio/raw_io.py

def arrays(
    self,
    n_blocks: int = -1,
    n_block_per_batch: int = 1000,
    sub_detectors: Optional[list[str]] = None,
    max_workers: Optional[int] = None,
) -> ak.Array:
    """
    Read and return arrays of data from the BES raw file.

    Parameters:
        n_blocks (int, optional): The number of blocks to read. Defaults to -1, which means read all blocks.
        n_block_per_batch (int, optional): The number of blocks to read per batch. Defaults to 1000.
        sub_detectors (Optional[list[str]]): List of sub-detectors to read. Defaults to `None`, which means read all sub-detectors.
        max_workers (Optional[int]): The maximum number of worker threads to use for reading the data. Defaults to `None`, which means use the default number of worker threads.

    Returns:
        An Awkward Array containing the read data.
    """

    self._reset_cursor()

    if sub_detectors is None:
        sub_detectors = []

    executor = ThreadPoolExecutor(max_workers=max_workers)

    n_total_blocks_read = 0

    futures: list[Future] = []
    while n_total_blocks_read < n_blocks or (
        n_blocks == -1 and self._file.tell() < self.data_end
    ):
        n_tlock_to_read = (
            min(n_blocks - n_total_blocks_read, n_block_per_batch)
            if n_blocks != -1
            else n_block_per_batch
        )

        batch_data, n_read = self._read_batch(n_tlock_to_read)
        futures.append(executor.submit(read_bes_raw, batch_data, sub_detectors))
        n_total_blocks_read += n_read

    res = []
    for future in futures:
        org_dict = future.result()
        res.append(_raw_dict_to_ak(org_dict))

    return ak.concatenate(res)

concatenate(files, n_block_per_batch=10000, sub_detectors=None, max_workers=None, verbose=False)

Concatenate multiple raw binary files into ak.Array

Parameters:

Name	Type	Description	Default
files	Union[Union[str, Path], list[Union[str, Path]]]	files to be read.	required
n_block_per_batch	int	The number of blocks to read per batch. Defaults to 1000.	10000
sub_detectors	Optional[list[str]]	List of sub-detectors to read. Defaults to None, which means read all sub-detectors.	None
max_workers	Optional[int]	The maximum number of worker threads to use for reading the data. Defaults to None, which means use the default number of worker threads.	None
verbose	bool	Show reading process.	False

Returns:

Type	Description
Array	Concatenated raw data array.

Source code in src/pybes3/besio/raw_io.py

def concatenate(
    files: Union[Union[str, Path], list[Union[str, Path]]],
    n_block_per_batch: int = 10000,
    sub_detectors: Optional[list[str]] = None,
    max_workers: Optional[int] = None,
    verbose: bool = False,
) -> ak.Array:
    """
    Concatenate multiple raw binary files into `ak.Array`

    Parameters:
        files (Union[Union[str, Path], list[Union[str, Path]]]): files to be read.
        n_block_per_batch (int, optional): The number of blocks to read per batch. Defaults to 1000.
        sub_detectors (Optional[list[str]]): List of sub-detectors to read. Defaults to `None`, which means read all sub-detectors.
        max_workers (Optional[int]): The maximum number of worker threads to use for reading the data. Defaults to `None`, which means use the default number of worker threads.
        verbose (bool): Show reading process.

    Returns:
        Concatenated raw data array.
    """

    if not isinstance(files, list):
        files = glob.glob(files)

    files = [str(Path(file).resolve()) for file in files if _is_raw(file)]

    if len(files) == 0:
        raise ValueError("No valid raw files found")

    res = []
    for i, f in enumerate(files):
        if verbose:
            print(f"\rreading file {i+1}/{len(files)} ...", end="")

        res.append(
            RawBinaryReader(f).arrays(-1, n_block_per_batch, sub_detectors, max_workers)
        )

    if verbose:
        print()

    return ak.concatenate(res)