main298-python

Back to index.

# main298.py is a part of the PYTHIA event generator.
# Copyright (C) 2025 Artem Havryliuk and Torbjorn Sjostrand.
# PYTHIA is licenced under the GNU GPL v2 or later, see COPYING for details.
# Please respect the MCnet Guidelines, see GUIDELINES for details.

# Authors:
#            Philip Ilten
#            Artem Havryliuk
#            Jim Pivarski

# Keywords:
#            Python

# Example usage of the 'nextBatch' function in the PYTHIA event
# generator.  This script demonstrates the usage of different
# 'errorMode' values ('skip', 'fail', 'none', and a float) to generate
# events and handle failures.

# The 'nextBatch' function generates a batch of particle physics events,
# returning the results as an Awkward Array, with behavior determined
# by the specified error mode.

# To set the path to the Pythia 8 Python interface do either
# (in a shell prompt):
#      export PYTHONPATH=$(PREFIX_LIB):$PYTHONPATH
# or the following which sets the path from within Python.
import sys
cfg = open("Makefile.inc")
lib = "../lib"
for line in cfg:
    if line.startswith("PREFIX_LIB="): lib = line[11:-1]; break
sys.path.insert(0, lib)

# Import the Pythia module.
import pythia8

# Import the Awkward module.
import awkward as ak

#==========================================================================

def setup():
    """
    Basic Pythia setup for generation examples.
    """
    # Create the Pythia instance.
    py = pythia8.Pythia("", False)
    # Configure hard QCD.
    py.readString("HardQCD:all = on")
    # Turn on some weights.
    py.readString("VariationFrag:List = {var0 frag:aLund=0.6, "
                  "var1 frag:aLund=0.4}")
    # Turn off most of Pythia's messages.
    py.readString("Print:quiet = on")
    # Initialize and return.
    py.init()
    return py

#==========================================================================

def generateSkip():
    """
    Demonstrates event generation with errorMode='skip'. In this
    mode, failed events are simply skipped, and no record of them is
    included in the result. Only successful events are returned.
    """
    py = setup()
    nEvents = 10
    errorMode = "skip"
    run = py.nextBatch(nEvents, errorMode)
    print("Generated events (skip mode):", len(run))

#==========================================================================

def generateFail():
    """
    Demonstrates event generation with errorMode='fail'.
    In this mode, the process stops immediately if an event fails,
    raising a runtime exception.
    """
    py = setup()
    nEvents = 10
    errorMode = "fail"
    try:
        run = py.nextBatch(nEvents, errorMode)
        print("All events generated successfully!")
    except RuntimeError as e:
        print(f"Event generation failed: {e}")

#==========================================================================

def generateNone():
    """
    Demonstrates event generation with errorMode='none'. In this
    mode, failed events are included in the output array, but they are
    marked as invalid (e.g., None). The total number of events in the
    output matches the requested number.
    """
    py = setup()
    nEvents = 10
    errorMode = "none"
    run = py.nextBatch(nEvents, errorMode)
    print("Generated events (including invalid):", len(run))

#==========================================================================

def generateFloat():
    """
    Demonstrates event generation with errorMode as a float (e.g.,
    1.5). In this mode, additional attempts are allowed to ensure the
    desired number of successful events. The float value specifies a
    multiplier for the number of trials relative to the number of
    events.
    """
    py = setup()
    nEvents = 10
    errorMode = 1.5
    run = py.nextBatch(nEvents, errorMode)
    print("Generated events (float mode):", len(run))
    
#==========================================================================

def demoEvent():
    """
    Demonstrates how to work with the event record generated by
    'nextBatch'. The output is an Awkward Array containing particle
    data and event info. This function shows how to access various
    elements, such as 'id', 'x1', 'x2', 'mother1', and others, from
    the event record.
    """
    py = setup()
    nEvents = 5
    # Generate events, including failed ones as 'None'.
    errorMode = "none" 
    # Generate a batch of events.
    run = py.nextBatch(nEvents, errorMode)

    # Print the structure of the Awkward array.
    print("\nAwkward Array structure:")
    print(ak.fields(run))

    # Access particle-level information.
    print("\nParticle-level information:")
    prts = run.prt
    print(f"IDs of particles: {prts.id}")
    print(f"Mother 1: {prts.mother1}")
    print(f"Momentum px: {prts.p.px}")
    print(f"Momentum py: {prts.p.py}")

    # Accessing event-level information.
    print("\nEvent-level information:")
    info = run.info
    print(f"x1: {info.x1}")
    print(f"x2: {info.x2}")
    print(f"alphaS: {info.alphaS}")
    print(f"Q2Fac: {info.Q2Fac}")

    # Access weights.
    print("\nEvent weights:")
    print(info.weights)

    # Example: Accessing the first event's particle data.
    print("\nFirst event particle data:")
    print(prts[0])

#==========================================================================

def demoSlice():
    """
    Demonstrates how to work with and slice the particle arrays in the
    event record.  Shows examples of accessing specific particles,
    slicing arrays, and filtering data.
    """
    # Generate events with failed events as None.
    py = setup()
    nEvents = 5
    errorMode = "none"
    run = py.nextBatch(nEvents, errorMode)

    # Particle-level information.
    prts = run.prt
    print("\n### Particle Array Slicing and Manipulation ###")

    # Example: accessing particles from the first event.
    print("\nParticles from the first event:")
    firstEventPrts = prts[0]
    print(firstEventPrts)

    # Example: accessing only the first 5 particles of the first event.
    print("\nFirst 5 particles from the first event:")
    print(firstEventPrts[:5])

    # Example: slicing all events to retrieve the first particle of
    # each event.
    print("\nFirst particle of each event:")
    firstPrtEachEvent = prts[:, 0]
    print(firstPrtEachEvent)

    # Example: filtering particles with px > 50.0 across all events.
    print("\nParticles with px > 0.0182:")
    filterPrts = prts[prts.p.px > 0.0182]
    print(filterPrts)

    # Example: combining slices and conditions.
    print("\nParticles from the first event with px > 0.2:")
    filtered_first_event = firstEventPrts[firstEventPrts.p.px > 0.2]
    print(filtered_first_event)
    
#==========================================================================

def demoWeights():
    """
    Demonstrates how weights from the events are accessed and used.
    Shows examples of retrieving weights for each event and calculating
    weighted statistics.
    """
    # Generate events with failed events as None.
    py = setup()
    nEvents = 5
    errorMode = "none"
    run = py.nextBatch(nEvents, errorMode)

    # Access the weights for each event.
    print("\n### Demonstrating Event Weights ###")
    weights = run.info.weights

    # Print the weights for all events
    print("\nEvent weights:")
    for i, weight in enumerate(weights):
        print(f"Event {i + 1} weights: {weight}")

    # Example: sum over all events for each variation.
    sumWeights = ak.sum(weights, axis=0)
    print("\nSum over all events for each variation:")
    for i, sumWeight in enumerate(sumWeights):
        print(f"Variation {i + 1}: {sumWeight}")
    
#==========================================================================

def demoFilter():
    """
    Demonstrates how to filter out events where particles ('prt') are
    'None'.  Uses 'ak.is_none' to identify invalid events and provides
    examples of handling such cases.
    
    Important: We only get 'None' events in the output when
    'errorMode='none'' is used during event generation. In this mode,
    failed events are included in the result but marked as 'None',
    while successful events contain valid particle and event data.

    Example workflow:
    - Identify invalid events using 'ak.is_none'.
    - Count the number of invalid events.
    - Filter out invalid events to retain only valid events.
    - Optionally handle or analyze invalid events separately.
    """
    # Generate events with failed events as None.
    py = setup()
    nEvents = 10
    errorMode = "none"
    run = py.nextBatch(nEvents, errorMode)

    # Identify events where particles ('prt') are 'None'.
    print("\n### Filtering Invalid Events ###")
    invalid = ak.is_none(run.prt)
    print("\nInvalid events (where particles are None):")
    print(invalid)

    # Count the number of invalid events.
    nInvalid = ak.sum(invalid)
    print(f"\nNumber of invalid events: {nInvalid}")

    # Filter out invalid events.
    valid = run[~invalid]
    print(f"\nNumber of valid events after filtering: {len(valid)}")

    # Example: accessing valid particles from filtered events.
    print("\nParticles from the first valid event:")
    if len(valid) > 0:
        print(valid.prt[0])
    else:
        print("No valid events found.")

    # Example: handling invalid events (optional use case).
    if nInvalid > 0:
        print("\nInvalid events can be further analyzed or logged if needed.")
        # You can access invalid events like this.
        invalidEvents = run[invalid]
        print("Example of an invalid event:")
        print(invalidEvents)

# Run the examples to demonstrate different error modes.
if __name__ == "__main__":
    
    print("Running event generation with errorMode='skip'...")
    generateSkip()
    
    print("\nRunning event generation with errorMode='fail'...")
    generateFail()
    
    print("\nRunning event generation with errorMode='none'...")
    generateNone()
    
    print("\nRunning event generation with errorMode=1.5...")
    generateFloat()
    
    print("\nDemonstrating event record analysis...")
    demoEvent()
    
    print("\nDemonstrating event record analysis with slicing...")
    demoSlice()
    
    print("\nDemonstrating event record analysis with weights...")
    demoWeights()

    print("\nDemonstrating filtering invalid events...")
    demoFilter()