Script to take predictions, add some fixed shoebox size and then extract to shoeboxes.refl

make_shoeboxes.py

# to get to the starting point:
#
# dials.import ~/data/i04_bag_training/*gz
# dials.find_spots imported.expt
# dials.index imported.expt strong.refl
# dials.refine indexed.*
# dials.predict refined.expt
#
# then
#
# dials.python make_shoeboxes.py predicted.refl imported.expt nx=2 ny=2 nz=2

import logging

logger = logging.getLogger("make_shoeboxes")

from libtbx.phil import parse

phil_scope = parse(
    """
  nx = 1
    .type = int(value_min=1)
    .help = "+/- x around centroid"

  ny = 1
    .type = int(value_min=1)
    .help = "+/- y around centroid"

  nz = 1
    .type = int(value_min=1)
    .help = "+/- z around centroid"

  output {
    reflections = 'shoeboxes.refl'
      .type = str
      .help = "The integrated output filename"
  }
"""
)


class Script(object):
    """Class to run the script."""

    def __init__(self):
        """Initialise the script."""

        from dials.util.options import ArgumentParser

        # The script usage
        usage = (
            "usage: dials.extract_shoeboxes [options] experiment.expt observations.refl"
        )

        # Create the parser
        self.parser = ArgumentParser(
            usage=usage,
            phil=phil_scope,
            read_experiments=True,
            read_reflections=True,
        )

    def run(self):
        """Extract the shoeboxes."""
        from dials.util.options import flatten_reflections
        from dials.util.options import flatten_experiments
        from dials.util.options import flatten_experiments
        from dials.util import log
        from dials.array_family import flex
        from dials.util import Sorry

        # Parse the command line
        params, options = self.parser.parse_args(show_diff_phil=False)

        # Configure logging
        log.config()

        # Log the diff phil
        diff_phil = self.parser.diff_phil.as_str()
        if diff_phil:
            logger.info("The following parameters have been modified:\n")
            logger.info(diff_phil)

        # Get the data
        reflections = flatten_reflections(params.input.reflections)
        experiments = flatten_experiments(params.input.experiments)
        if not any([experiments, reflections]):
            self.parser.print_help()
            exit(0)
        elif len(experiments) > 1:
            raise Sorry("More than 1 experiment set")
        elif len(experiments) == 1:
            imageset = experiments[0].imageset
        if len(reflections) != 1:
            raise Sorry("Need 1 reflection table, got %d" % len(reflections))
        else:
            reflections = reflections[0]

        # Check the reflections contain the necessary stuff
        assert "bbox" in reflections
        assert "panel" in reflections

        # Get some models
        detector = imageset.get_detector()
        scan = imageset.get_scan()
        frame0, frame1 = scan.get_array_range()

        x, y, z = reflections["xyzcal.px"].parts()

        x = flex.floor(x).iround()
        y = flex.floor(y).iround()
        z = flex.floor(z).iround()

        bbox = flex.int6(len(reflections))

        dx, dy = detector[0].get_image_size()

        for j, (_x, _y, _z) in enumerate(zip(x, y, z)):
            x0 = _x - params.nx
            x1 = _x + params.nx + 1
            y0 = _y - params.ny
            y1 = _y + params.ny + 1
            z0 = _z - params.nz
            z1 = _z + params.nz + 1
            if x0 < 0:
                x0 = 0
            if x1 >= dx:
                x1 = dx - 1
            if y0 < 0:
                y0 = 0
            if y1 >= dy:
                y1 = dy - 1
            if z0 < frame0:
                z0 = frame0
            if z1 >= frame1:
                z1 = frame1 - 1
            bbox[j] = (x0, x1, y0, y1, z0, z1)

        reflections["bbox"] = bbox

        # beware change of variables - removing those which were a different shape because boundary conditions
        x0, x1, y0, y1, z0, z1 = bbox.parts()

        dx, dy, dz = (x1 - x0), (y1 - y0), (z1 - z0)

        good = (
            (dx == 2 * params.nx + 1)
            & (dy == 2 * params.ny + 1)
            & (dz == 2 * params.nz + 1)
        )

        print(f"{good.count(True)} / {good.size()} kept")

        reflections = reflections.select(good)

        reflections["shoebox"] = flex.shoebox(
            reflections["panel"], reflections["bbox"], allocate=True
        )

        reflections.extract_shoeboxes(imageset)

        filename = params.output.reflections
        logger.info("Saving %d reflections to %s" % (len(reflections), filename))
        reflections.as_file(filename)


if __name__ == "__main__":
    import dials.util

    with dials.util.show_mail_on_error():
        script = Script()
        script.run()