lheutils 0.0.2__py3-none-any.whl

lheutils/cli/lhefilter.py

@@ -0,0 +1,371 @@
+#!/usr/bin/env python3
+# ruff: noqa: SIM103
+"""
+CLI tool to filter LHE files based on various criteria.
+
+This tool filters Les Houches Event (LHE) files based on process ID,
+particle PDG IDs (incoming/outgoing), and event numbers.
+"""
+
+import argparse
+import sys
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Optional
+
+import pylhe
+
+from lheutils.cli.util import create_base_parser
+
+
+def matches_process_filter(
+    event: pylhe.LHEEvent,
+    process_ids: Optional[set[int]],
+    exclude_process_ids: Optional[set[int]],
+) -> bool:
+    """Check if event matches process ID filters."""
+    if process_ids is not None and event.eventinfo.pid not in process_ids:
+        return False
+    if exclude_process_ids is not None and event.eventinfo.pid in exclude_process_ids:
+        return False
+    return True
+
+
+def matches_particle_filter(
+    event: pylhe.LHEEvent,
+    incoming_pdgids: Optional[set[int]],
+    exclude_incoming_pdgids: Optional[set[int]],
+    outgoing_pdgids: Optional[set[int]],
+    exclude_outgoing_pdgids: Optional[set[int]],
+) -> bool:
+    """Check if event matches particle PDG ID filters."""
+    # Get incoming particles (status -1)
+    incoming_particles = [p for p in event.particles if p.status == -1]
+    incoming_ids = {p.id for p in incoming_particles}
+
+    # Get outgoing particles (status 1)
+    outgoing_particles = [p for p in event.particles if p.status == 1]
+    outgoing_ids = {p.id for p in outgoing_particles}
+
+    # Check incoming particle filters
+    if incoming_pdgids is not None and not incoming_ids.intersection(incoming_pdgids):
+        return False
+
+    if exclude_incoming_pdgids is not None and incoming_ids.intersection(
+        exclude_incoming_pdgids
+    ):
+        return False
+
+    # Check outgoing particle filters
+    if outgoing_pdgids is not None and not outgoing_ids.intersection(outgoing_pdgids):
+        return False
+
+    if exclude_outgoing_pdgids is not None and outgoing_ids.intersection(
+        exclude_outgoing_pdgids
+    ):
+        return False
+    return True
+
+
+def matches_event_filter(
+    event_index: int,
+    include_event_ranges: Optional[list[tuple[int, int]]],
+    exclude_event_ranges: Optional[list[tuple[int, int]]],
+) -> bool:
+    """Check if event matches event number filters."""
+    # event_index is 0-based, but user input is 1-based
+    event_number = event_index + 1
+
+    # Check include ranges
+    if include_event_ranges is not None:
+        matches_include = False
+        for start, end in include_event_ranges:
+            if end == -1:  # Open upper bound
+                if event_number >= start:
+                    matches_include = True
+                    break
+            elif start <= event_number <= end:
+                matches_include = True
+                break
+        if not matches_include:
+            return False
+
+    # Check exclude ranges
+    if exclude_event_ranges is not None:
+        for start, end in exclude_event_ranges:
+            if end == -1:  # Open upper bound
+                if event_number >= start:
+                    return False
+            elif start <= event_number <= end:
+                return False
+
+    return True
+
+
+def filter_lhe_file(
+    input_file: str,
+    rwgt: bool,
+    weights: bool,
+    output_file: Optional[str] = None,
+    process_ids: Optional[set[int]] = None,
+    exclude_process_ids: Optional[set[int]] = None,
+    incoming_pdgids: Optional[set[int]] = None,
+    exclude_incoming_pdgids: Optional[set[int]] = None,
+    outgoing_pdgids: Optional[set[int]] = None,
+    exclude_outgoing_pdgids: Optional[set[int]] = None,
+    include_event_ranges: Optional[list[tuple[int, int]]] = None,
+    exclude_event_ranges: Optional[list[tuple[int, int]]] = None,
+) -> None:
+    """Filter an LHE file based on the given criteria."""
+    try:
+        # Read the input LHE file
+        if input_file == "-":
+            lhefile = pylhe.LHEFile.frombuffer(sys.stdin)
+        else:
+            lhefile = pylhe.LHEFile.fromfile(input_file)
+
+        # Filter events
+        def _generator() -> Iterable[pylhe.LHEEvent]:
+            for event_index, event in enumerate(lhefile.events):
+                # Apply all filters
+                if (
+                    matches_process_filter(event, process_ids, exclude_process_ids)
+                    and matches_particle_filter(
+                        event,
+                        incoming_pdgids,
+                        exclude_incoming_pdgids,
+                        outgoing_pdgids,
+                        exclude_outgoing_pdgids,
+                    )
+                    and matches_event_filter(
+                        event_index, include_event_ranges, exclude_event_ranges
+                    )
+                ):
+                    yield event
+
+        # Create filtered LHE file
+        filtered_lhefile = pylhe.LHEFile(init=lhefile.init, events=_generator())
+
+        # Output the result
+        if output_file:
+            filtered_lhefile.tofile(output_file, rwgt=rwgt, weights=weights)
+        else:
+            # Write to stdout
+            filtered_lhefile.write(sys.stdout, rwgt=rwgt, weights=weights)
+
+    except FileNotFoundError:
+        if input_file == "-":
+            print("Error: Unable to read from stdin", file=sys.stderr)
+        else:
+            print(f"Error: File '{input_file}' not found", file=sys.stderr)
+        sys.exit(1)
+    except Exception as e:
+        source = "stdin" if input_file == "-" else f"file '{input_file}'"
+        print(f"Error processing {source}: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+def parse_int_list(value: str) -> set[int]:
+    """Parse comma-separated list of integers."""
+    try:
+        return {int(x.strip()) for x in value.split(",")}
+    except ValueError as e:
+        err = f"Invalid integer list: {value}"
+        raise argparse.ArgumentTypeError(err) from e
+
+
+def parse_range_list(value: str) -> list[tuple[int, int]]:
+    """Parse comma-separated list of integers and ranges.
+
+    Returns list of (start, end) tuples where -1 indicates open bound.
+
+    Supports:
+    - Individual numbers: 5 -> [(5, 5)]
+    - Ranges: 5-10 -> [(5, 10)] (inclusive)
+    - Lower bound: 5- -> [(5, -1)] (from 5 to end)
+    - Upper bound: -10 -> [(1, 10)] (from start to 10)
+    - Mixed: 1,5-10,15-,20,-25 -> [(1, 1), (5, 10), (15, -1), (20, 20), (1, 25)]
+    """
+    result = []
+
+    try:
+        for vitem in value.split(","):
+            item = vitem.strip()
+
+            if "-" not in item:
+                # Single number
+                num = int(item)
+                result.append((num, num))
+            elif item.startswith("-"):
+                # Upper bound: -N
+                upper = int(item[1:])
+                result.append((1, upper))
+            elif item.endswith("-"):
+                # Lower bound: N- (open range)
+                lower = int(item[:-1])
+                result.append((lower, -1))  # -1 indicates open upper bound
+            else:
+                # Range: N-M
+                parts = item.split("-")
+                if len(parts) == 2:
+                    lower, upper = int(parts[0]), int(parts[1])
+                    if lower > upper:
+                        err = f"Invalid range: {item} (start > end)"
+                        raise ValueError(err)
+                    result.append((lower, upper))
+                else:
+                    err = f"Invalid range format: {item}"
+                    raise ValueError(err)
+
+    except ValueError as e:
+        err = f"Invalid range specification: {value} ({e})"
+        raise argparse.ArgumentTypeError(err) from e
+
+    return result
+
+
+def main() -> None:
+    """Main CLI function."""
+    parser = create_base_parser(
+        description="Filter LHE files based on process ID, particle PDG IDs, and event numbers",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  lhefilter input.lhe -o filtered.lhe --process-p 81,82
+  lhefilter input.lhe --PROCESS 91 --incoming 21 --outgoing 11,-11
+  lhefilter input.lhe --events 1,5,10 --outgoing 13,-13
+  lhefilter input.lhe --EVENTS 7 --incoming 2,-2
+  lhefilter input.lhe.gz --out 6,-6 | gzip > filtered.lhe.gz
+  lhefilter input.lhe --events 10-20 --outgoing 11,-11
+  lhefilter input.lhe --events 50- --EVENTS 55-60
+  cat input.lhe | lhefilter --outgoing 11,-11
+  zcat input.lhe.gz | lhefilter --process-p 81 --outgoing 13,-13
+
+Process ID filters:
+  --process-p ID[,ID...]    Include only events with these process IDs
+  --PROCESS/-P ID[,ID...]   Exclude events with these process IDs
+
+Particle PDG ID filters:
+  --incoming/-in ID[,ID...] Include events containing these incoming particles
+  --INCOMING/-IN ID[,ID...] Exclude events containing these incoming particles
+  --outgoing/-out ID[,ID...] Include events containing these outgoing particles
+  --OUTGOING/-OUT ID[,ID...] Exclude events containing these outgoing particles
+
+Event filters:
+  --events RANGE[,RANGE...] Include events in these ranges (1-indexed)
+                            Supports: N (single), N-M (range), N- (from N to end), -M (up to M)
+  --EVENTS RANGE[,RANGE...] Exclude events in these ranges (1-indexed)
+                            Supports: N (single), N-M (range), N- (from N to end), -M (up to M)
+
+Note: Multiple filters are combined with AND logic.
+      PDG ID 0 can be used as wildcard for any particle.
+      Open ranges (N-) use -1 internally to represent no upper bound.
+        """,
+    )
+
+    parser.add_argument(
+        "input", nargs="?", default="-", help="Input LHE file (default: stdin)"
+    )
+    parser.add_argument("-o", "--output", help="Output file (default: write to stdout)")
+
+    # Process ID filters
+    parser.add_argument(
+        "--process-p",
+        type=parse_int_list,
+        metavar="ID[,ID...]",
+        help="Include only events with these process IDs",
+    )
+    parser.add_argument(
+        "--PROCESS",
+        "-P",
+        type=parse_int_list,
+        metavar="ID[,ID...]",
+        help="Exclude events with these process IDs",
+    )
+
+    # Incoming particle filters
+    parser.add_argument(
+        "--incoming",
+        "-in",
+        type=parse_int_list,
+        metavar="PDGID[,PDGID...]",
+        help="Include events containing these incoming particles",
+    )
+    parser.add_argument(
+        "--INCOMING",
+        "-IN",
+        type=parse_int_list,
+        metavar="PDGID[,PDGID...]",
+        help="Exclude events containing these incoming particles",
+    )
+
+    # Outgoing particle filters
+    parser.add_argument(
+        "--outgoing",
+        "-out",
+        type=parse_int_list,
+        metavar="PDGID[,PDGID...]",
+        help="Include events containing these outgoing particles",
+    )
+    parser.add_argument(
+        "--OUTGOING",
+        "-OUT",
+        type=parse_int_list,
+        metavar="PDGID[,PDGID...]",
+        help="Exclude events containing these outgoing particles",
+    )
+
+    # Event range filters
+    parser.add_argument(
+        "--events",
+        type=parse_range_list,
+        metavar="RANGE[,RANGE...]",
+        help="Include events in these ranges (1-indexed). Supports: N (single), N-M (range), N- (from N), -M (up to M)",
+    )
+    parser.add_argument(
+        "--EVENTS",
+        type=parse_range_list,
+        metavar="RANGE[,RANGE...]",
+        help="Exclude events in these ranges (1-indexed). Supports: N (single), N-M (range), N- (from N), -M (up to M)",
+    )
+
+    parser.add_argument(
+        "--rwgt",
+        action="store_true",
+        help="Use rwgt section if present in the input file",
+    )
+
+    parser.add_argument(
+        "--no-weights",
+        action="store_true",
+        help="Do not preserve event weights in output file",
+    )
+
+    args = parser.parse_args()
+
+    # Validate input file (skip validation for stdin)
+    if args.input != "-":
+        input_path = Path(args.input)
+        if not input_path.exists():
+            print(f"Error: Input file '{args.input}' does not exist", file=sys.stderr)
+            sys.exit(1)
+
+    # Call the filtering function
+    filter_lhe_file(
+        rwgt=args.rwgt,
+        weights=not args.no_weights,
+        input_file=args.input,
+        output_file=args.output,
+        process_ids=args.process_p,
+        exclude_process_ids=args.PROCESS,
+        incoming_pdgids=args.incoming,
+        exclude_incoming_pdgids=args.INCOMING,
+        outgoing_pdgids=args.outgoing,
+        exclude_outgoing_pdgids=args.OUTGOING,
+        include_event_ranges=args.events,
+        exclude_event_ranges=args.EVENTS,
+    )
+
+
+if __name__ == "__main__":
+    main()

lheutils/cli/lheinfo.py

@@ -0,0 +1,306 @@
+#!/usr/bin/env python3
+"""
+CLI tool to display information about LHE files.
+
+This tool analyzes Les Houches Event (LHE) files and displays relevant information
+including number of events, process information, particle combinations, etc.
+"""
+
+import argparse
+import json
+import sys
+import warnings
+from collections import Counter
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TextIO, Union
+
+import yaml  # type: ignore[import-untyped]
+
+import pylhe
+from pylhe.cli.util import dataclass_with_properties_to_dict
+
+from lheutils.cli.util import create_base_parser
+
+
+@dataclass
+class LHEChannel:
+    """Information about a single channel in an LHE file."""
+
+    incoming_pdgid: list[int]
+    outgoing_pdgid: list[int]
+    num_events: int
+
+
+@dataclass
+class LHEProcess:
+    """Information about a single process in an LHE file."""
+
+    procId: int
+    xSection: float
+    error: float
+    channels: list[LHEChannel]
+
+
+@dataclass
+class LHEInfo:
+    """Information about a single LHE file."""
+
+    filepath: str
+    beamA: int
+    energyA: float
+    beamB: int
+    energyB: float
+    weight_groups: dict[str, int]
+    num_events: int
+    negative_weighted_events: int
+    process_info: list[LHEProcess]
+
+    @property
+    def negative_weighted_events_ratio(self) -> float:
+        """Ratio of negative weighted events to total events."""
+        return (
+            self.negative_weighted_events / self.num_events
+            if self.num_events > 0
+            else 0.0
+        )
+
+    def __str__(self) -> str:
+        lines = []
+        lines.append("-" * 60)
+        lines.append(f"File: {self.filepath}")
+
+        # Beam information
+        lines.append(f"Beam A: {self.beamA} @ {self.energyA} GeV")
+        lines.append(f"Beam B: {self.beamB} @ {self.energyB} GeV")
+        # Weight groups
+        if self.weight_groups:
+            lines.append("  Weight Groups:")
+            for name, count in self.weight_groups.items():
+                lines.append(f"    {name}: {count} weights")
+        # Number of events
+        lines.append(
+            f"Number of events: {self.num_events} (negative: {self.negative_weighted_events_ratio:.2%})"
+        )
+
+        # Process information
+        processes = self.process_info
+        if processes:
+            for proc in processes:
+                lines.append(
+                    f"Process {proc.procId} cross-section: ({proc.xSection:.3e} +- {proc.error:.3e}) pb"
+                )
+
+                channels = proc.channels
+                if channels:
+                    # Sort channels by num_events in descending order
+                    sorted_channels = sorted(
+                        channels, key=lambda ch: ch.num_events, reverse=True
+                    )
+                    for channel in sorted_channels:
+                        percentage = 100 * channel.num_events / self.num_events
+                        lines.append(
+                            f"  {channel.incoming_pdgid} -> {channel.outgoing_pdgid}: {channel.num_events:,} events ({percentage:.1f}%)"
+                        )
+
+        return "\n".join(lines)
+
+
+def get_lheinfo(filepath_or_fileobj: Union[str, TextIO]) -> LHEInfo:
+    # Read LHE file
+    if isinstance(filepath_or_fileobj, str):
+        lhefile = pylhe.LHEFile.fromfile(filepath_or_fileobj)
+        file_display_name = filepath_or_fileobj
+    else:
+        lhefile = pylhe.LHEFile.frombuffer(filepath_or_fileobj)
+        file_display_name = "<stdin>"
+    init_info = lhefile.init.initInfo
+
+    initial_final_combinations: Counter[
+        tuple[int, tuple[int, ...], tuple[int, ...]]
+    ] = Counter()
+
+    num_events = 0
+    num_negative_weighted_events = 0
+    for event in lhefile.events:
+        num_events += 1
+        if event.eventinfo.weight < 0:
+            num_negative_weighted_events += 1
+        initial = []
+        final = []
+
+        for particle in event.particles:
+            if particle.status == -1:  # Incoming particles
+                initial.append(particle.id)
+            elif particle.status == 1:  # Outgoing particles
+                final.append(particle.id)
+
+        # Count initial particles
+        initial_tuple = tuple(sorted(initial))
+
+        # Count final particles
+        final_tuple = tuple(sorted(final))
+
+        # Count initial -> final combinations
+        combination = (event.eventinfo.pid, initial_tuple, final_tuple)
+        initial_final_combinations[combination] += 1
+
+    return LHEInfo(
+        filepath=file_display_name,
+        beamA=init_info.beamA,
+        energyA=init_info.energyA,
+        beamB=init_info.beamB,
+        energyB=init_info.energyB,
+        weight_groups={
+            name: len(wg.weights) for name, wg in lhefile.init.weightgroup.items()
+        },
+        num_events=num_events,
+        negative_weighted_events=num_negative_weighted_events,
+        process_info=[
+            LHEProcess(
+                procId=proc.procId,
+                xSection=proc.xSection,
+                error=proc.error,
+                channels=[
+                    LHEChannel(
+                        incoming_pdgid=list(incoming_pdgid),
+                        outgoing_pdgid=list(outgoing_pdgid),
+                        num_events=count,
+                    )
+                    for (
+                        pid,
+                        incoming_pdgid,
+                        outgoing_pdgid,
+                    ), count in initial_final_combinations.items()
+                    if pid == proc.procId
+                ],
+            )
+            for proc in lhefile.init.procInfo
+        ],
+    )
+
+
+@dataclass
+class LHESummary:
+    """Summary information from multiple LHE files."""
+
+    files: list[LHEInfo]
+
+    @property
+    def total_events(self) -> int:
+        total_events = 0
+        for lheinfo in self.files:
+            total_events += lheinfo.num_events
+        return total_events
+
+    @property
+    def total_negative_weighted_events(self) -> int:
+        total_negative_weighted_events = 0
+        for lheinfo in self.files:
+            total_negative_weighted_events += lheinfo.negative_weighted_events
+        return total_negative_weighted_events
+
+    @property
+    def negative_weighted_events_ratio(self) -> float:
+        """Ratio of negative weighted events to total events."""
+        return (
+            self.total_negative_weighted_events / self.total_events
+            if self.total_events > 0
+            else 0.0
+        )
+
+    def __str__(self) -> str:
+        lines = []
+        for lheinfo in self.files:
+            lines.append(str(lheinfo))
+        lines.append("=" * 60)
+        lines.append(
+            f"Total number of events: {self.total_events} (negative: {self.negative_weighted_events_ratio:.2%})"
+        )
+        lines.append("=" * 60)
+        return "\n".join(lines)
+
+
+def get_lhesummary(
+    filepaths_or_fileobjs: list[Union[str, TextIO]],
+) -> LHESummary:
+    lheinfos = []
+    # Analyze all files
+    for filepath_or_fileobj in filepaths_or_fileobjs:
+        lheinfo = get_lheinfo(filepath_or_fileobj)
+        lheinfos.append(lheinfo)
+
+    return LHESummary(files=lheinfos)
+
+
+def main() -> None:
+    """Main CLI function."""
+    parser = create_base_parser(
+        description="Display information about LHE files",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  lheinfo file.lhe                      # Analyze single file (plain format)
+  cat file.lhe | lheinfo                # Read from stdin
+  lheinfo *.lhe                         # Analyze multiple files
+  lheinfo file1.lhe --format=json       # Output results in JSON format
+  lheinfo file1.lhe --format=yaml       # Output results in YAML format
+        """,
+    )
+
+    parser.add_argument(
+        "files",
+        nargs="*",
+        help="LHE file(s) to analyze (or read from stdin if not provided)",
+    )
+
+    parser.add_argument(
+        "--format",
+        choices=["plain", "json", "yaml"],
+        default="plain",
+        help="Output format (default: plain)",
+    )
+
+    args = parser.parse_args()
+
+    # Check if reading from stdin
+    use_stdin = not args.files and not sys.stdin.isatty()
+
+    file_inputs: list[Union[str, TextIO]] = []
+    if use_stdin:
+        # Read from stdin
+        file_inputs += [sys.stdin]
+    else:
+        # Expand file paths
+        for pattern in args.files:
+            path = Path(pattern)
+            if path.exists():
+                if path.is_file():
+                    file_inputs.append(str(path))
+                else:
+                    warnings.warn(f"{pattern} is not a file", UserWarning, stacklevel=2)
+        if not file_inputs:
+            print("Error: No valid files found and no stdin data", file=sys.stderr)
+            sys.exit(1)
+
+    summary = get_lhesummary(file_inputs)
+    if args.format == "json":
+        print(
+            json.dumps(
+                dataclass_with_properties_to_dict(summary),
+                indent=2,
+            )
+        )
+    elif args.format == "yaml":
+        print(
+            yaml.dump(
+                dataclass_with_properties_to_dict(summary),
+                sort_keys=False,
+            )
+        )
+    else:
+        print(str(summary))
+
+
+if __name__ == "__main__":
+    main()