benchmark_disori_grid.py

import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from matplotlib.colors import LinearSegmentedColormap

# Batlow colormap data
batlow_data = [
    [0.005193, 0.098238, 0.349842],
    [0.009065, 0.104487, 0.350933],
    [0.012963, 0.110779, 0.351992],
    [0.01653, 0.116913, 0.35307],
    [0.019936, 0.122985, 0.35412],
    [0.023189, 0.129035, 0.355182],
    [0.026291, 0.135044, 0.35621],
    [0.029245, 0.140964, 0.357239],
    [0.032053, 0.146774, 0.358239],
    [0.034853, 0.152558, 0.359233],
    [0.037449, 0.158313, 0.360216],
    [0.039845, 0.163978, 0.361187],
    [0.042104, 0.169557, 0.362151],
    [0.044069, 0.175053, 0.363084],
    [0.045905, 0.18046, 0.364007],
    [0.047665, 0.185844, 0.364915],
    [0.049378, 0.191076, 0.36581],
    [0.050795, 0.196274, 0.366684],
    [0.052164, 0.201323, 0.367524],
    [0.053471, 0.206357, 0.36837],
    [0.054721, 0.211234, 0.369184],
    [0.055928, 0.216046, 0.369974],
    [0.057033, 0.220754, 0.37075],
    [0.058032, 0.22534, 0.371509],
    [0.059164, 0.229842, 0.372252],
    [0.060167, 0.234299, 0.372978],
    [0.061052, 0.238625, 0.373691],
    [0.06206, 0.242888, 0.374386],
    [0.063071, 0.247085, 0.37505],
    [0.063982, 0.251213, 0.375709],
    [0.064936, 0.255264, 0.376362],
    [0.065903, 0.259257, 0.376987],
    [0.066899, 0.263188, 0.377594],
    [0.067921, 0.267056, 0.378191],
    [0.069002, 0.270922, 0.378774],
    [0.070001, 0.274713, 0.379342],
    [0.071115, 0.278497, 0.379895],
    [0.072192, 0.282249, 0.380434],
    [0.07344, 0.285942, 0.380957],
    [0.074595, 0.289653, 0.381452],
    [0.075833, 0.293321, 0.381922],
    [0.077136, 0.296996, 0.382376],
    [0.078517, 0.300622, 0.382814],
    [0.079984, 0.304252, 0.383224],
    [0.081553, 0.307858, 0.383598],
    [0.083082, 0.311461, 0.383936],
    [0.084778, 0.315043, 0.38424],
    [0.086503, 0.318615, 0.384506],
    [0.088353, 0.322167, 0.384731],
    [0.090281, 0.325685, 0.38491],
    [0.092304, 0.32922, 0.38504],
    [0.094462, 0.332712, 0.385116],
    [0.096618, 0.336161, 0.385134],
    [0.099015, 0.339621, 0.38509],
    [0.101481, 0.343036, 0.384981],
    [0.104078, 0.34641, 0.384801],
    [0.106842, 0.349774, 0.384548],
    [0.109695, 0.353098, 0.384217],
    [0.112655, 0.356391, 0.383807],
    [0.115748, 0.359638, 0.38331],
    [0.118992, 0.362849, 0.382713],
    [0.12232, 0.36603, 0.382026],
    [0.125889, 0.36916, 0.381259],
    [0.129519, 0.372238, 0.380378],
    [0.133298, 0.375282, 0.379395],
    [0.137212, 0.378282, 0.378315],
    [0.14126, 0.38124, 0.377135],
    [0.145432, 0.38413, 0.37584],
    [0.149706, 0.386975, 0.374449],
    [0.154073, 0.389777, 0.372934],
    [0.15862, 0.392531, 0.37132],
    [0.163246, 0.395237, 0.369609],
    [0.167952, 0.397889, 0.367784],
    [0.172788, 0.400496, 0.365867],
    [0.177752, 0.403041, 0.363833],
    [0.182732, 0.405551, 0.361714],
    [0.187886, 0.408003, 0.359484],
    [0.19305, 0.410427, 0.357177],
    [0.19831, 0.412798, 0.354767],
    [0.203676, 0.415116, 0.352253],
    [0.209075, 0.417412, 0.349677],
    [0.214555, 0.419661, 0.347019],
    [0.220112, 0.421864, 0.344261],
    [0.225707, 0.424049, 0.341459],
    [0.231362, 0.426197, 0.338572],
    [0.237075, 0.428325, 0.335634],
    [0.242795, 0.430418, 0.332635],
    [0.248617, 0.432493, 0.329571],
    [0.254452, 0.434529, 0.326434],
    [0.26032, 0.436556, 0.323285],
    [0.266241, 0.438555, 0.320085],
    [0.272168, 0.440541, 0.316831],
    [0.278171, 0.442524, 0.313552],
    [0.284175, 0.444484, 0.310243],
    [0.290214, 0.44642, 0.306889],
    [0.296294, 0.448357, 0.303509],
    [0.302379, 0.450282, 0.300122],
    [0.308517, 0.452205, 0.296721],
    [0.314648, 0.454107, 0.293279],
    [0.320834, 0.456006, 0.289841],
    [0.327007, 0.4579, 0.286377],
    [0.333235, 0.459794, 0.282937],
    [0.339469, 0.461685, 0.279468],
    [0.345703, 0.463563, 0.275998],
    [0.351976, 0.46544, 0.272492],
    [0.358277, 0.467331, 0.269037],
    [0.364589, 0.469213, 0.265543],
    [0.370922, 0.471085, 0.262064],
    [0.377291, 0.472952, 0.258588],
    [0.383675, 0.474842, 0.255131],
    [0.39007, 0.476711, 0.251665],
    [0.396505, 0.478587, 0.248212],
    [0.402968, 0.480466, 0.244731],
    [0.409455, 0.482351, 0.241314],
    [0.415967, 0.484225, 0.237895],
    [0.422507, 0.486113, 0.234493],
    [0.429094, 0.488011, 0.231096],
    [0.435714, 0.48989, 0.227728],
    [0.442365, 0.491795, 0.224354],
    [0.449052, 0.493684, 0.221074],
    [0.455774, 0.495585, 0.217774],
    [0.462539, 0.497497, 0.214518],
    [0.469368, 0.499393, 0.211318],
    [0.476221, 0.501314, 0.208148],
    [0.483123, 0.503216, 0.205037],
    [0.490081, 0.505137, 0.201976],
    [0.497089, 0.507058, 0.198994],
    [0.504153, 0.508984, 0.196118],
    [0.511253, 0.510898, 0.193296],
    [0.518425, 0.512822, 0.190566],
    [0.525637, 0.514746, 0.18799],
    [0.532907, 0.516662, 0.185497],
    [0.540225, 0.518584, 0.183099],
    [0.547599, 0.520486, 0.180884],
    [0.555024, 0.522391, 0.178854],
    [0.562506, 0.524293, 0.176964],
    [0.570016, 0.526186, 0.175273],
    [0.577582, 0.528058, 0.173775],
    [0.585199, 0.529927, 0.172493],
    [0.592846, 0.531777, 0.171449],
    [0.60052, 0.533605, 0.170648],
    [0.60824, 0.535423, 0.170104],
    [0.615972, 0.537231, 0.169826],
    [0.623739, 0.539002, 0.169814],
    [0.631513, 0.540752, 0.170075],
    [0.639301, 0.542484, 0.170622],
    [0.647098, 0.544183, 0.171465],
    [0.654889, 0.545863, 0.172603],
    [0.662691, 0.547503, 0.174044],
    [0.670477, 0.549127, 0.175747],
    [0.678244, 0.550712, 0.177803],
    [0.685995, 0.552274, 0.180056],
    [0.69372, 0.553797, 0.18261],
    [0.701421, 0.555294, 0.185478],
    [0.709098, 0.556772, 0.188546],
    [0.716731, 0.558205, 0.191851],
    [0.724322, 0.559628, 0.195408],
    [0.731878, 0.561011, 0.199174],
    [0.739393, 0.562386, 0.203179],
    [0.74685, 0.563725, 0.207375],
    [0.754268, 0.565033, 0.211761],
    [0.761629, 0.566344, 0.216322],
    [0.768942, 0.56763, 0.221045],
    [0.776208, 0.568899, 0.22593],
    [0.783416, 0.570162, 0.230962],
    [0.790568, 0.571421, 0.23616],
    [0.797665, 0.572682, 0.24149],
    [0.804709, 0.573928, 0.246955],
    [0.811692, 0.575187, 0.252572],
    [0.81861, 0.576462, 0.258303],
    [0.825472, 0.577725, 0.264197],
    [0.832272, 0.579026, 0.270211],
    [0.838999, 0.580339, 0.276353],
    [0.845657, 0.581672, 0.282631],
    [0.852247, 0.583037, 0.289036],
    [0.858747, 0.58444, 0.295572],
    [0.865168, 0.585882, 0.302255],
    [0.871505, 0.587352, 0.309112],
    [0.877741, 0.588873, 0.316081],
    [0.883878, 0.59045, 0.323195],
    [0.8899, 0.592087, 0.330454],
    [0.895809, 0.593765, 0.337865],
    [0.90159, 0.595507, 0.345429],
    [0.907242, 0.597319, 0.353142],
    [0.912746, 0.599191, 0.360986],
    [0.918103, 0.601126, 0.368999],
    [0.9233, 0.603137, 0.377139],
    [0.928323, 0.605212, 0.385404],
    [0.933176, 0.607369, 0.393817],
    [0.93785, 0.609582, 0.402345],
    [0.942332, 0.611867, 0.411006],
    [0.946612, 0.614218, 0.419767],
    [0.950697, 0.616649, 0.428624],
    [0.954574, 0.619137, 0.437582],
    [0.958244, 0.621671, 0.446604],
    [0.961696, 0.624282, 0.455702],
    [0.964943, 0.626934, 0.46486],
    [0.967983, 0.629639, 0.474057],
    [0.970804, 0.632394, 0.48329],
    [0.973424, 0.635183, 0.492547],
    [0.975835, 0.638012, 0.501826],
    [0.978052, 0.640868, 0.51109],
    [0.980079, 0.643752, 0.52035],
    [0.981918, 0.646664, 0.529602],
    [0.983574, 0.64959, 0.538819],
    [0.985066, 0.652522, 0.547998],
    [0.986392, 0.65547, 0.557142],
    [0.987567, 0.658422, 0.566226],
    [0.988596, 0.661378, 0.575265],
    [0.989496, 0.664329, 0.584246],
    [0.990268, 0.66728, 0.593174],
    [0.990926, 0.67023, 0.602031],
    [0.991479, 0.673165, 0.610835],
    [0.991935, 0.676091, 0.619575],
    [0.992305, 0.679007, 0.628251],
    [0.992595, 0.681914, 0.636869],
    [0.992813, 0.684815, 0.645423],
    [0.992967, 0.687705, 0.653934],
    [0.993064, 0.690579, 0.662398],
    [0.993111, 0.693451, 0.67081],
    [0.993112, 0.696314, 0.679177],
    [0.993074, 0.699161, 0.687519],
    [0.993002, 0.702006, 0.695831],
    [0.9929, 0.704852, 0.704114],
    [0.992771, 0.707689, 0.71238],
    [0.992619, 0.71053, 0.720639],
    [0.992447, 0.713366, 0.728892],
    [0.992258, 0.71621, 0.737146],
    [0.992054, 0.719049, 0.745403],
    [0.991837, 0.721893, 0.753673],
    [0.991607, 0.724754, 0.761959],
    [0.991367, 0.727614, 0.77027],
    [0.991116, 0.730489, 0.778606],
    [0.990855, 0.733373, 0.786976],
    [0.990586, 0.736265, 0.795371],
    [0.990307, 0.739184, 0.80381],
    [0.990018, 0.742102, 0.812285],
    [0.98972, 0.745039, 0.820804],
    [0.989411, 0.747997, 0.829372],
    [0.989089, 0.750968, 0.837979],
    [0.988754, 0.753949, 0.846627],
    [0.988406, 0.756949, 0.855332],
    [0.988046, 0.759964, 0.864078],
    [0.987672, 0.762996, 0.872864],
    [0.98728, 0.766047, 0.881699],
    [0.986868, 0.769105, 0.890573],
    [0.986435, 0.772184, 0.899493],
    [0.98598, 0.775272, 0.908448],
    [0.985503, 0.778378, 0.917444],
    [0.985002, 0.781495, 0.926468],
    [0.984473, 0.784624, 0.935531],
    [0.983913, 0.787757, 0.944626],
    [0.983322, 0.790905, 0.953748],
    [0.982703, 0.794068, 0.962895],
    [0.982048, 0.797228, 0.97207],
    [0.981354, 0.800406, 0.981267],
]

# Create batlow colormap from data
batlow_cmap = LinearSegmentedColormap.from_list("batlow", batlow_data)

# Publication-quality settings
DPI = 300
plt.rcParams.update(
    {
        "figure.dpi": DPI,
        "savefig.dpi": DPI,
        "font.family": "sans-serif",
        "font.sans-serif": ["Arial", "Helvetica", "DejaVu Sans"],
        "font.size": 14,
        "axes.linewidth": 0.5,
    }
)


def create_disorientation_grid(
    data_file,
    grid_shape=(149, 200),
    dtype_filter="FP32",
    colormap_name="gray_r",
    log_scale=False,
):
    """Create 3x3 grid of disorientation heatmaps for different noise levels and methods."""

    # Load benchmark data
    data = np.load(data_file, allow_pickle=True).item()

    # Filter for specific dtype
    mask = np.array(data["dtype"]) == dtype_filter

    methods = np.array(data["method"])[mask]
    resolutions = np.array(data["dict_resolution"])[mask]
    pca_comps = np.array(data["pca_components"])[mask]
    raw_disorientations = np.array(data["raw_disorientations"], dtype=object)[mask]
    noise_ids = np.array(data["dataset_id"])[mask]

    # Use smallest resolution only
    smallest_resolution = min(np.unique(resolutions))
    res_mask = resolutions == smallest_resolution

    methods = methods[res_mask]
    pca_comps = pca_comps[res_mask]
    raw_disorientations = raw_disorientations[res_mask]
    noise_ids = noise_ids[res_mask]

    # Unique noise levels
    unique_noise_ids = sorted(np.unique(noise_ids))
    n_noise_levels = len(unique_noise_ids)

    # Noise labels
    noise_labels = {1: "Low", 5: "Medium", 10: "High"}

    # Method ordering
    method_order = ["DI", "PCA-512", "PCA-1024"]

    # Select colormap
    if colormap_name == "gray_r":
        cmap = plt.cm.gray_r  # Reversed grayscale (0=dark, 1=light)
    elif colormap_name == "batlow":
        cmap = batlow_cmap
    else:
        cmap = plt.cm.gray_r

    # Calculate figure size
    img_aspect = grid_shape[1] / grid_shape[0]  # width / height
    subplot_height = 3.0
    subplot_width = subplot_height * img_aspect

    # Create figure with 3 columns (no 4th column needed)
    fig, axes = plt.subplots(
        n_noise_levels,
        3,
        figsize=(3 * subplot_width + 1, n_noise_levels * subplot_height),
        squeeze=False,
    )

    # Set colorbar range to maximum cubic disorientation
    # Max cubic disorientation = 2*arccos((2 + sqrt(2)) / 4) ≈ 62.8°
    vmax_orig = 2 * np.arccos((2 + np.sqrt(2)) / 4) * 180 / np.pi  # 62.8 degrees
    if log_scale:
        vmin = np.log10(0.1)  # Minimum value for log scale (0.1 degrees)
        vmax = np.log10(vmax_orig)
    else:
        vmin = 0.0
        vmax = vmax_orig

    # Populate grid
    for row_idx, noise_id in enumerate(unique_noise_ids):
        for col_idx, method_label in enumerate(method_order):
            ax = axes[row_idx][col_idx]

            # Find matching data
            if method_label == "DI":
                method_mask = (noise_ids == noise_id) & (methods == "DI")
            else:
                n_comp = int(method_label.split("-")[1])
                method_mask = (
                    (noise_ids == noise_id) & (methods == "PCA") & (pca_comps == n_comp)
                )

            if not np.any(method_mask):
                ax.text(
                    0.5,
                    0.5,
                    "No Data",
                    ha="center",
                    va="center",
                    transform=ax.transAxes,
                )
                ax.axis("off")
                continue

            # Get the disorientations
            idx = np.where(method_mask)[0][0]
            disoris = raw_disorientations[idx]

            if len(disoris) == 0:
                ax.text(
                    0.5, 0.5, "Empty", ha="center", va="center", transform=ax.transAxes
                )
                ax.axis("off")
                continue

            # Reshape to grid
            H, W = grid_shape
            if len(disoris) != H * W:
                print(
                    f"Warning: Expected {H*W} values, got {len(disoris)} for {method_label} at noise_id {noise_id}"
                )
                ax.text(
                    0.5,
                    0.5,
                    f"Size Mismatch\n{len(disoris)} vs {H*W}",
                    ha="center",
                    va="center",
                    transform=ax.transAxes,
                )
                ax.axis("off")
                continue

            # Convert to heatmap
            try:
                disori_map = disoris.reshape(H, W)
                if log_scale:
                    # Apply log transform, clipping to avoid log(0)
                    disori_map = np.log10(np.clip(disori_map, 0.1, None))
                im = ax.imshow(
                    disori_map, cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest"
                )

                # Add border around the image to distinguish from background
                for spine in ax.spines.values():
                    spine.set_visible(True)
                    spine.set_edgecolor("black")
                    spine.set_linewidth(1.5)

            except Exception as e:
                print(f"Error processing {method_label} at noise_id {noise_id}: {e}")
                ax.text(
                    0.5, 0.5, f"Error", ha="center", va="center", transform=ax.transAxes
                )
                ax.axis("off")
                continue

            ax.set_xticks([])
            ax.set_yticks([])

            # Add column labels on top row
            if row_idx == 0:
                ax.set_title(method_label, fontsize=16, fontweight="bold", pad=10)

            # Add row labels on left column
            if col_idx == 0:
                noise_label = noise_labels.get(noise_id, f"ID {noise_id}")
                ax.text(
                    -0.05,
                    0.5,
                    f"{noise_label} Noise",
                    transform=ax.transAxes,
                    fontsize=16,
                    fontweight="bold",
                    rotation=90,
                    va="center",
                    ha="right",
                )

    # Add colorbar to the right of the grid
    cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
    cbar = fig.colorbar(im, cax=cbar_ax)
    if log_scale:
        # Set colorbar ticks to show actual values, not log values
        # Use standard log scale sequence plus additional labels at 3, 30, 60
        major_tick_values = [0.1, 1, 3, 10, 30, 60]
        major_tick_positions = [
            np.log10(v) for v in major_tick_values if 0.1 <= v <= vmax_orig
        ]
        cbar.set_ticks(major_tick_positions)
        cbar.set_ticklabels(
            [
                f"{v:.1f}" if v < 1 else f"{int(v)}"
                for v in major_tick_values
                if 0.1 <= v <= vmax_orig
            ]
        )

        # Add minor ticks at intermediate log scale values
        minor_tick_values = [
            0.2,
            0.3,
            0.4,
            0.5,
            0.6,
            0.7,
            0.8,
            0.9,
            2,
            3,
            4,
            5,
            6,
            7,
            8,
            9,
            20,
            30,
            40,
            50,
            60,
        ]
        minor_tick_positions = [
            np.log10(v) for v in minor_tick_values if 0.1 <= v <= vmax_orig
        ]
        cbar.set_ticks(minor_tick_positions, minor=True)

        cbar.set_label(
            "Indexing Disorientation (degrees) (logscale)",
            rotation=270,
            labelpad=20,
            fontsize=14,
        )
    else:
        cbar.set_label(
            "Indexing Disorientation (degrees)", rotation=270, labelpad=20, fontsize=14
        )

    # # Add overall title
    # fig.suptitle(
    plt.tight_layout(rect=[0, 0, 0.9, 1])

    # Save figure
    output_dir = Path("benchmark_results")
    cmap_suffix = f"_{colormap_name}" if colormap_name != "gray_r" else ""
    log_suffix = "_log" if log_scale else ""
    output_file = output_dir / f"figure_disori_grid{cmap_suffix}{log_suffix}.png"
    plt.savefig(output_file, dpi=DPI, bbox_inches="tight", facecolor="white")
    print(f"✓ Disorientation grid saved to: {output_file}")
    print(f"  Using {smallest_resolution}° dictionary resolution")
    print(f"  Colorbar range: {vmin:.2f}° to {vmax:.2f}°")
    print(f"  Colormap: {colormap_name}")

    output_pdf = output_dir / f"figure_disori_grid{cmap_suffix}{log_suffix}.pdf"
    plt.savefig(output_pdf, bbox_inches="tight", facecolor="white")
    print(f"✓ Disorientation grid PDF saved to: {output_pdf}")

    plt.close(fig)

    # Print summary statistics
    print("\n" + "=" * 70)
    print("DISORIENTATION STATISTICS")
    print("=" * 70)
    for row_idx, noise_id in enumerate(unique_noise_ids):
        noise_label = noise_labels.get(noise_id, f"ID {noise_id}")
        print(f"\n{noise_label} Noise (ID={noise_id}):")
        for method_label in method_order:
            # Find matching data
            if method_label == "DI":
                method_mask = (noise_ids == noise_id) & (methods == "DI")
            else:
                n_comp = int(method_label.split("-")[1])
                method_mask = (
                    (noise_ids == noise_id) & (methods == "PCA") & (pca_comps == n_comp)
                )

            if np.any(method_mask):
                idx = np.where(method_mask)[0][0]
                disoris = raw_disorientations[idx]
                if len(disoris) > 0:
                    mean_disori = np.mean(disoris)
                    median_disori = np.median(disoris)
                    max_disori = np.max(disoris)
                    frac_above_3 = (disoris > 3.0).mean() * 100
                    print(
                        f"  {method_label:12s}: mean={mean_disori:5.2f}°, median={median_disori:5.2f}°, max={max_disori:5.2f}°, >3°={frac_above_3:5.1f}%"
                    )


if __name__ == "__main__":
    data_file = Path("benchmark_results/benchmark_dictionary.npy")

    # Generate grids with different colormaps
    for cmap_name in ["gray_r", "batlow"]:
        print(f"\nGenerating grid with {cmap_name} colormap...")
        create_disorientation_grid(
            data_file=data_file,
            grid_shape=(149, 200),  # Adjust based on your scan dimensions
            dtype_filter="FP32",
            colormap_name=cmap_name,
        )

    # Generate batlow with log scale
    print(f"\nGenerating grid with batlow colormap (log scale)...")
    create_disorientation_grid(
        data_file=data_file,
        grid_shape=(149, 200),
        dtype_filter="FP32",
        colormap_name="batlow",
        log_scale=True,
    )

    print("\n" + "=" * 70)
    print("Disorientation grid generation complete!")
    print("=" * 70)