A prototype tool for creating a mask for a pair of coaxial apex-sharing conical shadows, which rotate with the sample, using the metadata in a NXmx NeXus file. Use cases include shadowing the body of a diamond anvil pressure cell.

shadow_mask.py

"""
Create a conical shadow mask for NXmx-flavour NeXus image data.

From a NXmx NeXus-format HDF5 data file, compute a stack of mask frames corresponding
to the shadow on the detector of a circular object or aperture that rotates with the
sample.  This might be used to mask the anvils of a diamond anvil pressure cell or
the sample mount assembly in a complex goniometer apparatus.

It is assumed that the shadowed region is bounded by a double cone, with the apex at
the lab frame origin.  You must specify the cones' axis as a vector, and aperture (twice
the cone opening angle) in degrees.  If you want the two cones to have different
apertures, the 'aperture' parameter is presumed to apply to the downstream cone (i.e.
the one in the direction specified by 'axis') and the '--upstream-aperture' parameter
is used to specify the aperture of the other cone.

By default, the region to be masked is taken to be the volume outside the cones,
as appropriate for a diamond anvil cell or other annulus.  To mask the volume inside
the cones instead, use the option '--inside'.

To specify the axis vector, use the NeXus McStas coordinate system.
See https://doi.org/10.1107/S2052252520008672

If the input NeXus file is named like <my-nexus-file>.<extension>, the new mask will be
stored in a file named <my-nexus-file>-mask.h5.  The mask is encoded using the NeXus
mask standard, see
https://manual.nexusformat.org/classes/base_classes/NXdetector.html#nxdetector-pixel-mask-field
"""

from __future__ import annotations

import argparse
import math
import re
from pathlib import Path
from typing import Sequence

import h5py
import hdf5plugin
import numpy as np
from dask import array as da
from dask.diagnostics import ProgressBar
from numpy.typing import ArrayLike

import nxmx

examples = """Examples:

To mask the shadow of a diamond anvil cell, with an aperture of 100° and an axis that
is parallel to the beam at the goniometer zero datum, use the following
    %(prog)s my-nexus-file.h5 100

If the cone upstream of the apex (i.e. in the direction anti-parallel to the axis 
vector) has a different aperture, for example 60°, you can specify it with 
--upstream-aperture, like this
    %(prog)s my-nexus-file.h5 100 --upstream-aperture 60 

If the axis of the cell is rotated at 90° to the beam axis at the goniometer zero datum,
we can specify the axis of the shadowed cones with --axis
    %(prog)s my-nexus-file.h5 100 --axis "0,1,0"

If the object to be masked is a circular obstruction downstream of the sample, 
so that we wish to mask the inside of the downstream cone and have no upstream mask 
at all, we can achieve this by specifying '--inside' and 'upstream-aperture 0', 
like this
    %(prog)s  my-nexus-file.h5 100 --upstream-aperture 0 --inside
"""


class _AngleAction(argparse.Action):
    """Store in radians an angle that was input in degrees."""

    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, math.radians(float(values)))


class _VectorAction(argparse.Action):
    """
    Store a comma-separated sequence of three numbers as a tuple of floats.

    The values describe a normalised vector.
    """

    def __call__(self, parser, namespace, values, option_string=None):
        try:
            x, y, z = map(float, values.split(","))
            norm = math.hypot(x, y, z)
        except ValueError:
            raise argparse.ArgumentError(
                self,
                "Please specify the axis vector as a comma-separated sequence of three "
                "numbers.",
            )
        setattr(namespace, self.dest, (x / norm, y / norm, z / norm))


parser = argparse.ArgumentParser(
    description=__doc__,
    epilog=examples,
    formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
    "nexus_file", help="NXmx format NeXus file.", metavar="nexus-file", type=Path
)
parser.add_argument(
    "aperture",
    help="Aperture of the shadowed cone (twice its opening angle) in degrees.  For a "
    "diamond anvil cell, this is sometimes known as the 4θ angle.",
    action=_AngleAction,
)
parser.add_argument(
    "--upstream-aperture",
    "-u",
    help="Aperture of the upstream cone in degrees, if it differs from that of the "
    "downstream cone.",
    action=_AngleAction,
)
parser.add_argument(
    "--axis",
    "-a",
    help="A vector describing the axis at zero datum of the shadowed cone.  Specify "
    "the axis as comma-separated values with no spaces, like '--axis 0,0,1'.  McStas "
    "coordinates are used, as per the NeXus standard.  By default, the axis is aligned "
    "with the beam — (0, 0, 1) in the NeXus McStas coordinate system.",
    action=_VectorAction,
    default=(0, 0, 1),
)
parser.add_argument(
    "--inside",
    "-i",
    action="store_true",
    help="By default, the region outside the specified cone is masked as appropriate "
    "for a diamond anvil cell, for example.  Use this option to mask the region inside "
    "the cone instead.",
)
parser.add_argument(
    "--make-link",
    "--link",
    "-l",
    action="store_true",
    help="Modify the NeXus file to include a link to the new mask file, under the "
    "NXdetector HDF5 group, with the name 'pixel_mask' or 'pixel_mask_n', where n is "
    "an integer.  The key is chosen so that it doesn't clobber any existing keys.",
)

# Match 'pixel_mask' with an optional suffix '_n' where n is an integer of any length.
regex = re.compile(r"pixel_mask(?:_(\d+))?")


def cumulative_transform(t: nxmx.NXtransformationsAxis) -> ArrayLike:
    return nxmx.get_cumulative_transformation(nxmx.get_dependency_chain(t))


def make_shadow_mask(
    nexus_file: str | Path,
    axis: Sequence[int],
    angle: float,
    angle_upstream: float | None,
    inside: bool,
) -> da.Array:
    """
    Compute the shadow mask of a pair of apex-sharing cones.

    The cones rotate with the sample, as described in the NXmx metadata in
    'nexus_file'.  It is assumed that the cones are never translated — their apex is
    never offset from the lab frame origin.

    Args:
        nexus_file:      The path to a NXmx NeXus file.
        axis:            The three-vector describing the normalised axis vector of the
                           cones.
        angle:           The opening angle (half-aperture) of the downstream cone in
                           radians.
        angle_upstream:  The opening angle (half-aperture) of the upstream cone in
                           radians.  If none, this defaults to the value of 'angle'.
        inside:          If true, mask the region inside the specified cones.  If
                           false, mask the region outside the cones.
    Returns:
        The stack of mask frames as a Dask array.
    """
    with h5py.File(nexus_file) as f:
        metadata = nxmx.NXmx(f)

        entry, *_ = metadata.entries
        sample, *_ = entry.samples
        instrument, *_ = entry.instruments
        detector, *_ = instrument.detectors
        module, *_ = detector.modules

        # Image size.
        slow_size, fast_size = module.data_size

        # Transformations from lab frame to basis of slow and fast axes.
        slow, fast = module.slow_pixel_direction, module.fast_pixel_direction
        slow_dep, fast_dep = slow.depends_on, fast.depends_on

        orientation, slow_dep, fast_dep = map(
            cumulative_transform, (sample.depends_on, slow_dep, fast_dep)
        )

        slow = (slow[()] * slow.vector).to("mm").magnitude
        fast = (fast[()] * fast.vector).to("mm").magnitude

    # Reshape for easy matrix multiplication.
    slow = slow.reshape(-1, 1)
    fast = fast.reshape(-1, 1)

    # Pixel centre offsets in number of pixels.
    slow_offset = np.arange(0.5, slow_size + 0.5).reshape(-1, 1, 1, 1)
    fast_offset = np.arange(0.5, fast_size + 0.5).reshape(1, -1, 1, 1)

    # Slow axis component of pixel locs in mm in homogeneous coords in slow axis frame.
    slow_position = np.append(slow_offset * slow, np.ones(slow_offset.shape), axis=-2)
    # Fast axis component of pixel offsets in mm in homogen coords in fast axis frame.
    fast_offset = np.append(fast_offset * fast, np.zeros(fast_offset.shape), axis=-2)

    # Pixel centre positions in mm in homogenous coords in lab frame.
    pixel_centres = slow_dep @ slow_position + fast_dep @ fast_offset
    # Pixel centres in mm in cartesian coords in lab frame.  Drop convenience axis.
    pixel_centres = np.squeeze(pixel_centres, axis=-1)[..., :3]

    # Normalise the pixel position vectors.
    pixel_centres /= np.linalg.norm(pixel_centres, axis=-1, keepdims=True)

    # Get the cone axis orientation in the lab frame in cartesian coordinates.
    # For ease of computation, it is assumed that 'orientation' is a pure rotation
    # operation, with no translation.  In other words, it is assumed that the cone
    # axis always passes through the sample position at the lab frame origin.
    orientation = orientation[..., :3, :3]  # Convert homogeneous to cartesian coords.
    axis = orientation @ axis  # Find orientation of cone axis for each image.
    axis = da.from_array(axis, chunks=(1, -1))

    # Mark the pixels to be masked.
    cosines = da.einsum("ij,klj->ikl", axis, pixel_centres)  # Broadcast scalar product.
    if angle_upstream:
        if inside:
            mask = (cosines > math.cos(angle)) | (cosines < -math.cos(angle_upstream))
        else:
            mask = (cosines < math.cos(angle)) & (cosines > -math.cos(angle_upstream))
    elif inside:
        mask = da.abs(cosines) > math.cos(angle)
    else:
        mask = da.abs(cosines) < math.cos(angle)

    # Create the mask using the NeXus standard.
    # https://manual.nexusformat.org/classes/base_classes/NXdetector.html#nxdetector-pixel-mask-field
    mask = da.where(mask, np.uint32(1 << 8), np.uint32(0))

    return mask


def save_mask_file(mask: da.Array, nexus_file: Path | str, make_link: bool = False):
    mask_file = nexus_file.parent / (nexus_file.stem + "-mask.h5")
    with ProgressBar():
        mask.to_hdf5(mask_file, "mask", **hdf5plugin.Bitshuffle())

    if make_link:
        with h5py.File(nexus_file, "r+") as f:
            metadata = nxmx.NXmx(f)

            entry, *_ = metadata.entries
            instrument, *_ = entry.instruments
            detector, *_ = instrument.detectors
            detector_path = detector.path.strip("/")

            # Find an appropriate mask key.  If no 'pixel_mask' exists, use that as 
            # the key. If 'pixel_mask_n' exists, use 'pixel_mask_(n+1)', else if 
            # 'pixel_mask' exists, use 'pixel_mask_1'. 
            keys = f[detector_path].keys()
            masks = [m.group(1) or 0 for m in filter(None, map(regex.match, keys))]
            new_key = f"pixel_mask_{max(map(int, masks))}" if masks else "pixel_mask"

            f["/".join([detector_path, new_key])] = h5py.ExternalLink(mask_file, "mask")


if __name__ == "__main__":
    args = parser.parse_args()
    opening_angle = args.aperture / 2
    opening_angle_upstream = (
        args.upstream_aperture / 2 if args.upstream_aperture else None
    )

    mask = make_shadow_mask(
        args.nexus_file,
        axis=args.axis,
        angle=opening_angle,
        angle_upstream=opening_angle_upstream,
        inside=args.inside,
    )

    save_mask_file(mask, nexus_file=args.nexus_file, make_link=args.make_link)