AFM_script.py

# analysis_script.py
import numpy as np
from pathlib import Path
from skimage.transform import resize
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import matplotlib.patheffects as path_effects

# ----------------------------
# Unit scaling helpers
# ----------------------------
# Multipliers to convert FROM base SI to target display units:
#   - length_map: meters -> target length units
#   - volt_map:   volts  -> target voltage units
LENGTH_MAP = {
    "pm": 1e12,
    "nm": 1e9,
    "um": 1e6,   # allow "um" as ASCII micro
    "µm": 1e6,   # and "µm" (micro sign)
    "mm": 1e3,
    "cm": 1e2,
    "m": 1.0,
    "km": 1e-3,
}
VOLT_MAP = {
    "pv": 1e12,
    "nv": 1e9,
    "uv": 1e6,   # "uv" and "µv"
    "µv": 1e6,
    "mv": 1e3,
    "v": 1.0,
    "kv": 1e-3,
}

def _normalize_units(units: str) -> str:
    """Normalize a units string for lookup (lowercase, unify µ)."""
    if units is None:
        return ""
    u = units.strip()
    # treat Greek mu and ASCII 'u' as the same
    u = u.replace("μ", "µ")
    return u.lower()

def _unit_multiplier(units: str, log=None) -> float:
    """
    Return the scale factor to convert FROM base SI (meters or volts) TO the requested units.
    If unknown, warn and return 1.0.
    """
    _log = log or print
    key = _normalize_units(units)
    if key in LENGTH_MAP:
        return LENGTH_MAP[key]
    if key in VOLT_MAP:
        return VOLT_MAP[key]
    if key == "":
        _log("No units provided; using raw base SI values (factor=1).")
    else:
        _log(f"Warning: units '{units}' not recognized; using raw values (factor=1).")
    return 1.0

def _units_suffix(units: str) -> str:
    """Make a filesystem-friendly suffix for filenames (replace µ with u)."""
    if not units:
        return "units"
    return units.replace("μ", "µ").replace("µ", "u").replace("/", "-").replace(" ", "")

# ----------------------------
# Main analysis
# ----------------------------
def analyze2(
    input_dir: str | Path,
    output_dir: str | Path,
    expected_shape=(256, 256),
    target_size=(1024, 1024),
    contrast_percentiles=(2, 98),
    scalebar_length: float = 4.0,
    colormap: str = "viridis",
    colorscale_label: str = "nm",
    show_true_scale: bool = False,
    show_high_contrast: bool = True,
    save_figures: bool = True,
    save_shifted_data: bool = False,
    save_parameters: bool = False,       # <-- NEW
    real_width_um: float = 20.0,         # <-- needed for surface area calc
    log=None,
    **_,
):
    """
    Read AFM .txt matrices (Z in base SI: meters or volts), rescale Z into user-chosen units,
    then render/save either:
      - True scale panel (absolute min..max),
      - High contrast panel (robust percentiles),
      - or both side-by-side.

    - colorscale_label: text shown on the colorbar (also used to determine the multiplier)
    - colormap: Matplotlib palette name
    """
    input_dir = Path(input_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    def _log(msg: str):
        (log or print)(msg)

    if not (show_true_scale or show_high_contrast):
        raise ValueError("At least one of 'show_true_scale' or 'show_high_contrast' must be True.")

    # Resolve scaling factor based on desired units label
    scale_units = colorscale_label or "nm"
    mult = _unit_multiplier(scale_units, log=_log)  # e.g., "nm" => 1e9; "mV" => 1e3

    files = sorted([f for f in input_dir.glob("*.txt")])
    if not files:
        _log(f"No HeightRetrace .txt files found in: {input_dir}")
        return {"processed": 0, "skipped": 0, "saved": 0}

    processed = 0
    skipped = 0
    saved = 0
    label = f"{scalebar_length:g} µm"
    unit_suffix = _units_suffix(scale_units)

    for file in files:
        try:
            data_raw = np.loadtxt(file)  # Z in base SI (meters or volts)
        except Exception as e:
            _log(f"Skipping {file.name}: could not read ({e})")
            skipped += 1
            continue

        # Rotate + flip to keep original orientation
        data_rf = np.fliplr(np.rot90(data_raw, k=1))

        # Enforce expected size
        if data_rf.shape != tuple(expected_shape):
            _log(f"Skipping {file.name} due to unexpected shape: {data_rf.shape}")
            skipped += 1
            continue

        # Convert from base SI to requested units, then shift so min = 0 in those units
        data_units = data_rf * mult
        zmin = float(np.min(data_units))
        data_units = data_units - zmin

        # Upscale for nicer viewing (preserve numeric range)
        data_up = resize(
            data_units, tuple(target_size), order=3, anti_aliasing=True, preserve_range=True
        )

        # Compute ranges in chosen units
        real_min = float(np.min(data_units))  # should be 0.0
        real_max = float(np.max(data_units))
        if real_max == real_min:
            real_max = real_min + 1.0  # avoid zero span

        p_lo, p_hi = np.percentile(data_units, contrast_percentiles)
        if p_hi <= p_lo:
            p_lo, p_hi = real_min, real_max

        # Decide which panels to draw
        panels = []
        suffix_bits = []
        if show_true_scale:
            panels.append(("True scale", real_min, real_max))
            suffix_bits.append("true")
        if show_high_contrast:
            panels.append(("High contrast", p_lo, p_hi))
            suffix_bits.append("high")

        n_panels = len(panels)
        figsize = (5.7 * n_panels, 5) if n_panels == 1 else (11.5, 5)
        fig, axes = plt.subplots(1, n_panels, figsize=figsize, constrained_layout=True)
        if n_panels == 1:
            axes = [axes]

        for ax, (title, vmin, vmax) in zip(axes, panels):
            im = ax.imshow(
                data_up,
                cmap=colormap,
                vmin=vmin,
                vmax=vmax,
                extent=[0, 20, 0, 20],  # 20×20 µm field of view
                origin="upper",
            )
            ax.set_title(f"{file.stem} — {title}", fontsize=12)
            cbar = plt.colorbar(im, ax=ax, orientation="vertical", pad=0.02)
            cbar.ax.tick_params(direction="in", labelsize=12)
            cbar.outline.set_edgecolor("black")
            cbar.outline.set_linewidth(1.0)
            cbar.ax.set_title(scale_units, fontsize=14, pad=4)  # units label here

            # Remove axes
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_xlabel("")
            ax.set_ylabel("")

            # Scale bar rectangle
            scalebar_start = 1     # µm
            scalebar_height = 0.2  # µm
            rect = Rectangle(
                (scalebar_start, 1), scalebar_length, scalebar_height,
                linewidth=1.0, edgecolor="black", facecolor="white"
            )
            ax.add_patch(rect)
            text = ax.text(
                scalebar_start + scalebar_length / 2,
                1.3,
                label,
                color="white",
                ha="center",
                va="bottom",
                fontsize=11,
            )
            text.set_path_effects([
                path_effects.Stroke(linewidth=1.5, foreground="black"),
                path_effects.Normal()
            ])

        processed += 1

        # Save outputs
        if save_figures:
            suffix = "_and_".join(suffix_bits) if suffix_bits else "image"
            out_png = output_dir / f"{file.stem}_{suffix}.png"
            fig.savefig(out_png, dpi=300, bbox_inches="tight")
            saved += 1

        if save_shifted_data:
            out_txt = output_dir / f"{file.stem}_shifted_{unit_suffix}.txt"
            np.savetxt(out_txt, data_units, fmt="%.6f")
            saved += 1
               
                # --- Optional: save calculated parameters ---
        if save_parameters:
            try:
                params = compute_topography_parameters(data_units, real_width_um)
                out_param = output_dir / f"{file.stem}_parameters.txt"
                with open(out_param, "w", encoding="utf-8") as f:
                    f.write(f"Topography parameters for {file.name}\n")
                    f.write(f"Units: {scale_units}\n")
                    f.write(f"Real width: {real_width_um} µm\n\n")
                    for k, v in params.items():
                        f.write(f"{k:25s}: {v:.6f}\n")
                saved += 1
                _log(f"Saved parameters → {out_param.name}")
            except Exception as e:
                _log(f"Warning: could not compute/save parameters for {file.name} ({e})")

        plt.close(fig)

    _log(f"Done. Processed={processed}, Skipped={skipped}, Saved files={saved}")
    return {"processed": processed, "skipped": skipped, "saved": saved}




from skimage import measure
import pyvista as pv

def compute_topography_parameters(Z: np.ndarray, real_width_um: float):
    """
    Compute statistical and geometric parameters from a topography/voltage map.

    Parameters
    ----------
    Z : np.ndarray
        2D array of Z values (in chosen display units, e.g. nm or µm).
    real_width_um : float
        The real width of the image in micrometers (same for height).

    Returns
    -------
    dict
        Dictionary with computed parameters:
        - z_max
        - z_min
        - std_dev
        - avg_dev
        - rms
        - total_area_um2
        - area_increase_percent
    """
    if Z.ndim != 2:
        raise ValueError("Z must be a 2D array.")

    # --- Basic statistics ---
    z_max = float(np.max(Z))
    z_min = float(np.min(Z))
    std_dev = float(np.std(Z))
    mean_val = float(np.mean(Z))
    avg_dev = float(np.mean(np.abs(Z - mean_val)))
    rms = float(np.sqrt(np.mean(Z ** 2)))

    # --- Geometric area calculation ---
    H, W = Z.shape
    pixel_size_um = real_width_um / (W-1)  # µm per pixel

    # Generate X, Y grids in µm
    y, x = np.mgrid[0:H, 0:W]
    X = x * pixel_size_um
    Y = y * pixel_size_um

    # Compute 3D surface area (µm²)
    Z_um = Z * 1e6

    surf = pv.StructuredGrid(X, Y, Z_um)
    total_area_um2 = surf.area

    # Flat (projected) area
    flat_area_um2 = (real_width_um ** 2)

    # Area increase from flat (%)
    area_increase_percent = (total_area_um2 / flat_area_um2 - 1.0) * 100.0

    return {
        "z_max": z_max,
        "z_min": z_min,
        "std_dev": std_dev,
        "avg_dev": avg_dev,
        "rms": rms,
        "total_area_um2": total_area_um2,
        "area_increase_percent": area_increase_percent,
    }