Add crystfel geometry and detector plotting functions

src/geomtools/__init__.py

@@ -2,5 +2,6 @@
 from .sfx import (
     parse_crystfel_streamfile, read_crystfel_streamfile, extract_geometry,
     plot_center_shift, plot_cell_parameters, plot_peakogram, plot_powder,
+    plot_detector, ph_en_to_lambda, get_q_from_xyz, get_min_bragg_dist,
     spacing, gauss2d_fit, ellipse, parse_xwiz_summary, get_peak_position,
 )

src/geomtools/detector/__init__.py

+# flake8: noqa E401
+from .geom import (
+    read_crystfel_geom, get_pixel_positions, assemble_data
+)
+from .draw import plot_detector_layout, plot_data_on_detector

src/geomtools/detector/draw.py

+import numpy as np
+
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from .geom import assemble_data, get_pixel_positions
+
+
+def plot_detector_layout(panels, figsize=(16, 16), **kwargs):
+    """Draw the detector panels according to the detector geometry."""
+    rect = []
+    xmn, xmx, ymn, ymx = np.inf, -np.inf, np.inf, -np.inf
+    for name, p in panels.iterrows():
+        x0, y0 = p.cnx, p.cny
+        x1 = p.cnx + p.w * p.fsx + p.h * p.ssx
+        y1 = p.cny + p.w * p.fsy + p.h * p.ssy
+        rect.append(np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]))
+        xmn, xmx = min(xmn, x0, x1), max(xmx, x1, x0)
+        ymn, ymx = min(ymn, y0, y1), max(ymx, y1, y0)
+
+    width = xmx - xmn + 4
+    height = ymx - ymn + 4
+    figsize = (16, 16 / width * height)
+    fig, ax = plt.subplots(1, 1, figsize=figsize, **kwargs)
+    ax.set_xlim(xmn - 2, xmx + 2)
+    ax.set_ylim(ymn - 2, ymx + 2)
+    ax.axis(False)
+    ax.set_aspect('equal')
+    for r in rect:
+        ax.add_patch(plt.Polygon(r, fill=False, color='lightgray'))
+
+    return fig, ax
+
+
+def plot_data_on_detector(data, panels, figwidth=16, cmap=plt.cm.magma,
+                          **kwargs):
+    """Plots data according to the detector geometry."""
+
+    if isinstance(panels, np.ndarray):
+        pos = panels
+    else:
+        pos = get_pixel_positions(panels)
+
+    img, (xc, yc) = assemble_data(data, pos)
+    msk, _ = assemble_data(np.ones(data.shape, bool), pos)
+    img = np.ma.masked_array(img, ~msk)
+    ny, nx = img.shape
+
+    figsize = (figwidth, figwidth / nx * ny)
+    fig, ax = plt.subplots(1, 1, figsize=figsize, **kwargs)
+
+    extent = (nx - xc + 0.5, -xc - 0.5, -yc - 0.5, ny - yc + 0.5)
+    im = ax.imshow(np.flip(img, axis=1), origin='lower',
+                   extent=extent, vmin=0, vmax=1, cmap=cmap)
+    ax.plot([-50, 50], [0, 0], 'r')
+    ax.plot([0, 0], [-50, 50], 'r')
+
+    ax.set_aspect('equal')
+    ax.axis(False)
+
+    ax_divider = make_axes_locatable(ax)
+    cax = ax_divider.append_axes("right", size="5%", pad="2%")
+
+    fig.colorbar(im, ax=ax, cax=cax)
+    return fig, ax

src/geomtools/detector/geom.py

+import numpy as np
+import pandas as pd
+
+from cfelpyutils.geometry import load_crystfel_geometry
+
+
+def append_rigid_group(panels, geom, group_name):
+    """Appends the rigid group name to the panels."""
+    groups = dict(
+        (rg, geom.detector['rigid_groups'][rg])
+        for rg in geom.detector['rigid_group_collections'][group_name]
+    )
+
+    ix, pa = [], []
+    for m, p in groups.items():
+        ix += [m] * len(p)
+        pa += p
+
+    panels = panels.join(
+        pd.DataFrame(data={group_name: ix, 'panel': pa}).set_index('panel'))
+
+    return panels
+
+
+def read_crystfel_geom(filename, indexes=dict()):
+    """Reads Crystfel geometry file."""
+    geom = load_crystfel_geometry(filename)
+
+    panels = pd.DataFrame(
+        geom.detector['panels'].values(),
+        index=geom.detector['panels'].keys()
+    )
+    panels.index.name = 'panel'
+
+    for group_name in geom.detector['rigid_group_collections']:
+        panels = append_rigid_group(panels, geom, group_name)
+
+    for column, dimno in indexes.items():
+        panels[column] = panels.dim_structure.apply(lambda x: x[dimno])
+
+    return panels
+
+
+def get_pixel_positions(panels):
+    """Returns pixel positions according to the detector geometry."""
+    lo = [1 << 63, 1 << 63, 1 << 63]
+    hi = [-1 << 63, -1 << 63, -1 << 63]
+    for pn, pv in panels.iterrows():
+        sh = pv.dim_structure[:]
+        sh.remove('%')
+
+        repl = {'fs': pv.orig_min_fs, 'ss': pv.orig_min_ss}
+        lo = min(lo, list(map(lambda x: repl.get(x, x), sh)))
+
+        repl = {'fs': pv.orig_max_fs, 'ss': pv.orig_max_ss}
+        hi = max(hi, list(map(lambda x: repl.get(x, x) + 1, sh)))
+
+    shape = [b - a for a, b in zip(lo, hi)]
+    crd = np.zeros(shape + [3], float)
+
+    for pn, pv, in panels.iterrows():
+        sh = pv.dim_structure[:]
+        sh.remove('%')
+
+        ss, fs = np.meshgrid(
+            range(pv.orig_min_ss, pv.orig_max_ss + 1),
+            range(pv.orig_min_fs, pv.orig_max_fs + 1),
+            indexing="ij"
+        )
+
+        repl = {'fs': fs, 'ss': ss}
+        ix = tuple(map(lambda x: repl.get(x[0], x[0]) - x[1], zip(sh, lo)))
+
+        ss -= pv.orig_min_ss
+        fs -= pv.orig_min_fs
+
+        crd[ix + (0,)] = ss * pv.ssx + fs * pv.fsx + pv.cnx
+        crd[ix + (1,)] = ss * pv.ssy + fs * pv.fsy + pv.cny
+        crd[ix + (2,)] = ss * pv.ssz + fs * pv.fsz
+
+    return crd
+
+
+def assemble_data(data, pos):
+    """Assembles data in image according to the detector geometry."""
+    x = (pos[..., 0] + 0.5).astype(int)
+    y = (pos[..., 1] + 0.5).astype(int)
+
+    xmin, xmax = np.min(x), np.max(x)
+    ymin, ymax = np.min(y), np.max(y)
+
+    nx, ny = xmax - xmin + 1, ymax - ymin + 1
+
+    img = np.zeros([ny, nx], dtype=data.dtype)
+    img[y - ymin, x - xmin] = data
+
+    return img, (-xmin, -ymin)

src/geomtools/sfx/__init__.py

@@ -4,7 +4,10 @@ from .crystfelio import (
 )
 from .draw import (
     plot_center_shift, plot_cell_parameters, plot_peakogram, plot_powder,
+    plot_detector,
+)
+from .lattice import (
+    spacing, ph_en_to_lambda, get_q_from_xyz, get_min_bragg_dist,
 )
-from .lattice import spacing
 from .misc import gauss2d_fit, ellipse, get_peak_position
 from .xwizio import parse_xwiz_summary

src/geomtools/sfx/draw.py

@@ -126,3 +126,28 @@ def plot_powder(peaks, counts=True, figwidth=14, frameon=False, **kwargs):
     ax.axis(False)
 
     return fig, ax
+
+
+def plot_detector(panels, figsize=(16, 16), **kwargs):
+    rect = []
+    xmn, xmx, ymn, ymx = np.inf, -np.inf, np.inf, -np.inf
+    for name, p in panels.iterrows():
+        x0, y0 = p.cnx, p.cny
+        x1 = p.cnx + p.w * p.fsx + p.h * p.ssx
+        y1 = p.cny + p.w * p.fsy + p.h * p.ssy
+        rect.append(np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]))
+        xmn, xmx = min(xmn, x0, x1), max(xmx, x1, x0)
+        ymn, ymx = min(ymn, y0, y1), max(ymx, y1, y0)
+
+    width = xmx - xmn + 4
+    height = ymx - ymn + 4
+    figsize = (16, 16 / width * height)
+    fig, ax = plt.subplots(1, 1, figsize=figsize, **kwargs)
+    ax.set_xlim(xmn - 2, xmx + 2)
+    ax.set_ylim(ymn - 2, ymx + 2)
+    ax.axis(False)
+    ax.set_aspect('equal')
+    for r in rect:
+        ax.add_patch(plt.Polygon(r, fill=False, color='lightgray'))
+
+    return fig, ax