PythonToolHandle.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/

 
#include <AsgMessaging/MessageCheck.h>
#include <AsgMessaging/AsgMessaging.h>
#include <AsgTools/AsgTool.h>
#include <AsgTools/AsgToolConfig.h>
#include <AsgTools/ToolHandle.h>
#include <ColumnarInterfaces/IColumnarTool.h>
#include <ColumnarInterfaces/ColumnInfo.h>
#include <ColumnarToolWrapper/ToolColumnVectorMap.h>
#include <PATInterfaces/ISystematicsTool.h>
#include <PATInterfaces/SystematicSet.h>
#include <PATInterfaces/SystematicsUtil.h>
#include <span>
 
namespace columnar
{
 
  class PythonToolHandle final : public asg::AsgMessaging
  {
 
  public:

    PythonToolHandle() : asg::AsgMessaging("PythonToolHandle") {}

    void setTypeAndName (const std::string& typeAndName)
    {
      m_config.setTypeAndName (typeAndName);
    }

    template<typename T>
    void setProperty (const std::string& key, T&& value)
    {
      if (m_config.setProperty (key, std::forward<T> (value)).isFailure())
        throw std::runtime_error ("failed to set property: " + key);
    }

    [[nodiscard]] const asg::AsgToolConfig& getConfig () const
    {
      return m_config;
    }

    void preinitialize ()
    {
      ANA_MSG_DEBUG("preinitializing with " << m_toolHandle << " and cleanup " << m_cleanup);
      if (!m_config.makeTool (m_toolHandle, m_cleanup).isSuccess())
        throw std::runtime_error ("failed to create tool");
 
      ANA_MSG_DEBUG("m_config created tool" << m_toolHandle);
 
      m_tool = dynamic_cast<IColumnarTool*> (&*m_toolHandle);
 
      ANA_MSG_DEBUG("attempting to dynamically cast to IColumnarTool* gives " << m_tool);
      if (m_tool == nullptr)
        throw std::runtime_error ("The tool does not implement IColumnarTool. First, check to make sure you're in the ColumnarAnalysis release. Then, check to see if the tool inherits from ColumnarTool.");
      m_systTool = dynamic_cast<CP::ISystematicsTool*> (m_tool);
    }

    void renameContainers (const std::vector<std::pair<std::string,std::string>>& renames)
    {
      if (m_tool == nullptr)
        preinitialize ();
      {
        auto columnInfo = m_tool->getColumnInfo ();
        for (auto& [from, to] : renames)
        {
          for (auto& column : columnInfo)
          {
            if (column.name.starts_with (from) && (column.name.size() == from.size() || column.name[from.size()] == '.'))
            {
              std::string newName = to + column.name.substr (from.size());
              m_tool->renameColumn (column.name, newName);
            }
          }
        }
      }
      if (m_toolWrapper)
      {
        auto columnHeader = std::make_shared<ColumnVectorHeader> ();
        m_columnHeader = columnHeader;
        m_toolWrapper = std::make_shared<ToolColumnVectorMap> (*columnHeader, *m_tool);
        m_columns = std::make_unique<ColumnVectorData> (m_columnHeader.get());
      }
    }

    void initialize ()
    {
      if (m_tool == nullptr)
        preinitialize ();
 
      auto columnHeader = std::make_shared<ColumnVectorHeader> ();
      m_columnHeader = columnHeader;
      m_toolWrapper = std::make_shared<ToolColumnVectorMap> (*columnHeader, *m_tool);
      m_columns = std::make_unique<ColumnVectorData> (m_columnHeader.get());
    }

    void applySystematicVariation (const std::string& sysName)
    {
      // by convention setting a systematic on a non-systematics tool
      // will do nothing
      if (m_systTool == nullptr)
        return;
      if (!m_systTool->applySystematicVariation (CP::SystematicSet (sysName)).isSuccess())
        throw std::runtime_error ("failed to apply systematic variation");
    }

    template<typename CT>
    void setColumn (const std::string& key, std::size_t size, CT* dataPtr)
    {
      if (!m_columns)
        throw std::runtime_error ("tool not initialized");
      m_columns->setColumn (m_toolWrapper->getColumnIndex (key), size, dataPtr);
    }

    void setColumnVoid (const std::string& name, std::size_t size, const void *dataPtr, const std::type_info& type, bool isConst) {
      if (!m_columns)
        throw std::runtime_error ("tool not initialized");
      m_columns->setColumnVoid (m_toolWrapper->getColumnIndex (name), size, dataPtr, type, isConst);
    }

    void call ()
    {
      if (!m_columns)
        throw std::runtime_error ("no columns set");
      m_columns->checkData ();
      m_columns->callNoCheck (*m_tool);
      m_columns = std::make_unique<ColumnVectorData> (m_columnHeader.get());
    }

    [[nodiscard]] std::vector<ColumnInfo> getColumnInfo () const
    {
      if (!m_toolWrapper)
        throw std::runtime_error ("tool not initialized");
      return m_tool->getColumnInfo ();
    }

    std::vector<std::string> getColumnNames () const
    {
      if (!m_toolWrapper)
        throw std::runtime_error ("tool not initialized");
      return m_toolWrapper->getColumnNames ();
    }

    std::vector<std::string> getRecommendedSystematics () const
    {
      if (!m_systTool)
        return {""};
      std::vector<std::string> result;
      for (auto& sys : CP::make_systematics_vector (m_systTool->recommendedSystematics()))
        result.push_back (sys.name());
      return result;
    }
 
 

 
  private:
 
    asg::AsgToolConfig m_config;
    ToolHandle<asg::AsgTool> m_toolHandle;
    std::shared_ptr<void> m_cleanup;
 
    IColumnarTool* m_tool = nullptr;
    CP::ISystematicsTool* m_systTool = nullptr;
 
    std::shared_ptr<const ColumnVectorHeader> m_columnHeader;
    std::shared_ptr<const ToolColumnVectorMap> m_toolWrapper;
    std::unique_ptr<ColumnVectorData> m_columns;
  };
}