Skip to content

Commit

Permalink
Python bindings for Frame (#343)
Browse files Browse the repository at this point in the history
* Make python code a proper module

* Reorganize CMake config for shared libraries

* Refactor library and dict generation in cmake config

* Add basic version of python bindings for working with Frames
  • Loading branch information
tmadlener authored Nov 10, 2022
1 parent 77babd8 commit e8ae38b
Show file tree
Hide file tree
Showing 38 changed files with 909 additions and 245 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -61,3 +61,4 @@ spack*

# Populated by cmake before build
/include/podio/podioVersion.h
/python/podio/__init__.py
2 changes: 2 additions & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,8 @@ configure_file(${CMAKE_CURRENT_SOURCE_DIR}/podioVersion.in.h
${CMAKE_CURRENT_SOURCE_DIR}/include/podio/podioVersion.h )
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/include/podio/podioVersion.h
DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/podio )
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/python/__init__.py.in
${CMAKE_CURRENT_SOURCE_DIR}/python/podio/__init__.py)

#--- add license files ---------------------------------------------------------
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/LICENSE
Expand Down
25 changes: 25 additions & 0 deletions include/podio/Frame.h
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,10 @@ using EnableIfCollection = typename std::enable_if_t<isCollection<T>>;
template <typename T>
using EnableIfCollectionRValue = typename std::enable_if_t<isCollection<T> && !std::is_lvalue_reference_v<T>>;

/// Alias template for enabling overloads for r-values
template <typename T>
using EnableIfRValue = typename std::enable_if_t<!std ::is_lvalue_reference_v<T>>;

namespace detail {
/** The minimal interface for raw data types
*/
Expand Down Expand Up @@ -152,6 +156,14 @@ class Frame {
template <typename FrameDataT>
Frame(std::unique_ptr<FrameDataT>);

/** Frame constructor from (almost) arbitrary raw data.
*
* This r-value overload is mainly present for enabling the python bindings,
* where cppyy seems to strip the std::unique_ptr somewhere in the process
*/
template <typename FrameDataT, typename = EnableIfRValue<FrameDataT>>
Frame(FrameDataT&&);

// The frame is a non-copyable type
Frame(const Frame&) = delete;
Frame& operator=(const Frame&) = delete;
Expand All @@ -167,6 +179,11 @@ class Frame {
template <typename CollT, typename = EnableIfCollection<CollT>>
const CollT& get(const std::string& name) const;

/** Get a collection from the Frame. This is the pointer-to-base version for
* type-erased access (e.g. python interface)
*/
const podio::CollectionBase* get(const std::string& name) const;

/** (Destructively) move a collection into the Frame and get a const reference
* back for further use
*/
Expand Down Expand Up @@ -259,6 +276,10 @@ template <typename FrameDataT>
Frame::Frame(std::unique_ptr<FrameDataT> data) : m_self(std::make_unique<FrameModel<FrameDataT>>(std::move(data))) {
}

template <typename FrameDataT, typename>
Frame::Frame(FrameDataT&& data) : Frame(std::make_unique<FrameDataT>(std::move(data))) {
}

template <typename CollT, typename>
const CollT& Frame::get(const std::string& name) const {
const auto* coll = dynamic_cast<const CollT*>(m_self->get(name));
Expand All @@ -270,6 +291,10 @@ const CollT& Frame::get(const std::string& name) const {
return emptyColl;
}

const podio::CollectionBase* Frame::get(const std::string& name) const {
return m_self->get(name);
}

void Frame::put(std::unique_ptr<podio::CollectionBase> coll, const std::string& name) {
const auto* retColl = m_self->put(std::move(coll), name);
if (!retColl) {
Expand Down
5 changes: 5 additions & 0 deletions include/podio/ROOTFrameReader.h
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include <iostream>
#include <memory>
#include <string>
#include <string_view>
#include <tuple>
#include <utility>
#include <vector>
Expand Down Expand Up @@ -70,10 +71,14 @@ class ROOTFrameReader {
/// Returns number of entries for the given name
unsigned getEntries(const std::string& name) const;

/// Get the build version of podio that has been used to write the current file
podio::version::Version currentFileVersion() const {
return m_fileVersion;
}

/// Get the names of all the availalable Frame categories in the current file(s)
std::vector<std::string_view> getAvailableCategories() const;

private:
/**
* Helper struct to group together all the necessary state to read / process a
Expand Down
5 changes: 5 additions & 0 deletions include/podio/SIOBlock.h
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#include <map>
#include <memory>
#include <string>
#include <string_view>

namespace podio {

Expand Down Expand Up @@ -230,6 +231,10 @@ class SIOFileTOCRecord {
*/
PositionType getPosition(const std::string& name, unsigned iEntry = 0) const;

/** Get all the record names that are stored in this TOC record
*/
std::vector<std::string_view> getRecordNames() const;

private:
friend struct SIOFileTOCRecordBlock;

Expand Down
5 changes: 5 additions & 0 deletions include/podio/SIOFrameReader.h
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

#include <memory>
#include <string>
#include <string_view>
#include <unordered_map>

namespace podio {
Expand Down Expand Up @@ -44,10 +45,14 @@ class SIOFrameReader {

void openFile(const std::string& filename);

/// Get the build version of podio that has been used to write the current file
podio::version::Version currentFileVersion() const {
return m_fileVersion;
}

/// Get the names of all the availalable Frame categories in the current file(s)
std::vector<std::string_view> getAvailableCategories() const;

private:
void readPodioHeader();

Expand Down
29 changes: 19 additions & 10 deletions python/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,28 @@ SET(podio_PYTHON_INSTALLDIR python)
SET(podio_PYTHON_INSTALLDIR ${podio_PYTHON_INSTALLDIR} PARENT_SCOPE)
SET(podio_PYTHON_DIR ${CMAKE_CURRENT_LIST_DIR} PARENT_SCOPE)

file(GLOB to_install *.py figure.txt)

# remove test_*.py file from being installed
foreach(file_path ${to_install})
get_filename_component(file_name ${file_path} NAME)
string(REGEX MATCH test_.*\\.py$ FOUND_PY_TEST ${file_name})
if (NOT "${FOUND_PY_TEST}" STREQUAL "")
list(REMOVE_ITEM to_install "${file_path}")
endif()
endforeach()
set(to_install
podio_class_generator.py
figure.txt
EventStore.py)

install(FILES ${to_install} DESTINATION ${podio_PYTHON_INSTALLDIR})

if(ENABLE_SIO)
install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/podio
DESTINATION ${podio_PYTHON_INSTALLDIR}
REGEX test_.*\\.py$ EXCLUDE # Do not install test files
PATTERN __pycache__ EXCLUDE # Or pythons caches
)
else()
install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/podio
DESTINATION ${podio_PYTHON_INSTALLDIR}
REGEX test_.*\\.py$ EXCLUDE # Do not install test files
PATTERN __pycache__ EXCLUDE # Or pythons caches
REGEX .*sio.*\\.py$ EXCLUDE # All things sio related
)
endif()

#--- install templates ---------------------------------------------------------
install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/templates
DESTINATION ${podio_PYTHON_INSTALLDIR})
139 changes: 4 additions & 135 deletions python/EventStore.py
Original file line number Diff line number Diff line change
@@ -1,137 +1,6 @@
"""Python EventStore for reading files with podio generated datamodels"""
"""Legacy import wrapper for EventStore."""

import warnings
warnings.warn("You are using the legacy EventStore import. Switch to 'from podio import EventStore'", FutureWarning)

from ROOT import gSystem
gSystem.Load("libpodioPythonStore") # noqa: E402
from ROOT import podio # noqa: E402 # pylint: disable=wrong-import-position


def size(self):
"""Override size function that can be attached as __len__ method to
collections"""
return self.size()


def getitem(self, key):
"""Override getitem function that can be attached as __getitem__ method to
collections (see below why this is necessary sometimes)"""
return self.at(key)


class EventStore:
'''Interface to events in an podio root file.
Example of use:
events = EventStore(["example.root", "example1.root"])
for iev, store in islice(enumerate(events), 0, 2):
particles = store.get("GenParticle");
for i, p in islice(enumerate(particles), 0, 5):
print "particle ", i, p.ID(), p.P4().Pt
'''

def __init__(self, filenames):
'''Create an event list from the podio root file.
Parameters:
filenames: list of root files
you can of course provide a list containing a single
root file. you could use the glob module to get all
files matching a wildcard pattern.
'''
if isinstance(filenames, str):
filenames = (filenames,)
self.files = filenames
self.stores = []
self.current_store = None
for fname in self.files:
store = podio.PythonEventStore(fname)
if store.isZombie():
raise ValueError(fname + ' does not exist.')
store.name = fname
if self.current_store is None:
self.current_store = store
self.stores.append((store.getEntries(), store))

def __str__(self):
result = "Content:"
for item in self.current_store.getCollectionNames():
result += f"\n\t{item}"
return result

def get(self, name):
'''Returns a collection.
Parameters:
name: name of the collection in the podio root file.
'''
coll = self.current_store.get(name)
# adding length function
coll.__len__ = size
# enabling the use of [] notation on the collection
# cppyy defines the __getitem__ method if the underlying c++ class has an operator[]
# method. For some reason they do not conform to the usual signature and only
# pass one argument to the function they call. Here we simply check if we have to
# define the __getitem__ for the collection.
if not hasattr(coll, '__getitem__'):
coll.__getitem__ = getitem
return coll

def collections(self):
"""Return list of all collection names."""
return [str(c) for c in self.current_store.getCollectionNames()]

def metadata(self):
"""Get the metadata of the current event as GenericParameters"""
return self.current_store.getEventMetaData()

def isValid(self):
"""Check if the EventStore is in a valid state"""
return self.current_store is not None and self.current_store.isValid()

# def __getattr__(self, name):
# '''missing attributes are taken from self.current_store'''
# if name != 'current_store':
# return getattr(self.current_store, name)
# else:
# return None

def current_filename(self):
'''Returns the name of the current file.'''
if self.current_store is None:
return None
return self.current_store.fname

def __enter__(self):
return self

def __exit__(self, exception_type, exception_val, trace):
for store in self.stores:
store[1].close()

def __iter__(self):
'''iterate on events in the tree.
'''
for _, store in self.stores:
self.current_store = store
for _ in range(store.getEntries()):
yield store
store.endOfEvent()

def __getitem__(self, evnum):
'''Get event number evnum'''
current_store = None
rel_evnum = evnum
for nev, store in self.stores:
if rel_evnum < nev:
current_store = store
break
rel_evnum -= nev
if current_store is None:
raise ValueError('event number too large: ' + str(evnum))
self.current_store = current_store
self.current_store.goToEvent(rel_evnum)
return self

def __len__(self):
'''Returns the total number of events in all files.'''
nevts_all_files = 0
for nev, _ in self.stores:
nevts_all_files += nev
return nevts_all_files
from podio import EventStore # noqa: F401 # pylint: disable=wrong-import-position, unused-import
1 change: 1 addition & 0 deletions python/__init__.py.in
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
__version__ = '@podio_VERSION@'
Loading

0 comments on commit e8ae38b

Please sign in to comment.