Traffic Flow on a 1D Ring: NaSch and LWR Models

!pip install numpy>=1.23 matplotlib>=3.7 jupyter>=1.0 pillow>=9.4

# Core imports and paths
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple, Dict, List

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import display, Image

FIGURES_DIR = Path("figures")
DATA_DIR = Path("data")
FIGURES_DIR.mkdir(parents=True, exist_ok=True)
DATA_DIR.mkdir(parents=True, exist_ok=True)

print(f"Figures directory: {FIGURES_DIR.resolve()}")
print(f"Data directory:    {DATA_DIR.resolve()}")

Figures directory: /content/figures
Data directory:    /content/data

@dataclass
class NaSchState:
    """State of the NaSch system at a given time on a 1D ring.

    Attributes
    -----------
    L : int
        Number of cells on the ring.
    occ : np.ndarray[bool]
        Occupation array, True if the cell is occupied by a vehicle.
    vel : np.ndarray[int]
        Velocity (in cells per time step) of the vehicle in each cell,
        or -1 if the cell is empty.
    """

    L: int
    occ: np.ndarray
    vel: np.ndarray


def initialize_nasch(
    L: int,
    density: float,
    vmax: int,
    rng: Optional[np.random.Generator] = None,
) -> NaSchState:
    """Randomly initialize a NaSch configuration on a 1D ring.

    Parameters
    ----------
    L : int
        Number of cells.
    density : float
        Target density (0 < density < 1), approximately N/L.
    vmax : int
        Maximum velocity (cells per time step).
    rng : np.random.Generator, optional
        Random number generator for reproducibility.
    """
    if rng is None:
        rng = np.random.default_rng()

    N = int(round(density * L))
    if N > L:
        raise ValueError("Density too high: N > L.")

    occ = np.zeros(L, dtype=bool)
    vel = -np.ones(L, dtype=int)

    positions = rng.choice(L, size=N, replace=False)
    occ[positions] = True
    vel[positions] = rng.integers(low=0, high=vmax + 1, size=N)

    return NaSchState(L=L, occ=occ, vel=vel)


def nasch_step(
    state: NaSchState,
    vmax: int,
    p_slow: float,
    rng: Optional[np.random.Generator] = None,
) -> Tuple[NaSchState, float]:
    """Perform **one time step** of the NaSch model with parallel update.

    Rules (per vehicle)
    -------------------
    1) Acceleration: v -> min(v+1, vmax)
    2) Safe braking: v -> min(v, gap)  (gap = number of free cells ahead)
    3) Random braking: v -> max(v-1, 0) with probability p_slow
    4) Movement: x -> x + v (mod L)

    Returns
    -------
    new_state : NaSchState
    flow_step : float
        Mean flow in this time step: sum(v_i) / L.
    """
    if rng is None:
        rng = np.random.default_rng()

    L = state.L
    occ = state.occ
    vel = state.vel

    new_vel = vel.copy()

    # 1) Acceleration
    car_cells = np.where(occ)[0]
    new_vel[car_cells] = np.minimum(new_vel[car_cells] + 1, vmax)

    # 2) Safe braking (respecting gaps)
    car_cells_sorted = np.sort(car_cells)
    n_cars = car_cells_sorted.size

    if n_cars > 0:
        for idx_k in range(n_cars):
            j = car_cells_sorted[idx_k]
            j_next = car_cells_sorted[(idx_k + 1) % n_cars]
            gap = (j_next - j - 1) % L
            new_vel[j] = min(new_vel[j], gap)

        # 3) Random braking
        random_vals = rng.random(size=n_cars)
        to_slow = random_vals < p_slow

        for idx_k, slow_flag in enumerate(to_slow):
            if not slow_flag:
                continue
            j = car_cells_sorted[idx_k]
            if new_vel[j] > 0:
                new_vel[j] -= 1

    # 4) Movement
    new_occ = np.zeros_like(occ, dtype=bool)
    final_vel = -np.ones_like(vel, dtype=int)

    total_distance = 0

    for j in car_cells_sorted:
        vj = new_vel[j]
        if vj < 0:
            continue
        new_pos = (j + vj) % L
        if new_occ[new_pos]:
            raise RuntimeError("Collision detected: two cars in the same cell.")
        new_occ[new_pos] = True
        final_vel[new_pos] = vj
        total_distance += vj

    flow_step = total_distance / L
    new_state = NaSchState(L=L, occ=new_occ, vel=final_vel)
    return new_state, flow_step


def simulate_nasch(
    L: int,
    density: float,
    vmax: int,
    p_slow: float,
    n_steps: int,
    warmup: int = 0,
    save_history: bool = True,
    rng_seed: Optional[int] = None,
) -> Dict[str, np.ndarray]:
    """Run a full NaSch simulation.

    Returns a dictionary with:
    - ``occ_history``: (n_steps, L) occupation history (if ``save_history``).
    - ``vel_history``: (n_steps, L) velocity history (if ``save_history``).
    - ``flow_per_step``: (n_steps,) flow per time step.
    - ``mean_flow_stationary``: scalar array with mean flow after warmup.
    """
    rng = np.random.default_rng(rng_seed)
    state = initialize_nasch(L=L, density=density, vmax=vmax, rng=rng)

    if save_history:
        occ_history = np.zeros((n_steps, L), dtype=bool)
        vel_history = np.zeros((n_steps, L), dtype=int)
    else:
        occ_history = None
        vel_history = None

    flow_per_step = np.zeros(n_steps, dtype=float)

    for n in range(n_steps):
        if save_history:
            occ_history[n, :] = state.occ
            vel_history[n, :] = state.vel
        state, flow = nasch_step(state, vmax=vmax, p_slow=p_slow, rng=rng)
        flow_per_step[n] = flow

    if warmup >= n_steps:
        mean_flow_stationary = np.nan
    else:
        mean_flow_stationary = flow_per_step[warmup:].mean()

    results: Dict[str, np.ndarray] = {
        "flow_per_step": flow_per_step,
        "mean_flow_stationary": np.array([mean_flow_stationary]),
    }
    if save_history:
        results["occ_history"] = occ_history
        results["vel_history"] = vel_history

    return results

def plot_nasch_spacetime(occ_history: np.ndarray, title: str) -> plt.Figure:
    """Space–time diagram from a Boolean occupation history (T, L)."""
    T, L = occ_history.shape
    data = occ_history.astype(float)

    fig, ax = plt.subplots(figsize=(10, 4))
    im = ax.imshow(
        data,
        aspect="auto",
        interpolation="nearest",
        origin="lower",
        cmap="binary",
    )
    cbar = fig.colorbar(im, ax=ax)
    cbar.set_label("Occupation (0 empty, 1 car)")
    ax.set_xlabel("Position (cell)")
    ax.set_ylabel("Time step")
    ax.set_title(title)
    fig.tight_layout()
    return fig


def plot_nasch_flow(flow_per_step: np.ndarray, mean_flow: float, title: str) -> plt.Figure:
    """Plot flow vs time and mark the stationary mean."""
    n_steps = flow_per_step.size
    t = np.arange(n_steps)

    fig, ax = plt.subplots(figsize=(8, 3))
    ax.plot(t, flow_per_step, lw=1)
    ax.axhline(mean_flow, linestyle="--", label=f"Mean flow = {mean_flow:.3f}")
    ax.set_xlabel("Time step")
    ax.set_ylabel("Flow per cell")
    ax.set_title(f"Flow vs time — {title}")
    ax.legend()
    fig.tight_layout()
    return fig


def make_nasch_ring_animation(
    occ_history: np.ndarray,
    filename: str,
    fps: int = 10,
) -> None:
    """GIF of NaSch traffic on a 1D ring represented as a circle in 2D."""
    T, L = occ_history.shape

    theta = 2.0 * np.pi * np.arange(L) / L
    x_all = np.cos(theta)
    y_all = np.sin(theta)

    fig, ax = plt.subplots(figsize=(5, 5))
    ax.set_aspect("equal")
    ax.set_axis_off()
    ax.plot(x_all, y_all, lw=1, alpha=0.3)
    (points,) = ax.plot([], [], "o", markersize=4)

    def init():
        points.set_data([], [])
        ax.set_title("t = 0")
        return (points,)

    def update(frame: int):
        occ = occ_history[frame]
        x_cars = x_all[occ]
        y_cars = y_all[occ]
        points.set_data(x_cars, y_cars)
        ax.set_title(f"t = {frame}")
        return (points,)

    anim = animation.FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=T,
        interval=1000 / fps,
        blit=True,
    )

    gif_path = FIGURES_DIR / filename
    writer = animation.PillowWriter(fps=fps)
    anim.save(gif_path, writer=writer)
    plt.close(fig)
    print(f"NaSch ring GIF saved to: {gif_path}")


def make_nasch_concentric_rings_animation(
    occ_histories: List[np.ndarray],
    filename: str,
    fps: int = 10,
) -> None:
    """GIF with multiple NaSch regimes as concentric rings.

    All histories must share the same number of cells L; the effective
    number of frames is the minimum T among them.
    """
    n_cases = len(occ_histories)
    if n_cases == 0:
        raise ValueError("At least one occ_history is required.")

    L_list = [occ.shape[1] for occ in occ_histories]
    if len(set(L_list)) != 1:
        raise ValueError("All occ_history arrays must share the same L.")
    L = L_list[0]

    T_list = [occ.shape[0] for occ in occ_histories]
    T = min(T_list)

    base_r = 1.0
    dr = 0.4
    radii = [base_r + i * dr for i in range(n_cases)]

    theta = 2.0 * np.pi * np.arange(L) / L
    x_all_rings = [r * np.cos(theta) for r in radii]
    y_all_rings = [r * np.sin(theta) for r in radii]

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_aspect("equal")
    ax.set_axis_off()

    for x_all, y_all in zip(x_all_rings, y_all_rings):
        ax.plot(x_all, y_all, lw=1, alpha=0.3)

    points_list = []
    for _ in range(n_cases):
        (points,) = ax.plot([], [], "o", markersize=3)
        points_list.append(points)

    max_r = max(radii)
    ax.set_xlim(-max_r - 0.2, max_r + 0.2)
    ax.set_ylim(-max_r - 0.2, max_r + 0.2)

    def init():
        for pts in points_list:
            pts.set_data([], [])
        ax.set_title("t = 0")
        return tuple(points_list)

    def update(frame: int):
        for k in range(n_cases):
            occ = occ_histories[k][frame]
            x_all = x_all_rings[k]
            y_all = y_all_rings[k]
            x_cars = x_all[occ]
            y_cars = y_all[occ]
            points_list[k].set_data(x_cars, y_cars)
        ax.set_title(f"Concentric rings — t = {frame}")
        return tuple(points_list)

    anim = animation.FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=T,
        interval=1000 / fps,
        blit=True,
    )

    gif_path = FIGURES_DIR / filename
    writer = animation.PillowWriter(fps=fps)
    anim.save(gif_path, writer=writer)
    plt.close(fig)
    print(f"NaSch concentric rings GIF saved to: {gif_path}")


def run_nasch_examples(
    make_animations: bool = True,
    n_steps: int = 400,
    L: int = 200,
    vmax: int = 5,
) -> List[np.ndarray]:
    """Reproduce the three canonical NaSch regimes on the ring.

    Returns the list of occupation histories [free, stop_go, jammed].
    """
    cases = [
        {
            "name": "free",
            "density": 0.05,
            "p_slow": 0.1,
            "seed": 1,
            "title": "Free flow (low density, p_slow=0.1)",
        },
        {
            "name": "stop_go",
            "density": 0.2,
            "p_slow": 0.3,
            "seed": 2,
            "title": "Stop-and-go waves (intermediate density, p_slow=0.3)",
        },
        {
            "name": "jammed",
            "density": 0.5,
            "p_slow": 0.3,
            "seed": 3,
            "title": "Dense jam (high density, p_slow=0.3)",
        },
    ]

    occ_histories: List[np.ndarray] = []

    for cfg in cases:
        results = simulate_nasch(
            L=L,
            density=cfg["density"],
            vmax=vmax,
            p_slow=cfg["p_slow"],
            n_steps=n_steps,
            warmup=200,
            save_history=True,
            rng_seed=cfg["seed"],
        )
        occ_history = results["occ_history"]
        flow_per_step = results["flow_per_step"]
        mean_flow = float(results["mean_flow_stationary"][0])

        fig_st = plot_nasch_spacetime(occ_history, title=cfg["title"])
        fig_st_path = FIGURES_DIR / f"spacetime_{cfg['name']}.png"
        fig_st.savefig(fig_st_path, dpi=200)
        display(fig_st)
        plt.close(fig_st)
        print(f"Saved NaSch space–time plot: {fig_st_path}")

        fig_flow = plot_nasch_flow(flow_per_step, mean_flow, title=cfg["title"])
        fig_flow_path = FIGURES_DIR / f"flow_{cfg['name']}.png"
        fig_flow.savefig(fig_flow_path, dpi=200)
        display(fig_flow)
        plt.close(fig_flow)
        print(f"Saved NaSch flow plot: {fig_flow_path}")

        if make_animations:
            gif_name = f"ring_{cfg['name']}.gif"
            make_nasch_ring_animation(occ_history, filename=gif_name, fps=10)
            # Show GIF inline in the notebook
            display(Image(filename=str(FIGURES_DIR / gif_name)))

        occ_histories.append(occ_history)

    if make_animations and len(occ_histories) == 3:
        concentric_name = "ring_concentric_all.gif"
        make_nasch_concentric_rings_animation(
            occ_histories=occ_histories,
            filename=concentric_name,
            fps=10,
        )
        # Show concentric GIF inline
        display(Image(filename=str(FIGURES_DIR / concentric_name)))

    return occ_histories

_ = run_nasch_examples(make_animations=True)  # set True to also generate GIFs, but it takes a while

Out [9]:

Saved NaSch space–time plot: figures/spacetime_free.png

Saved NaSch flow plot: figures/flow_free.png
NaSch ring GIF saved to: figures/ring_free.gif

Saved NaSch space–time plot: figures/spacetime_stop_go.png

Saved NaSch flow plot: figures/flow_stop_go.png
NaSch ring GIF saved to: figures/ring_stop_go.gif

Saved NaSch space–time plot: figures/spacetime_jammed.png

Saved NaSch flow plot: figures/flow_jammed.png
NaSch ring GIF saved to: figures/ring_jammed.gif

NaSch concentric rings GIF saved to: figures/ring_concentric_all.gif

from typing import Optional, Tuple
import numpy as np

def estimate_flow_for_density(
    density: float,
    L: int,
    vmax: int,
    p_slow: float,
    n_steps: int,
    warmup: int,
    n_runs: int,
    base_seed: int = 1234,
) -> Tuple[float, float]:
    """Estimate stationary mean flow for a given density, averaging over n_runs independent simulations."""
    flows = []

    for k in range(n_runs):
        seed = base_seed + k
        results = simulate_nasch(
            L=L,
            density=density,
            vmax=vmax,
            p_slow=p_slow,
            n_steps=n_steps,
            warmup=warmup,
            save_history=False,
            rng_seed=seed,
        )
        flow_k = float(results["mean_flow_stationary"][0])
        flows.append(flow_k)

    flows = np.array(flows)
    mean_flow = flows.mean()
    std_flow = flows.std(ddof=1) if n_runs > 1 else 0.0

    return mean_flow, std_flow


def build_fundamental_diagram(
    L: int = 200,
    vmax: int = 5,
    p_slow: float = 0.3,
    n_steps: int = 4000,
    warmup: int = 1000,
    n_runs: int = 5,
    rho_min: float = 0.02,
    rho_max: float = 0.80,
    n_rho: int = 20,
    custom_densities: Optional[np.ndarray] = None,
) -> None:
    """Build the fundamental diagram Q(ρ) by sweeping densities and generate plots/data."""
    if custom_densities is not None:
        densities = np.asarray(custom_densities)
        actual_rho_min = densities.min()
        actual_rho_max = densities.max()
        actual_n_rho = densities.size
    else:
        densities = np.linspace(rho_min, rho_max, n_rho)
        actual_rho_min = rho_min
        actual_rho_max = rho_max
        actual_n_rho = n_rho

    mean_flows = np.zeros_like(densities)
    std_flows = np.zeros_like(densities)

    print("Building fundamental diagram Q(ρ)...")
    print(f"L={L}, vmax={vmax}, p_slow={p_slow}")
    print(f"n_steps={n_steps}, warmup={warmup}, n_runs={n_runs}")
    print(f"densities: {actual_rho_min:.3f} -> {actual_rho_max:.3f} ({actual_n_rho} points)")
    print()

    for i, rho in enumerate(densities):
        mean_flow, std_flow = estimate_flow_for_density(
            density=rho,
            L=L,
            vmax=vmax,
            p_slow=p_slow,
            n_steps=n_steps,
            warmup=warmup,
            n_runs=n_runs,
            base_seed=1000,
        )
        mean_flows[i] = mean_flow
        std_flows[i] = std_flow
        print(f"ρ = {rho:.3f}  ->  Q = {mean_flow:.4f}  (std = {std_flow:.4f})")

    # Mean speed u(ρ) = Q/ρ
    with np.errstate(divide="ignore", invalid="ignore"):
        mean_speeds = np.where(densities > 0, mean_flows / densities, 0.0)

    # Save to .npz
    npz_path = DATA_DIR / "fundamental_diagram.npz"
    np.savez(
        npz_path,
        densities=densities,
        mean_flows=mean_flows,
        std_flows=std_flows,
        mean_speeds=mean_speeds,
        L=L,
        vmax=vmax,
        p_slow=p_slow,
        n_steps=n_steps,
        warmup=warmup,
        n_runs=n_runs,
    )
    print(f"\nData saved to: {npz_path}")

    # Save CSV
    csv_path = DATA_DIR / "fundamental_diagram.csv"
    header = "rho,mean_flow,std_flow,mean_speed"
    data_for_csv = np.column_stack([densities, mean_flows, std_flows, mean_speeds])
    np.savetxt(csv_path, data_for_csv, delimiter=",", header=header, comments="")
    print(f"CSV saved to: {csv_path}")

    # Plot Q(ρ)
    fig1, ax1 = plt.subplots(figsize=(6, 4))
    ax1.errorbar(
        densities,
        mean_flows,
        yerr=std_flows,
        fmt="o-",
        markersize=4,
        capsize=3,
        label="NaSch (average over runs)",
    )
    ax1.set_xlabel(r"Density $\rho$ (N/L)")
    ax1.set_ylabel(r"Mean flow $Q(\rho)$ (cars/step/cell)")
    ax1.set_title(r"Fundamental diagram $Q(\rho)$ for NaSch") # Changed to raw string
    ax1.grid(True, alpha=0.3)
    ax1.legend()
    fig1.tight_layout()
    fig1_path = FIGURES_DIR / "fundamental_diagram_Q_vs_rho.png"
    fig1.savefig(fig1_path, dpi=200)
    display(fig1)
    plt.close(fig1)
    print(f"Figure Q(ρ) saved to: {fig1_path}")

    # Plot u(ρ) = Q/ρ
    fig2, ax2 = plt.subplots(figsize=(6, 4))
    ax2.plot(densities, mean_speeds, "o-", markersize=4, label="NaSch")
    ax2.set_xlabel(r"Density $\rho__STASH_1790932053694_3__quot;)
    ax2.set_ylabel(r"Mean speed $u(\rho)$ (cells/step)")
    ax2.set_title(r"Mean speed $u(\rho)$ for NaSch") # Changed to raw string
    ax2.grid(True, alpha=0.3)
    ax2.legend()
    fig2.tight_layout()
    fig2_path = FIGURES_DIR / "fundamental_diagram_u_vs_rho.png"
    fig2.savefig(fig2_path, dpi=200)
    display(fig2)
    plt.close(fig2)
    print(f"Figure u(ρ) saved to: {fig2_path}")

import numpy as np

# Define parameters for dense and sparse segments
rho_dense_start = 0.0
rho_dense_end = 0.2
n_points_dense = 100  # More points in the 0-0.2 range

rho_sparse_start = 0.2  # Start from the end of the dense segment
rho_sparse_end = 1.0  # Original rho_max
n_points_sparse = 20  # Fewer points for the rest of the range

# Generate the two segments
densities_dense = np.linspace(rho_dense_start, rho_dense_end, n_points_dense)
densities_sparse = np.linspace(rho_sparse_start, rho_sparse_end, n_points_sparse)

# Combine, remove duplicates, and sort to ensure a clean, monotonically increasing array
custom_densities = np.unique(np.concatenate((densities_dense, densities_sparse)))
custom_densities = np.sort(custom_densities)

build_fundamental_diagram(custom_densities=custom_densities, rho_min=0.0, rho_max=1.0)

Out [13]:

Building fundamental diagram Q(ρ)...
L=200, vmax=5, p_slow=0.3
n_steps=4000, warmup=1000, n_runs=5
densities: 0.000 -> 1.000 (119 points)

ρ = 0.000  ->  Q = 0.0000  (std = 0.0000)
ρ = 0.002  ->  Q = 0.0000  (std = 0.0000)
ρ = 0.004  ->  Q = 0.0235  (std = 0.0000)
ρ = 0.006  ->  Q = 0.0235  (std = 0.0000)
ρ = 0.008  ->  Q = 0.0470  (std = 0.0001)
ρ = 0.010  ->  Q = 0.0470  (std = 0.0001)
ρ = 0.012  ->  Q = 0.0470  (std = 0.0001)
ρ = 0.014  ->  Q = 0.0705  (std = 0.0001)
ρ = 0.016  ->  Q = 0.0705  (std = 0.0001)
ρ = 0.018  ->  Q = 0.0940  (std = 0.0001)
ρ = 0.020  ->  Q = 0.0940  (std = 0.0001)
ρ = 0.022  ->  Q = 0.0940  (std = 0.0001)
ρ = 0.024  ->  Q = 0.1174  (std = 0.0002)
ρ = 0.026  ->  Q = 0.1174  (std = 0.0002)
ρ = 0.028  ->  Q = 0.1408  (std = 0.0003)
ρ = 0.030  ->  Q = 0.1408  (std = 0.0003)
ρ = 0.032  ->  Q = 0.1408  (std = 0.0003)
ρ = 0.034  ->  Q = 0.1643  (std = 0.0002)
ρ = 0.036  ->  Q = 0.1643  (std = 0.0002)
ρ = 0.038  ->  Q = 0.1876  (std = 0.0002)
ρ = 0.040  ->  Q = 0.1876  (std = 0.0002)
ρ = 0.042  ->  Q = 0.1876  (std = 0.0002)
ρ = 0.044  ->  Q = 0.2110  (std = 0.0001)
ρ = 0.046  ->  Q = 0.2110  (std = 0.0001)
ρ = 0.048  ->  Q = 0.2343  (std = 0.0003)
ρ = 0.051  ->  Q = 0.2343  (std = 0.0003)
ρ = 0.053  ->  Q = 0.2575  (std = 0.0003)
ρ = 0.055  ->  Q = 0.2575  (std = 0.0003)
ρ = 0.057  ->  Q = 0.2575  (std = 0.0003)
ρ = 0.059  ->  Q = 0.2809  (std = 0.0002)
ρ = 0.061  ->  Q = 0.2809  (std = 0.0002)
ρ = 0.063  ->  Q = 0.3040  (std = 0.0001)
ρ = 0.065  ->  Q = 0.3040  (std = 0.0001)
ρ = 0.067  ->  Q = 0.3040  (std = 0.0001)
ρ = 0.069  ->  Q = 0.3272  (std = 0.0004)
ρ = 0.071  ->  Q = 0.3272  (std = 0.0004)
ρ = 0.073  ->  Q = 0.3503  (std = 0.0002)
ρ = 0.075  ->  Q = 0.3503  (std = 0.0002)
ρ = 0.077  ->  Q = 0.3503  (std = 0.0002)
ρ = 0.079  ->  Q = 0.3731  (std = 0.0002)
ρ = 0.081  ->  Q = 0.3731  (std = 0.0002)
ρ = 0.083  ->  Q = 0.3960  (std = 0.0003)
ρ = 0.085  ->  Q = 0.3960  (std = 0.0003)
ρ = 0.087  ->  Q = 0.3960  (std = 0.0003)
ρ = 0.089  ->  Q = 0.4184  (std = 0.0008)
ρ = 0.091  ->  Q = 0.4184  (std = 0.0008)
ρ = 0.093  ->  Q = 0.4412  (std = 0.0008)
ρ = 0.095  ->  Q = 0.4412  (std = 0.0008)
ρ = 0.097  ->  Q = 0.4412  (std = 0.0008)
ρ = 0.099  ->  Q = 0.4623  (std = 0.0009)
ρ = 0.101  ->  Q = 0.4623  (std = 0.0009)
ρ = 0.103  ->  Q = 0.4824  (std = 0.0019)
ρ = 0.105  ->  Q = 0.4824  (std = 0.0019)
ρ = 0.107  ->  Q = 0.4824  (std = 0.0019)
ρ = 0.109  ->  Q = 0.5007  (std = 0.0048)
ρ = 0.111  ->  Q = 0.5007  (std = 0.0048)
ρ = 0.113  ->  Q = 0.5102  (std = 0.0080)
ρ = 0.115  ->  Q = 0.5102  (std = 0.0080)
ρ = 0.117  ->  Q = 0.5102  (std = 0.0080)
ρ = 0.119  ->  Q = 0.5000  (std = 0.0115)
ρ = 0.121  ->  Q = 0.5000  (std = 0.0115)
ρ = 0.123  ->  Q = 0.4874  (std = 0.0045)
ρ = 0.125  ->  Q = 0.4874  (std = 0.0045)
ρ = 0.127  ->  Q = 0.4874  (std = 0.0045)
ρ = 0.129  ->  Q = 0.4728  (std = 0.0114)
ρ = 0.131  ->  Q = 0.4728  (std = 0.0114)
ρ = 0.133  ->  Q = 0.4696  (std = 0.0087)
ρ = 0.135  ->  Q = 0.4696  (std = 0.0087)
ρ = 0.137  ->  Q = 0.4696  (std = 0.0087)
ρ = 0.139  ->  Q = 0.4625  (std = 0.0050)
ρ = 0.141  ->  Q = 0.4625  (std = 0.0050)
ρ = 0.143  ->  Q = 0.4657  (std = 0.0058)
ρ = 0.145  ->  Q = 0.4657  (std = 0.0058)
ρ = 0.147  ->  Q = 0.4657  (std = 0.0058)
ρ = 0.149  ->  Q = 0.4563  (std = 0.0064)
ρ = 0.152  ->  Q = 0.4563  (std = 0.0064)
ρ = 0.154  ->  Q = 0.4568  (std = 0.0047)
ρ = 0.156  ->  Q = 0.4568  (std = 0.0047)
ρ = 0.158  ->  Q = 0.4552  (std = 0.0018)
ρ = 0.160  ->  Q = 0.4552  (std = 0.0018)
ρ = 0.162  ->  Q = 0.4552  (std = 0.0018)
ρ = 0.164  ->  Q = 0.4546  (std = 0.0033)
ρ = 0.166  ->  Q = 0.4546  (std = 0.0033)
ρ = 0.168  ->  Q = 0.4533  (std = 0.0030)
ρ = 0.170  ->  Q = 0.4533  (std = 0.0030)
ρ = 0.172  ->  Q = 0.4533  (std = 0.0030)
ρ = 0.174  ->  Q = 0.4505  (std = 0.0064)
ρ = 0.176  ->  Q = 0.4505  (std = 0.0064)
ρ = 0.178  ->  Q = 0.4468  (std = 0.0044)
ρ = 0.180  ->  Q = 0.4468  (std = 0.0044)
ρ = 0.182  ->  Q = 0.4468  (std = 0.0044)
ρ = 0.184  ->  Q = 0.4471  (std = 0.0049)
ρ = 0.186  ->  Q = 0.4471  (std = 0.0049)
ρ = 0.188  ->  Q = 0.4433  (std = 0.0067)
ρ = 0.190  ->  Q = 0.4433  (std = 0.0067)
ρ = 0.192  ->  Q = 0.4433  (std = 0.0067)
ρ = 0.194  ->  Q = 0.4411  (std = 0.0037)
ρ = 0.196  ->  Q = 0.4411  (std = 0.0037)
ρ = 0.198  ->  Q = 0.4355  (std = 0.0057)
ρ = 0.200  ->  Q = 0.4355  (std = 0.0057)
ρ = 0.242  ->  Q = 0.4187  (std = 0.0030)
ρ = 0.284  ->  Q = 0.4003  (std = 0.0031)
ρ = 0.326  ->  Q = 0.3819  (std = 0.0023)
ρ = 0.368  ->  Q = 0.3613  (std = 0.0016)
ρ = 0.411  ->  Q = 0.3415  (std = 0.0009)
ρ = 0.453  ->  Q = 0.3186  (std = 0.0016)
ρ = 0.495  ->  Q = 0.2995  (std = 0.0018)
ρ = 0.537  ->  Q = 0.2789  (std = 0.0009)
ρ = 0.579  ->  Q = 0.2544  (std = 0.0010)
ρ = 0.621  ->  Q = 0.2330  (std = 0.0009)
ρ = 0.663  ->  Q = 0.2085  (std = 0.0003)
ρ = 0.705  ->  Q = 0.1863  (std = 0.0004)
ρ = 0.747  ->  Q = 0.1625  (std = 0.0003)
ρ = 0.789  ->  Q = 0.1366  (std = 0.0003)
ρ = 0.832  ->  Q = 0.1124  (std = 0.0004)
ρ = 0.874  ->  Q = 0.0839  (std = 0.0002)
ρ = 0.916  ->  Q = 0.0577  (std = 0.0004)
ρ = 0.958  ->  Q = 0.0276  (std = 0.0001)
ρ = 1.000  ->  Q = 0.0000  (std = 0.0000)

Data saved to: data/fundamental_diagram.npz
CSV saved to: data/fundamental_diagram.csv

Figure Q(ρ) saved to: figures/fundamental_diagram_Q_vs_rho.png

Figure u(ρ) saved to: figures/fundamental_diagram_u_vs_rho.png

@dataclass
class FundamentalDiagram:
    """Discrete representation of a fundamental diagram Q(ρ).

    Attributes
    ----------
    rho_grid : np.ndarray
        1D array of densities where Q was measured.
    Q_grid : np.ndarray
        1D array of corresponding Q(ρ) values.
    c_max : float
        Maximum wave speed |Q'(ρ)| on this grid.
    """

    rho_grid: np.ndarray
    Q_grid: np.ndarray
    c_max: float

    @property
    def rho_min(self) -> float:
        return float(self.rho_grid[0])

    @property
    def rho_max(self) -> float:
        return float(self.rho_grid[-1])

    def Q(self, rho: np.ndarray) -> np.ndarray:
        """Evaluate Q(ρ) by linear interpolation over rho_grid.

        Any ρ outside [rho_min, rho_max] is clipped to avoid extrapolation.
        """
        rho_arr = np.asarray(rho, dtype=float)
        rho_clip = np.clip(rho_arr, self.rho_min, self.rho_max)
        return np.interp(rho_clip, self.rho_grid, self.Q_grid)


def load_fundamental_diagram(
    npz_path: Path | str = "data/fundamental_diagram.npz",
) -> FundamentalDiagram:
    """Load the fundamental diagram from a .npz file and build a FundamentalDiagram object."""
    npz_path = Path(npz_path)
    data = np.load(npz_path)
    rho_grid = data["densities"]
    Q_grid = data["mean_flows"]

    # Estimate Q'(ρ) by finite differences and take c_max = max |Q'|
    drho = np.diff(rho_grid)
    dQ = np.diff(Q_grid)
    slopes = dQ / drho
    c_max = float(np.max(np.abs(slopes)))

    return FundamentalDiagram(rho_grid=rho_grid, Q_grid=Q_grid, c_max=c_max)


def simulate_lwr_ring(
    fd: FundamentalDiagram,
    L: int,
    T_final: float,
    CFL: float = 0.9,
    rho_init: np.ndarray | None = None,
) -> Tuple[np.ndarray, float]:
    """Simulate the LWR equation on a 1D ring using Lax–Friedrichs.

    The PDE is: ∂_t ρ + ∂_x Q(ρ) = 0, with periodic boundary conditions.

    Domain: x = 0,...,L-1 with Δx = 1 (consistent with NaSch model).

    Returns
    -------
    rho_history : (n_steps+1, L) array with the evolution of ρ.
    dt : time step used (computed from CFL condition).
    """
    dx = 1.0

    if CFL <= 0.0 or CFL > 1.0:
        raise ValueError("CFL must be in (0, 1].")

    if fd.c_max <= 0:
        raise ValueError("c_max must be positive; check the fundamental diagram.")

    # Time step from CFL condition
    dt = CFL * dx / fd.c_max

    # Number of time steps (round up to cover T_final)
    n_steps = int(np.ceil(T_final / dt))

    # Initialize ρ
    if rho_init is None:
        rho0 = 0.5 * (fd.rho_min + fd.rho_max)
        rng = np.random.default_rng(123)
        rho = rho0 * np.ones(L)
        rho += 0.01 * rho0 * rng.standard_normal(size=L)
        rho = np.clip(rho, 0.0, fd.rho_max)
    else:
        rho = np.asarray(rho_init, dtype=float)
        if rho.shape != (L,):
            raise ValueError(f"rho_init must have shape ({L},), but has {rho.shape}.")
        rho = np.clip(rho, 0.0, fd.rho_max)

    # History array
    rho_history = np.zeros((n_steps + 1, L), dtype=float)
    rho_history[0, :] = rho

    Q_vals = np.zeros(L, dtype=float)

    for n in range(n_steps):
        # Evaluate Q(ρ) at current points
        Q_vals[:] = fd.Q(rho)

        # Periodicity: neighbor indices i-1, i+1
        rho_left = np.roll(rho, 1)
        rho_right = np.roll(rho, -1)
        Q_left = np.roll(Q_vals, 1)
        Q_right = np.roll(Q_vals, -1)

        # Lax–Friedrichs scheme (conservative form)
        rho_new = 0.5 * (rho_right + rho_left) - (dt / (2.0 * dx)) * (Q_right - Q_left)

        # Clipping for safety
        rho_new = np.clip(rho_new, 0.0, fd.rho_max)

        rho = rho_new
        rho_history[n + 1, :] = rho

    return rho_history, dt


def simulate_lwr_ring_nsteps(
    fd: FundamentalDiagram,
    L: int,
    n_steps: int,
    CFL: float = 0.9,
    rho_init: np.ndarray | None = None,
) -> Tuple[np.ndarray, float]:
    """Variant of simulate_lwr_ring that fixes the number of time steps n_steps directly."""
    if n_steps < 1:
        raise ValueError("n_steps must be an integer >= 1.")

    dx = 1.0

    if CFL <= 0.0 or CFL > 1.0:
        raise ValueError("CFL must be in (0, 1].")

    if fd.c_max <= 0:
        raise ValueError("c_max must be positive; check the fundamental diagram.")

    dt = CFL * dx / fd.c_max

    # Initialize ρ
    if rho_init is None:
        rho0 = 0.5 * (fd.rho_min + fd.rho_max)
        rng = np.random.default_rng(123)
        rho = rho0 * np.ones(L)
        rho += 0.01 * rho0 * rng.standard_normal(size=L)
        rho = np.clip(rho, 0.0, fd.rho_max)
    else:
        rho = np.asarray(rho_init, dtype=float)
        if rho.shape != (L,):
            raise ValueError(f"rho_init must have shape ({L},), but has {rho.shape}.")
        rho = np.clip(rho, 0.0, fd.rho_max)

    # History
    rho_history = np.zeros((n_steps + 1, L), dtype=float)
    rho_history[0, :] = rho

    Q_vals = np.zeros(L, dtype=float)

    for n in range(n_steps):
        Q_vals[:] = fd.Q(rho)

        rho_left = np.roll(rho, 1)
        rho_right = np.roll(rho, -1)
        Q_left = np.roll(Q_vals, 1)
        Q_right = np.roll(Q_vals, -1)

        rho_new = 0.5 * (rho_right + rho_left) - (dt / (2.0 * dx)) * (Q_right - Q_left)
        rho_new = np.clip(rho_new, 0.0, fd.rho_max)

        rho = rho_new
        rho_history[n + 1, :] = rho

    return rho_history, dt

def plot_lwr_spacetime(rho_history: np.ndarray, title: str) -> plt.Figure:
    """Plot a space–time diagram for the LWR solution."""
    T, L = rho_history.shape

    fig, ax = plt.subplots(figsize=(10, 4))
    im = ax.imshow(
        rho_history,
        aspect="auto",
        interpolation="nearest",
        origin="lower",
        cmap="viridis",
    )
    cbar = fig.colorbar(im, ax=ax)
    cbar.set_label(r"Density $\rho(x,t)__STASH_1790932053694_4__quot;)
    ax.set_xlabel("Position (cell)")
    ax.set_ylabel("Time step index (LWR)")
    ax.set_title(title)
    fig.tight_layout()
    return fig


def make_lwr_ring_animation(
    rho_history: np.ndarray,
    rho_min: float,
    rho_max: float,
    filename: str,
    fps: int = 10,
    max_frames: int = 400,
) -> None:
    """Create a GIF where LWR density evolves on a 1D ring.

    Frames are subsampled to avoid huge files (max_frames maximum).
    """
    T, L = rho_history.shape

    if T <= max_frames:
        frame_indices = np.arange(T)
    else:
        step = int(np.ceil(T / max_frames))
        frame_indices = np.arange(0, T, step, dtype=int)

    print(f"Generating GIF '{filename}' with {len(frame_indices)} frames (original T={T}).")

    theta = 2.0 * np.pi * np.arange(L) / L
    x_all = np.cos(theta)
    y_all = np.sin(theta)

    fig, ax = plt.subplots(figsize=(5, 5))
    ax.set_aspect("equal")
    ax.set_axis_off()

    scat = ax.scatter(
        x_all,
        y_all,
        c=rho_history[frame_indices[0]],
        cmap="viridis",
        vmin=rho_min,
        vmax=rho_max,
        s=25,
    )

    cbar = fig.colorbar(scat, ax=ax)
    cbar.set_label(r"Density $\rho__STASH_1790932053694_4__quot;)
    ax.plot(x_all, y_all, lw=1, alpha=0.3)

    def init():
        scat.set_array(rho_history[frame_indices[0]])
        ax.set_title(f"t = {frame_indices[0]}")
        return (scat,)

    def update(frame_idx: int):
        rho_frame = rho_history[frame_idx]
        scat.set_array(rho_frame)
        ax.set_title(f"t = {frame_idx}")
        return (scat,)

    anim = animation.FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=frame_indices,
        interval=1000 / fps,
        blit=False,
    )

    gif_path = FIGURES_DIR / filename
    writer = animation.PillowWriter(fps=fps)
    anim.save(gif_path, writer=writer)
    plt.close(fig)
    print(f"LWR ring GIF saved to: {gif_path}")


def make_lwr_concentric_rings_animation(
    rho_histories: List[np.ndarray],
    filename: str,
    fps: int = 10,
    max_frames: int = 400,
) -> None:
    """Create a GIF with multiple concentric rings (one per rho_history).

    Each ring has its own color scale and colormap for better contrast.
    """
    n_cases = len(rho_histories)
    if n_cases == 0:
        raise ValueError("At least one rho_history is required.")

    L_list = [rho.shape[1] for rho in rho_histories]
    if len(set(L_list)) != 1:
        raise ValueError("All rho_history arrays must share the same L.")
    L = L_list[0]

    T_list = [rho.shape[0] for rho in rho_histories]
    T = min(T_list)

    if T <= max_frames:
        frame_indices = np.arange(T)
    else:
        step = int(np.ceil(T / max_frames))
        frame_indices = np.arange(0, T, step, dtype=int)

    print(f"Generating concentric GIF '{filename}' with {len(frame_indices)} frames (min T={T}).")

    vmin_list = []
    vmax_list = []
    for k in range(n_cases):
        vals = rho_histories[k][:T].ravel()
        vmin_list.append(float(vals.min()))
        vmax_list.append(float(vals.max()))

    cmap_list = ["Blues", "Greens", "Reds", "Purples", "Oranges"]

    base_r = 1.0
    dr = 0.4
    radii = [base_r + i * dr for i in range(n_cases)]

    theta = 2.0 * np.pi * np.arange(L) / L
    x_all_rings = [r * np.cos(theta) for r in radii]
    y_all_rings = [r * np.sin(theta) for r in radii]

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_aspect("equal")
    ax.set_axis_off()

    for x_all, y_all in zip(x_all_rings, y_all_rings):
        ax.plot(x_all, y_all, lw=1, alpha=0.3)

    scat_list = []
    for k in range(n_cases):
        cmap = cmap_list[k % len(cmap_list)]
        scat = ax.scatter(
            x_all_rings[k],
            y_all_rings[k],
            c=rho_histories[k][frame_indices[0]],
            cmap=cmap,
            vmin=vmin_list[k],
            vmax=vmax_list[k],
            s=20,
        )
        scat_list.append(scat)

    labels_default = ["low ρ", "mid ρ", "high ρ"]
    for k in range(min(n_cases, len(labels_default))):
        ax.text(0.0, radii[k], labels_default[k], ha="center", va="bottom", fontsize=9)

    max_r = max(radii)
    ax.set_xlim(-max_r - 0.3, max_r + 0.3)
    ax.set_ylim(-max_r - 0.3, max_r + 0.3)

    def init():
        for k in range(n_cases):
            scat_list[k].set_array(rho_histories[k][frame_indices[0]])
        ax.set_title(f"t = {frame_indices[0]}")
        return tuple(scat_list)

    def update(frame_idx: int):
        for k in range(n_cases):
            rho_frame = rho_histories[k][frame_idx]
            scat_list[k].set_array(rho_frame)
        ax.set_title(f"t = {frame_idx}")
        return tuple(scat_list)

    anim = animation.FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=frame_indices,
        interval=1000 / fps,
        blit=False,
    )

    gif_path = FIGURES_DIR / filename
    writer = animation.PillowWriter(fps=fps)
    anim.save(gif_path, writer=writer)
    plt.close(fig)
    print(f"LWR concentric rings GIF saved to: {gif_path}")


def make_initial_profile(L: int, rho0: float, perturbation: str = "bump") -> np.ndarray:
    """Build an initial profile for LWR on a ring of L cells.

    Parameters
    ----------
    L : int
        Number of cells.
    rho0 : float
        Background mean density.
    perturbation : str
        Type of initial perturbation:
        - "bump": local Gaussian bump of higher density.
        - "step": step (half with high ρ, half with low ρ).
        - "noise": small white noise around rho0.
    """
    x = np.arange(L)
    rho = rho0 * np.ones(L)

    if perturbation == "bump":
        center = L // 2
        width = L // 10
        bump = np.exp(-0.5 * ((x - center) / width) ** 2)
        rho += 0.5 * rho0 * bump
    elif perturbation == "step":
        rho[: L // 2] = rho0 * 0.5
        rho[L // 2 :] = rho0 * 1.5
    elif perturbation == "noise":
        rng = np.random.default_rng(42)
        rho += 0.1 * rho0 * rng.standard_normal(size=L)
    else:
        raise ValueError(f"Unknown perturbation: {perturbation}")

    rho = np.clip(rho, 0.0, 1.0)
    return rho


def run_lwr_examples(
    make_animations: bool = True,
    L: int = 200,
    T_final: float = 400.0,
    CFL: float = 0.9,
) -> List[np.ndarray]:
    """Run LWR examples for three density regimes (low, intermediate, high).

    Returns the list of density histories [low, mid, high].
    """
    fd = load_fundamental_diagram("data/fundamental_diagram.npz")
    print(
        f"Fundamental diagram loaded: ρ in [{fd.rho_min:.3f}, {fd.rho_max:.3f}], "
        f"c_max ≈ {fd.c_max:.3f}"
    )

    cases = [
        {"rho0": 0.05, "name": "low", "title": rf"LWR — low density $\rho_0 \approx 0.05__STASH_1790932053694_4__quot;},
        {"rho0": 0.2, "name": "mid", "title": rf"LWR — intermediate density $\rho_0 \approx 0.2__STASH_1790932053694_4__quot;},
        {"rho0": 0.5, "name": "high", "title": rf"LWR — high density $\rho_0 \approx 0.5__STASH_1790932053694_4__quot;},
    ]

    rho_histories: List[np.ndarray] = []

    for cfg in cases:
        rho_init = make_initial_profile(L, rho0=cfg["rho0"], perturbation="bump")
        rho_hist, dt = simulate_lwr_ring(
            fd, L=L, T_final=T_final, CFL=CFL, rho_init=rho_init
        )
        print(f"Case {cfg['name']} density: dt={dt:.4f}, n_steps={rho_hist.shape[0]-1}")

        fig = plot_lwr_spacetime(rho_hist, title=cfg["title"])
        fig_path = FIGURES_DIR / f"lwr_spacetime_{cfg['name']}_density.png"
        fig.savefig(fig_path, dpi=200)
        display(fig)
        plt.close(fig)
        print(f"Saved LWR space–time plot: {fig_path}")

        if make_animations:
            gif_name = f"lwr_ring_{cfg['name']}_density.gif"
            make_lwr_ring_animation(
                rho_history=rho_hist,
                rho_min=fd.rho_min,
                rho_max=fd.rho_max,
                filename=gif_name,
                fps=10,
            )
            display(Image(filename=str(FIGURES_DIR / gif_name)))

        rho_histories.append(rho_hist)

    if make_animations and len(rho_histories) == 3:
        concentric_name = "lwr_ring_concentric_all.gif"
        make_lwr_concentric_rings_animation(
            rho_histories=rho_histories,
            filename=concentric_name,
            fps=10,
        )
        display(Image(filename=str(FIGURES_DIR / concentric_name)))

    return rho_histories

_ = run_lwr_examples(make_animations=True)

Out [19]:

Fundamental diagram loaded: ρ in [0.000, 1.000], c_max ≈ 11.641
Case low density: dt=0.0773, n_steps=5174

Saved LWR space–time plot: figures/lwr_spacetime_low_density.png
Generating GIF 'lwr_ring_low_density.gif' with 399 frames (original T=5175).
LWR ring GIF saved to: figures/lwr_ring_low_density.gif

Case mid density: dt=0.0773, n_steps=5174

Saved LWR space–time plot: figures/lwr_spacetime_mid_density.png
Generating GIF 'lwr_ring_mid_density.gif' with 399 frames (original T=5175).
LWR ring GIF saved to: figures/lwr_ring_mid_density.gif

Case high density: dt=0.0773, n_steps=5174

Saved LWR space–time plot: figures/lwr_spacetime_high_density.png
Generating GIF 'lwr_ring_high_density.gif' with 399 frames (original T=5175).
LWR ring GIF saved to: figures/lwr_ring_high_density.gif

Generating concentric GIF 'lwr_ring_concentric_all.gif' with 399 frames (min T=5175).
LWR concentric rings GIF saved to: figures/lwr_ring_concentric_all.gif

def coarse_grain_space(
    occ_history: np.ndarray,
    window_radius: int = 3,
) -> np.ndarray:
    """Spatial coarse-graining of binary occupation on a 1D ring.

    For each time t and cell i, compute the local mean density:
        ρ_micro(t, i) = average of occ(t, j) for j in [i-R, ..., i+R]

    using periodic boundaries. Window size is W = 2*R+1.
    """
    T, L = occ_history.shape
    R = window_radius
    W = 2 * R + 1

    kernel = np.ones(W, dtype=float) / W
    rho_micro = np.zeros((T, L), dtype=float)

    for n in range(T):
        occ_float = occ_history[n].astype(float)
        # Periodic padding
        padded = np.concatenate([occ_float[-R:], occ_float, occ_float[:R]])
        conv = np.convolve(padded, kernel, mode="valid")
        rho_micro[n, :] = conv

    return rho_micro


def run_micro_macro_case(
    density: float,
    name_tag: str,
    L: int = 200,
    vmax: int = 5,
    p_slow: float = 0.3,
    n_warmup: int = 500,
    T_compare: int = 400,
    window_radius: int = 3,
    CFL: float = 0.9,
    rng_seed: int = 1234,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Run a NaSch vs LWR comparison for a fixed density.

    Steps:
      1) Simulate NaSch for n_warmup + (T_compare-1) steps, saving history.
      2) Spatial coarse-graining -> ρ_micro_all(t,x).
      3) Extract comparison window of size T_compare:
            t = n_warmup, ..., n_warmup + T_compare - 1
         and call it ρ_micro_cmp(t,x).
      4) Use ρ_micro_cmp[0] as initial condition for LWR.
      5) Simulate LWR for T_compare-1 steps with that rho_init.
      6) Compute error E(t) at each time (including t=0).

    Returns
    -------
    rho_micro_cmp : (T_compare, L) array, coarse-grained micro density.
    rho_macro : (T_compare, L) array, LWR solution.
    E_t : (T_compare,) array, L2 error at each time.
    """
    n_steps_total = n_warmup + T_compare

    print(f"\n=== Case {name_tag} ===")
    print(f"NaSch: density={density:.3f}, L={L}, vmax={vmax}, p_slow={p_slow}")
    print(f"  warmup={n_warmup}, T_compare={T_compare}, n_steps_total={n_steps_total}")

    # 1) NaSch simulation
    results_nasch = simulate_nasch(
        L=L,
        density=density,
        vmax=vmax,
        p_slow=p_slow,
        n_steps=n_steps_total,
        warmup=0,
        save_history=True,
        rng_seed=rng_seed,
    )

    occ_history = results_nasch["occ_history"]

    # 2) Spatial coarse-graining
    rho_micro_all = coarse_grain_space(occ_history, window_radius=window_radius)

    # 3) Comparison window
    rho_micro_cmp = rho_micro_all[n_warmup : n_warmup + T_compare, :]
    rho_init = rho_micro_cmp[0, :]

    # 4) LWR with same initial condition
    fd = load_fundamental_diagram("data/fundamental_diagram.npz")
    rho_macro, dt = simulate_lwr_ring_nsteps(
        fd,
        L=L,
        n_steps=T_compare - 1,
        CFL=CFL,
        rho_init=rho_init,
    )

    print(
        f"LWR: dt={dt:.4f}, steps={rho_macro.shape[0]-1}, "
        f"T_physical≈{(rho_macro.shape[0]-1)*dt:.1f}"
    )

    # 5) Error E(t)
    diff = rho_macro - rho_micro_cmp
    E_t = np.sqrt((diff**2).mean(axis=1))

    # Save data
    out_npz = DATA_DIR / f"micro_macro_{name_tag}.npz"
    np.savez(
        out_npz,
        rho_micro=rho_micro_cmp,
        rho_macro=rho_macro,
        E_t=E_t,
        density=density,
        L=L,
        vmax=vmax,
        p_slow=p_slow,
        n_warmup=n_warmup,
        T_compare=T_compare,
        window_radius=window_radius,
        CFL=CFL,
        dt_LWR=dt,
    )
    print(f"NaSch/LWR data saved to: {out_npz}")

    # 6) Comparative plots

    vmin = min(rho_micro_cmp.min(), rho_macro.min())
    vmax_plot = max(rho_micro_cmp.max(), rho_macro.max())

    # (a) Space–time NaSch vs LWR
    fig, axes = plt.subplots(2, 1, figsize=(10, 6), sharex=True)

    im0 = axes[0].imshow(
        rho_micro_cmp,
        aspect="auto",
        origin="lower",
        cmap="viridis",
        vmin=vmin,
        vmax=vmax_plot,
    )
    axes[0].set_ylabel("t (NaSch)")
    axes[0].set_title(rf"NaSch (micro, coarse-grained) — $\rho \approx {density:.2f}__STASH_1790932053694_5__quot;)

    im1 = axes[1].imshow(
        rho_macro,
        aspect="auto",
        origin="lower",
        cmap="viridis",
        vmin=vmin,
        vmax=vmax_plot,
    )
    axes[1].set_ylabel("t (LWR)")
    axes[1].set_xlabel("Position (cell)")
    axes[1].set_title("LWR (macro)")

    fig.tight_layout(rect=[0.0, 0.0, 0.9, 1.0])
    cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
    cbar = fig.colorbar(im1, cax=cbar_ax)
    cbar.set_label(r"Density $\rho__STASH_1790932053694_5__quot;)

    fig_path = FIGURES_DIR / f"micro_macro_spacetime_{name_tag}.png"
    fig.savefig(fig_path, dpi=200)
    display(fig)
    plt.close(fig)
    print(f"Space–time NaSch vs LWR figure saved to: {fig_path}")

    # (b) Error E(t)
    fig2, ax2 = plt.subplots(figsize=(6, 4))
    t_indices = np.arange(T_compare)
    ax2.plot(t_indices, E_t, "-o", markersize=3)
    ax2.set_xlabel("Time step (comparison)")
    ax2.set_ylabel(r"$E(t)__STASH_1790932053694_5__quot;)
    ax2.set_title(rf"Error $E(t)$ NaSch vs LWR — $\rho \approx {density:.2f}__STASH_1790932053694_5__quot;)
    ax2.grid(alpha=0.3)
    fig2.tight_layout()
    fig2_path = FIGURES_DIR / f"micro_macro_error_{name_tag}.png"
    fig2.savefig(fig2_path, dpi=200)
    display(fig2)
    plt.close(fig2)
    print(f"Error E(t) figure saved to: {fig2_path}")

    return rho_micro_cmp, rho_macro, E_t

# Common parameters
L = 200
vmax = 5
p_slow = 0.3
n_warmup = 500
T_compare = 400
window_radius = 3
CFL = 0.9

densities = [0.0, 0.015, 0.025, 0.05, 0.10, 0.12, 0.15, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.8, 0.90, 0.95, 1.0]
seeds = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 160, 170] # Extend seeds to match the number of densities

for density, seed in zip(densities, seeds):
    name_tag = f"ρ≈{density:.3f}"
    run_micro_macro_case(
        density=density,
        name_tag=name_tag,
        L=L,
        vmax=vmax,
        p_slow=p_slow,
        n_warmup=n_warmup,
        T_compare=T_compare,
        window_radius=window_radius,
        CFL=CFL,
        rng_seed=seed,
    )

Out [23]:

=== Case ρ≈0.000 ===
NaSch: density=0.000, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.000.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.000.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.000.png

=== Case ρ≈0.015 ===
NaSch: density=0.015, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.015.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.015.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.015.png

=== Case ρ≈0.025 ===
NaSch: density=0.025, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.025.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.025.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.025.png

=== Case ρ≈0.050 ===
NaSch: density=0.050, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.050.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.050.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.050.png

=== Case ρ≈0.100 ===
NaSch: density=0.100, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.100.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.100.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.100.png

=== Case ρ≈0.120 ===
NaSch: density=0.120, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.120.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.120.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.120.png

=== Case ρ≈0.150 ===
NaSch: density=0.150, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.150.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.150.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.150.png

=== Case ρ≈0.200 ===
NaSch: density=0.200, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.200.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.200.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.200.png

=== Case ρ≈0.300 ===
NaSch: density=0.300, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.300.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.300.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.300.png

=== Case ρ≈0.400 ===
NaSch: density=0.400, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.400.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.400.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.400.png

=== Case ρ≈0.500 ===
NaSch: density=0.500, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.500.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.500.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.500.png

=== Case ρ≈0.600 ===
NaSch: density=0.600, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.600.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.600.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.600.png

=== Case ρ≈0.700 ===
NaSch: density=0.700, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.700.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.700.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.700.png

=== Case ρ≈0.800 ===
NaSch: density=0.800, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.800.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.800.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.800.png

=== Case ρ≈0.900 ===
NaSch: density=0.900, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.900.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.900.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.900.png

=== Case ρ≈0.950 ===
NaSch: density=0.950, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈0.950.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈0.950.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈0.950.png

=== Case ρ≈1.000 ===
NaSch: density=1.000, L=200, vmax=5, p_slow=0.3
  warmup=500, T_compare=400, n_steps_total=900
LWR: dt=0.0773, steps=399, T_physical≈30.8
NaSch/LWR data saved to: data/micro_macro_ρ≈1.000.npz

Space–time NaSch vs LWR figure saved to: figures/micro_macro_spacetime_ρ≈1.000.png

Error E(t) figure saved to: figures/micro_macro_error_ρ≈1.000.png

from pathlib import Path
import matplotlib.colors as mcolors

densities = [0.0, 0.015, 0.025, 0.05, 0.10, 0.12, 0.15, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.8, 0.90, 0.95, 1.0]
mm_paths = [Path(f"data/micro_macro_ρ≈{d:.3f}.npz") for d in densities] # Changed to .3f for consistency

fig, ax = plt.subplots(figsize=(12, 6)) # Aumentar el tamaño de la figura para acomodar la leyenda a un lado

# Custom list of contrasting colors as requested
# These are chosen to be distinct and vibrant
cust_colors = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#FF00FF', '#00FFFF', '#FFA500', '#800080', '#008000', '#FFD700', '#DA70D6', '#20B2AA', '#F08080', '#00BFFF', '#BA55D3', '#FF4500', '#7FFF00'] # Added more distinct colors
colors = mcolors.to_rgba_array(cust_colors[:len(densities)]) # Use only as many colors as densities

for p, density, color in zip(mm_paths, densities, colors):
    if not p.exists():
        raise FileNotFoundError(
            f"Missing {p}. Run the micro–macro comparison cell above first."
        )
    data = np.load(p)
    E_t = data["E_t"]
    t = np.arange(E_t.shape[0])

    label = f"ρ ≈ {density:.3f}" # Changed to .3f for consistency
    if density == 0.0: # Special style for density 0 to make it visible
        ax.plot(t, E_t, label=label, color=color, lw=3.0, linestyle='--', marker='o', markersize=5, alpha=0.9)
    else:
        ax.plot(t, E_t, label=label, color=color, lw=2.0, alpha=0.8)

# Oформación del gráfico
ax.set_xlabel("Time step (comparison)")
ax.set_ylabel(r"$E(t)__STASH_1790932053694_6__quot;)
ax.set_title("NaSch vs LWR error over time (Combined Densities)")
ax.grid(alpha=0.3)
ax.legend(frameon=False, fontsize=10, loc="center left", bbox_to_anchor=(1, 0.5), title="Density") # Mover la leyenda fuera del gráfico

fig.tight_layout()

out_path = Path("figures/micro_macro_error_combined.png")
out_path.parent.mkdir(parents=True, exist_ok=True)

fig.savefig(out_path, dpi=200)
display(fig)
plt.close(fig)
print(f"Combined error figure saved to: {out_path}")

Out [24]:

Combined error figure saved to: figures/micro_macro_error_combined.png

# Parameters for enhanced GIF visualization
T_GIF = 600          # NaSch/LWR steps for the GIF (longer than before)
GIF_MAX_FRAMES = 400  # Maximum number of frames to display
GIF_FPS = 15          # Frames per second (faster than before)
MACRO_CONTRAST = 3.0  # Factor to amplify contrast in macro (visual only)


def build_char_speed_from_fd(fd: FundamentalDiagram) -> Tuple[np.ndarray, np.ndarray]:
    """Build a table (rho_grid, dQdrho_grid) from a FundamentalDiagram.

    Computes dQ/drho using numerical gradient.
    """
    rho_grid = np.asarray(fd.rho_grid)
    Q_grid = np.asarray(fd.Q_grid)
    dQdrho_grid = np.gradient(Q_grid, rho_grid)
    return rho_grid, dQdrho_grid


def compute_characteristics(
    rho_history: np.ndarray,
    dt: float,
    rho_grid: np.ndarray,
    dQdrho_grid: np.ndarray,
    n_chars: int = 20,
    dx: float = 1.0,
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute approximate characteristic curves for LWR:

        dx/dt = c(ρ) = dQ/dρ,

    using the local density from rho_history as the wave speed.

    Parameters
    ----------
    rho_history : (T, L) array
        Density evolution from LWR.
    dt : float
        LWR time step.
    rho_grid : np.ndarray
        Density grid from fundamental diagram.
    dQdrho_grid : np.ndarray
        dQ/dρ values on that grid.
    n_chars : int
        Number of characteristic curves to compute.
    dx : float
        Spatial cell size.

    Returns
    -------
    t_vals : (T,) array
        Physical time values.
    x_chars : (n_chars, T) array
        Characteristic trajectories x(t).
    """
    T, L = rho_history.shape
    L_phys = L * dx

    def char_speed(rho_val: np.ndarray) -> np.ndarray:
        rho_clip = np.clip(rho_val, rho_grid[0], rho_grid[-1])
        return np.interp(rho_clip, rho_grid, dQdrho_grid)

    x0 = np.linspace(0.0, L_phys, n_chars, endpoint=False)
    x_chars = np.zeros((n_chars, T), dtype=float)
    x_chars[:, 0] = x0

    for n in range(T - 1):
        idx = (x_chars[:, n] / dx).astype(int) % L
        rho_local = rho_history[n, idx]
        c_local = char_speed(rho_local)

        x_next = x_chars[:, n] + c_local * dt
        x_next = np.mod(x_next, L_phys)
        x_chars[:, n + 1] = x_next

    t_vals = dt * np.arange(T)
    return t_vals, x_chars


def run_lwr_characteristics_for_density(
    rho0: float,
    name_tag: str,
    L: int = 200,
    CFL: float = 0.9,
    T_steps: int = 400,
) -> None:
    """Simulate LWR for mean density rho0 and generate a figure with
    space–time diagram and characteristic curves overlaid.
    """
    fd = load_fundamental_diagram(DATA_DIR / "fundamental_diagram.npz")
    rho_grid, dQdrho_grid = build_char_speed_from_fd(fd)

    x = np.arange(L)
    center = 0.25 * L
    width = 0.08 * L
    # Perturbation amplitude adapted to mean density
    amp = 0.15 * (1.0 - rho0)

    rho_init = rho0 + amp * np.exp(-0.5 * ((x - center) / width) ** 2)
    rho_init = np.clip(rho_init, fd.rho_min, fd.rho_max)

    rho_hist, dt = simulate_lwr_ring_nsteps(
        fd=fd,
        L=L,
        n_steps=T_steps - 1,
        CFL=CFL,
        rho_init=rho_init,
    )

    t_vals, x_chars = compute_characteristics(
        rho_history=rho_hist,
        dt=dt,
        rho_grid=rho_grid,
        dQdrho_grid=dQdrho_grid,
        n_chars=20,
        dx=1.0,
    )

    fig, ax = plt.subplots(figsize=(10, 5))

    im = ax.imshow(
        rho_hist,
        aspect="auto",
        origin="lower",
        cmap="viridis",
        extent=[0, L, t_vals[0], t_vals[-1]],
    )
    ax.set_xlabel("Position (cell)")
    ax.set_ylabel("Physical time $t__STASH_1790932053694_7__quot;)
    ax.set_title(
        f"LWR: space–time and characteristics ($\\rho_0 \\approx {rho0:.2f}$)"
    )

    # Overlay characteristic curves
    for k in range(x_chars.shape[0]):
        ax.plot(
            x_chars[k, :],
            t_vals,
            color="w",
            lw=0.8,
            alpha=0.9,
        )

    cbar = fig.colorbar(im, ax=ax)
    cbar.set_label(r"Density $\rho(x,t)__STASH_1790932053694_7__quot;)

    fig.tight_layout()
    out_path = FIGURES_DIR / f"lwr_characteristics_{name_tag}_density.png"
    fig.savefig(out_path, dpi=200)
    display(fig)
    plt.close(fig)
    print(f"Characteristics figure ({name_tag}) saved to: {out_path}")


def simulate_micro_macro_for_density(
    density: float,
    name_tag: str,
    L: int = 200,
    vmax: int = 5,
    p_slow: float = 0.3,
    n_warmup: int = 500,
    T_gif: int = T_GIF,
    window_radius: int = 3,
    CFL: float = 0.9,
    rng_seed: int = 1234,
) -> Tuple[np.ndarray, np.ndarray]:
    """Run a coherent NaSch + LWR simulation for a given density.

    Steps:
      - Simulate NaSch with warm-up + window for GIF.
      - Apply spatial coarse-graining -> ρ_micro(t,x).
      - Use ρ_micro(0,x) as initial condition for LWR.
      - Simulate LWR for the same number of steps (T_gif-1).
      - Return ρ_micro(t,x) and ρ_macro(t,x) with shape (T_gif, L).

    These results are used to build the 6-ring GIF.
    """
    n_steps_total = n_warmup + T_gif

    print(f"\n=== Case {name_tag} (density={density:.3f}) ===")
    print(f"NaSch: L={L}, vmax={vmax}, p_slow={p_slow}")
    print(f"  warmup={n_warmup}, T_gif={T_gif}, n_steps_total={n_steps_total}")

    results_nasch = simulate_nasch(
        L=L,
        density=density,
        vmax=vmax,
        p_slow=p_slow,
        n_steps=n_steps_total,
        warmup=0,
        save_history=True,
        rng_seed=rng_seed,
    )

    occ_history = results_nasch["occ_history"]

    # Window for GIF
    occ_cmp = occ_history[n_warmup : n_warmup + T_gif, :]
    rho_micro = coarse_grain_space(occ_cmp, window_radius=window_radius)

    # LWR with same initial condition
    fd = load_fundamental_diagram(DATA_DIR / "fundamental_diagram.npz")
    rho_init = rho_micro[0, :]

    rho_macro, dt = simulate_lwr_ring_nsteps(
        fd=fd,
        L=L,
        n_steps=T_gif - 1,
        CFL=CFL,
        rho_init=rho_init,
    )

    print(
        f"LWR: dt={dt:.4f}, steps={rho_macro.shape[0]-1}, "
        f"T_physical ≈ {(rho_macro.shape[0]-1)*dt:.1f}"
    )

    return rho_micro, rho_macro


def make_six_ring_micro_macro_gif(
    rho_micro_list: List[np.ndarray],
    rho_macro_list: List[np.ndarray],
    out_name: str = "six_ring_micro_macro_comparison.gif",
    fps: int = GIF_FPS,
    max_frames: int = GIF_MAX_FRAMES,
) -> None:
    """Build a GIF with 6 concentric rings:

      low_ρ micro, low_ρ macro,
      mid_ρ micro, mid_ρ macro,
      high_ρ micro, high_ρ macro.

    Improvements:
      - Longer simulation (T_GIF)
      - Temporal subsampling for smooth animation
      - Amplified visual contrast in macro rings.
    """
    assert len(rho_micro_list) == 3
    assert len(rho_macro_list) == 3

    # Check all have the same shape
    T_list = [arr.shape[0] for arr in rho_micro_list]
    L_list = [arr.shape[1] for arr in rho_micro_list]
    for arr in rho_macro_list:
        T_list.append(arr.shape[0])
        L_list.append(arr.shape[1])
    T = min(T_list)
    L = L_list[0]
    if not all(Li == L for Li in L_list):
        raise ValueError("All simulations must have the same L.")

    # Temporal subsampling for GIF
    if T <= max_frames:
        frame_indices = np.arange(T)
    else:
        step = int(np.ceil(T / max_frames))
        frame_indices = np.arange(0, T, step, dtype=int)

    print(
        f"Generating 6-ring GIF with {len(frame_indices)} frames "
        f"(original T={T})."
    )

    # Prepare "visible" versions of macro with more contrast
    rho_macro_vis_list: List[np.ndarray] = []
    for k in range(3):
        rho_macro = rho_macro_list[k][:T]
        mean_macro = rho_macro.mean()
        rho_macro_vis = mean_macro + MACRO_CONTRAST * (rho_macro - mean_macro)
        rho_macro_vis = np.clip(rho_macro_vis, 0.0, 1.0)
        rho_macro_vis_list.append(rho_macro_vis)

    # Density ranges and colormaps per density (shared micro/macro_vis)
    vmin_list = []
    vmax_list = []
    cmap_list = ["Blues", "Greens", "Reds"]

    for k in range(3):
        vals = np.concatenate(
            [rho_micro_list[k][:T].ravel(), rho_macro_vis_list[k][:T].ravel()]
        )
        # Use percentiles to avoid outliers and increase contrast
        vmin = float(np.percentile(vals, 1))
        vmax = float(np.percentile(vals, 99))
        vmin_list.append(vmin)
        vmax_list.append(vmax)
        print(
            f"  Density {k}: visual range ρ ≈ [{vmin_list[-1]:.3f}, "
            f"{vmax_list[-1]:.3f}]"
        )

    # Define radii for 6 rings
    base_r = 1.0
    dr = 0.35
    n_rings = 6
    radii = [base_r + i * dr for i in range(n_rings)]

    theta = 2.0 * np.pi * np.arange(L) / L
    x_all_rings = [r * np.cos(theta) for r in radii]
    y_all_rings = [r * np.sin(theta) for r in radii]

    fig, ax = plt.subplots(figsize=(7, 7))
    ax.set_aspect("equal")
    ax.set_axis_off()

    # Very faint guide circles
    for x_all, y_all in zip(x_all_rings, y_all_rings):
        ax.plot(x_all, y_all, lw=0.8, alpha=0.25, color="lightgray")

    # Scatter for each ring
    # Order:
    #   0: low micro
    #   1: low macro
    #   2: mid micro
    #   3: mid macro
    #   4: high micro
    #   5: high macro
    ring_info = []
    scat_list = []

    frame0 = frame_indices[0]

    for ring_idx in range(n_rings):
        dens_idx = ring_idx // 2
        is_micro = (ring_idx % 2 == 0)
        cmap = cmap_list[dens_idx]
        vmin = vmin_list[dens_idx]
        vmax = vmax_list[dens_idx]

        if is_micro:
            data0 = rho_micro_list[dens_idx][frame0]
        else:
            data0 = rho_macro_vis_list[dens_idx][frame0]

        scat = ax.scatter(
            x_all_rings[ring_idx],
            y_all_rings[ring_idx],
            c=data0,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            s=20,
            edgecolors="none",
        )
        scat_list.append(scat)
        ring_info.append((dens_idx, is_micro))

    labels = [
        r"low $\rho$ (micro)",
        r"low $\rho$ (macro)",
        r"mid $\rho$ (micro)",
        r"mid $\rho$ (macro)",
        r"high $\rho$ (micro)",
        r"high $\rho$ (macro)",
    ]
    for ring_idx, label in enumerate(labels):
        r = radii[ring_idx]
        ax.text(
            0.0,
            r + 0.05,
            label,
            ha="center",
            va="bottom",
            fontsize=8,
        )

    max_r = max(radii)
    ax.set_xlim(-max_r - 0.5, max_r + 0.5)
    ax.set_ylim(-max_r - 0.5, max_r + 0.5)

    def init():
        t0 = frame_indices[0]
        for scat, (dens_idx, is_micro) in zip(scat_list, ring_info):
            if is_micro:
                arr0 = rho_micro_list[dens_idx][t0]
            else:
                arr0 = rho_macro_vis_list[dens_idx][t0]
            scat.set_array(arr0)
        ax.set_title(f"t = {t0}")
        return tuple(scat_list)

    def update(frame_idx: int):
        for scat, (dens_idx, is_micro) in zip(scat_list, ring_info):
            if is_micro:
                arr = rho_micro_list[dens_idx][frame_idx]
            else:
                arr = rho_macro_vis_list[dens_idx][frame_idx]
            scat.set_array(arr)
        ax.set_title(f"t = {frame_idx}")
        return tuple(scat_list)

    anim = animation.FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=frame_indices,
        interval=1000 / fps,
        blit=False,
    )

    out_path = FIGURES_DIR / out_name
    writer = animation.PillowWriter(fps=fps)
    anim.save(out_path, writer=writer)

    plt.close(fig)
    print(f"6-ring micro/macro GIF saved to: {out_path}")
    # Display the GIF inline
    display(Image(filename=str(out_path)))

# Coherent micro-macro simulations for three densities
L = 200
vmax = 5
p_slow = 0.3
n_warmup = 500
window_radius = 3
CFL = 0.9

densities = [0.05, 0.20, 0.50]
seeds = [10, 20, 30]

rho_micro_list = []
rho_macro_list = []

for density, seed, tag in zip(densities, seeds, ["low", "mid", "high"]):
    rho_micro, rho_macro = simulate_micro_macro_for_density(
        density=density,
        name_tag=tag,
        L=L,
        vmax=vmax,
        p_slow=p_slow,
        n_warmup=n_warmup,
        T_gif=T_GIF,
        window_radius=window_radius,
        CFL=CFL,
        rng_seed=seed,
    )
    rho_micro_list.append(rho_micro)
    rho_macro_list.append(rho_macro)

# Generate the 6-ring GIF
make_six_ring_micro_macro_gif(
    rho_micro_list=rho_micro_list,
    rho_macro_list=rho_macro_list,
    out_name="six_ring_micro_macro_comparison.gif",
    fps=GIF_FPS,
    max_frames=GIF_MAX_FRAMES,
)

Out [28]:

=== Case low (density=0.050) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3

=== Case mid (density=0.200) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3

=== Case high (density=0.500) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
Generating 6-ring GIF with 300 frames (original T=600).
  Density 0: visual range ρ ≈ [0.000, 0.152]
  Density 1: visual range ρ ≈ [0.000, 0.990]
  Density 2: visual range ρ ≈ [0.000, 1.000]
6-ring micro/macro GIF saved to: figures/six_ring_micro_macro_comparison.gif

def make_multi_ring_micro_macro_gif(
    rho_micro_list: list[np.ndarray],
    rho_macro_list: list[np.ndarray],
    densities_for_labels: list[float],
    out_name: str = "multi_ring_micro_macro_comparison.gif",
    fps: int = GIF_FPS,
    max_frames: int = GIF_MAX_FRAMES,
) -> None:
    """Build a GIF with multiple concentric rings interleaving micro (NaSch) and macro (LWR).

    Each pair of rings corresponds to a density regime.
    """
    import matplotlib.pyplot as plt
    import matplotlib.lines as mlines # Added mlines for custom legend handles
    from matplotlib import animation
    import matplotlib.colors as mcolors # Added mcolors for colormap handling

    n_cases = len(rho_micro_list)
    if n_cases == 0:
        raise ValueError("At least one micro/macro pair is required.")
    if n_cases != len(rho_macro_list):
        raise ValueError("rho_micro_list and rho_macro_list must have the same length.")
    if n_cases != len(densities_for_labels):
        raise ValueError("densities_for_labels must have the same length as rho_micro_list.")

    # Check for consistent L and find minimum T
    T_list = []
    L_list = []
    for k in range(n_cases):
        T_list.append(rho_micro_list[k].shape[0])
        L_list.append(rho_micro_list[k].shape[1])
        T_list.append(rho_macro_list[k].shape[0])
        L_list.append(rho_macro_list[k].shape[1])

    T = min(T_list)
    L = L_list[0]
    if not all(li == L for li in L_list):
        raise ValueError("All simulations must have the same L.")

    # Temporal subsampling for GIF
    if T <= max_frames:
        frame_indices = np.arange(T)
    else:
        step = int(np.ceil(T / max_frames))
        frame_indices = np.arange(0, T, step, dtype=int)

    print(
        f"Generando {2 * n_cases}-anillos GIF con {len(frame_indices)} frames "
        f"(original T={T})."
    )

    # Prepare "visible" versions of macro with more contrast
    rho_macro_vis_list: list[np.ndarray] = []
    for k in range(n_cases):
        rho_macro = rho_macro_list[k][:T]
        mean_macro = rho_macro.mean()
        rho_macro_vis = mean_macro + MACRO_CONTRAST * (rho_macro - mean_macro)
        rho_macro_vis = np.clip(rho_macro_vis, 0.0, 1.0)
        rho_macro_vis_list.append(rho_macro_vis)

    # Density ranges and colormaps per density (shared micro/macro_vis)
    vmin_list = []
    vmax_list = []
    # Extended list of diverse colormaps
    cmap_cycle = ["Blues", "Greens", "Reds", "Oranges", "Purples", "Greys", "YlOrBr", "PuRd", "GnBu", "BuGn"]

    for k in range(n_cases):
        vals = np.concatenate(
            [rho_micro_list[k][:T].ravel(), rho_macro_vis_list[k][:T].ravel()]
        )
        # Use percentiles to avoid outliers and increase contrast
        vmin = float(np.percentile(vals, 1))
        vmax = float(np.percentile(vals, 99))
        vmin_list.append(vmin)
        vmax_list.append(vmax)
        print(
            f"  Densidad {densities_for_labels[k]:.3f}: rango visual ρ ≈ [{vmin_list[-1]:.3f}, "
            f"{vmax_list[-1]:.3f}]"
        )

    # Define radii for all rings
    base_r = 1.0
    dr = 0.45 # Increased spacing between rings
    n_total_rings = 2 * n_cases
    radii = [base_r + i * dr for i in range(n_total_rings)]

    theta = 2.0 * np.pi * np.arange(L) / L
    x_all_rings = [r * np.cos(theta) for r in radii]
    y_all_rings = [r * np.sin(theta) for r in radii]

    # Adjust figsize to make room for legend
    fig, ax = plt.subplots(figsize=(10, 8)) # Increased width
    ax.set_aspect("equal")
    ax.set_axis_off()

    # Very faint guide circles
    for x_all, y_all in zip(x_all_rings, y_all_rings):
        ax.plot(x_all, y_all, lw=0.8, alpha=0.25, color="lightgray")

    # Scatter for each ring
    scat_list = []
    ring_data_map = [] # To map ring_idx to (dens_idx, is_micro)

    frame0 = frame_indices[0]

    # Prepare legend handles and labels
    legend_handles = []
    legend_labels = []

    for ring_idx in range(n_total_rings):
        dens_idx = ring_idx // 2
        is_micro = (ring_idx % 2 == 0)
        cmap_name = cmap_cycle[dens_idx % len(cmap_cycle)]
        cmap_obj = plt.cm.get_cmap(cmap_name)
        vmin = vmin_list[dens_idx]
        vmax = vmax_list[dens_idx]

        if is_micro:
            data0 = rho_micro_list[dens_idx][frame0]
        else:
            data0 = rho_macro_vis_list[dens_idx][frame0]

        scat = ax.scatter(
            x_all_rings[ring_idx],
            y_all_rings[ring_idx],
            c=data0,
            cmap=cmap_obj,
            norm=mcolors.Normalize(vmin=vmin, vmax=vmax), # Use Normalize for consistent color mapping
            s=20,
            edgecolors="none",
        )
        scat_list.append(scat)
        ring_data_map.append((dens_idx, is_micro))

        # Create legend handles and labels dynamically
        label_type = "micro" if is_micro else "macro"
        label_text = fr"$\rho \approx {densities_for_labels[dens_idx]:.2f}$ ({label_type})"
        # Use a Line2D object as a proxy for the scatter points in the legend
        legend_handles.append(mlines.Line2D([], [], color=cmap_obj(0.5), marker='o', linestyle='None', markersize=5))
        legend_labels.append(label_text)

    max_r = max(radii)
    ax.set_xlim(-max_r - 0.5, max_r + 0.5)
    ax.set_ylim(-max_r - 0.5, max_r + 0.5)

    def init():
        t0 = frame_indices[0]
        for scat, (dens_idx, is_micro) in zip(scat_list, ring_data_map):
            if is_micro:
                arr0 = rho_micro_list[dens_idx][t0]
            else:
                arr0 = rho_macro_vis_list[dens_idx][t0]
            scat.set_array(arr0)
        ax.set_title(f"t = {t0}")
        return tuple(scat_list)

    def update(frame_idx_in_indices: int): # frame_idx_in_indices is an index into frame_indices
        current_frame = frame_indices[frame_idx_in_indices]
        for scat, (dens_idx, is_micro) in zip(scat_list, ring_data_map):
            if is_micro:
                arr = rho_micro_list[dens_idx][current_frame]
            else:
                arr = rho_macro_vis_list[dens_idx][current_frame]
            scat.set_array(arr)
        ax.set_title(f"t = {current_frame}")
        return tuple(scat_list)

    anim = animation.FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=len(frame_indices), # Iterate over the length of frame_indices
        interval=1000 / fps,
        blit=False,
    )

    # Add the legend outside the plot
    ax.legend(handles=legend_handles, labels=legend_labels, loc='center left', bbox_to_anchor=(1, 0.5), title="Densities (Micro/Macro)", frameon=False, fontsize=8)

    fig.tight_layout(rect=[0, 0, 0.85, 1]) # Adjust tight_layout to make space for the legend

    out_path = FIGURES_DIR / out_name
    writer = animation.PillowWriter(fps=fps)
    anim.save(out_path, writer=writer)

    plt.close(fig)
    print(f"{n_total_rings}-ring micro/macro GIF saved to: {out_path}")
    # Display the GIF inline
    display(Image(filename=str(out_path)))

expanded_densities = [0.01, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 0.9]
expanded_seeds = [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]

# Common parameters for simulation
L = 200
vmax = 5
p_slow = 0.3
n_warmup = 500
window_radius = 3
CFL = 0.9

rho_micro_histories_for_gif = []
rho_macro_histories_for_gif = []

print("Starting simulations for multi-ring GIF...")
for i, density in enumerate(expanded_densities):
    seed = expanded_seeds[i % len(expanded_seeds)] # Cycle through seeds if fewer seeds than densities
    name_tag = f"multi_ring_rho_{density:.2f}"

    print(f"Processing density: {density:.2f} with seed: {seed}")
    rho_micro, rho_macro = simulate_micro_macro_for_density(
        density=density,
        name_tag=name_tag,
        L=L,
        vmax=vmax,
        p_slow=p_slow,
        n_warmup=n_warmup,
        T_gif=T_GIF, # Using the global T_GIF parameter
        window_radius=window_radius,
        CFL=CFL,
        rng_seed=seed,
    )
    rho_micro_histories_for_gif.append(rho_micro)
    rho_macro_histories_for_gif.append(rho_macro)

print("All simulations complete. Generating multi-ring GIF...")

# Generate the multi-ring GIF
make_multi_ring_micro_macro_gif(
    rho_micro_list=rho_micro_histories_for_gif,
    rho_macro_list=rho_macro_histories_for_gif,
    densities_for_labels=expanded_densities,
    out_name="multi_ring_micro_macro_comparison_expanded.gif",
    fps=GIF_FPS,
    max_frames=GIF_MAX_FRAMES,
)

print("Multi-ring GIF generation process finished.")

Out [30]:

Starting simulations for multi-ring GIF...
Processing density: 0.01 with seed: 1001

=== Case multi_ring_rho_0.01 (density=0.010) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
Processing density: 0.05 with seed: 1002

=== Case multi_ring_rho_0.05 (density=0.050) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
Processing density: 0.10 with seed: 1003

=== Case multi_ring_rho_0.10 (density=0.100) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
Processing density: 0.20 with seed: 1004

=== Case multi_ring_rho_0.20 (density=0.200) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
Processing density: 0.30 with seed: 1005

=== Case multi_ring_rho_0.30 (density=0.300) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
Processing density: 0.50 with seed: 1006

=== Case multi_ring_rho_0.50 (density=0.500) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
Processing density: 0.70 with seed: 1007

=== Case multi_ring_rho_0.70 (density=0.700) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
Processing density: 0.90 with seed: 1008

=== Case multi_ring_rho_0.90 (density=0.900) ===
NaSch: L=200, vmax=5, p_slow=0.3
  warmup=500, T_gif=600, n_steps_total=1100
LWR: dt=0.0773, steps=599, T_physical ≈ 46.3
All simulations complete. Generating multi-ring GIF...
Generando 16-anillos GIF con 300 frames (original T=600).
  Densidad 0.010: rango visual ρ ≈ [0.000, 0.143]
  Densidad 0.050: rango visual ρ ≈ [0.000, 0.212]
  Densidad 0.100: rango visual ρ ≈ [0.000, 0.229]
  Densidad 0.200: rango visual ρ ≈ [0.000, 0.917]
  Densidad 0.300: rango visual ρ ≈ [0.000, 1.000]
  Densidad 0.500: rango visual ρ ≈ [0.000, 1.000]
  Densidad 0.700: rango visual ρ ≈ [0.286, 1.000]
  Densidad 0.900: rango visual ρ ≈ [0.571, 1.000]

/tmp/ipython-input-2254860142.py:114: MatplotlibDeprecationWarning: The get_cmap function was deprecated in Matplotlib 3.7 and will be removed in 3.11. Use ``matplotlib.colormaps[name]`` or ``matplotlib.colormaps.get_cmap()`` or ``pyplot.get_cmap()`` instead.
  cmap_obj = plt.cm.get_cmap(cmap_name)

16-ring micro/macro GIF saved to: figures/multi_ring_micro_macro_comparison_expanded.gif

Multi-ring GIF generation process finished.