triggerflow 0.1.4__py3-none-any.whl → 0.2.4__py3-none-any.whl

triggerflow/core.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 import json
+import yaml
 import numpy as np
 import tarfile
 import importlib
@@ -8,6 +9,7 @@ from typing import Optional, Dict, Any, Union
 import shutil, warnings
 import importlib.resources as pkg_resources
 import triggerflow.templates 
+from triggerflow.interfaces.uGT import build_ugt_model
 
 
 class ModelConverter(ABC):
@@ -73,37 +75,31 @@ class NoOpConverter(ModelConverter):
 
 
 class HLS4MLStrategy(CompilerStrategy):
-    """HLS4ML compilation strategy for Keras models"""
-    
-    def compile(self, model, workspace: Path, config: Optional[Dict] = None, **kwargs) -> Any:
+    def compile(self, model, workspace: Path, config: Optional[Dict] = None) -> Any:
         import hls4ml
-        
+
         firmware_dir = workspace / "firmware"
         firmware_dir.mkdir(exist_ok=True)
-        
-        cfg = config or hls4ml.utils.config_from_keras_model(model, granularity="name")
 
-        hls_kwargs = {
-            "hls_config": cfg,
-            "output_dir": str(firmware_dir),
-            "io_type": "io_stream",
-            "backend": "Vitis"
-        }
-        hls_kwargs.update(kwargs)
+        hls_config = hls4ml.utils.config_from_keras_model(model, granularity="name")
+        hls_kwargs = {}
+
+        for key in ["project_name", "namespace", "io_type", "backend", "write_weights_txt"]:
+            if key in config:
+                hls_kwargs[key] = config[key]
 
         firmware_model = hls4ml.converters.convert_from_keras_model(
             model,
+            hls_config=hls_config,
+            output_dir=str(firmware_dir),
             **hls_kwargs
         )
 
         firmware_model.compile()
-        if shutil.which("vivado") is not None:
-             firmware_model.build()
-        else:
-            warnings.warn("Vivado not found in PATH. Firmware build failed.", UserWarning)
         firmware_model.save(workspace / "firmware_model.fml")
         return firmware_model
-    
+
+        
     def load_compiled_model(self, workspace: Path) -> Any:
         from hls4ml.converters import link_existing_project
         
@@ -113,27 +109,31 @@ class HLS4MLStrategy(CompilerStrategy):
 
 
 class ConiferStrategy(CompilerStrategy):
-    """Conifer compilation strategy for XGBoost models"""
-    
+    """Conifer compilation strategy for XGBoost models, unified config/workspace handling."""
+
     def compile(self, model, workspace: Path, config: Optional[Dict] = None) -> Any:
         import conifer
-        
+        import os
+
         firmware_dir = workspace / "firmware"
         firmware_dir.mkdir(exist_ok=True)
-        
-        cfg = config or conifer.backends.xilinxhls.auto_config()
-        firmware_model = conifer.converters.convert_from_xgboost(
-            model,
-            config=cfg,
-            output_dir=str(firmware_dir)
-        )
+
+        cfg = conifer.backends.xilinxhls.auto_config()
+        cfg['OutputDir'] = str(firmware_dir)
+        cfg['ProjectName'] = config['project_name']
+        cfg['XilinxPart'] = config['fpga_part']
+        cfg['ClockPeriod'] = config['clock_period']
+
+        if config:
+            for key, value in config.items():
+                cfg[key] = value
+
+        firmware_model = conifer.converters.convert_from_xgboost(model, config=cfg)
         firmware_model.compile()
-        if shutil.which("vivado") is not None:
-             firmware_model.build()
-        else:
-            warnings.warn("Vivado not found in PATH. Firmware build failed.", UserWarning)
-        firmware_model.save(workspace / "firmware_model.fml")
+        firmware_model.save(firmware_dir / "firmware_model.fml")
+
         return firmware_model
+
     
     def load_compiled_model(self, workspace: Path) -> Any:
         from conifer import load_model
@@ -202,13 +202,18 @@ class QONNXPredictor(ModelPredictor):
 class FirmwarePredictor(ModelPredictor):
     """Firmware-based model predictor"""
     
-    def __init__(self, firmware_model):
+    def __init__(self, firmware_model, compiler):
         if firmware_model is None:
             raise RuntimeError("Firmware model not built.")
         self.firmware_model = firmware_model
+        self.compiler = compiler
+    
     
     def predict(self, input_data: np.ndarray) -> np.ndarray:
-        return self.firmware_model.predict(input_data)
+        if self.compiler == "conifer":
+            return self.firmware_model.decision_function(input_data)
+        else:
+            return self.firmware_model.predict(input_data)
 
 
 class ConverterFactory:
@@ -360,33 +365,45 @@ class ModelSerializer:
             return model, input_name
         return None, None
 
+
 class TriggerModel:
-    """Main facade class that orchestrates model conversion, compilation, and inference"""
-    
-    def __init__(self, name: str, ml_backend: str, scales: dict, n_outputs:int, compiler: str,
-                 native_model: object, dataset_object: object, compiler_config: dict = None):
-        
-        if ml_backend.lower() not in ("keras", "xgboost"):
-            raise ValueError("Only Keras or XGBoost backends are currently supported.")
-        
-        self.name = name
-        self.ml_backend = ml_backend.lower()
-        self.scales = scales
-        self.n_outputs = n_outputs
-        self.compiler = compiler.lower()
+    def __init__(self, config: Union[str, Path, Dict], native_model, scales):
+        if isinstance(config, (str, Path)):
+            with open(config, "r") as f:
+                config = yaml.safe_load(f)
+        elif not isinstance(config, dict):
+            raise TypeError("config must be a dict or path to a YAML file")
+
         self.native_model = native_model
-        self.dataset_object = dataset_object
-        self.compiler_conifg = compiler_config
-        
+        self.scales = scales
+
+        self.compiler_cfg  = config.get("compiler", {})
+        self.subsystem_cfg = config.get("subsystem", {})
+
+        self.name = self.compiler_cfg.get("name", "model")
+        self.ml_backend = self.compiler_cfg.get("ml_backend", "").lower()
+        self.compiler = self.compiler_cfg.get("compiler", "").lower()
+
+        self.n_outputs = self.compiler_cfg.get("n_outputs")
+        self.unscaled_type = self.subsystem_cfg.get("unscaled_type", "ap_fixed<16,6>")
+
+        if self.ml_backend not in ("keras", "xgboost"):
+            raise ValueError("Unsupported backend")
+
         self.workspace_manager = WorkspaceManager()
-        self.converter = ConverterFactory.create_converter(ml_backend, compiler)
-        self.compiler_strategy = CompilerFactory.create_compiler(ml_backend, compiler)
-        
+        self.converter = ConverterFactory.create_converter(self.ml_backend, self.compiler)
+        self.compiler_strategy = CompilerFactory.create_compiler(self.ml_backend, self.compiler)
+
         self.firmware_model = None
         self.model_qonnx = None
         self.input_name = None
         
-        self.workspace_manager.setup_workspace(name, self.ml_backend, self.compiler)
+        
+        self.workspace_manager.setup_workspace(
+            self.name, 
+            self.ml_backend, 
+            self.compiler
+        )
     
     @property
     def workspace(self) -> Path:
@@ -403,8 +420,8 @@ class TriggerModel:
         """Get metadata dictionary"""
         return self.workspace_manager.metadata
     
-    def __call__(self, **compiler_kwargs):
-        """Execute the full model conversion and compilation pipeline"""
+    def __call__(self):
+        """Execute full model conversion and compilation pipeline using YAML config"""
         self.parse_dataset_object()
         
         # Save native model
@@ -421,21 +438,57 @@ class TriggerModel:
             self.input_name = self.model_qonnx.graph.input[0].name
             self.workspace_manager.add_artifact("qonnx", qonnx_path)
             self.workspace_manager.add_version({"qonnx": str(qonnx_path)})
+
         
         # Compile model
         self.firmware_model = self.compiler_strategy.compile(
             self.native_model,
             self.workspace_manager.workspace,
-            self.compiler_conifg, 
-            **compiler_kwargs
+            self.compiler_cfg,
+            **self.compiler_cfg.get("kwargs", {})
         )
         
         self.workspace_manager.add_artifact("firmware", self.workspace_manager.workspace / "firmware")
+        if self.compiler != "conifer" and self.scales is not None:
+            self.build_emulator(
+                self.scales['shifts'],
+                self.scales['offsets'],
+                self.n_outputs,
+                self.unscaled_type
+            )
 
-        self.build_emulator(self.scales['shifts'], self.scales['offsets'], self.n_outputs)
         
+        if shutil.which("vivado") is not None:
+            build_ugt_model(
+                templates_dir=self.subsystem_cfg.get("templates_dir", Path("templates")),
+                firmware_dir=self.workspace_manager.workspace / "firmware",
+                compiler = self.compiler,
+                model_name=self.name,
+                n_inputs=self.subsystem_cfg["n_inputs"],
+                n_outputs=self.subsystem_cfg.get("n_outputs", self.n_outputs),
+                nn_offsets=self.scales["offsets"],
+                nn_shifts=self.scales["shifts"],
+                muon_size=self.subsystem_cfg.get("muon_size", 0),
+                jet_size=self.subsystem_cfg.get("jet_size", 0),
+                egamma_size=self.subsystem_cfg.get("egamma_size", 0),
+                tau_size=self.subsystem_cfg.get("tau_size", 0),
+                output_type=self.subsystem_cfg.get("output_type", "result_t"),
+                offset_type=self.subsystem_cfg.get("offset_type", "ap_fixed<10,10>"),
+                shift_type=self.subsystem_cfg.get("shift_type", "ap_fixed<10,10>"),
+                object_features=self.subsystem_cfg.get("object_features"),
+                global_features=self.subsystem_cfg.get("global_features")
+            )
+        else:
+            warnings.warn(
+                "Vivado executable not found on the system PATH. "
+                "Skipping FW build. ",
+                UserWarning
+            )
+
+
         self.workspace_manager.add_artifact("firmware", self.workspace_manager.workspace / "firmware")
         self.workspace_manager.save_metadata()
+
     
     @staticmethod
     def parse_dataset_object():
@@ -457,23 +510,27 @@ class TriggerModel:
         predictor = SoftwarePredictor(self.native_model, self.ml_backend)
         return predictor.predict(input_data)
     
-    def qonnx_predict(self, input_data: np.ndarray) -> np.ndarray:
+    def qonnx_predict(self, input_data: np.ndarray) -> np.ndarray | None:
         """Make predictions using QONNX model"""
+        
         if self.model_qonnx is None:
-            raise RuntimeError("QONNX model not available")
+            warnings.warn(
+                "QONNX model is not available. Prediction skipped.",
+                UserWarning
+            )
+            return None 
+            
         predictor = QONNXPredictor(self.model_qonnx, self.input_name)
         return predictor.predict(input_data)
     
     def firmware_predict(self, input_data: np.ndarray) -> np.ndarray:
         """Make predictions using firmware model"""
-        predictor = FirmwarePredictor(self.firmware_model)
+        predictor = FirmwarePredictor(self.firmware_model, self.compiler)
         return predictor.predict(input_data)
     
-    def build_emulator(self, ad_shift: list, ad_offsets: list, n_outputs: int):
-        """
-        Create an emulator directory for this model.
-        Copies HLS sources and generates emulator scaffolding.
-        """
+    def build_emulator(self, ad_shift: list, ad_offsets: list, n_outputs: int, unscaled_type: str = "ap_fixed<16,6>"):
+        """Builds CMSSW emulator"""
+    
         emulator_dir = self.workspace / "emulator"
         emulator_dir.mkdir(exist_ok=True)
 
@@ -493,12 +550,13 @@ class TriggerModel:
                 "N_OUTPUTS": n_outputs,
                 "AD_SHIFT": ", ".join(map(str, ad_shift)),
                 "AD_OFFSETS": ", ".join(map(str, ad_offsets)),
+                "UNSCALED_TYPE": unscaled_type,  
             }
             self._render_template(scales_template_path, scales_out_path, context)
 
         with pkg_resources.path(triggerflow.templates, "model_template.cpp") as emulator_template_path:
             emulator_out_path = model_dir / "emulator.cpp"
-            self._render_template(emulator_template_path, emulator_out_path, context)
+            self._render_template(emulator_template_path, emulator_out_path, context) 
         
         with pkg_resources.path(triggerflow.templates, "makefile_version") as makefile_template_path:
             makefile_out_path = model_dir / "Makefile"

triggerflow/interfaces/uGT.py

@@ -0,0 +1,127 @@
+from pathlib import Path
+import shutil
+import pkg_resources
+from jinja2 import Template
+import re
+
+def _render_template(template_file: str, output_file: Path, context: dict):
+    with open(template_file, "r") as f:
+        template_text = f.read()
+
+    template = Template(template_text)
+    rendered = template.render(**context)
+
+    with open(output_file, "w") as f:
+        f.write(rendered)
+
+def build_ugt_model(
+    templates_dir: Path,
+    firmware_dir: Path,
+    compiler: str,
+    model_name: str,
+    n_inputs: int,
+    n_outputs: int,
+    nn_offsets: list,
+    nn_shifts: list,
+    muon_size: int,
+    jet_size: int,
+    egamma_size: int,
+    tau_size: int,
+    output_type: str = "result_t",
+    offset_type: str = "ap_fixed<10,10>",
+    shift_type: str = "ap_fixed<10,10>",
+    object_features: dict = None,
+    global_features: list = None
+):
+    """
+    Render uGT top func.
+    """
+
+
+    if object_features is None:
+        object_features = {
+            "muons": ["pt", "eta_extrapolated", "phi_extrapolated"],
+            "jets": ["et", "eta", "phi"],
+            "egammas": ["et", "eta", "phi"],
+            "taus": ["et", "eta", "phi"]
+        }
+
+    if global_features is None:
+        global_features = [
+            "et.et",
+            "ht.et",
+            "etmiss.et", "etmiss.phi",
+            "htmiss.et", "htmiss.phi",
+            "ethfmiss.et", "ethfmiss.phi",
+            "hthfmiss.et", "hthfmiss.phi"
+        ]
+
+    header_path = firmware_dir / "firmware" / f"{model_name}_project.h"
+    if compiler.lower() == "conifer":
+        output_layer = "score"
+        output_type = "score_arr_t"
+        header_path = firmware_dir / "firmware" / f"{model_name}_project.h"
+        removal_pattern = re.compile(
+            r',\s*score_t\s+tree_scores\[BDT::fn_classes\(n_classes\)\s*\*\s*n_trees\]', 
+            re.DOTALL
+        )
+        modified_content = removal_pattern.sub('', header_path.read_text(encoding='utf-8'))
+        header_path.write_text(modified_content, encoding='utf-8')
+        out = output_layer
+    else:
+        header_content = header_path.read_text(encoding='utf-8')
+        layer_pattern = re.compile(
+            r'result_t\s+(\w+)\[\d+\]\s*\)', 
+            re.DOTALL
+        )
+        match = layer_pattern.search(header_content)
+        layer_name = match.group(1) 
+        output_layer = f"{layer_name}[{n_outputs}]"
+        out = layer_name
+
+        
+    context = {
+    "MODEL_NAME": model_name,
+    "N_INPUTS": n_inputs,
+    "N_OUTPUTS": n_outputs,
+    "NN_OFFSETS": ", ".join(map(str, nn_offsets)),
+    "NN_SHIFTS": ", ".join(map(str, nn_shifts)),
+    "MUON_SIZE": muon_size,
+    "JET_SIZE": jet_size,
+    "EGAMMA_SIZE": egamma_size,
+    "TAU_SIZE": tau_size,
+    "OUTPUT_TYPE": output_type,
+    "OUTPUT_LAYER": output_layer,
+    "OUT": out,
+    "OFFSET_TYPE": offset_type,
+    "SHIFT_TYPE": shift_type,
+    "MUON_FEATURES": object_features["muons"],
+    "JET_FEATURES": object_features["jets"],
+    "EGAMMA_FEATURES": object_features["egammas"],
+    "TAU_FEATURES": object_features["taus"],
+    "GLOBAL_FEATURES": global_features
+    }
+
+    context_tcl = {
+    "MODEL_NAME": model_name,
+    }
+
+    out_path = firmware_dir / "firmware/model-gt.cpp"
+    
+    _render_template(f"{templates_dir}/model-gt.cpp", out_path, context)
+
+    out_path = firmware_dir / "firmware/build_ugt.tcl"
+    _render_template(f"{templates_dir}/build_ugt.tcl", out_path, context_tcl)
+
+    shutil.copy(f"{templates_dir}/data_types.h", firmware_dir / "firmware")
+ 
+    
+    subprocess.run(
+        ["vitis_hls", "-f", "build_ugt.tcl"],
+        cwd=firmware_dir/"firmware",
+        check=True
+    )
+
+
+
+