From d3cdc33358492b64272a91851d4eb55fe2c6d12f Mon Sep 17 00:00:00 2001
From: Egor Sobolev <egor.sobolev@xfel.eu>
Date: Thu, 15 Jun 2023 19:39:14 +0200
Subject: [PATCH] Add crystfel geometry and detector plotting functions
---
src/geomtools/__init__.py | 1 +
src/geomtools/detector/__init__.py | 5 ++
src/geomtools/detector/draw.py | 65 ++++++++++++++++++++
src/geomtools/detector/geom.py | 97 ++++++++++++++++++++++++++++++
src/geomtools/sfx/__init__.py | 5 +-
src/geomtools/sfx/draw.py | 25 ++++++++
6 files changed, 197 insertions(+), 1 deletion(-)
create mode 100644 src/geomtools/detector/__init__.py
create mode 100644 src/geomtools/detector/draw.py
create mode 100644 src/geomtools/detector/geom.py
diff --git a/src/geomtools/__init__.py b/src/geomtools/__init__.py
index cc94950..abf1c68 100644
--- a/src/geomtools/__init__.py
+++ b/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,
)
diff --git a/src/geomtools/detector/__init__.py b/src/geomtools/detector/__init__.py
new file mode 100644
index 0000000..3e87147
--- /dev/null
+++ b/src/geomtools/detector/__init__.py
@@ -0,0 +1,5 @@
+# flake8: noqa E401
+from .geom import (
+ read_crystfel_geom, get_pixel_positions, assemble_data
+)
+from .draw import plot_detector_layout, plot_data_on_detector
diff --git a/src/geomtools/detector/draw.py b/src/geomtools/detector/draw.py
new file mode 100644
index 0000000..dbc7846
--- /dev/null
+++ b/src/geomtools/detector/draw.py
@@ -0,0 +1,65 @@
+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
diff --git a/src/geomtools/detector/geom.py b/src/geomtools/detector/geom.py
new file mode 100644
index 0000000..11e6129
--- /dev/null
+++ b/src/geomtools/detector/geom.py
@@ -0,0 +1,97 @@
+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)
diff --git a/src/geomtools/sfx/__init__.py b/src/geomtools/sfx/__init__.py
index 76fa95b..966ba9b 100644
--- a/src/geomtools/sfx/__init__.py
+++ b/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
diff --git a/src/geomtools/sfx/draw.py b/src/geomtools/sfx/draw.py
index c9bd277..07eadd4 100644
--- a/src/geomtools/sfx/draw.py
+++ b/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
--
GitLab