fixing phase map plots

Author: cophus

py4DSTEM/process/diffraction/crystal_phase.py

@@ -966,8 +966,8 @@ def plot_phase_weights(
         if total_intensity_normalize:
             sub = self.int_total > 0.0
             for a0 in range(self.num_fits):
-                phase_weights[:, :, a0][sub] /= self.int_total[sub]
-            phase_residuals[sub] /= self.int_total[sub]
+                phase_weights[:, :, a0][sub] /= (self.int_total[sub]+1e-12)
+            phase_residuals[sub] /= (self.int_total[sub]+1e-12)
 
         # intensity range for plotting
         if weight_normalize:
@@ -1112,7 +1112,7 @@ def plot_phase_maps(
         if total_intensity_normalize:
             sub = self.int_total > 0.0
             for a0 in range(self.num_fits):
-                phase_weights[:, :, a0][sub] /= self.int_total[sub]
+                phase_weights[:, :, a0][sub] /= (self.int_total[sub]+1e-12)
 
         # intensity range for plotting
         if weight_normalize:
@@ -1228,7 +1228,8 @@ def plot_phase_maps(
     def plot_dominant_phase(
         self,
         use_correlation_scores=False,
-        reliability_range=(0.0, 1.0),
+        reliability_range=None,
+        correlation_range=None,
         normalize_exp_intensity=True,
         sigma=0.0,
         phase_colors=None,
@@ -1248,7 +1249,9 @@ def plot_dominant_phase(
         use_correlation_scores: bool
             Set to True to use correlation scores instead of reliabiltiy from intensity residuals.
         reliability_range: (float, float)
-            Plotting intensity range
+            Plotting intensity range from reliability
+        correlation_range: (float, float)
+            Plotting intensity range from correlation score
         sigma: float
             Smoothing in units of probe position.
         phase_colors: np.array
@@ -1273,6 +1276,7 @@ def plot_dominant_phase(
 
         """
 
+
         if phase_colors is None:
             phase_colors = np.array(
                 [
@@ -1323,16 +1327,38 @@ def plot_dominant_phase(
             sub = phase_sig[a0] > phase_corr
             phase_map[sub] = a0
             phase_corr[sub] = phase_sig[a0][sub]
+        self.phase_corr_total = np.sum(phase_corr,axis=0)
+
+        phase_scale = np.ones((
+            self.phase_sig.shape[1],
+            self.phase_sig.shape[2],
+        ))
+        # if self.single_phase:
+            # if reliability_range is not None:
+            #     phase_scale *= np.clip(
+            #         (self.phase_reliability - reliability_range[0])
+            #         / (reliability_range[1] - reliability_range[0]),
+            #         0,
+            #         1,
+            #     )
+            # if correlation_range is not None:
+            #     phase_scale *= np.clip(
+            #         (self.phase_corr_total - correlation_range[0])
+            #         / (correlation_range[1] - correlation_range[0]),
+            #         0,
+            #         1,
+            #     )
+            
+            # phase_scale = np.clip(
+            #     (self.phase_reliability - reliability_range[0])
+            #     / (reliability_range[1] - reliability_range[0]),
+            #     0,
+            #     1,
+            # )
+
+        # else:
+        if not self.single_phase:
 
-        if self.single_phase:
-            phase_scale = np.clip(
-                (self.phase_reliability - reliability_range[0])
-                / (reliability_range[1] - reliability_range[0]),
-                0,
-                1,
-            )
-
-        else:
             # find the second correlation score for each crystal and match index
             for a0 in range(self.num_crystals):
                 corr = phase_sig[a0].copy()
@@ -1345,14 +1371,28 @@ def plot_dominant_phase(
 
             # normalize the reliability by the intensity of each experimental pattern
             if normalize_exp_intensity:
-                phase_rel /= self.int_total
-
-            phase_scale = np.clip(
-                (phase_rel - reliability_range[0])
+                phase_rel /= (self.int_total+1e-12)
+
+            # phase_scale = np.clip(
+            #     (phase_rel - reliability_range[0])
+            #     / (reliability_range[1] - reliability_range[0]),
+            #     0,
+            #     1,
+            # )
+        if reliability_range is not None:
+            phase_scale *= np.clip(
+                (self.phase_reliability - reliability_range[0])
                 / (reliability_range[1] - reliability_range[0]),
                 0,
                 1,
             )
+        if correlation_range is not None:
+            phase_scale *= np.clip(
+                (self.phase_corr_total - correlation_range[0])
+                / (correlation_range[1] - correlation_range[0]),
+                0,
+                1,
+            )
 
         # Print the total area of fraction of each phase
         if print_fractions: