gratinglab 0.1.0__py3-none-any.whl

gratinglab/metrology/core/processing.py

@@ -0,0 +1,344 @@
+"""
+Core data processing functions for AFM analysis
+"""
+import numpy as np
+from scipy.signal import find_peaks
+
+
+def load_afm_data(filename, default_scan_size=2.0, settings=None):
+    """
+    Load AFM data, whatever format it arrived in.
+
+    Dispatches on content rather than on the caller knowing: raw Nanoscope files
+    (.spm, and the extensionless companions Nanoscope also writes) are read
+    directly; everything else takes the text path used for Gwyddion exports.
+
+    Dispatching here rather than at each call site is deliberate - there are six
+    callers, and none of them should have to care.
+
+    Parameters:
+        settings: optional AnalysisSettings, consulted for which .spm channel and
+            scan direction to read. Ignored for text files.
+
+    Returns:
+        data: 2D array of height data, in metres
+        scan_x_size: scan width in microns
+    """
+    from .. import io as _io  # noqa: F401 - keeps the package import explicit
+    from ..io.spm import (DEFAULT_CHANNEL, DEFAULT_DIRECTION, is_nanoscope_file,
+                          read_spm)
+
+    if is_nanoscope_file(filename):
+        channel = getattr(settings, 'spm_channel', None) or DEFAULT_CHANNEL
+        direction = getattr(settings, 'spm_direction', None) or DEFAULT_DIRECTION
+        data, scan_x_size = read_spm(filename, channel=channel,
+                                     direction=direction)
+        print(f"  Nanoscope file: {channel} / {direction}, "
+              f"{data.shape[0]}x{data.shape[1]} px, {scan_x_size:g} µm")
+        return data, scan_x_size
+
+    return _load_text(filename, default_scan_size)
+
+
+def _load_text(filename, default_scan_size=2.0):
+    """
+    Load a text export (Gwyddion and similar), reading the scan size from the
+    header comments when they carry it.
+
+    Unchanged behaviour - this is the original load_afm_data body, moved so the
+    dispatcher above stays readable. The stored baselines depend on it.
+    """
+    import re
+    
+    scan_x_size = default_scan_size
+    
+    # Try to read scan size from header
+    try:
+        with open(filename, 'r', encoding="utf-8", errors="replace") as f:
+            # Read first 10 lines to look for scan size info
+            header_lines = [f.readline() for _ in range(10)]
+            
+            for line in header_lines:
+                # Common AFM file formats include scan size info
+                # Look for patterns like "Width: 2.000 µm"
+                if 'width' in line.lower():
+                    match = re.search(r'(\d+\.?\d*)\s*(um|µm|micron)', line, re.IGNORECASE)
+                    if match:
+                        scan_x_size = float(match.group(1))
+                        print(f"  Detected scan width from header: {scan_x_size} µm")
+                        break
+                    
+                    match = re.search(r'(\d+\.?\d*)\s*nm', line, re.IGNORECASE)
+                    if match:
+                        scan_x_size = float(match.group(1)) / 1000  # Convert nm to µm
+                        print(f"  Detected scan width from header: {scan_x_size} µm ({match.group(1)} nm)")
+                        break
+                        
+                elif 'scan' in line.lower() and 'size' in line.lower():
+                    # Extract number
+                    match = re.search(r'(\d+\.?\d*)\s*(um|µm|micron)', line, re.IGNORECASE)
+                    if match:
+                        scan_x_size = float(match.group(1))
+                        print(f"  Detected scan size from header: {scan_x_size} µm")
+                        break
+    except Exception:
+        pass
+    
+    # Load the data
+    try:
+        # Another program's file: only the numbers are parsed, so an
+        # undecodable byte in the skipped header must not stop the load.
+        with open(filename, encoding="utf-8", errors="replace") as handle:
+            data = np.genfromtxt(handle, skip_header=4, skip_footer=0)
+    except Exception as e:
+        raise Exception(f"Error loading {filename}: {e}") from e
+    
+    if scan_x_size == default_scan_size:
+        print(f"  Using default scan size: {scan_x_size} µm")
+    
+    return data, scan_x_size
+
+
+def raw_data(data, scan_x_size):
+    """Extract displacement and height profile from 2D AFM data (averages all rows)"""
+    disp_um = scan_x_size * np.arange(0, len(data[0])) / (len(data[0]) - 1)
+    groov_nm = (10**9) * (np.mean(data, axis=0) - np.min(np.mean(data, axis=0)))
+    return disp_um, groov_nm
+
+
+def raw_data_multi_group(data, scan_x_size, n_groups=10):
+    """
+    Extract multiple height profiles from different row groups in 2D AFM data
+    
+    This function divides the AFM image into n_groups horizontal bands and extracts
+    an averaged profile from each band. This allows for many more measurements per
+    image and enables assessment of within-image variation.
+    
+    Parameters:
+        data: 2D array of AFM height data (rows × columns)
+        scan_x_size: scan width in microns
+        n_groups: number of row groups to create (default: 10)
+        
+    Returns:
+        disp_um: 1D array of x-positions in microns (same for all groups)
+        profiles_nm: list of n_groups 1D arrays, each containing a height profile in nm
+        group_info: dict with metadata about the grouping
+    """
+    n_rows = data.shape[0]
+    n_cols = data.shape[1]
+    
+    # Calculate x-displacement (same for all profiles)
+    disp_um = scan_x_size * np.arange(0, n_cols) / (n_cols - 1)
+    
+    # Determine row ranges for each group
+    rows_per_group = n_rows // n_groups
+    
+    if rows_per_group < 3:
+        print(f"Warning: Only {rows_per_group} rows per group with n_groups={n_groups}")
+        print(f"  Reducing to n_groups={n_rows // 3} to ensure at least 3 rows per group")
+        n_groups = max(1, n_rows // 3)
+        rows_per_group = n_rows // n_groups
+    
+    profiles_nm = []
+    group_ranges = []
+    
+    for i in range(n_groups):
+        # Calculate row range for this group
+        start_row = i * rows_per_group
+        
+        # For the last group, include any remaining rows
+        if i == n_groups - 1:
+            end_row = n_rows
+        else:
+            end_row = (i + 1) * rows_per_group
+        
+        # Extract and average rows in this group
+        group_data = data[start_row:end_row, :]
+        profile_avg = np.mean(group_data, axis=0)
+        
+        # Convert to nm and subtract minimum
+        profile_nm = (10**9) * (profile_avg - np.min(profile_avg))
+        
+        profiles_nm.append(profile_nm)
+        group_ranges.append((start_row, end_row))
+    
+    # Package metadata
+    group_info = {
+        'n_groups': n_groups,
+        'n_rows': n_rows,
+        'rows_per_group': rows_per_group,
+        'group_ranges': group_ranges
+    }
+    
+    print(f"  Extracted {n_groups} row-group profiles")
+    print(f"    Rows per group: {rows_per_group}")
+    print(f"    Total rows used: {n_groups * rows_per_group} of {n_rows}")
+    
+    return disp_um, profiles_nm, group_info
+
+
+def flatten_profile(x, y, method='linear', poly_order=2, exclude_edges=0.0, 
+                   feature='both', period_nm=None):
+    """
+    Remove background tilt/curvature from AFM profile
+    
+    Parameters:
+        method: 'linear', 'polynomial', 'groove_peaks', or 'level_grooves'
+        poly_order: polynomial order (for 'polynomial' and 'level_grooves' methods)
+        exclude_edges: fraction of data to exclude from edges (0.0 to 0.5)
+        feature: for 'level_grooves': 'peaks', 'troughs', or 'both'
+        period_nm: estimated period for finding grooves (for 'level_grooves')
+    
+    Returns:
+        y_flat: flattened profile
+        background: the removed background
+    """
+    
+    if method == 'linear':
+        poly_coeffs = np.polyfit(x, y, 1)
+        background = np.polyval(poly_coeffs, x)
+        
+    elif method == 'polynomial':
+        if exclude_edges > 0:
+            n_exclude = int(len(x) * exclude_edges)
+            x_fit = x[n_exclude:-n_exclude]
+            y_fit = y[n_exclude:-n_exclude]
+            poly_coeffs = np.polyfit(x_fit, y_fit, poly_order)
+        else:
+            poly_coeffs = np.polyfit(x, y, poly_order)
+        
+        background = np.polyval(poly_coeffs, x)
+        
+    elif method == 'groove_peaks':
+        peaks, _ = find_peaks(y, distance=len(y)//20)
+        
+        if len(peaks) < 3:
+            print("Warning: Not enough peaks found for groove_peaks method, using polynomial")
+            poly_coeffs = np.polyfit(x, y, poly_order)
+            background = np.polyval(poly_coeffs, x)
+        else:
+            poly_coeffs = np.polyfit(x[peaks], y[peaks], poly_order)
+            background = np.polyval(poly_coeffs, x)
+            
+    elif method == 'level_grooves':
+        if period_nm is None:
+            raise ValueError("period_nm must be provided for level_grooves method")
+        
+        feature_x = []
+        feature_y = []
+        
+        if feature in ['peaks', 'both']:
+            dx_nm = (x[1] - x[0]) * 1000
+            min_distance = int(0.5 * period_nm / dx_nm)  # Changed from 0.3 to 0.5 for more reliable detection
+            # Use a lower prominence to catch more peaks
+            height_range = np.max(y) - np.min(y)
+            min_prominence = 0.2 * height_range  # Require peaks to be at least 20% of height range
+            peaks, _ = find_peaks(y, distance=min_distance, prominence=min_prominence)
+            if len(peaks) > 0:
+                feature_x.extend(x[peaks])
+                feature_y.extend(y[peaks])
+        
+        if feature in ['troughs', 'both']:
+            y_inv = -y
+            dx_nm = (x[1] - x[0]) * 1000
+            min_distance = int(0.5 * period_nm / dx_nm)  # Changed from 0.3 to 0.5
+            height_range = np.max(y) - np.min(y)
+            min_prominence = 0.2 * height_range
+            troughs, _ = find_peaks(y_inv, distance=min_distance, prominence=min_prominence)
+            if len(troughs) > 0:
+                feature_x.extend(x[troughs])
+                feature_y.extend(y[troughs])
+        
+        if len(feature_x) < 2:
+            print("    Warning: Not enough features found for level_grooves, using linear")
+            poly_coeffs = np.polyfit(x, y, 1)
+            background = np.polyval(poly_coeffs, x)
+        else:
+            feature_x = np.array(feature_x)
+            feature_y = np.array(feature_y)
+            
+            if exclude_edges > 0:
+                x_min = x[0] + exclude_edges * (x[-1] - x[0])
+                x_max = x[-1] - exclude_edges * (x[-1] - x[0])
+                mask = (feature_x >= x_min) & (feature_x <= x_max)
+                feature_x = feature_x[mask]
+                feature_y = feature_y[mask]
+            
+            if len(feature_x) < 2:
+                print("    Warning: Not enough features after edge exclusion, using all")
+            
+            poly_coeffs = np.polyfit(feature_x, feature_y, poly_order)
+            background = np.polyval(poly_coeffs, x)
+    else:
+        raise ValueError(f"Unknown flattening method: {method}")
+    
+    y_flat = y - background
+    
+    return y_flat, background
+
+
+def find_groove_positions(x, y, period_nm, prominence_factor=0.3, distance_factor=0.7,
+                          edge_exclusion=0.0, return_n_edge_rejected=False):
+    """
+    Find positions of groove centers (minima)
+
+    Parameters:
+        edge_exclusion: reject centers lying within this many periods of either
+            end of the profile. A groove that close to the edge has no complete
+            facet in the data - the scan starts or stops part-way through it -
+            so any angle fitted to it is meaningless. 0.0 disables the check.
+        return_n_edge_rejected: also return how many centers were rejected
+
+    Returns:
+        peaks, or (peaks, n_rejected) if return_n_edge_rejected
+    """
+    y_inv = -y  # Invert so grooves become peaks
+    dx_nm = (x[1] - x[0]) * 1000
+    period_idx = int(period_nm / dx_nm)
+    height_range = np.max(y) - np.min(y)
+    min_prominence = prominence_factor * height_range
+
+    peaks, properties = find_peaks(y_inv,
+                                   distance=int(distance_factor * period_idx),
+                                   prominence=min_prominence)
+
+    n_rejected = 0
+    if edge_exclusion > 0 and len(peaks) > 0:
+        margin_idx = edge_exclusion * period_idx
+        last = len(x) - 1
+        keep = (peaks >= margin_idx) & (peaks <= last - margin_idx)
+        n_rejected = int((~keep).sum())
+        peaks = peaks[keep]
+
+    if return_n_edge_rejected:
+        return peaks, n_rejected
+    return peaks
+
+
+def extract_single_groove(x, y, center_idx, period_nm, margin=0.1, allow_asymmetric=True):
+    """
+    Extract a single groove around center_idx
+    
+    Parameters:
+        allow_asymmetric: If True, allows asymmetric extraction near edges
+    """
+    dx_nm = (x[1] - x[0]) * 1000
+    half_width = int(round((period_nm / 2) * (1 + margin) / dx_nm))
+    
+    if allow_asymmetric:
+        # Allow asymmetric extraction for edge grooves
+        left_extent = min(half_width, center_idx)
+        right_extent = min(half_width, len(x) - 1 - center_idx)
+        
+        start = center_idx - left_extent
+        end = center_idx + right_extent + 1
+    else:
+        # Original symmetric extraction
+        L = min(half_width, center_idx, len(x) - 1 - center_idx)
+        start = center_idx - L
+        end = center_idx + L + 1
+    
+    x_groove = x[start:end]
+    y_groove = y[start:end]
+    
+    return x_groove, y_groove

gratinglab/metrology/core/tip.py

@@ -0,0 +1,233 @@
+"""
+Tip convolution: simulate it, and undo what can be undone.
+
+An AFM image is not the surface. It is the surface **dilated** by the tip
+shape -- at every pixel the tip descends until it touches, and the recorded
+height is where the *apex* stopped, which is above the true surface anywhere
+the contact happened on the tip's flank instead of its apex. Villarrubia
+(1997, J. Res. NIST 102, 425) gives the exact morphological statement and the
+two operations this module implements:
+
+- ``dilate`` -- image formation. What a tip of the given shape would report
+  for a known surface. Exact, and used to *make* test data: dilating the
+  synthetic scan produces an image whose true surface is known, which no real
+  scan allows.
+- ``erode`` -- reconstruction. The **least upper bound** on the true surface
+  consistent with the image and the tip: everywhere the apex made contact the
+  bound *is* the surface, and everywhere else the tip physically could not
+  reach, so no algorithm can recover what is there. The certainty map marks
+  the difference, which is the entire point: a reconstruction without one
+  manufactures exactly the false confidence it exists to remove.
+
+The tip is a cone with a spherical apex cap -- ``radius_nm`` and
+``half_angle_deg`` (measured from the tip axis, so the flank rises at
+``90 - half_angle`` degrees from the surface). The consequence worth knowing
+before trusting any corrected groove: a facet steeper than the flank is
+touched only at its top corner and is **unrecoverable**, and a facet close to
+the flank angle is the marginal case -- an 18 deg tip has a 72 deg flank,
+almost parallel to a 70.5 deg anti-blaze facet. That near-parallelism, not
+the nanometre apex, is what rounds the troughs of blazed gratings.
+
+**One dimension, deliberately.** The operations run along the fast-scan axis
+(rows), not over the full 2-D image. For a grating the surface is invariant
+along the grooves, and dilating an invariant surface with a cone reduces
+exactly to dilating each cross-section with the cone's cross-section -- the
+groove-parallel extent of the tip touches nothing the transverse extent does
+not. What a 2-D pass would add is sensitivity to along-groove features, and
+it would need the slow-axis pixel pitch, which the text-export format does
+not record.
+
+All heights are in metres (the unit the loaders promise), tip parameters in
+nm, and the pixel pitch comes from the scan width exactly as
+``core.processing.raw_data`` computes it.
+"""
+from dataclasses import dataclass
+
+import numpy as np
+
+#: Values ``AnalysisSettings.tip_correction`` may take. 'none' is the default:
+#: the correction changes measured numbers and must be asked for, not applied
+#: silently.
+VALID_TIP_CORRECTIONS = ('none', 'erosion')
+
+#: Height differences below this (metres) are numerically meaningless: 1e-12 m
+#: is a thousandth of a nanometre, far under any AFM noise floor. Used to call
+#: two candidate contacts a tie, which marks the pixel uncertain.
+_TIE_TOL_M = 1e-12
+
+
+def tip_cross_section(dx_nm, radius_nm, half_angle_deg, max_height_nm):
+    """The tip's transverse profile, sampled at the pixel pitch.
+
+    Returns heights in nm above the apex at lateral offsets
+    ``k * dx_nm`` for ``k in [-K, K]`` (odd length, apex at the centre,
+    ``t[K] == 0``). ``K`` is chosen so the profile just exceeds
+    ``max_height_nm`` -- parts of the tip higher than the tallest feature can
+    never make contact and would only widen every window.
+
+    Sphere of radius ``R`` capping a cone of half-angle ``theta``: the two
+    meet where the sphere's tangent matches the cone flank, at lateral offset
+    ``R cos(theta)`` and height ``R (1 - sin(theta))``.
+    """
+    if radius_nm <= 0:
+        raise ValueError(f"tip radius must be positive, got {radius_nm} nm")
+    if not 0.0 < half_angle_deg < 90.0:
+        raise ValueError(
+            f"tip half angle must lie in (0, 90) degrees, got {half_angle_deg}")
+
+    theta = np.radians(half_angle_deg)
+    r_tangent = radius_nm * np.cos(theta)
+    t_tangent = radius_nm * (1.0 - np.sin(theta))
+
+    # Lateral extent at which the tip profile reaches max_height_nm.
+    if max_height_nm <= t_tangent:
+        r_max = float(np.sqrt(max(radius_nm**2 - (radius_nm - max_height_nm)**2,
+                                  0.0)))
+    else:
+        r_max = float(r_tangent + (max_height_nm - t_tangent) * np.tan(theta))
+
+    # At least the nearest neighbours, so a sub-pixel tip degrades to a
+    # near-no-op rather than an empty window.
+    k = max(int(np.ceil(r_max / dx_nm)), 1)
+    r = np.abs(np.arange(-k, k + 1)) * dx_nm
+
+    sphere = radius_nm - np.sqrt(np.maximum(radius_nm**2 - r**2, 0.0))
+    cone = t_tangent + (r - r_tangent) / np.tan(theta)
+    return np.where(r <= r_tangent, sphere, cone)
+
+
+def _pixel_pitch_nm(data, scan_x_size):
+    """nm per column, matching ``raw_data``'s displacement axis exactly."""
+    return scan_x_size * 1000.0 / (data.shape[1] - 1)
+
+
+def _shifted(data, k, fill):
+    """``out[:, j] = data[:, j + k]``, with ``fill`` where that runs off."""
+    out = np.full_like(data, fill)
+    if k == 0:
+        return data.copy()
+    if k > 0:
+        out[:, :-k] = data[:, k:]
+    else:
+        out[:, -k:] = data[:, :k]
+    return out
+
+
+def dilate(data, scan_x_size, *, radius_nm, half_angle_deg):
+    """Image formation: what this tip reports for the surface ``data``.
+
+    ``img(x) = max_u [ surface(x + u) - tip(u) ]`` -- the apex height when the
+    tip, centred at ``x``, rests on whichever point it touches first. Always
+    ``>= data`` pointwise: a tip can hide a trough but never dig one.
+
+    Exists for synthesis and testing; the pipeline itself only erodes.
+    """
+    dx_nm = _pixel_pitch_nm(data, scan_x_size)
+    relief_nm = float(data.max() - data.min()) * 1e9
+    tip_m = tip_cross_section(dx_nm, radius_nm, half_angle_deg, relief_nm) * 1e-9
+    k_max = len(tip_m) // 2
+
+    img = np.full_like(data, -np.inf)
+    for k in range(-k_max, k_max + 1):
+        np.maximum(img, _shifted(data, k, -np.inf) - tip_m[k + k_max], out=img)
+    return img
+
+
+@dataclass(frozen=True)
+class TipCorrection:
+    """An eroded image, and how much of it is actually the surface.
+
+    ``data`` is the least upper bound on the true surface: equal to it at
+    every certain pixel, above it (by an unknowable amount) elsewhere. The
+    parameters are carried so downstream records can state what was assumed
+    rather than that something was.
+    """
+
+    data: np.ndarray       #: corrected heights, metres, same shape as input
+    certain: np.ndarray    #: bool; True where the apex provably made contact
+    radius_nm: float
+    half_angle_deg: float
+    footprint_px: int      #: window width used; 3 means the tip is sub-pixel
+
+    @property
+    def certain_fraction(self) -> float:
+        """Fraction of pixels where the correction recovered the surface."""
+        return float(np.mean(self.certain))
+
+    @property
+    def summary(self) -> str:
+        """One line for a log or a metrics sidecar."""
+        note = (" -- tip is sub-pixel at this pitch, correction is a near-no-op"
+                if self.footprint_px <= 3 else "")
+        return (f"erosion, R = {self.radius_nm:g} nm, "
+                f"half angle = {self.half_angle_deg:g} deg: "
+                f"{100.0 * self.certain_fraction:.1f}% of pixels certain, "
+                f"the rest upper bounds{note}")
+
+
+def erode(data, scan_x_size, *, radius_nm, half_angle_deg) -> TipCorrection:
+    """Reconstruct the least upper bound on the surface behind an image.
+
+    ``rec(x) = min_u [ img(x + u) + tip(u) ]`` -- the lowest the surface can
+    be at ``x`` given that the tip, wherever it stood, stopped where the image
+    says. Sits between the truth and the image: ``surface <= rec <= img``.
+
+    The certainty map re-runs the contact argument: re-dilating ``rec``
+    reproduces the image exactly, and for each image pixel the surface point
+    achieving that maximum is where the tip touched. A pixel is certain when
+    it is such a contact point for some image pixel **and** the contact is
+    unique -- a tie means the contact cannot be localised (the flank lying
+    flush along a facet, a wedge gripped at both walls), and every tied point
+    stays an upper bound. Window-truncated pixels at the scan edges are never
+    marked certain.
+    """
+    dx_nm = _pixel_pitch_nm(data, scan_x_size)
+    relief_nm = float(data.max() - data.min()) * 1e9
+    tip_m = tip_cross_section(dx_nm, radius_nm, half_angle_deg, relief_nm) * 1e-9
+    k_max = len(tip_m) // 2
+
+    rec = np.full_like(data, np.inf)
+    for k in range(-k_max, k_max + 1):
+        np.minimum(rec, _shifted(data, k, np.inf) + tip_m[k + k_max], out=rec)
+
+    # Contact scan: for each image pixel, the best (highest) candidate
+    # rec(x+u) - tip(u) is the redilated image; its argmax u is the contact.
+    best = np.full_like(data, -np.inf)
+    second = np.full_like(data, -np.inf)
+    best_k = np.zeros(data.shape, dtype=np.int64)
+    for k in range(-k_max, k_max + 1):
+        candidate = _shifted(rec, k, -np.inf) - tip_m[k + k_max]
+        improves = candidate > best
+        second = np.where(improves, best, np.maximum(second, candidate))
+        best = np.where(improves, candidate, best)
+        best_k = np.where(improves, k, best_k)
+
+    unique = np.isfinite(best) & (best - second > _TIE_TOL_M)
+    certain = np.zeros(data.shape, dtype=bool)
+    rows, cols = np.nonzero(unique)
+    certain[rows, cols + best_k[rows, cols]] = True
+
+    return TipCorrection(
+        data=rec,
+        certain=certain,
+        radius_nm=radius_nm,
+        half_angle_deg=half_angle_deg,
+        footprint_px=2 * k_max + 1,
+    )
+
+
+def apply_tip_correction(data, scan_x_size, settings):
+    """The pipeline entry point: dispatch on ``settings.tip_correction``.
+
+    Returns ``(data, correction)`` -- the array to analyse and the
+    :class:`TipCorrection` behind it, or ``(data, None)`` untouched when the
+    setting is ``'none'``. Callers record ``correction.summary`` wherever
+    their output lands, because a corrected depth and an uncorrected one are
+    different measurements and must not be filed under the same description.
+    """
+    if settings.tip_correction == 'none':
+        return data, None
+    correction = erode(data, scan_x_size,
+                       radius_nm=settings.tip_radius_nm,
+                       half_angle_deg=settings.tip_half_angle_deg)
+    return correction.data, correction

gratinglab/metrology/data/__init__.py

@@ -0,0 +1,16 @@
+"""
+Data management and aggregation
+Combining scans and grouping by temperature
+"""
+
+from .aggregation import (
+    group_by_temperature,
+    extract_temperatures_for_output,
+    combine_scans
+)
+
+__all__ = [
+    'group_by_temperature',
+    'extract_temperatures_for_output',
+    'combine_scans'
+]