Smack My Higgs Up

The initial prompt:

Can you write Python script that can run in a Google Colab notebook that generates an image that contains all 24-bit RGB colors exactly once, and that represents the discovery of the Higgs boson. If possible, I want you to use actual scattering data/imagery of the Large Hadron Collider.

After some back-and-forth, resulted in the following script:

import os
from io import BytesIO

import numpy as np
import pandas as pd
from PIL import Image
import matplotlib.pyplot as plt
import requests

# ---------------- Global configuration ----------------

CANVAS_SIZE  = 4096                  # 4096 x 4096 = 2^24
TOTAL_COLORS = 1 << 24               # 16,777,216

CSV_URL   = "https://zenodo.org/records/8034556/files/higgs2011.csv?download=1"
CSV_PATH  = "higgs2011.csv"

# Optional CMS event-display image (press-style)
EVENT_IMG_URL = (
    "https://upload.wikimedia.org/wikipedia/commons/3/32/"
    "CMS-PHO-EVENTS-2012-007-1.png"
)
USE_EVENT_DISPLAY = False  # set to True if you want to blend this in

OUTPUT_ALLRGB = "higgs_allrgb_4096_notext.png"


# ---------------- Data loading ----------------

def load_higgs_dataframe(local_path=CSV_PATH, url=CSV_URL):
    """
    Load the CMS Higgs 4-lepton CSV.
    If the file doesn't exist locally, download it from Zenodo.
    """
    if not os.path.exists(local_path):
        print(f"Downloading Higgs CSV from\n  {url}")
        resp = requests.get(url, timeout=120)
        resp.raise_for_status()
        with open(local_path, "wb") as f:
            f.write(resp.content)
        print(f"Saved as {local_path}")

    df = pd.read_csv(local_path)
    print("Loaded CSV with columns:")
    print(df.columns.tolist())
    print("Number of events:", len(df))
    return df


# ---------------- Higgs data → η–φ maps ----------------

def build_higgs_event_maps_from_df(df, mass_target=125.0, grid_size=1024):
    """
    Use the CMS Higgs 4-lepton CSV (E, px, py, pz, eta, phi for 4 leptons)
    to build:

      - intensity_map: 2D float32 in [0,1]
      - colour_map:    3D float32 in [0,1], RGB

    Steps:
      1) Compute the 4-lepton invariant mass per event.
      2) Pick the event with mass closest to `mass_target` (GeV).
      3) Paint each lepton as a Gaussian blob in (eta, phi).
         Intensity ~ pT, colour hue ~ phi, brightness ~ pT.
    """
    H = W = grid_size

    # Energies & momenta
    E_cols  = [f"E{i}"   for i in range(1, 5)]
    px_cols = [f"px{i}"  for i in range(1, 5)]
    py_cols = [f"py{i}"  for i in range(1, 5)]
    pz_cols = [f"pz{i}"  for i in range(1, 5)]

    for cols in (E_cols, px_cols, py_cols, pz_cols):
        for c in cols:
            if c not in df.columns:
                raise RuntimeError(f"Expected column '{c}' in dataframe, but it's missing.")

    # Total 4-vector per event
    E_tot  = df[E_cols].sum(axis=1).to_numpy(dtype=np.float64)
    px_tot = df[px_cols].sum(axis=1).to_numpy(dtype=np.float64)
    py_tot = df[py_cols].sum(axis=1).to_numpy(dtype=np.float64)
    pz_tot = df[pz_cols].sum(axis=1).to_numpy(dtype=np.float64)

    p2_tot = px_tot**2 + py_tot**2 + pz_tot**2
    m2_4l  = E_tot**2 - p2_tot
    m2_4l  = np.clip(m2_4l, 0.0, None)
    m4l    = np.sqrt(m2_4l)  # GeV

    # Event closest to target mass (≈125 GeV = Higgs)
    idx = np.abs(m4l - mass_target).argmin()
    ev  = df.iloc[idx]

    print(f"Selected event index: {idx}")
    print(f"Run: {ev['Run']}, Event: {ev['Event']}")
    print(f"4-lepton mass ≈ {m4l[idx]:.2f} GeV")

    # Kinematics for the 4 leptons
    eta_vals = np.array([ev[f"eta{i}"] for i in range(1, 5)], dtype=np.float32)
    phi_vals = np.array([ev[f"phi{i}"] for i in range(1, 5)], dtype=np.float32)
    px_vals  = np.array([ev[f"px{i}"]  for i in range(1, 5)], dtype=np.float32)
    py_vals  = np.array([ev[f"py{i}"]  for i in range(1, 5)], dtype=np.float32)
    pt_vals  = np.sqrt(px_vals**2 + py_vals**2)  # transverse momentum

    # (eta, phi) grid
    eta_min = float(eta_vals.min() - 0.5)
    eta_max = float(eta_vals.max() + 0.5)
    if eta_max - eta_min < 1.0:
        mid = 0.5 * (eta_min + eta_max)
        eta_min, eta_max = mid - 0.5, mid + 0.5

    phi_min = -np.pi
    phi_max =  np.pi

    y, x = np.mgrid[0:H, 0:W].astype(np.float32)
    eta_grid = eta_min + (eta_max - eta_min) * (y / (H - 1))
    phi_grid = phi_min + (phi_max - phi_min) * (x / (W - 1))

    intensity = np.zeros((H, W), dtype=np.float32)
    rgb       = np.zeros((H, W, 3), dtype=np.float32)
    weight    = np.zeros((H, W), dtype=np.float32)

    sigma_eta = 0.3 * max(eta_max - eta_min, 1.0)
    sigma_phi = 0.3 * (phi_max - phi_min)

    pt_max = float(pt_vals.max())
    if pt_max <= 0:
        pt_max = 1.0

    # Paint each lepton as a Gaussian blob
    for eta, phi, pt in zip(eta_vals, phi_vals, pt_vals):
        d_eta = eta_grid - eta
        # periodic wrap in phi
        d_phi = np.arctan2(np.sin(phi_grid - phi), np.cos(phi_grid - phi))

        g = np.exp(-0.5 * ((d_eta / sigma_eta) ** 2 + (d_phi / sigma_phi) ** 2))

        # Intensity contribution ~ pT * Gaussian
        w = float(pt)
        intensity += w * g

        # Colour: hue from phi, brightness from pT
        hue = (float(phi) + np.pi) / (2 * np.pi)  # 0..1
        val = float(pt) / pt_max                  # 0..1
        sat = 1.0

        # Simple HSV -> RGB (scalar h, choose 1 branch)
        h = hue * 6.0
        i = int(np.floor(h)) % 6
        f = h - np.floor(h)
        p = val * (1 - sat)
        q = val * (1 - sat * f)
        t = val * (1 - sat * (1 - f))

        r_c = np.zeros_like(g, dtype=np.float32)
        g_c = np.zeros_like(g, dtype=np.float32)
        b_c = np.zeros_like(g, dtype=np.float32)

        if i == 0:
            r_c[:, :], g_c[:, :], b_c[:, :] = val, t,   p
        elif i == 1:
            r_c[:, :], g_c[:, :], b_c[:, :] = q,   val, p
        elif i == 2:
            r_c[:, :], g_c[:, :], b_c[:, :] = p,   val, t
        elif i == 3:
            r_c[:, :], g_c[:, :], b_c[:, :] = p,   q,   val
        elif i == 4:
            r_c[:, :], g_c[:, :], b_c[:, :] = t,   p,   val
        else:
            r_c[:, :], g_c[:, :], b_c[:, :] = val, p,   q

        rgb[..., 0] += g * r_c
        rgb[..., 1] += g * g_c
        rgb[..., 2] += g * b_c
        weight      += g

    # Normalise intensity to [0,1]
    inten_min = float(intensity.min())
    inten_max = float(intensity.max())
    if inten_max > inten_min:
        intensity = (intensity - inten_min) / (inten_max - inten_min)
    else:
        intensity[:] = 0.0

    # Normalise RGB to [0,1]
    mask = weight > 0
    rgb[mask] /= weight[mask, None]

    rgb_min = float(rgb.min())
    rgb_max = float(rgb.max())
    if rgb_max > rgb_min:
        rgb = (rgb - rgb_min) / (rgb_max - rgb_min)
    else:
        rgb[:] = 0.0

    return intensity.astype(np.float32), rgb.astype(np.float32)


# ---------------- Analytic LHC pattern ----------------

def make_analytic_lhc_intensity(size=CANVAS_SIZE):
    """
    Analytic pattern: central glow, LHC ring, two beams.
    """
    H = W = size
    y, x = np.mgrid[0:H, 0:W].astype(np.float32)

    cx = (W - 1) / 2.0
    cy = (H - 1) / 2.0

    dx = x - cx
    dy = y - cy
    r  = np.sqrt(dx * dx + dy * dy)
    r_norm = r / r.max()

    central = np.clip(1.0 - r_norm**2, 0.0, 1.0)
    ring    = np.exp(-0.5 * ((r_norm - 0.7) / 0.05)**2)

    beam1 = np.exp(-0.5 * (dx / (0.01 * W))**2)
    beam2 = np.exp(-0.5 * (dy / (0.01 * H))**2)

    intensity = (
        0.35 * central +
        0.45 * ring +
        0.10 * beam1 +
        0.10 * beam2
    )
    intensity = np.clip(intensity, 0.0, 1.0).astype(np.float32)
    return intensity


# ---------------- Optional event-display image overlay ----------------

def load_event_display(max_side=None, url=EVENT_IMG_URL):
    """
    Load a (greyscale) event-display image if available, resize to fit max_side.
    """
    if max_side is None:
        max_side = int(CANVAS_SIZE * 0.8)

    try:
        print(f"Downloading event-display image from\n  {url}")
        resp = requests.get(url, timeout=60)
        resp.raise_for_status()
        im = Image.open(BytesIO(resp.content)).convert("L")
        w, h = im.size
        scale = min(max_side / float(w), max_side / float(h), 1.0)
        new_size = (max(1, int(w * scale)), max(1, int(h * scale)))
        if new_size != (w, h):
            im = im.resize(new_size, Image.BICUBIC)
        print("Event-display size:", im.size)
        return im
    except Exception as e:
        print("Warning: could not load event-display image:", e)
        return None


def blend_event_into_intensity(intensity, event_image, weight=0.5):
    """
    Blend an event-display image into the intensity map at the centre.
    intensity: 2D float32 [0,1]
    event_image: PIL L image
    """
    if event_image is None:
        return intensity

    H, W = intensity.shape
    ev = np.array(event_image, dtype=np.float32) / 255.0
    eh, ew = ev.shape

    x0 = (W - ew) // 2
    y0 = (H - eh) // 2

    sub = intensity[y0:y0 + eh, x0:x0 + ew]
    blended = np.clip((1.0 - weight) * sub + weight * ev, 0.0, 1.0)
    intensity[y0:y0 + eh, x0:x0 + ew] = blended
    return intensity


# ---------------- Combine intensity sources (NO TEXT) ----------------

def build_combined_intensity_map(df,
                                 size=CANVAS_SIZE,
                                 grid_size_data=1024,
                                 use_event_display=USE_EVENT_DISPLAY):
    """
    Build a 4096x4096 intensity map combining:
      - Analytic LHC ring+beams pattern
      - Data-driven Higgs event η–φ blobs
      - Optional event-display image
    (No text overlay.)
    Returns:
      combined_4k:  4096x4096 float32 [0,1]
      data_int:     grid_size_data x grid_size_data
      data_colour:  grid_size_data x grid_size_data x 3
    """
    # 1) Data-driven Higgs η–φ maps
    data_int, data_colour = build_higgs_event_maps_from_df(
        df,
        mass_target=125.0,
        grid_size=grid_size_data,
    )

    # Upscale data intensity to 4K
    data_int_u8 = Image.fromarray((np.clip(data_int, 0.0, 1.0) * 255).astype("uint8"))
    data_int_4k = np.array(
        data_int_u8.resize((size, size), resample=Image.BICUBIC),
        dtype=np.float32,
    ) / 255.0

    # 2) Analytic LHC pattern at 4K
    analytic_4k = make_analytic_lhc_intensity(size)

    # 3) Combine analytic + data (weights can be tuned)
    combined = 0.6 * analytic_4k + 0.8 * data_int_4k
    combined = np.clip(combined, 0.0, 1.0)

    # 4) Optional event-display overlay
    if use_event_display:
        event_img = load_event_display(max_side=int(size * 0.8))
        if event_img is not None:
            combined = blend_event_into_intensity(combined, event_img, weight=0.4)

    # 5) Final normalisation to [0,1]
    mn = float(combined.min())
    mx = float(combined.max())
    if mx > mn:
        combined = (combined - mn) / (mx - mn)
    else:
        combined[:] = 0.0

    return combined.astype(np.float32), data_int.astype(np.float32), data_colour.astype(np.float32)


# ---------------- AllRGB generator (simple luminance logic) ----------------

def generate_allrgb_from_intensity(intensity_4k, output_path=OUTPUT_ALLRGB,
                                   verify_permutation=True):
    """
    Given a 4096x4096 intensity field in [0,1], produce a 4096x4096 RGB image
    that uses every 24-bit color exactly once, assigning darker colors to
    darker pixels and vice versa.
    """
    H = W = CANVAS_SIZE
    assert intensity_4k.shape == (H, W), "intensity_4k must be 4096x4096"

    flat_intensity = intensity_4k.ravel().astype(np.float32)

    # Generate all 24-bit colors as integers 0xRRGGBB
    print("Generating full 24-bit RGB set...")
    colors = np.arange(TOTAL_COLORS, dtype=np.uint32)

    # Compute approximate sRGB luminance
    print("Computing luminance for all colors...")
    r = ((colors >> 16) & 0xFF).astype(np.float32)
    g = ((colors >> 8)  & 0xFF).astype(np.float32)
    b = (colors         & 0xFF).astype(np.float32)
    lum = 0.2126 * r + 0.7152 * g + 0.0722 * b
    del r, g, b

    # Sort pixels by target intensity (dark -> bright)
    print("Sorting pixels by intensity...")
    pixel_order = np.argsort(flat_intensity)
    del flat_intensity

    # Sort colors by luminance (dark -> bright)
    print("Sorting colors by luminance...")
    color_order = np.argsort(lum)
    del lum

    # Assign each color exactly once
    print("Assigning colors to pixels...")
    flat_result = np.empty(TOTAL_COLORS, dtype=np.uint32)
    flat_result[pixel_order] = colors[color_order]
    del pixel_order, color_order, colors

    # Optional permutation check
    if verify_permutation:
        print("Verifying that every 24-bit color is used exactly once...")
        sorted_result = np.sort(flat_result)
        reference     = np.arange(TOTAL_COLORS, dtype=np.uint32)
        assert np.array_equal(sorted_result, reference), "Color permutation check FAILED"
        del sorted_result, reference
        print("Permutation check passed ✅")

    # Reshape to RGB image and save
    print("Reshaping and saving PNG...")
    img_r = ((flat_result >> 16) & 0xFF).astype(np.uint8).reshape(H, W)
    img_g = ((flat_result >> 8)  & 0xFF).astype(np.uint8).reshape(H, W)
    img_b = (flat_result         & 0xFF).astype(np.uint8).reshape(H, W)
    del flat_result

    img = np.stack([img_r, img_g, img_b], axis=-1)
    out = Image.fromarray(img)
    out.save(output_path, format="PNG")
    print(f"Saved allRGB Higgs image as: {output_path}")


# ---------------- Main orchestration ----------------

def main():
    # 1) Load real CMS Higgs data
    df = load_higgs_dataframe()

    # 2) Build combined 4K intensity map + 1024² data previews
    combined_4k, data_int, data_colour = build_combined_intensity_map(
        df,
        size=CANVAS_SIZE,
        grid_size_data=1024,
        use_event_display=USE_EVENT_DISPLAY,
    )

    print("Combined 4K intensity shape:", combined_4k.shape)
    print("Data intensity shape        :", data_int.shape)
    print("Data colour shape          :", data_colour.shape)

    # 3) Save some previews
    gray_data   = Image.fromarray((np.clip(data_int, 0.0, 1.0) * 255).astype("uint8"))
    colour_data = Image.fromarray((np.clip(data_colour, 0.0, 1.0) * 255).astype("uint8"))
    gray_data.save("higgs_event_intensity_gray_1024.png")
    colour_data.save("higgs_event_intensity_colour_1024.png")

    combined_preview = Image.fromarray((combined_4k * 255).astype("uint8"))
    combined_preview_1024 = combined_preview.resize((1024, 1024), resample=Image.BICUBIC)
    combined_preview_1024.save("combined_intensity_preview_1024.png")

    # 4) Quick inline preview in Colab
    plt.figure(figsize=(15, 5))
    plt.subplot(1, 3, 1)
    plt.title("Higgs η–φ intensity (1024²)")
    plt.imshow(data_int, cmap="inferno", origin="lower")
    plt.axis("off")

    plt.subplot(1, 3, 2)
    plt.title("Higgs η–φ coloured (1024²)")
    plt.imshow(data_colour, origin="lower")
    plt.axis("off")

    plt.subplot(1, 3, 3)
    plt.title("Combined intensity (4K → 1024 preview, no text)")
    plt.imshow(np.array(combined_preview_1024, dtype=np.uint8),
               cmap="gray", origin="lower")
    plt.axis("off")

    plt.tight_layout()
    plt.show()

    # 5) Generate the 4096x4096 allRGB image
    generate_allrgb_from_intensity(combined_4k, output_path=OUTPUT_ALLRGB)


if __name__ == "__main__":
    main()