NXmx writer in simtbx

simtbx/command_line/integrate.py

@@ -61,6 +61,7 @@ def print0(*args, **kwargs):
 from copy import deepcopy
 from collections import Counter
 
+from simtbx.diffBragg.device import DeviceWrapper
 
 def filter_refls(R):
     vec3_dbl_keys = 'xyzcal.px', 'xyzcal.mm', 'xyzobs.px.value', 'xyzobs.px.value', 'rlp', 's1'
@@ -338,6 +339,7 @@ def refls_from_sims(panel_imgs, detector, beam, thresh=0, filter=None, panel_ids
             os.makedirs(args.outdir)
     COMM.barrier()
 
+
     params = utils.get_extracted_params_from_phil_sources(args.predPhil, args.cmdlinePhil)
     if os.path.isfile(args.inputGlob):
         df_all = pandas.read_pickle(args.inputGlob)
@@ -369,117 +371,116 @@ def df_iter():
     dev = COMM.rank % args.numdev
 
     print0("Found %d input files" % Nf)
-
-    all_dfs = []
-    all_pred_names = []
-    exp_ref_spec_lines = []
-    for i_f, df in df_iter():
-        printR("Shot %d / %d" % (i_f+1, Nf), flush=True)
-
-        expt_name = df.opt_exp_name.values[0]
-        tag = os.path.splitext(os.path.basename(expt_name))[0]
-        new_expt_name = "%s/%s_%d_predicted.expt" % (args.outdir,tag,  i_f)
-        new_expt_name = os.path.abspath(new_expt_name)
-        df["opt_exp_name"] = new_expt_name
-
-        shutil.copyfile(expt_name,  new_expt_name)
-
-        data_exptList = ExperimentList.from_file(expt_name)
-        data_expt = data_exptList[0]
-
-        try:
-            spectrum_override = None
-            if params.spectrum_from_imageset:
-                spectrum_override = downsamp_spec_from_params(params, data_expt)
-            pred = predictions.get_predicted_from_pandas(
-                df, params, strong=None, device_Id=dev, spectrum_override=spectrum_override)
-            if args.filterDupes:
-                pred = filter_refls(pred)
-        except ValueError:
-            os.remove(new_expt_name)
-            continue
-
-        data = utils.image_data_from_expt(data_expt)
-        Rstrong = refls_from_sims(data, data_expt.detector, data_expt.beam, phil_file=args.procPhil )
-        Rstrong['id'] = flex.int(len(Rstrong), 0)
-        num_panels = len(data_expt.detector)
-        if num_panels > 1:
-            assert params.predictions.label_weak_col == "rlp"
-
-        Rstrong.centroid_px_to_mm(data_exptList)
-        Rstrong.map_centroids_to_reciprocal_space(data_exptList)
-        predictions.label_weak_predictions(pred, Rstrong, q_cutoff=params.predictions.qcut, col=params.predictions.label_weak_col )
-
-        pred['is_strong'] = flex.bool(np.logical_not(pred['is_weak']))
-        strong_sel = np.logical_not(pred['is_weak'])
-
-        pred["refl_idx"] = flex.int(np.arange(len(pred)))
-        weaks = pred.select(pred['is_weak'])
-        weaks_sorted = np.argsort(weaks["scatter"])[::-1]
-        nweak = len(weaks)
-        num_keep = int(nweak*params.predictions.weak_fraction)
-        weak_refl_inds_keep = set(np.array(weaks["refl_idx"])[weaks_sorted[:num_keep]])
-
-        weak_sel = flex.bool([i in weak_refl_inds_keep for i in pred['refl_idx']])
-        keeps = np.logical_or( pred['is_strong'], weak_sel)
-        #printR("Sum keeps=%d; num_strong=%d, num_kept_weak=%d" % (sum(keeps), sum(strong_sel), sum(weak_sel)))
-        pred = pred.select(flex.bool(keeps))
-        nstrong = np.sum(strong_sel)
-        printR("Will save %d refls (%d strong, %d weak)" % (len(pred), np.sum(strong_sel), np.sum(weak_sel)))
-        pred_file = os.path.abspath("%s/%s_%d_predicted.refl" % ( args.outdir, tag, i_f))
-        pred.as_file(pred_file)
-
-        Rindexed = Rstrong.select(Rstrong['indexed'])
-        if len(Rindexed)==0:
-            print("No strong indexed refls for shot %s" % new_expt_name)
-            continue
-
-        utils.refls_to_hkl(Rindexed, data_expt.detector, data_expt.beam, data_expt.crystal, update_table=True)
-        try:
-            int_expt, int_refl = integrate(args.procPhil, data_exptList, Rindexed, pred)
+    with DeviceWrapper(dev) as _:
+        all_dfs = []
+        all_pred_names = []
+        exp_ref_spec_lines = []
+        for i_f, df in df_iter():
+            printR("Shot %d / %d" % (i_f+1, Nf), flush=True)
+
+            expt_name = df.opt_exp_name.values[0]
+            tag = os.path.splitext(os.path.basename(expt_name))[0]
+            new_expt_name = "%s/%s_%d_predicted.expt" % (args.outdir,tag,  i_f)
+            new_expt_name = os.path.abspath(new_expt_name)
+            df["opt_exp_name"] = new_expt_name
+
+            shutil.copyfile(expt_name,  new_expt_name)
+
+            data_exptList = ExperimentList.from_file(expt_name)
+            data_expt = data_exptList[0]
+
+            try:
+                spectrum_override = None
+                if params.spectrum_from_imageset:
+                    spectrum_override = downsamp_spec_from_params(params, data_expt)
+                pred = predictions.get_predicted_from_pandas(
+                    df, params, strong=None, device_Id=dev, spectrum_override=spectrum_override)
+                if args.filterDupes:
+                    pred = filter_refls(pred)
+            except ValueError:
+                os.remove(new_expt_name)
+                continue
+
+            data = utils.image_data_from_expt(data_expt)
+            Rstrong = refls_from_sims(data, data_expt.detector, data_expt.beam, phil_file=args.procPhil )
+            Rstrong['id'] = flex.int(len(Rstrong), 0)
+            num_panels = len(data_expt.detector)
+            if num_panels > 1:
+                assert params.predictions.label_weak_col == "rlp"
+
+            Rstrong.centroid_px_to_mm(data_exptList)
+            Rstrong.map_centroids_to_reciprocal_space(data_exptList)
+            predictions.label_weak_predictions(pred, Rstrong, q_cutoff=params.predictions.qcut, col=params.predictions.label_weak_col )
+
+            pred['is_strong'] = flex.bool(np.logical_not(pred['is_weak']))
+            strong_sel = np.logical_not(pred['is_weak'])
+
+            pred["refl_idx"] = flex.int(np.arange(len(pred)))
+            weaks = pred.select(pred['is_weak'])
+            weaks_sorted = np.argsort(weaks["scatter"])[::-1]
+            nweak = len(weaks)
+            num_keep = int(nweak*params.predictions.weak_fraction)
+            weak_refl_inds_keep = set(np.array(weaks["refl_idx"])[weaks_sorted[:num_keep]])
+
+            weak_sel = flex.bool([i in weak_refl_inds_keep for i in pred['refl_idx']])
+            keeps = np.logical_or( pred['is_strong'], weak_sel)
+            #printR("Sum keeps=%d; num_strong=%d, num_kept_weak=%d" % (sum(keeps), sum(strong_sel), sum(weak_sel)))
+            pred = pred.select(flex.bool(keeps))
+            nstrong = np.sum(strong_sel)
+            printR("Will save %d refls (%d strong, %d weak)" % (len(pred), np.sum(strong_sel), np.sum(weak_sel)))
+            pred_file = os.path.abspath("%s/%s_%d_predicted.refl" % ( args.outdir, tag, i_f))
+            pred.as_file(pred_file)
+
+            Rindexed = Rstrong.select(Rstrong['indexed'])
+            if len(Rindexed)==0:
+                print("No strong indexed refls for shot %s" % new_expt_name)
+                continue
+
+            utils.refls_to_hkl(Rindexed, data_expt.detector, data_expt.beam, data_expt.crystal, update_table=True)
+            try:
+                int_expt, int_refl = integrate(args.procPhil, data_exptList, Rindexed, pred)
+            except RuntimeError:
+                print("Integration failed for %s" % pred_file)
+                continue
             int_expt_name = "%s/%s_%d_integrated.expt" % ( args.outdir,tag, i_f)
             int_expt.as_file(int_expt_name)
             int_refl['bbox'] = int_refl['shoebox'].bounding_boxes()
             int_refl_name = int_expt_name.replace(".expt", ".refl")
             int_refl.as_file(int_refl_name)
-        except RuntimeError as err:
-            print("Integration failed for %s because: '%s'" % (pred_file, str(err)))
-
-        all_dfs.append(df)
-        all_pred_names.append(pred_file)
-        spec_name = df.spectrum_filename.values[0]
-        if spec_name is None:
-            spec_name = ""
-        exp_ref_spec_lines.append("%s %s %s\n" % (new_expt_name, pred_file, spec_name))
-
-    if all_dfs:
-        all_dfs = pandas.concat(all_dfs)
-        all_dfs["predicted_refls"] = all_pred_names
-        all_dfs["predictions"] = all_pred_names
-    else:
-        all_dfs = None
-
-    all_dfs = COMM.gather(all_dfs)
-    exp_ref_spec_lines = COMM.reduce(exp_ref_spec_lines)
-    print0("\nReflections written to folder %s.\n" % args.outdir)
-    if COMM.rank==0:
-        hopper_input_name = os.path.abspath(os.path.join(args.outdir , "%s.txt" % args.hopInputName))
-        o = open(hopper_input_name, "w")
-        for l in exp_ref_spec_lines:
-            o.write(l)
-        o.close()
-        all_dfs = [df for df in all_dfs if df is not None]
-        if not all_dfs:
-            raise ValueError("No dataframes to concat: prediction/integration failed for all shots..")
-
-        all_dfs = pandas.concat([df for df in all_dfs if df is not None])
-        all_dfs.reset_index(inplace=True, drop=True)
-        best_pkl_name = os.path.abspath(os.path.join(args.outdir , "%s.pkl" % args.hopInputName))
-        all_dfs.to_pickle(best_pkl_name)
-        print("Wrote %s (best_pickle option for simtbx.diffBragg.hopper) and %s (exp_ref_spec option for simtbx.diffBragg.hopper). Use them to run the predictions through hopper. Use the centroid=cal option to specify the predictions" % (best_pkl_name, hopper_input_name))
-
-        with open(args.outdir +".exectution.txt", "w") as o:
-            o.write("integrate was run from folder: %s\n" % os.getcwd())
-            o.write("The command line input was:\n")
-            o.write(" ".join(sys.argv))
-            #TODO: write the diff phils here:
+            all_dfs.append(df)
+            all_pred_names.append(pred_file)
+            spec_name = df.spectrum_filename.values[0]
+            if spec_name is None:
+                spec_name = ""
+            exp_ref_spec_lines.append("%s %s %s\n" % (new_expt_name, int_refl_name, spec_name))
+
+        if all_dfs:
+            all_dfs = pandas.concat(all_dfs)
+            all_dfs["predicted_refls"] = all_pred_names
+        else:
+            all_dfs = None
+
+        all_dfs = COMM.gather(all_dfs)
+        exp_ref_spec_lines = COMM.reduce(exp_ref_spec_lines)
+        print0("\nReflections written to folder %s.\n" % args.outdir)
+        if COMM.rank==0:
+            hopper_input_name = os.path.abspath(os.path.join(args.outdir , "%s.txt" % args.hopInputName))
+            o = open(hopper_input_name, "w")
+            for l in exp_ref_spec_lines:
+                o.write(l)
+            o.close()
+            all_dfs = [df for df in all_dfs if df is not None]
+            if not all_dfs:
+                raise ValueError("No dataframes to concat: prediction/integration failed for all shots..")
+
+            all_dfs = pandas.concat([df for df in all_dfs if df is not None])
+            all_dfs.reset_index(inplace=True, drop=True)
+            best_pkl_name = os.path.abspath(os.path.join(args.outdir , "%s.pkl" % args.hopInputName))
+            all_dfs.to_pickle(best_pkl_name)
+            print("Wrote %s (best_pickle option for simtbx.diffBragg.hopper) and %s (exp_ref_spec option for simtbx.diffBragg.hopper). Use them to run the predictions through hopper. Use the centroid=cal option to specify the predictions" % (best_pkl_name, hopper_input_name))
+
+            with open(args.outdir +".exectution.txt", "w") as o:
+                o.write("integrate was run from folder: %s\n" % os.getcwd())
+                o.write("The command line input was:\n")
+                o.write(" ".join(sys.argv))
+                #TODO: write the diff phils here:

simtbx/nanoBragg/__init__.py

@@ -2,6 +2,7 @@
 import boost_adaptbx.boost.python as bp
 import cctbx.uctbx # possibly implicit
 ext = bp.import_ext("simtbx_nanoBragg_ext")
+from libtbx.phil import parse
 from scitbx.array_family import flex
 from simtbx_nanoBragg_ext import *
 from scitbx.matrix import col, sqr
@@ -12,7 +13,8 @@
 from dxtbx.model import CrystalFactory
 from dxtbx.model import BeamFactory
 from dxtbx.model import DetectorFactory
-from dxtbx.format import cbf_writer
+from dxtbx.format.Format import Format
+from dxtbx.format import cbf_writer, nxmx_writer
 
 
 @bp.inject_into(ext.nanoBragg)
@@ -152,15 +154,15 @@ def detector(self):
                      'identifier': '',
                      'image_size': im_shape,
                      'mask': [],
-                     'material': '',  # TODO
+                     'material': 'Si',
                      'mu': 0.0,  # TODO
                      'name': 'Panel',
                      'origin': origin,
                      'pedestal': 0.0,
                      'pixel_size': (pixsize, pixsize),
                      'px_mm_strategy': {'type': 'SimplePxMmStrategy'},
-                     'raw_image_offset': (0, 0),  # TODO
-                     'thickness': 0.0,  # TODO
+                     'raw_image_offset': (0, 0),
+                     'thickness': 0.00001,  # TODO
                      'trusted_range': (-1e3, 1e10),  # TODO
                      'type': ''}]}
     detector = DetectorFactory.from_dict(det_descr)
@@ -171,7 +173,7 @@ def imageset(self):
     format_class = FormatBraggInMemory(self.raw_pixels)
     reader = MemReaderNamedPath("virtual_Bragg_path", [format_class])
     reader.format_class = FormatBraggInMemory
-    imageset_data = ImageSetData(reader, None)
+    imageset_data = ImageSetData(reader, None, vendor="", params={'raw_pixels': self.raw_pixels}, format=FormatBraggInMemory)
     imageset = ImageSet(imageset_data)
     imageset.set_beam(self.beam)
     imageset.set_detector(self.detector)
@@ -242,6 +244,60 @@ def to_cbf(self, cbf_filename, toggle_conventions=False, intfile_scale=1.0):
       self.raw_pixels = cache_pixels
       # print("switch back to cached")
 
+  def to_nexus_nxmx(self, nxmx_filename, toggle_conventions=False, intfile_scale=1.0):
+    """write a NeXus NXmx-format image file to disk from the raw pixel array
+    intfile_scale: multiplicative factor applied to raw pixels before output
+         intfile_scale > 0 : value of the multiplicative factor
+         intfile_scale = 1 (default): do not apply a factor
+         intfile_scale = 0 : compute a reasonable scale factor to set max pixel to 55000; given by get_intfile_scale()"""
+
+    if intfile_scale != 1.0:
+      cache_pixels = self.raw_pixels
+      if intfile_scale > 0: self.raw_pixels = self.raw_pixels * intfile_scale
+      else: self.raw_pixels = self.raw_pixels * self.get_intfile_scale()
+      # print("switch to scaled")
+
+    if toggle_conventions:
+      # switch to DIALS convention before writing CBF
+      CURRENT_CONV = self.beamcenter_convention
+      self.beamcenter_convention=DIALS
+
+    params = nxmx_writer.phil_scope.fetch(parse("""
+    output_file=%s
+    nexus_details {
+      instrument_name=nanoBragg
+      source_name=nanoBragg
+      start_time=NA
+      end_time_estimated=NA
+      sample_name=nanoBragg
+    }
+    """%nxmx_filename)).extract()
+    writer = nxmx_writer.NXmxWriter(params)
+    writer(imageset=self.imageset)
+
+    if toggle_conventions:
+      self.beamcenter_convention=CURRENT_CONV
+
+    if intfile_scale != 1.0:
+      self.raw_pixels = cache_pixels
+      # print("switch back to cached")
+
+def nexus_factory(nxmx_filename):
+    params = nxmx_writer.phil_scope.fetch(parse("""
+    output_file=%s
+    nexus_details {
+      instrument_name=nanoBragg
+      source_name=nanoBragg
+      start_time=NA
+      end_time_estimated=NA
+      sample_name=nanoBragg
+    }
+    dtype=int32
+    """%nxmx_filename)).extract()
+    writer = nxmx_writer.NXmxWriter(params)
+    return writer
+
+
 def make_imageset(data, beam, detector):
   format_class = FormatBraggInMemoryMultiPanel(data)
   reader = MemReaderNamedPath("virtual_Bragg_path", [format_class])
@@ -279,7 +335,7 @@ def get_mask(self, goniometer=None):
     """dummie place holder for mask, consider using internal nanoBragg mask"""
     return self.mask
 
-class FormatBraggInMemory:
+class FormatBraggInMemory(Format):
 
   def __init__(self, raw_pixels):
     self.raw_pixels = raw_pixels
@@ -305,6 +361,10 @@ def get_mask(self, goniometer=None):
     """dummie place holder for mask, consider using internal nanoBragg mask"""
     return self.mask,
 
+  @classmethod
+  def get_instance(Class, filename, **kwargs):
+    return Class(raw_pixels = kwargs.pop('raw_pixels'), **kwargs)
+
   #def paths(self):
   #  return ["InMemoryBraggPath"]  # TODO: CBFLib complains if no datablock path provided which comes from path
 

simtbx/nanoBragg/tst_gauss_argchk.py

@@ -231,6 +231,7 @@ def simple_monochromatic_case(bragg_engine, BEAM, DETECTOR, CRYSTAL, SF_model, a
   SIM.to_smv_format(fileout="test_full_001.img", intfile_scale=output_scale)
   assert approx_equal(SIM.raw_pixels, SIM2.raw_pixels)
   SIM.to_cbf("test_full_001.cbf", intfile_scale=output_scale)
+  SIM.to_nexus_nxmx("test_full_001.h5", intfile_scale=output_scale)
 
   if runmode=="GPU":
     bragg_engine = nanoBragg.add_nanoBragg_spots_cuda