Merge pull request #306 from Kvieta1990/master

Bragg plot bugs fixed.

Author: marshallmcdonnell

addie/rietveld/braggtree.py

@@ -61,6 +61,45 @@ def __init__(self, parent):
         # reset
         self.reset_bragg_tree()
 
+    def process_selected_nodes(self, selected_nodes):
+        """
+        Process selected nodes so that returned nodes will only
+        contain main leaf but not any children.
+
+        If both parent and children are selected from the tree, we want
+        to get rid of children from the selected wks list. Meanwhile, we want
+        to check the box corresponding to the selected children.
+
+        Arguments:
+            selected_nodes {list} -- List of selected nodes
+        Return:
+            selected_nodes_temp {list} -- List of nodes with all children removed.
+        """
+
+        leaf_dict_temp = self._main_window.rietveld_ui.treeWidget_braggWSList._leafDict
+        for item_temp in selected_nodes:
+            if item_temp not in list(leaf_dict_temp.keys()):
+                parent_found = False
+                for key, list_temp in leaf_dict_temp.items():
+                    for item_temp_1 in list_temp:
+                        if item_temp in item_temp_1:
+                            bank_temp = list_temp.index(item_temp_1) + 1
+                            self._main_window._braggBankWidgets[bank_temp].setChecked(True)
+                            parent_found = True
+                            parent_in_tree = key
+                            break
+                    if parent_found:
+                        break
+                if parent_in_tree not in selected_nodes:
+                    selected_nodes.append(parent_in_tree)
+
+        selected_nodes_temp = []
+        for item_temp in selected_nodes:
+            if item_temp in list(leaf_dict_temp.keys()):
+                selected_nodes_temp.append(item_temp)
+
+        return selected_nodes_temp
+
     def _get_bank_id(self, bank_wksp):
         """Get bank ID from a workspace name with the structure:
              Bank 1 - <float for theta angle>
@@ -234,6 +273,9 @@ def do_plot_ws(self):
         # get the selected items of tree and sort them alphabetically
         item_list = self.get_selected_items()
         item_list = [str(item.text()) for item in item_list]
+        # item_list.sort()
+
+        item_list = self.process_selected_nodes(item_list)
         item_list.sort()
 
         # FIXME/LATER - replace this by signal
@@ -317,11 +359,15 @@ def do_select_gss_node(self):
         """
         # get selected nodes
         selected_nodes = self.get_selected_items()
+        selected_nodes = [str(item.text()) for item in selected_nodes]
+
+        selected_nodes = self.process_selected_nodes(selected_nodes)
 
         # set to plot
         for gss_group_node in selected_nodes:
-            gss_group_name = str(gss_group_node.text())
-            self._main_window.set_bragg_ws_to_plot(gss_group_name)
+            # gss_group_name = str(gss_group_node.text())
+            # self._main_window.set_bragg_ws_to_plot(gss_group_name)
+            self._main_window.set_bragg_ws_to_plot(gss_group_node)
 
     def get_current_main_nodes(self):
         """

addie/rietveld/event_handler.py

@@ -261,16 +261,24 @@ def plot_bragg_bank(main_window):
         # multiple GSAS file/single bank mode: get GSAS group from tree
         gss_group_list = main_window.rietveld_ui.treeWidget_braggWSList.get_main_nodes()
         gss_group_list.remove('workspaces')
-
     else:
         # single GSAS file mode
         status, ret_obj = main_window.rietveld_ui.treeWidget_braggWSList.get_current_main_nodes()
+        # ZYP -> Sometimes, it can happen that the returned `status` here is
+        # `True`, but we have empty `ret_obj`. So we need to tackle such a
+        # situation. Since this will happen when multiple GSS files exist in wks,
+        # and therefore when this happens, we just simply roll into the multiple GSS mode.
         if status:
-            gss_group = ret_obj[0]
+            if len(ret_obj) > 0:
+                gss_group = ret_obj[0]
+            else:
+                gss_group_list = main_window.rietveld_ui.treeWidget_braggWSList.get_main_nodes()
+                gss_group_list.remove('workspaces')
         else:
             raise RuntimeError(
                 'Unable to get current selected main node(s) due to {0}.'.format(ret_obj))
-        gss_group_list = [gss_group]
+        if len(ret_obj) > 0:
+            gss_group_list = [gss_group]
 
     # remove banks from plot
     for gss_group_name in gss_group_list:
@@ -282,22 +290,42 @@ def plot_bragg_bank(main_window):
     for ws_group in gss_group_list:
         plot_data_dict[ws_group] = new_bank_list[:]
 
+    # ZYP -> Check whether selected bank (to plot) exists or not.
+    # ZYP -> If not, no action will be taken and msg will be printed out to terminal.
+    banks_checker_ok = True
+    leaf_dict_temp = main_window.rietveld_ui.treeWidget_braggWSList._leafDict
+    for ws_group in gss_group_list:
+        for bank_to_check in plot_data_dict[ws_group]:
+            if int(bank_to_check) > len(leaf_dict_temp[ws_group]):
+                print("Bank-{0} not existing in {1} and therefore not to be plotted.".format(bank_to_check, ws_group))
+                banks_checker_ok = False
+                break
+
     if plot_multi_gss:
         main_window.rietveld_ui.graphicsView_bragg.set_to_single_gss(False)
     else:
         main_window.rietveld_ui.graphicsView_bragg.set_to_single_gss(True)
         status, ws_name_list = main_window.rietveld_ui.treeWidget_braggWSList.get_current_main_nodes()
+        # ZYP -> Sometimes, it can happen that the returned `status` here is
+        # `True`, but we have empty `ret_obj`. So we need to tackle such a
+        # situation. Since this will happen when multiple GSS files exist in wks,
+        # and therefore when this happens, we just simply roll into the multiple GSS mode.
+        if len(ws_name_list) == 0:
+            ws_name_list = main_window.rietveld_ui.treeWidget_braggWSList.get_main_nodes()
+            ws_name_list.remove("workspaces")
         plot_data_dict = {ws_name_list[0]: plot_data_dict[ws_name_list[0]]}
 
     # plot new
-    main_window.rietveld_ui.graphicsView_bragg.plot_banks(
-        plot_data_dict, main_window._currBraggXUnit)
+    # ZYP -> Check whether selected bank (to plot) exists or not.
+    if banks_checker_ok:
+        main_window.rietveld_ui.graphicsView_bragg.plot_banks(
+            plot_data_dict, main_window._currBraggXUnit)
 
-    # reset
-    reset_bragg_data_range(main_window, main_window._currBraggXUnit)
+        # reset
+        reset_bragg_data_range(main_window, main_window._currBraggXUnit)
 
-    # rescale
-    do_rescale_bragg(main_window)
+        # rescale
+        do_rescale_bragg(main_window)
 
 
 def do_rescale_bragg(main_window):
@@ -526,10 +554,18 @@ def plot_bragg(main_window, ws_list, bankIds, clear_canvas=False):
 
     # plot all workspsaces
     plot_data_dict = dict()
+    leaf_dict_temp = main_window.rietveld_ui.treeWidget_braggWSList._leafDict
     for bragg_ws_name in ws_list:
         # construct dictionary for plotting
         # main_window._myController.get_bank_numbers(ws_group)
-        plot_data_dict[bragg_ws_name] = bankIds
+        to_plot = []
+        for bank_temp in bankIds:
+            if int(bank_temp) <= len(leaf_dict_temp[bragg_ws_name]):
+                to_plot.append(bank_temp)
+            else:
+                print("Bank-{0} not existing in {1} and therefore not to be plotted.".format(bank_temp, bragg_ws_name))
+        # plot_data_dict[bragg_ws_name] = bankIds
+        plot_data_dict[bragg_ws_name] = to_plot
 
     # plot
     main_window.rietveld_ui.graphicsView_bragg.plot_banks(