{
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "nbsphinx": {
   "execute": "never"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5,
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Measured frames against a map's forward model\n\nGive this notebook an ImageD11 dataset and a map of it (a TensorMap from the [indexer](indexing_parameters.ipynb), ImageD11 or a refinement), and it shows the measured frames next to the frames the map predicts, for any row (dty) and frame (ω), with sliders. It is the most direct check of a map: a map that is right reproduces the spots where they are, with their shapes and intensities, frame by frame.\n\nThe model simulates the dataset's own scan (`python -m anri.simulate prepare --dataset`): every frame's ω and dty from the DataSet, its detector and geometry, and the map's phases. Rows are rendered when first shown (a few seconds on a GPU, after the first compile) and kept. The measured side is the dataset's sparse file (what the indexer sees, above the segmentation cut) or the raw detector frames from the master file; both are moved by the DataSet's detector distortion maps to where an ideal detector would see them, as the model is. Each row's model is scaled onto the data by one factor (least squares over the pixels the data hold or the scaled model puts above the cut), so the overall brightness of the map does not matter.\n\nIt needs the dataset (not distributed with anri), so it is stored without outputs. With `interactive = True` (the default) the sliders redraw the plots; it needs `ipywidgets`, `ipympl` and `matplotlib`.",
   "id": "cell-00"
  },
  {
   "cell_type": "code",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "execution_count": null,
   "outputs": [],
   "source": "# Parameters\n\ndataset_file = None  # an ImageD11 <sample>_<dataset>_dataset.h5 (required)\nmap_file = None  # a TensorMap (.h5) or python -m anri.index's _entries.npz (required)\nremap = None  # {old prefix: new prefix} for the paths the DataSet records, e.g. data mounted elsewhere\npars_file = None  # geometry (.par or pars.json; None: the DataSet's)\ny0 = None  # where the rotation axis is in the beam (None: the DataSet's); use the y0 the map was made with\ncif = None  # --cif: one path, or {phase name: path}\noptics = \"al-crl\"  # ID11's focusing optics: \"al-crl\" or \"si-nanolens\"\nwindow, max_frames = (3, 7, 7), 31  # as python -m anri.simulate render\ncut = 1.0  # the segmentation's threshold: model pixels below it do not count against the data in the scale fit\noutdir = None  # where the prepare stage file goes (None: a temporary folder)\nn_cpu = 4\n\n# this notebook only\ninteractive = True  # sliders (ipympl); False: static figures at the starting values\nrow = None  # the row shown first (None: the middle row)\nframe = None  # the frame shown first (None: the model's brightest in that row)\nsource = \"sparse\"  # \"sparse\" (the segmented pixels) or \"raw\" (the detector frames from the master file)",
   "id": "cell-01"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "import anri.utils\n\nanri.utils.setup(n_cpu=n_cpu)  # before any JAX computation\n\nimport os\nimport tempfile\nimport time\nfrom types import SimpleNamespace\n\nimport ipywidgets as widgets\nimport numpy as np\nfrom IPython import get_ipython\nfrom IPython.display import display\nfrom matplotlib import pyplot as plt\n\nimport anri.io\nimport anri.simulate as sim\nfrom anri.simulate import diagnostics as dg\nfrom anri.simulate import plots as pl\n\nget_ipython().run_line_magic(\"matplotlib\", \"widget\" if interactive else \"inline\")\nT0 = time.time()\n\n\ndef log(msg):\n    print(f\"[{time.time() - T0:7.1f} s] {msg}\", flush=True)\n\n\ndef figure(name, **kwargs):\n    plt.close(name)\n    with plt.ioff():\n        fig = plt.figure(name, layout=\"constrained\", **kwargs)\n    if interactive:\n        fig.canvas.header_visible = False\n    return fig\n\n\ndef interact(fig, update, **controls):\n    # sliders that call a plot's update, which redraws it and returns a summary; their values are in .kwargs\n    if not interactive:\n        values = {k: c.value for k, c in controls.items()}\n        print(update(**values))\n        display(fig)\n        return SimpleNamespace(kwargs=values)\n    w = widgets.interactive(lambda **kw: print(update(**kw)), **controls)\n    display(w, fig.canvas)\n    return w",
   "id": "cell-02"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## The dataset and the model\n\nThe DataSet gives the scan, the sparse and raw files and the detector distortion; `stage_prepare(..., dataset=...)` prepares the map to simulate exactly that scan. A map made by `python -m anri.index` is centred on the y0 its fit used (`python -m anri.index show <file>` prints it): give that y0, or the model's spots move with the rotation.",
   "id": "cell-03"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "if dataset_file is None or map_file is None:\n    raise ValueError(\"set dataset_file and map_file in the parameters\")\nds = anri.io.imaged11.read_dataset(dataset_file, remap)\nspatial = anri.io.imaged11.read_spatial(ds)\noutdir = outdir or tempfile.mkdtemp()\nprepare_file = sim.stages.stage_prepare(map_file, pars_file, os.path.join(outdir, \"compare_prepare.h5\"), dataset=dataset_file,\n                                remap=remap, y0=y0, cif=cif, optics=optics, log=log)\np = sim.stages.load_prepare(prepare_file)\nn_rows, n_frames = p[\"omega\"].shape\nprint(f\"{n_rows} rows x {n_frames} frames; detector distortion: {'yes' if spatial is not None else 'none'}\")",
   "id": "cell-04"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Frames\n\nData, model and both overlaid (data red, model cyan: white where they agree) for `frames` frames around `frame` of `row`, summed. The panels share their zoom: use the toolbar's zoom on one, and the others follow. What to look for:\n\n- **Positions:** spots in the same place at the same frame. A shift that grows with the distance from the rotation axis points to y0; one that rotates with ω, to the map's position or orientation.\n- **Red spots** (data only): grains or phases the map lacks, or orientations it has wrong.\n- **Cyan spots** (model only): orientations or voxels the data do not support.\n- **Shapes:** a measured spot wider or longer than the model's is spread the model lacks (orientation spread within the voxels, the beam's convergence, the detector's point spread); smeared along the ring it is orientation, across it strain.\n- **Intensities:** the row's scale factor is printed; spot-to-spot differences within a row point to structure factors, absorption or the voxels' densities.",
   "id": "cell-05"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "rows_seen = {}  # row: (observed, model, scale), filled as rows are shown\n\n\ndef row_data(r):\n    if r not in rows_seen:\n        t = time.time()\n        obs, mod = dg.observed_row(ds, r, spatial), dg.model_row(p, r, window, max_frames)\n        rows_seen[r] = (obs, mod, dg.row_scale(obs, mod, cut))\n        log(f\"row {r}: {len(obs['value'])} measured and {len(mod['value'])} model pixels ({time.time() - t:.1f} s)\")\n    return rows_seen[r]\n\n\nstate = {\"source\": source}\n\n\ndef images(r, frames):\n    obs, mod, s = row_data(r)\n    data = dg.raw_image(ds, r, frames, spatial) if state[\"source\"] == \"raw\" else dg.frames_image(obs, frames)\n    return data, s * dg.frames_image(mod, frames), f\"{state['source']}; model scaled by {s:.3g}\"\n\n\nrow = n_rows // 2 if row is None else row\nif frame is None:  # the model's brightest pixel in that row\n    _, mod0, _ = row_data(row)\n    frame = int(mod0[\"frame\"][np.argmax(mod0[\"value\"])])\nfig = figure(\"frames\", figsize=(18, 6.5))\nupdate = pl.plot_frame_compare(fig, images, p[\"det_shape\"])\n\n\ndef show_frames(row, frame, frames, source):\n    state[\"source\"] = source\n    return update(row, frame, frames)\n\n\nw = interact(fig, show_frames,\n             row=widgets.IntSlider(value=row, min=0, max=n_rows - 1, description=\"row\", continuous_update=False),\n             frame=widgets.IntSlider(value=frame, min=0, max=n_frames - 1, description=\"frame\",\n                                     continuous_update=False),\n             frames=widgets.IntSlider(value=3, min=1, max=31, step=2, description=\"frames\", continuous_update=False),\n             source=widgets.Dropdown(options=[\"sparse\", \"raw\"], value=source, description=\"data\"))",
   "id": "cell-06"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "Zoom on the model's brightest spot in the row shown first (static: the figure above, cropped).",
   "id": "cell-07"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "obs, mod, s = row_data(row)\nk = int(np.argmax(mod[\"value\"]))\nr0, c0 = divmod(int(mod[\"pixel\"][k]), p[\"det_shape\"][1])\nfig = figure(\"zoom\", figsize=(18, 6.2))\nzoom = pl.plot_frame_compare(fig, images, p[\"det_shape\"])\nprint(zoom(row, int(mod[\"frame\"][k]), 3))\nfor a in fig.axes[:3]:\n    a.set_xlim(c0 - 60, c0 + 60)\n    a.set_ylim(r0 + 60, r0 - 60)\ndisplay(fig.canvas if interactive else fig)",
   "id": "cell-08"
  }
 ]
}