""" Application to concatenate several scans; each corresponding to a z-stage, into one nexus scan for nabu-helical """ import argparse import os import re import json import numpy as np from silx.io.url import DataUrl from tomoscan.esrf.scan.nxtomoscan import NXtomoScan from tomoscan.scanbase import ReducedFramesInfos from tomoscan.io import HDF5File, get_swmr_mode from nxtomo.application.nxtomo import NXtomo from nxtomo.nxobject.utils import concatenate as nx_concatenate from nxtomo.nxobject.nxdetector import ImageKey from nxtomomill.utils.hdf5 import DatasetReader def str2bool(v): if isinstance(v, bool): return v if v.lower() in ("yes", "true", "t", "y", "1"): return True elif v.lower() in ("no", "false", "f", "n", "0"): return False else: raise argparse.ArgumentTypeError("Boolean value expected.") def get_arguments(argv): parser = argparse.ArgumentParser(description="") parser.add_argument( "--filename_template", required=True, help="""The filename template. It must contain a segment equal to "X"*ndigits which will be replaced by the stage number """, ) parser.add_argument( "--target_file", required=True, help="The new nexus filename that we are going to create ", ) parser.add_argument( "--entry_name", required=False, help="Optional.Output data path aka entry_name. Its default is entry0000", default="entry0000", ) parser.add_argument( "--total_nstages", type=int, required=True, help="The total number of stages" ) parser.add_argument( "--first_stage", type=int, default=0, required=False, help="Optional. Defaults to zero. The number of the first considered stage. Use this to do a smaller sequence", ) parser.add_argument( "--last_stage", type=int, default=-1, required=False, help="Optional. Defaults to total_nstages-1. The number of the last considered stage. Use this to do a smaller sequence", ) parser.add_argument( "--cors_file", required=False, type=str, help="Optional. If given, it is a txt file with a column of centers of rotation. We expect one COR per stage. They are use to set the x_translation.\n" "This file can be either a one column text file, or a json file. In this case the key rotation_axis_position_list must be present and give a list", ) parser.add_argument( "--pixel_size_m", required=False, default=None, help="Optional. The pixel size in meters. If given it will overwrite the nexus one, in the final nexus file", ) parser.add_argument( "--neutral_flat", type=str2bool, required=False, default=False, help="Optional. Its default is False. If true then it will set flats to arrays filled with ones and darks with zeroes", ) parser.add_argument( "--flats_from_reduced", type=str2bool, required=False, default=False, help="Optional. Its default is False. If true then the flats will be set from the reduced flats of the original scans", ) parser.add_argument( "--flats_from_before_after", type=str2bool, required=False, default=False, help="""Optional. Its default is False. If true set flats from the reduced flats of "before and after" scans.\n" "If given then you have to provide also the names these two scans, using parameters '--scan_before' and '--scan_after'""", ) parser.add_argument( "--scan_before", required=False, type=str, help="""Optional. To be used with "flats_from_before_after selected". The name of the scan before all the scans. From this the flat/dark has been extracted""", ) parser.add_argument( "--scan_after", required=False, type=str, help="""Optional. To be used with "flats_from_before_after selected". The name of the scan after all the scans. From this the flat/dark has been extracted""", ) args = parser.parse_args(argv) sum_bool = ( args.neutral_flat + args.flats_from_reduced + args.flats_from_before_after ) if (sum_bool) not in [0, 1]: message = f""" Only one or none of neutral_flat, flats_from_reduced, flats_from_before_after options can be selected. You selected {sum_bool} of them """ raise ValueError(message) if args.last_stage == -1: args.last_stage = args.total_nstages - 1 args.nstages = args.last_stage + 1 - args.first_stage if args.cors_file is None: args.cors = np.zeros([args.nstages], "f") else: try: # check if it is json with open(args.cors_file, "r") as fj: json_dict = json.load(fj) except ValueError: args.cors = np.loadtxt(args.cors_file) else: args.cors = json_dict["rotation_axis_position_list"] pattern = re.compile("[X]+") # X represent the variable part of the 'template' # for example if we want to treat scans HA_2000_sample_0000.nx, ..., HA_2000_sample_9999.nx then # we expect the template to be HA_2000_sample_XXXX.nx # warning: If the dataset base names contains several X substrings the longest ones will be taken. ps = pattern.findall(args.filename_template) ls = list(map(len, ps)) if len(ls) < 1: args.name_template = args.filename_template if args.first_stage != args.last_stage: message = f" The argument for filename_template , which was '{args.filename_template}' does not seem to contain a pattern with multiple X for the numerical part" raise ValueError(message) else: idx = np.argmax(ls) if len(ps[idx]) < 2: message = f""" The argument filename_template should contain a substring formed by at least two 'X' The filename_template was {args.filename_template} """ raise ValueError(message) args.name_template = args.filename_template.replace( ps[idx], "{i_stage:" + "0" + str(ls[idx]) + "d}" ) args.file_list = [] for i_stage in range(args.first_stage, args.last_stage + 1): args.file_list.append(args.name_template.format(i_stage=i_stage)) args.used_current = None if args.pixel_size_m is not None: args.overwrite_pixel_size = True args.pixel_size_m = float(args.pixel_size_m) else: args.overwrite_pixel_size = False with HDF5File(args.file_list[0], "r", swmr=get_swmr_mode()) as h5f: args.pixel_size_m = h5f[ os.path.join(args.entry_name, "instrument", "detector", "x_pixel_size") ][()] return args def main(argv): args = get_arguments(argv[1:]) output_file = args.target_file dummy_frame = None """the dummy frame is a frame fill with ones. It will be insert between two stages to ensure series from the twos stages will be perceived as 'uncontiguous'. Separates flats at the end of a stage and at the beginning. """ nxt_z_list = [] npoints_list = [] for i_stage in range(args.first_stage, args.last_stage + 1): filename = args.file_list[i_stage - args.first_stage] # step 1: load the stage nxt = NXtomo() nxt.load(filename, data_path=args.entry_name, detector_data_as="as_data_url") nxt_z_list.append(nxt) scan = NXtomoScan(filename, args.entry_name) npoints_list.append(len(scan.image_key_control)) if dummy_frame is None: if len(nxt.instrument.detector.data) > 0: # create empty frame to give it to the 'invalid_nxt' NXtomo later with DatasetReader(nxt.instrument.detector.data[0]) as raw_data: data_type = raw_data.dtype # FIXME: avoid keeping some file open. not clear why this is needed raw_data = None args.img_shape = (scan.dim_2, scan.dim_1) with HDF5File(output_file, mode="w") as h5s: h5s["empty_frame"] = np.ones( (1, scan.dim_2, scan.dim_1), dtype=data_type ) dummy_frame = DataUrl( file_path=output_file, data_path="empty_frame" ) else: message = """The first nx file must have some data""" raise ValueError(message) # step 2: insert single 'empty' frame between the two stages (aka dummy frame) invalid_nxt = NXtomo(args.entry_name) invalid_nxt.instrument.detector.image_key_control = [ImageKey.INVALID] invalid_nxt.energy = nxt.energy.value invalid_nxt.control.data = np.ones(1, "f") invalid_nxt.sample.rotation_angle = [0.0] invalid_nxt.instrument.detector.field_of_view = ( nxt.instrument.detector.field_of_view.value ) invalid_nxt.instrument.detector.x_pixel_size = ( nxt.instrument.detector.x_pixel_size.value ) invalid_nxt.instrument.detector.y_pixel_size = ( nxt.instrument.detector.y_pixel_size.value ) invalid_nxt.instrument.detector.distance = ( nxt.instrument.detector.distance.value ) invalid_nxt.sample.x_translation = np.zeros([1], "f") invalid_nxt.sample.y_translation = np.zeros([1], "f") invalid_nxt.sample.z_translation = np.zeros([1], "f") invalid_nxt.instrument.detector.data = (dummy_frame,) nxt_z_list.append(invalid_nxt) nx_concatenated = nx_concatenate(nxt_z_list) nx_concatenated.save( file_path=output_file, data_path=args.entry_name, overwrite=True ) with HDF5File(output_file, "r+") as output_target: N_total = output_target[ os.path.join(args.entry_name, "sample", "x_translation") ].shape[0] if args.neutral_flat: # will set flats to arrays filled with ones and darks with zeroes original_currents = output_target[ os.path.join(args.entry_name, "control", "data") ][()] original_currents[:] = 1 del output_target[os.path.join(args.entry_name, "control", "data")] output_target[os.path.join(args.entry_name, "control", "data")] = ( original_currents ) do_flat = ( args.flats_from_reduced or args.flats_from_before_after or args.neutral_flat ) if do_flat or (args.cors_file is not None): one_for_invalid = 1 actual_pos = 0 old_pos = actual_pos i_stage = args.first_stage if do_flat: darks_dictionary = {} flats_dictionary = {} darks_infos_dictionary = {"count_time": []} flats_infos_dictionary = { "machine_electric_current": [], "count_time": [], } add_dark_to_dict( args, darks_dictionary, darks_infos_dictionary, i_stage, actual_pos ) for i_stage in range(args.first_stage, args.last_stage + 1): if do_flat: add_flat_to_dict( args, flats_dictionary, flats_infos_dictionary, i_stage, actual_pos, ) actual_pos += npoints_list[i_stage - args.first_stage] if do_flat: add_flat_to_dict( args, flats_dictionary, flats_infos_dictionary, i_stage, actual_pos, position_in_list=1, ) # if i_stage != args.last_stage: actual_pos += one_for_invalid if args.cors_file is not None: output_target[ os.path.join(args.entry_name, "sample", "x_translation") ][old_pos:actual_pos] = ( -args.cors[i_stage - args.first_stage] + args.cors[0] ) * args.pixel_size_m old_pos = actual_pos assert actual_pos == N_total if args.overwrite_pixel_size: g = output_target[ os.path.join(args.entry_name, "instrument", "detector") ] g["x_pixel_size"][()] = args.pixel_size_m g["y_pixel_size"][()] = args.pixel_size_m if do_flat: # darks metadata meta_darks = ReducedFramesInfos() meta_darks.count_time.extend(darks_infos_dictionary["count_time"][:1]) # flats metadata meta_flats = ReducedFramesInfos() meta_flats.machine_electric_current.extend( flats_infos_dictionary["machine_electric_current"] ) meta_flats.count_time.extend(flats_infos_dictionary["count_time"]) # save metadata scan = NXtomoScan(output_file, args.entry_name) scan.save_reduced_flats( flats_dictionary, flats_infos=meta_flats, overwrite=True ) scan.save_reduced_darks( darks_dictionary, darks_infos=meta_darks, overwrite=True ) def add_flat_to_dict( args, flats_dictionary, flats_infos_dictionary, i_stage, actual_pos, position_in_list=0, ): if args.flats_from_reduced: filename = args.file_list[i_stage - args.first_stage] scan = NXtomoScan(filename, args.entry_name) reduced_flats, metadata_flats = scan.load_reduced_flats(return_info=True) position_in_list = min(position_in_list, len(list(reduced_flats.keys())) - 1) my_flat = reduced_flats[list(reduced_flats.keys())[position_in_list]] my_current = metadata_flats.machine_electric_current[position_in_list] my_count_time = metadata_flats.count_time[position_in_list] elif args.neutral_flat: my_flat = np.ones(args.img_shape, "f") my_current = 1.0 my_count_time = 1.0 elif args.flats_from_before_after: scan = NXtomoScan(args.scan_before, args.entry_name) reduced_flats_b, metadata_flats_b = scan.load_reduced_flats(return_info=True) reduced_darks_b, _ = scan.load_reduced_darks(return_info=True) scan = NXtomoScan(args.scan_after, args.entry_name) reduced_flats_e, metadata_flats_e = scan.load_reduced_flats(return_info=True) reduced_darks_e, _ = scan.load_reduced_darks(return_info=True) flat_b = reduced_flats_b[list(reduced_flats_b.keys())[0]] flat_e = reduced_flats_e[list(reduced_flats_e.keys())[0]] dark_b = reduced_darks_b[list(reduced_darks_b.keys())[0]] dark_e = reduced_darks_b[list(reduced_darks_e.keys())[0]] current_b = metadata_flats_b.machine_electric_current[0] current_e = metadata_flats_e.machine_electric_current[0] if args.used_current is None: args.used_current = current_b factor = (i_stage) / (args.total_nstages) my_flat = args.used_dark + ( (flat_b - dark_b) * (args.used_current / current_b) * (1 - factor) + (flat_e - dark_e) * (args.used_current / current_e) * factor ) my_current = args.used_current my_count_time = args.used_count_time else: raise ValueError( "one of the option must be activated in[neutral_flat, flats_from_reduced, flats_from_before_after]" ) flats_dictionary[actual_pos] = my_flat flats_infos_dictionary["machine_electric_current"].append(my_current) flats_infos_dictionary["count_time"].append(my_count_time) def add_dark_to_dict( args, darks_dictionary, darks_infos_dictionary, i_stage, actual_pos ): if args.flats_from_reduced: filename = args.file_list[i_stage - args.first_stage] elif args.flats_from_before_after: filename = args.scan_before else: assert args.neutral_flat filename = None if filename is not None: scan = NXtomoScan(filename, args.entry_name) reduced_darks, metadata_darks = scan.load_reduced_darks(return_info=True) my_dark = reduced_darks[list(reduced_darks.keys())[0]] my_count_time = metadata_darks.count_time[0] else: my_dark = np.zeros(args.img_shape, "f") my_count_time = 1.0 darks_dictionary[actual_pos] = my_dark darks_infos_dictionary["count_time"].append(my_count_time) args.used_count_time = my_count_time args.used_dark = darks_dictionary[actual_pos]