Merge pull request FABLE-3DXRD#573 from haixing0a/master

haixing0a · web-flow · commit 32a72e8ced94 · 2026-04-15T09:16:56.000+02:00
a few clean ups and minor modifications for forward_projector etc.
diff --git a/ImageD11/forward_model/forward_projector.py b/ImageD11/forward_model/forward_projector.py
@@ -6,6 +6,7 @@
 # Jan 2nd, 2025
 # version 1.0: updated on January 30, 2025
 # version 1.1: fixed a bug with ray origin on December 10, 2025
+# version 1.2: some clean-ups and ensure all print-outs go to .out file instead of .log file, April 12, 2026
 
 import numpy as np
 from matplotlib import pyplot as plt
@@ -39,7 +40,11 @@
 from ImageD11.nbGui.nb_utils import slurm_submit_and_wait
 
 econst = 12.3984
-logging.basicConfig(level=logging.INFO, force=True)
+logging.basicConfig(
+    level=logging.INFO,
+    stream=sys.stdout,
+    force=True
+)
 os.environ["JOBLIB_TEMP_FOLDER"] = "/tmp_14_days/slurm/log"
 
 """
@@ -380,14 +385,14 @@ def set_all(self):
         self.check_halfy()
 
     def check_halfy(self):
-        #quick check to see if halfy is big enough
+        # quick check to see if halfy is big enough to include all sample voxels to be illuminated
         if self.sample_mask is None:
             mask_tmp = (self.sample_input.DS['labels'] > -1) & (~np.isnan(self.sample_input.DS['U'][:, :, :, 0, 0]))
             mask_tmp = np.transpose(mask_tmp, (2, 1, 0))
         else:
             mask_tmp = self.sample_mask
         center = (mask_tmp.shape[0]/2, mask_tmp.shape[1]/2)
-        radius = self.args["halfy"]/np.mean(self.sample_input.DS['voxel_size']) #radius covered by halfy, in pixels
+        radius = self.args["halfy"]/np.mean(self.sample_input.DS['voxel_size']) # radius covered by halfy, [pixel]
         x = np.linspace(0, mask_tmp.shape[0], mask_tmp.shape[1])
         y = np.linspace(0, mask_tmp.shape[1], mask_tmp.shape[1])
         xv, yv = np.meshgrid(x, y)
@@ -396,8 +401,9 @@ def check_halfy(self):
         combine_mask = mask_circle*mask_tmp[:,:,0]
         if self.verbose >= 1 and np.sum(combine_mask)>=1:
             logging.warning('****** WARNING ******')
-            logging.warning('half_y = {} is too small, you should increase it.'.format(self.args["halfy"]))
-            
+            logging.info('halfy = {} um is too small, you should increase to ~{:6f} um.'.format(self.args["halfy"],
+                         self.args["halfy"] + 2.5*np.sum(combine_mask)/(2*np.pi*radius)*np.mean(self.sample_input.DS['voxel_size'])))
+
     def set_beam(self):
         """
         set the X-ray beam source, including beam energy, FWHM, flux, shape, source-to-sample distance, name etc.
@@ -406,7 +412,7 @@ def set_beam(self):
         
     def set_sample(self):
         """
-        set the sample properties: including the grain map, voxel size, .
+        set the sample properties: including the grain map, voxel size, density, mass absorption.
         """
         self.sample_input = sample(self.sample_filename, min_misori = self.args['min_misori'], crystal_system = self.args['crystal_system'],
                                       remove_small_grains = self.args['remove_small_grains'], min_vol = self.args['min_vol'])    
@@ -620,19 +626,23 @@ def generate_slurm_script(self, script_name="fwd_projector_slurm.sh"):
 #SBATCH --mem-per-cpu=8000
 #SBATCH --time=01:00:00
 
+module load jupyter-slurm/2025.04.5
+echo "===== JOB START ====="
 date
-source /cvmfs/hpc.esrf.fr/software/packages/linux/x86_64/jupyter-slurm/latest/envs/jupyter-slurm/bin/activate
+
 log_path="{log_path}_${{SLURM_JOB_ID}}_${{SLURM_ARRAY_TASK_ID}}.log"
 dty_file="{dty_file_list}"
 dty_values=$(sed -n "${{SLURM_ARRAY_TASK_ID}}p" $dty_file)
-echo "OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 PYTHONPATH={id11_code_path}" >> $log_path 2>&1
-OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 PYTHONPATH={id11_code_path} \
+echo "OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 PYTHONPATH={id11_code_path}"
+OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 PYTHONPATH={id11_code_path}
+export PYTHONPATH={id11_code_path}
 python3 -c "from ImageD11.forward_model import forward_projector; \
 fp = forward_projector.forward_projector(sample_filename='{sample_filename}', pars_filename='{pars_filename}', \
 phase_name='{phase_name}', output_folder='{output_folder}', gids_mask={gids_mask}, \
 verbose={verbose}, auto_set={auto_set}, detector_mask={detector_mask}, to_sparse={to_sparse}, \
 read_fwd_peaks_from_file={read_fwd_peaks_from_file}, **{args}); \
-fp.run_series_dty([float(v) for v in '$dty_values'.split()])" >> $log_path 2>&1
+fp.run_series_dty([float(v) for v in '$dty_values'.split()])"
+echo "===== JOB END ====="
 date""".format(
             outfile_path=outfile_path,
             errfile_path=errfile_path,
@@ -780,6 +790,7 @@ def write_cf(self):
     def plot_cf(self, cf_type = '2d', m=None):
         """
         plot the sinogram of cf peaks, i.e. dty vs omega colored by peak intensity
+        m is a boolen mask
         """
         plt.figure(figsize=(10,8))
         if cf_type == '2d':
@@ -870,7 +881,7 @@ def forward_peaks_voxels(beam, sample, omega_angles, ucell, pars, dty=0.0,
     t0 = time.time()
     results = Parallel(n_jobs=-1, backend="loky")(delayed(process_omega)(
         omega, mask, DS, beam, sample, pars, ucell, rot_step, ds_max, Lsam2det, args, verbose
-    ) for omega in tqdm(omega_angles, desc="Processing omega angles"))
+    ) for omega in tqdm(omega_angles, desc="Processing omega angles", file=sys.stdout))
     # # Regarding multiprocessing.Pool, tried chunk, shared memory etc, it does not help to improve the speed
     # set_tmp_dir()
     # with multiprocessing.Pool(processes=min(multiprocessing.cpu_count(), 24)) as pool:
diff --git a/ImageD11/forward_model/grainmaps.py b/ImageD11/forward_model/grainmaps.py
@@ -51,12 +51,13 @@ class grainmap:
     example:
     filename = ds.grains_file
     gm = grainmap(filename)
-    gm.merge_and_identify_grains(FirstGrainID = 0, dis_tol = np.sqrt(2))
+    gm.merge_and_identify_grains(FirstGrainID = 0, dis_tol = np.sqrt(3))
     DS = gm.DS
     '''
-    def __init__(self, filename = None, outname = None, min_misori = 3, crystal_system = 'cubic', remove_small_grains = True, min_vol = 2):
+    def __init__(self, filename = None, outname = None, min_misori = 3, crystal_system = 'cubic', remove_small_grains = True, min_vol = 2, verbose = 0):
         self.filename = filename
         self.outname = outname
+        self.verbose = verbose
         if self.outname is None and self.filename is not None:
             self.outname = os.path.join(os.path.split(self.filename)[0], 'DS.h5')
         if filename is not None:
@@ -68,19 +69,21 @@ def __init__(self, filename = None, outname = None, min_misori = 3, crystal_syst
         self.crystal_system = crystal_system
         self.remove_small_grains = remove_small_grains
         self.min_vol = min_vol
-        print('****************************************** Parameters for operating the grain map: ')
-        print('Output file name: {}'.format(self.outname))
-        print('min_misori = {}'.format(min_misori))
-        print('crystal_system: {}'.format(crystal_system))
-        print('remove_small_grains = {}'.format(remove_small_grains))
-        print('min_vol = {} voxels'.format(min_vol))
-        
+        if self.verbose >= 1:
+            print('****************************************** Parameters for operating the grain map: ')
+            print('Output file name: {}'.format(self.outname))
+            print('min_misori = {}'.format(min_misori))
+            print('crystal_system: {}'.format(crystal_system))
+            print('remove_small_grains = {}'.format(remove_small_grains))
+            print('min_vol = {} voxels'.format(min_vol))
+            
     
     def read_tensor_map(self):
         if os.path.exists(self.filename):
             tensor_map = io.read_h5_file(self.filename)
             if isinstance(tensor_map, dict):
-                io.print_all_keys(tensor_map)
+                if self.verbose >= 1:
+                    io.print_all_keys(tensor_map)
             else:
                 print("The file did not produce a dictionary.")
             return tensor_map
@@ -97,7 +100,8 @@ def convert2DS(self):
         tensor_map_flag = False
         for key in tensor_map.keys():
             if 'TensorMap' in key:
-                print('Got the key name {}'.format(key))
+                if self.verbose >= 1:
+                    print('Got the key name {}'.format(key))
                 keyname = key
                 tensor_map_flag = True
         
@@ -108,7 +112,8 @@ def convert2DS(self):
             DS = {}
             for k in keys_list:
                 if k in tensor_map[keyname]['maps'].keys():
-                    print('Loading {} with a shape of {} to DS ...'.format(k, tensor_map[keyname]['maps'][k].shape))
+                    if self.verbose >= 1:
+                        print('Loading {} with a shape of {} to DS ...'.format(k, tensor_map[keyname]['maps'][k].shape))
                     DS[k] = tensor_map[keyname]['maps'][k]
             if 'step' in tensor_map[keyname].keys():
                 DS['voxel_size'] = tensor_map[keyname]['step'] # [um/pixel]
@@ -131,7 +136,7 @@ def convert2DS(self):
             print('Found a map (dimensions {}) with {} voxels to be labeled with grain ID'.format(DS['labels'].shape, np.max(DS['labels'])))
             print('Please note: current labels are only randomized without any region merging !!!')
 
-        # get U from UBI if no 'U' existed
+        # get U from UBI if no 'U' existed, this is usually the case with tensor map before refinement
         if 'U' not in DS.keys() and 'UBI' in DS.keys():
             try:
                 latticepar = np.array(tensor_map[keyname]['phases']['0'].decode('utf-8').split(), dtype = float)
@@ -141,11 +146,15 @@ def convert2DS(self):
                 
                 DS['U'] = np.empty((DS['UBI'].shape[0], DS['UBI'].shape[1], DS['UBI'].shape[2], 3, 3))
                 DS['U'].fill(np.nan)
+
+                DS['B'] = np.empty((DS['UBI'].shape[0], DS['UBI'].shape[1], DS['UBI'].shape[2], 3, 3))
+                DS['B'].fill(np.nan)
                 
                 indices = np.argwhere(~np.isnan(DS['UBI'][:, :, :, 0, 0]))
                 for (i, j, k) in indices:
                     DS['U'][i, j, k, :, :] = np.dot(B, DS['UBI'][i, j, k, :, :]).T
-                print('Done with deriving U matrix !')
+                    DS['B'][i, j, k, :, :] = B
+                print('Done with deriving U and B matrices !')
             except KeyError as e:
                 print("Key error: {}".format(e))
             except Exception as e:
@@ -284,7 +293,7 @@ def DS_remove_small_grains(DS, min_vol = 2):
     return DS_out
 
 
-def DS_merge_and_identify_grains_sub(DS, FirstGrainID = 0, min_misori = 3.0, dis_tol = np.sqrt(2), crystal_system = 'cubic'):
+def DS_merge_and_identify_grains_sub(DS, FirstGrainID = 0, min_misori = 3.0, dis_tol = np.sqrt(3), crystal_system = 'cubic'):
     
     """
     sub function:
@@ -296,7 +305,7 @@ def DS_merge_and_identify_grains_sub(DS, FirstGrainID = 0, min_misori = 3.0, dis
     DS                -- grain map dictionary, gm = grainmap(tensor_map_file); DS = gm.DS
     FirstGrainID      -- First grain ID to be considered for merging, by default = 0
     min_misori        -- misorientation for merging regions
-    dis_tol           -- maximum distance around the target voxel for identifying its neigboring grains, np.sqrt(2) corresponds to distance to the diagonal voxel in 2D
+    dis_tol           -- maximum distance around the target voxel for identifying its neigboring grains, np.sqrt(3) ensure the diagonal voxel in 2D included
     crystal_system    -- crystal system name one of ['cubic', 'hexagonal', 'orthorhombic', 'tetragonal', 'trigonal', 'monoclinic', 'triclinic']
 
     Returns:
diff --git a/ImageD11/forward_model/io.py b/ImageD11/forward_model/io.py
@@ -11,7 +11,7 @@
 # March 20, 2024
 # updated on January 30, 2025
 
-import os, shutil
+import os, sys, shutil
 import glob
 
 import numpy as np
@@ -394,7 +394,7 @@ def write_xdmf(h5name, xdmf_filename = None, ctrl = 'recon_mlem', attributes = [
 
     
 def camera_names_ID11():
-    return ['marana3', 'marana1', 'marana2', 'marana', 'frelon1', 'frelon3', 'frelon6', 'eiger', 'basler_eh32', 'frelon16']
+    return ['marana3', 'marana1', 'marana2', 'marana', 'frelon1', 'frelon3', 'frelon6', 'eiger', 'basler_eh32', 'frelon16', 'pco1']
 
 
 def guess_camera_name(h5name, scan = '1.1'):
@@ -630,6 +630,32 @@ def read_fsparse(h5_file, group_name = "/entry_0000/ESRF-ID11/eiger/data"):
     return fsparse_pks
 
 
+def deep_getsizeof(obj, seen=None):
+    """
+    Get memory info of a dict, list, tuple or set, it does so recursively so that it can collect memory info of all elements.
+    Args:
+        obj: an object, e.g. dict, list, tuple or set
+    Returns:
+        size: size of the object in bytes
+    Example to use:
+        size_bytes = deep_getsizeof(sparse_cuda_dict)
+        print(size_bytes / (1024**3), "GB")
+    """
+    if seen is None:
+        seen = set()
+    obj_id = id(obj)
+    if obj_id in seen:
+        return 0
+    seen.add(obj_id)
+    size = sys.getsizeof(obj)
+    if isinstance(obj, dict):
+        size += sum(deep_getsizeof(k, seen) + deep_getsizeof(v, seen)
+                    for k, v in obj.items())
+    elif isinstance(obj, (list, tuple, set)):
+        size += sum(deep_getsizeof(i, seen) for i in obj)
+    return size
+
+
 def copytree_with_progress(src, dst, skip_keys=None, create_subfolder = True, overwrite = False):
     """
     Copies a folder with all subfolders and files, displaying progress. 
