diff --git a/tutorials/diffraction/polymer_peak_detection.ipynb b/tutorials/diffraction/polymer_peak_detection.ipynb new file mode 100644 index 0000000..f60aff5 --- /dev/null +++ b/tutorials/diffraction/polymer_peak_detection.ipynb @@ -0,0 +1,224 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro", + "metadata": {}, + "source": [ + "# Polymer 4D-STEM peak detection with the pinned paper model\n", + "\n", + "This workflow resolves an immutable, checksum-verified model, reuses its local cache, loads a calibrated 4D-STEM scan, detects peaks, builds count maps and flowlines, and exports selected figures. The model record remains private during review, so set `QUANTEM_POLYMER_MODEL_DIR` to the supplied local archive. Once the immutable public record exists, remove that override; the same pinned call will download and cache it atomically." + ] + }, + { + "cell_type": "markdown", + "id": "install", + "metadata": {}, + "source": [ + "## Installation\n", + "\n", + "Install the released QuantEM version containing this workflow. During PR review, an isolated environment may instead install the `paper/polymers` branch." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from functools import partial\n", + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import quantem as em\n", + "from quantem.diffraction import BraggPeaksPolymer, resolve_polymer_model\n", + "\n", + "DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n", + "SCAN_PATH = Path(os.environ['POLYMER_4DSTEM_PATH'])\n", + "MODEL_DIR = os.environ.get('QUANTEM_POLYMER_MODEL_DIR')\n", + "OUTPUT_DIR = Path(os.environ.get('POLYMER_TUTORIAL_OUTPUT_DIR', 'polymer_tutorial_outputs'))\n", + "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "DEVICE, SCAN_PATH, OUTPUT_DIR" + ] + }, + { + "cell_type": "markdown", + "id": "model-resolution", + "metadata": {}, + "source": [ + "## Resolve the immutable model\n", + "\n", + "Omitting both `version` and `latest` deliberately selects the paper-pinned version. Use `latest=True` only when reproducibility against the paper is not required." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "resolve-model", + "metadata": {}, + "outputs": [], + "source": [ + "model = resolve_polymer_model(local_model_dir=MODEL_DIR)\n", + "cached = resolve_polymer_model(local_model_dir=MODEL_DIR)\n", + "assert cached.weights_path == model.weights_path\n", + "print(model.model_id, model.version, model.checksum)\n", + "print(model.weights_path)" + ] + }, + { + "cell_type": "markdown", + "id": "load-scan", + "metadata": {}, + "source": [ + "## Load and calibrate the scan\n", + "\n", + "The paper scan is a DigitalMicrograph file whose 4D signal is dataset 1. Detector binning is applied before inference. For another instrument, change the reader arguments and confirm the reciprocal or angular calibration metadata before interpreting radial positions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "scan", + "metadata": {}, + "outputs": [], + "source": [ + "scan = em.io.read_4dstem(SCAN_PATH, file_type='digitalmicrograph', dataset_index=1)\n", + "max_scan = int(os.environ.get('POLYMER_TUTORIAL_MAX_SCAN', '0'))\n", + "if max_scan > 0:\n", + " scan = scan[:max_scan, :max_scan]\n", + "scan = scan.bin(bin_factors=(1, 1, 4, 4))\n", + "print('shape:', scan.shape, 'sampling:', scan.sampling, 'units:', scan.units)" + ] + }, + { + "cell_type": "markdown", + "id": "inference", + "metadata": {}, + "source": [ + "## Pinned-model inference\n", + "\n", + "The experimental normalization uses scan-level percentiles 1 and 99, as recorded in the model specification. A circular sample mask keeps normalization and BatchNorm adaptation away from scan corners." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "normalize", + "metadata": {}, + "outputs": [], + "source": [ + "def compute_parameters(data, lower_percentile=1.0, upper_percentile=99.0):\n", + " if isinstance(data, torch.Tensor):\n", + " values = data.detach().float().flatten()\n", + " return tuple(torch.quantile(values, torch.tensor([lower_percentile, upper_percentile], device=values.device) / 100).cpu().tolist())\n", + " return tuple(np.percentile(np.asarray(data), [lower_percentile, upper_percentile]))\n", + "\n", + "def normalize_data(data, lower, upper):\n", + " if isinstance(data, torch.Tensor):\n", + " return torch.clamp(data, lower, upper).sub(lower).div(upper - lower + 1e-8)\n", + " if upper == lower:\n", + " return np.zeros_like(data, dtype=np.float32)\n", + " return (np.clip(data, lower, upper) - lower) / (upper - lower)\n", + "\n", + "yy, xx = np.ogrid[:scan.shape[0], :scan.shape[1]]\n", + "cy, cx = (np.array(scan.shape[:2]) - 1) / 2\n", + "sample_mask = (yy - cy) ** 2 + (xx - cx) ** 2 <= (0.47 * min(scan.shape[:2])) ** 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "detect", + "metadata": {}, + "outputs": [], + "source": [ + "bp = BraggPeaksPolymer.from_data(\n", + " scan,\n", + " device=DEVICE,\n", + " compute_parameters=partial(compute_parameters, lower_percentile=1.0, upper_percentile=99.0),\n", + " normalize_data=normalize_data,\n", + ")\n", + "bp.set_model_weights(\n", + " model_id=model.model_id, version=model.version, local_model_dir=MODEL_DIR\n", + ")\n", + "bp.find_peaks_model(\n", + " device=DEVICE, scan_mask=sample_mask, threshold_peak=0.5,\n", + " n_normalize_samples=1000, initial_chunk_size=100, accelerating_voltage_kv=300,\n", + ")\n", + "bp.process_polar(scan_mask=sample_mask, center_device=DEVICE)" + ] + }, + { + "cell_type": "markdown", + "id": "maps", + "metadata": {}, + "source": [ + "## Count maps, flowlines, and selected exports\n", + "\n", + "The radial ranges below are examples in inverse angstroms; choose scientifically justified windows from the radial peak profile for a different specimen. `orientation_offset_degrees=-6.8` is the paper scan's calibrated display rotation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "count-map", + "metadata": {}, + "outputs": [], + "source": [ + "q_ranges = np.array([[0.00, 0.10], [0.10, 0.20], [0.20, 0.30]])\n", + "fig, axes, count_maps = bp.plot_peak_count_map(q_ranges, return_values=True)\n", + "fig.savefig(OUTPUT_DIR / 'polymer_peak_count_maps.png', dpi=180, bbox_inches='tight')\n", + "plt.close(fig)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "flowlines", + "metadata": {}, + "outputs": [], + "source": [ + "orientation = bp.make_orientation_histogram(\n", + " radial_ranges=q_ranges, orientation_offset_degrees=-6.8, upsample_factor=1,\n", + " theta_step_deg=2, progress_bar=True,\n", + ")\n", + "flowlines = bp.make_flowline_map(orientation, progress_bar=True)\n", + "flowline_rgb = bp.make_flowline_rainbow_image(\n", + " flowlines, sum_radial_bins=True, plot_images=False, white_background=True\n", + ")\n", + "plt.imsave(OUTPUT_DIR / 'polymer_flowlines.png', np.asarray(flowline_rgb))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "figures", + "metadata": {}, + "outputs": [], + "source": [ + "bp.save_peak_figures(\n", + " ry=scan.shape[0] // 2, rx=scan.shape[1] // 2,\n", + " prefix='paper_model', save_dir=OUTPUT_DIR, show_polar=True,\n", + ")\n", + "sorted(path.name for path in OUTPUT_DIR.iterdir())" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}