Пример #1
0
class MainWindow():
    """
    Defines the main window for the envelope tool GUI
    """
    def __init__(self):
        """Initialise the main window"""
        self.window = Window(ROOT.gClient.GetRoot(), # pylint: disable=E1101
                             ROOT.gClient.GetRoot(), # pylint: disable=E1101
                             SHARE_DIR+"main_frame.json")
        self.lattice = Lattice()
        self.beam_setup = None
        self.magnet_setup = None
        self.plot_setup = None
        self.plot_setup_options = [{"variable_type":0,
                                    "first_var":0,
                                    "plot_apertures":True}]
        self.window.set_button_action("&Beam Setup", self.beam_button_action)
        self.window.set_button_action("&Magnet Setup",
                                                      self.magnet_button_action)
        self.window.set_button_action("&Plot Setup", self.plot_button_action)
        self.window.set_button_action("E&xit", self.exit_button_action)
        self.update_plot()

    def beam_button_action(self):
        """Handle a Beam Setup button press"""
        if self.beam_setup == None:
            self.beam_setup = BeamSetup(self, self.window.main_frame)
            ref, ellipse = self.lattice.get_beam()
            self.beam_setup.set_reference(ref)
            self.beam_setup.set_matrix(ellipse)

    def magnet_button_action(self):
        """Handle a Magnet Setup button press"""
        if self.magnet_setup == None:
            self.magnet_setup = MagnetSetup(self, self.window.main_frame)

    def plot_button_action(self):
        """Handle a Plot Setup button press"""
        if self.plot_setup == None:
            self.plot_setup = PlotSetup(self, self.window.main_frame,
                                        self.plot_setup_options)

    def exit_button_action(self):
        """Handle a Exit button press"""
        self.window.close_window()

    def update_plot(self):
        """Update the plot"""
        canvas_frame = self.window.get_frame("main_canvas", "canvas")
        canvas = canvas_frame.GetCanvas()
        Plotter(self.plot_setup_options, canvas, self.lattice.ref_list,
                self.lattice.ellipse_list, self.lattice.get_field_list())
Пример #2
0
class TwissSetup:
    """GUI window to setup a beam according to 2D Twiss parameterisation"""
    def __init__(self, beam_setup, root_frame):
        """Initialise the window"""
        self.window = Window(
            ROOT.gClient.GetRoot(),  # pylint: disable = E1101
            root_frame,
            SHARE_DIR + "twiss_setup.json")
        self.beam_setup = beam_setup
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)

    def okay_action(self):
        """Handle a click on the Okay button, building the beam matrix"""
        ref = self.beam_setup.get_reference()
        matrix = maus_cpp.covariance_matrix.create_from_twiss_parameters(
            mass=ref['mass'],
            momentum=ref['p'],
            emittance_x=self.window.get_text_entry("emittance_x", type(1.)),
            beta_x=self.window.get_text_entry("beta_x", type(1.)),
            emittance_y=self.window.get_text_entry("emittance_y", type(1.)),
            beta_y=self.window.get_text_entry("beta_y", type(1.)),
            emittance_l=self.window.get_text_entry("emittance_l", type(1.)),
            beta_l=self.window.get_text_entry("beta_l", type(1.)),
            alpha_x=self.window.get_text_entry("alpha_x", type(1.)),
            alpha_y=self.window.get_text_entry("alpha_y", type(1.)),
            alpha_l=self.window.get_text_entry("alpha_l", type(1.)),
            dispersion_x=self.window.get_text_entry("disp_x", type(1.)),
            dispersion_prime_x=self.window.get_text_entry(
                "disp_prime_x", type(1.)),
            dispersion_y=self.window.get_text_entry("disp_y", type(1.)),
            dispersion_prime_y=self.window.get_text_entry(
                "disp_prime_y", type(1.)))
        self.beam_setup.set_matrix(matrix)
        self.beam_setup.get_matrix()  # check for consistency
        self.window.close_window()
        self.beam_setup.matrix_select = None

    def cancel_action(self):
        """Handle a click on the Cancel button"""
        self.window.close_window()
        self.beam_setup.matrix_select = None
Пример #3
0
class TwissSetup:
    """GUI window to setup a beam according to 2D Twiss parameterisation"""
    def __init__(self, beam_setup, root_frame):
        """Initialise the window"""
        self.window = Window(ROOT.gClient.GetRoot(), # pylint: disable = E1101
                             root_frame,
                             SHARE_DIR+"twiss_setup.json")
        self.beam_setup = beam_setup
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)

    def okay_action(self):
        """Handle a click on the Okay button, building the beam matrix"""
        ref = self.beam_setup.get_reference()
        matrix = maus_cpp.covariance_matrix.create_from_twiss_parameters(
          mass = ref['mass'],
          momentum = ref['p'],
          emittance_x = self.window.get_text_entry("emittance_x", type(1.)),
          beta_x = self.window.get_text_entry("beta_x", type(1.)),
          emittance_y = self.window.get_text_entry("emittance_y", type(1.)),
          beta_y = self.window.get_text_entry("beta_y", type(1.)),
          emittance_l = self.window.get_text_entry("emittance_l", type(1.)),
          beta_l = self.window.get_text_entry("beta_l", type(1.)),
          alpha_x = self.window.get_text_entry("alpha_x", type(1.)),
          alpha_y = self.window.get_text_entry("alpha_y", type(1.)),
          alpha_l = self.window.get_text_entry("alpha_l", type(1.)),
          dispersion_x = self.window.get_text_entry("disp_x", type(1.)),
          dispersion_prime_x = self.window.get_text_entry("disp_prime_x",
                                                          type(1.)),
          dispersion_y = self.window.get_text_entry("disp_y", type(1.)),
          dispersion_prime_y = self.window.get_text_entry("disp_prime_y",
                                                          type(1.))
        )
        self.beam_setup.set_matrix(matrix)
        self.beam_setup.get_matrix() # check for consistency
        self.window.close_window()
        self.beam_setup.matrix_select = None

    def cancel_action(self):
        """Handle a click on the Cancel button"""
        self.window.close_window()
        self.beam_setup.matrix_select = None
Пример #4
0
class MagnetSetup:
    """
    GUI handler for field setup
    """
    def __init__(self, main_window, parent):
        """Initialise the GUI window"""
        self.main_window = main_window
        self.parent = parent
        manipulators = {"magnet_setup_action":self.magnet_setup_action}
        self.window = Window(ROOT.gClient.GetRoot(), # pylint: disable = E1101
                             parent,
                             data_file = self.share_dir+"magnet_setup.json",
                             manipulator_dict = manipulators)
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)

    def okay_action(self):
        """Handle okay button press"""
        field_list_out = []
        field_frame = self.window.get_frame_dict("field_list", "vertical_frame")
        for frame in field_frame["children"]:
            try:
                name = frame["name"]
                scale_factor = float(frame["text_entry"].text_entry.GetText())
                field_list_out.append({"field_name":name,
                                       "scale_factor":scale_factor})
            except KeyError:
                pass
        self.main_window.lattice.set_fields(field_list_out)
        self.main_window.lattice.run_lattice()
        self.main_window.update_plot()
        self.window.close_window()
        self.main_window.magnet_setup = None
    
    def cancel_action(self):
        """Handle cancel button press"""
        self.window.close_window()
        self.main_window.magnet_setup = None

    def magnet_setup_action(self, item): # pylint: disable=W0613
        """Iterate over magnets and add a line per magnet (special)"""
        field_list = self.main_window.lattice.get_field_list()
        geometry_list = {
        "type":"vertical_frame",
        "layout":"close",
        "name":"field_list",
        "children":[]}
        if len(field_list) == 0:
            geometry_list["children"].append({"type":"label",
                    "name":"No fields in lattice",
                    "label_length":30})
            return geometry_list
        for field in field_list:
            geometry_list["children"].append({
              "type":"named_text_entry",
              "default_text":str(field["scale_factor"]),
              "label_length":10,
              "name":str(field["field_name"]),
              "tool_tip":"Longitudinal beta function"
            })
        return geometry_list

    share_dir = ""
Пример #5
0
class OnlineGui(): # pylint: disable = R0902
    """Main frame for the online gui window"""
    def __init__(self, datacards):
        """Initialise the main window"""
        self.last = datetime.datetime.now()
        self.window = Window(ROOT.gClient.GetRoot(),
                             ROOT.gClient.GetRoot(),
                             SHARE_DIR+"online_gui.json")
        self._set_actions()
        self._docstore = None
        self._datacards = datacards
        self._draw_titles = Queue.Queue()
        for title in datacards["online_gui_default_canvases"]:
            self._draw_titles.put(title)
        self._canvases_read = Queue.Queue()
        self._canvases_draw = ThreadedValue([])
        self._redraw_target = None
        self._paused = ThreadedValue(False)
        rotate_tmp = self.window.get_text_entry("Rotate period", type(1.))
        self._rotate_period = ThreadedValue(rotate_tmp)
        reload_tmp = self.window.get_text_entry("Reload period", type(1.))
        self._reload_period = ThreadedValue(reload_tmp)
        self._poll_period = ThreadedValue(0.1)
        self._collections = ThreadedValue([])
        self._init_docstore()
        self._start_polling()

    def _set_actions(self):
        """
        Set up actions (responses to signals)
        """
        self.window.set_button_action("&Pause", self.pause_action)
        self.window.set_button_action("&Redraw", self.redraw_action)
        self.window.set_button_action("Re&load", self.reload_action)
        self.window.set_button_action("Re&connect", self.reconnect_action)
        self.window.set_button_action("E&xit", self.exit_action)
        self.window.set_button_action("&<", self.back_action)
        self.window.set_button_action("&>", self.forward_action)
        self.window.get_frame("canvas_options", "vertical_frame").\
                                                                Resize(200, 500)
        self.window.get_frame("canvas_select", "list_box").Resize(200, 200)
        self.window.set_action("canvas_select",
                               "list_box",
                               "SelectionChanged()",
                               self.canvas_select_action)
        self.window.set_action("Rotate period", "named_text_entry",
                               "ReturnPressed()", self.rotate_period_action)
        self.window.set_action("Reload period", "named_text_entry",
                               "ReturnPressed()", self.reload_period_action)

    def _init_docstore(self):
        """
        Initialise the docstore
        """
        datacards = self._datacards
        self._docstore = RootDocumentStore(
                  datacards["root_document_store_timeout"],
                  datacards["root_document_store_poll_time"], 10, 10)
        self._docstore.connect({"host":datacards["online_gui_host"],
                               "port":datacards["online_gui_port"]})

    def _start_polling(self):
        """
        Start polling the database and rotating the canvases
        """
        my_thread = threading.Thread(target=self._poll_canvas_draw,
                                     args=())
        my_thread.daemon = True
        my_thread.start()
        my_thread = threading.Thread(target=self._poll_docstore,
                                     args=())
        my_thread.daemon = True
        my_thread.start()

    def pause_action(self):
        """Handle a Pause/Unpause button press"""
        pause_button = self.window.get_frame("&Pause", "button")
        if self._paused.get_value():
            self._paused.set_value(False)
            pause_button.SetText("&Pause")
        else:
            self._paused.set_value(True)
            pause_button.SetText("&Play")

    def redraw_action(self):
        """
        Handle a Redraw button press

        Redraws self._redraw_target
        """
        if self._redraw_target == None:
            raise GuiError("No canvas to draw")
        self._do_redraw(self._redraw_target)

    def reload_action(self):
        """Handle a Reload button press"""
        self._reload()
        if self._redraw_target != None:
            go_back_fifo_queue(self._draw_titles)
            self._draw_next()

    def reconnect_action(self):
        """Handle a Reconnect button press"""
        self._docstore.disconnect()
        self._init_docstore()

    def exit_action(self):
        """Handle a Exit button press"""
        self._paused.set_value(True) # tell polling to wait
        time.sleep(self._poll_period.get_value()) # thinkapow
        self.window.close_window()

    def reload_period_action(self):
        """Handle a change to the reload period"""
        tmp_reload_period = self.window.get_text_entry("Reload period",
                                                       type(1.))
        tmp_reload_period = float(tmp_reload_period)
        self._reload_period.set_value(tmp_reload_period)

    def rotate_period_action(self):
        """Handle a change to the rotate period"""
        tmp_rotate_period = self.window.get_text_entry("Rotate period",
                                                       type(1.))
        tmp_rotate_period = float(tmp_rotate_period)
        self._rotate_period.set_value(tmp_rotate_period)

    def forward_action(self):
        """Handle a click on the > forward button"""
        self._draw_next()

    def back_action(self):
        """Handle a click on the < back button"""
        go_back_fifo_queue(self._draw_titles)
        go_back_fifo_queue(self._draw_titles)
        self._draw_next()

    def canvas_select_action(self):
        """Handle an update to the selected canvases"""
        selected = self._get_canvas_selected()
        new_draw_titles = Queue.Queue()
        for item in selected:
            new_draw_titles.put(item)
        self._draw_titles = new_draw_titles

    def _poll_canvas_draw(self):
        """
        Iterate over the canvas_draw queue, updating the GUI as appropriate

        * Update the queue from the canvas_read queue
        * Refresh the combo box with items in the canvases_draw queue
        """
        # poll_number = 0
        while True:
            # poll_number += 1
            self._draw_next()
            self._sleepish(self._rotate_period)
            # handle pause
            while self._paused.get_value():
                time.sleep(self._poll_period.get_value())

    def _draw_next(self):
        """
        Draw the next canvas in the selected queue
        """
        canvases_draw = self._filter_canvases(self._canvases_draw.get_value())
        if len(canvases_draw) > 0:
            try:
                title = self._draw_titles.get_nowait()
                new_title = copy.deepcopy(title)
                self._draw_titles.put(title)
                for wrap in canvases_draw:
                    if wrap.canvas_title() == new_title:
                        self._do_redraw(wrap)
                        break
            except Queue.Empty:
                pass # no canvas selected
            except ValueError:
                pass
        self._canvases_draw.set_value(canvases_draw)


    def _sleepish(self, sleep_length):
        """
        Sleep - but keep polling for changes to the sleep_length
        * sleep_length - threaded value containing numeric sleep time
        """
        time_slept = 0.
        poll_period = self._poll_period.get_value()
        while sleep_length.get_value() - time_slept > poll_period:
            time_slept += self._poll_period.get_value()
            time.sleep(poll_period)

    def _update_canvases_draw(self):
        """update canvases"""
        draw = [wrap.deepcopy() \
                           for wrap in generate_fifo_queue(self._canvases_read)]
        draw.sort(key = lambda wrap: wrap.sort_key())
        return draw

    def _do_redraw(self, wrap):
        """
        Draw the canvas wrap in the main_canvas
        """
        # wrap_copy is for ROOT
        wrap_copy = wrap.deepcopy()
        embedded_canvas = self.window.get_frame("main_canvas", "canvas")
        # this eats the wrap_copy TCanvas
        wrap_copy.canvas_wrapper.EmbedCanvas(embedded_canvas)
        embedded_canvas.GetCanvas().Update()
        # this _redraw_target is for the main thread
        self._redraw_target = wrap.deepcopy()

    def _poll_docstore(self):
        """
        Poll the doc store for documents
        """
        while True:
            print "poll_docstore"
            self._reload()
            self._sleepish(self._reload_period)
            # handle pause
            while self._paused.get_value():
                time.sleep(self._poll_period.get_value())

    def _reload(self):
        """
        Reload canvases; if the list of canvases changed, update the multiselect
        """
        self._reload_canvases()
        # now update GUI elements (each call iterates over the whole
        # _canvases_draw queue)
        self._update_canvas_select()

    def _update_canvas_select(self):
        """
        Update the list of canvases in the canvas_select frame
        """
        all_canvas_titles = [wrap.canvas_wrapper.GetCanvas().GetTitle() \
                           for wrap in generate_fifo_queue(self._canvases_read)]

        select_box = self.window.get_frame("canvas_select", "list_box")
        select_box.RemoveAll()
        for i, title in enumerate(all_canvas_titles):
            select_box.AddEntry(title, i)
        for title in generate_fifo_queue(self._draw_titles):
            try:
                index = all_canvas_titles.index(title)
                select_box.Select(index)
            except ValueError: # item was no longer in the select - ignore it
                pass
        go_back_fifo_queue(self._draw_titles)
        # ROOT doesn't like redrawing the histogram titles properly - need to
        # force it
        self.window.get_frame("canvas_select", "list_box").Resize(200, 10)
        self.window.get_frame("canvas_select", "list_box").Resize(200, 200)

    def _get_canvas_selected(self):
        """
        Get the list of canvases selected in the canvas_select frame
        """
        select_box = self.window.get_frame("canvas_select", "list_box")
        selected_root = ROOT.TList()
        selected = []
        select_box.GetSelectedEntries(selected_root)
        while selected_root.Last() != None:
            selected.insert(0, copy.deepcopy(selected_root.Last().GetText()))
            selected_root.RemoveLast()
        return selected

    def _reload_canvases(self):
        """
        Reload canvases from the docstore
        """
        # get all of the collections (one for each reducer)
        try:
            collections = self._docstore.collection_names()
        except (DocumentStoreException, SocketError):
            print "Failed to get collection names from the docstore"
            return
        self._collections.set_value(collections)
        temp_canvases = []
        # update "last accessed" timestamp
        self.last = datetime.datetime.now()
        # get new canvases
        for collection_name in collections:
            temp_canvases += self._get_new_canvases(collection_name)
        # purge any duplicate canvases
        filtered_canvases = self._filter_canvases(temp_canvases)
        # build a new queue and reassign _canvases_read
        canvases_read = Queue.Queue()
        for canvas in filtered_canvases:
            canvases_read.put_nowait(canvas)
        self._canvases_read = canvases_read
        self._collections.set_value(collections)

    def _get_new_canvases(self, collection_name):
        """
        Get a list of new canvases from the collection
        """
        try:
            doc_list = self._docstore.get_since(collection_name, self.last)
        except DocumentStoreException:
            print "Failed to get documents from the docstore"
            return []
        if len(doc_list) == 0:
            return []
        tree = doc_list[-1]["doc"]
        image_data = ROOT.MAUS.ImageData()
        tree.SetBranchAddress("data", image_data)
        if tree.GetEntries() == 0:
            return []
        tree.GetEntry()
        return CanvasRewrapped.new_list_from_image(collection_name,
                                                   image_data.GetImage())

    def _filter_canvases(self, new_canvases):
        """
        Enforce that we only have one copy with a given (collection, title)

        Note that we want to keep the item with the highest input_time if there
        is a choice
        """
        combined_list = [wrap.deepcopy() \
                           for wrap in generate_fifo_queue(self._canvases_read)]
        combined_list += new_canvases
        combined_list.sort(key=lambda wrap: wrap.sort_key())

        if len(combined_list) == 0:
            return []
        filtered_list = []
        for i, item_2 in enumerate(combined_list[1:]):
            item_1 = combined_list[i]
            if item_1.collection == item_2.collection and \
               item_1.canvas_title() == item_2.canvas_title():
                pass
            else:
                filtered_list.append(item_1)
        filtered_list.append(combined_list[-1])
        return filtered_list
Пример #6
0
class PlotSetup():
    """
    GUI handler for defining data to be plotted
    """
    def __init__(self, main_window, parent, selected):
        """
        Initialise the window
        """
        self.main_window = main_window
        self.parent = parent
        self.selected = selected
        self.window = Window(ROOT.gClient.GetRoot(), # pylint: disable = E1101
                             parent,
                             self.share_dir+"plot_setup.json")
        for i, item in enumerate(self.selected):
            self.window.set_action("variable_type_"+str(i), "drop_down",
                                  "Selected(Int_t)", self.select_action)
            type_frame = self.window.get_frame("variable_type_"+str(i),
                                               "drop_down")
            type_frame.Select(item["variable_type"])
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)

    def okay_action(self):
        """
        Handle Okay button press; get the plot selection and push it to the
        main window
        """
        type_int = self.window.get_frame_dict("variable_type_0",
                                             "drop_down")["frame"].GetSelected()
        first_int = self.window.get_frame_dict("first_var_0",
                                            "drop_down")["frame"].GetSelected()
        plot_apertures = self.window.get_frame("plot_apertures",
                                               "check_button").IsOn()
        if type_int != 0 and first_int == 0:
            raise GuiError("Please select plot variable")
        self.main_window.plot_setup_options = [{
            "variable_type":type_int,
            "first_var":first_int,
            "plot_apertures":plot_apertures
        }]
        self.main_window.update_plot()
        self.window.close_window()
        self.main_window.plot_setup = None
    
    def cancel_action(self):
        """
        Handle Cancel button press; just close the window
        """
        self.window.close_window()
        self.main_window.plot_setup = None

    def select_action(self):
        """
        Dynamically change the list of items in first_var depending on
        variable_type
        """
        for i, item in enumerate(self.selected):
            type_select = self.window.get_frame_dict("variable_type_"+str(i),
                                                     "drop_down")["frame"]
            first_var_select = self.window.get_frame_dict("first_var_"+str(i),
                                                    "drop_down")["frame"]
            selected_type_int = type_select.GetSelected()
            selected_type_str = self.type_list[selected_type_int]
            var_list_name = self.type_variables_dict[selected_type_str][0]
            var_list = self.select_lists[var_list_name]
            first_var_select.RemoveAll()
            for i, entry in enumerate(var_list):
                first_var_select.AddEntry(entry, i)
            first_var_select.Select(item["first_var"])

    @staticmethod
    def get_variable_type(options):
        """
        Return a string corresponding to integer variable_type selected
        """
        var_type_int = options[0]["variable_type"]
        return PlotSetup.type_list[var_type_int]

    @staticmethod
    def get_first_var(options):
        """
        Return a string corresponding to integer first_var selected
        """
        var_type_string = PlotSetup.get_variable_type(options)
        select_list_key = PlotSetup.type_variables_dict[var_type_string][0]
        select_list = PlotSetup.select_lists[select_list_key]
        first_var_int = options[0]["first_var"]
        return select_list[first_var_int]

    my_hit_get_variables = [var for var in Hit.get_variables() if len(var)]

    type_list = [
        "<Select plot type>",
        "mean",
        "envelope",
        "RMS",
        "beta",
        "alpha",
        "gamma",
        "emittance",
        "dispersion",
        "dispersion_prime"
    ]

    select_lists = {
        "no_var":[
            "",
        ], "optics_var":[
            "<Select plot variable>",
            "x",
            "y",
            "transverse",
            "longitudinal"
        ], "physics_var":[
            "<Select plot variable>"
        ]+my_hit_get_variables, "kinematic_var":[
            "<Select plot variable>",
            "x",
            "y",
            "px",
            "py",
            "t",
            "energy"
        ], "disp_var":[
            "<Select plot variable>",
            "x",
            "y"
        ]
    }

    type_variables_dict = {
        "<Select plot type>":["no_var"],
        "mean":["physics_var"],
        "envelope":["kinematic_var"],
        "RMS":["kinematic_var"],
        "beta":["optics_var"],
        "alpha":["optics_var"],
        "gamma":["optics_var"],
        "emittance":["optics_var"],
        "dispersion":["disp_var"],
        "dispersion_prime":["disp_var"],
    }

    share_dir = ""
Пример #7
0
class PennSetup:
    """GUI window to setup a beam according to 4D Penn parameterisation"""
    def __init__(self, beam_setup, root_frame):
        """Initialise the window"""
        self.window = Window(
            ROOT.gClient.GetRoot(),  # pylint: disable = E1101
            root_frame,
            SHARE_DIR + "penn_setup.json")
        self.beam_setup = beam_setup
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)
        self.window.set_button_action("Get &B0", self.get_b0_action)
        self.window.set_button_action("&Constant Beta",
                                      self.constant_beta_action)

    def okay_action(self):
        """Handle a click on the Okay button, building the beam matrix"""
        ref = self.beam_setup.get_reference()
        matrix = maus_cpp.covariance_matrix.create_from_penn_parameters(
            mass=ref['mass'],
            momentum=ref['p'],
            emittance_t=self.window.get_text_entry("emittance_trans",
                                                   type(1.)),
            beta_t=self.window.get_text_entry("beta_trans", type(1.)),
            emittance_l=self.window.get_text_entry("emittance_l", type(1.)),
            beta_l=self.window.get_text_entry("beta_l", type(1.)),
            alpha_t=self.window.get_text_entry("alpha_trans", type(1.)),
            alpha_l=self.window.get_text_entry("alpha_l", type(1.)),
            charge=ref["charge"],
            # BUG - multiply by -1 to fix BUG in source code
            bz=-self.window.get_text_entry("B0", type(1.)) * common.units['T'],
            ltwiddle=self.window.get_text_entry("Lc", type(1.)),
            dispersion_x=self.window.get_text_entry("disp_x", type(1.)),
            dispersion_prime_x=self.window.get_text_entry(
                "disp_prime_x", type(1.)),
            dispersion_y=self.window.get_text_entry("disp_y", type(1.)),
            dispersion_prime_y=self.window.get_text_entry(
                "disp_prime_y", type(1.)))
        self.beam_setup.set_matrix(matrix)
        self.beam_setup.get_matrix()  # check for consistency
        self.window.close_window()
        self.beam_setup.matrix_select = None

    def cancel_action(self):
        """Handle a click on the Cancel button"""
        self.window.close_window()
        self.beam_setup.matrix_select = None

    def get_b0_action(self):
        """
        Handle a click on the Get B0 button, getting B magnitude at the location
        of the reference particle
        """
        ref = self.beam_setup.get_reference()
        b_vec = maus_cpp.field.get_field_value(ref['x'], ref['y'], ref['z'],
                                               ref['t'])
        _b0 = (b_vec[0]**2 + b_vec[1]**2 + b_vec[2]**2)**0.5
        self.window.set_text_entry('B0', _b0 * common.units['T'])

    def constant_beta_action(self):
        """
        Handle a click on the Constant Beta button, getting beta function that
        would be required to get a constant beta function, given other settings
        (reference particle, B0, Lc)
        """
        _b0 = abs(self.window.get_text_entry("B0", type(1.)) * 1e-3)  # kT
        if abs(_b0) < 1e-9:
            raise ValueError(
              "Constant beta requires B0 is not 0. "+\
              "Modify the B0 field and try again."
            )
        _lc = self.window.get_text_entry("Lc", type(1.))  # kT
        mom = self.beam_setup.get_reference()['p']  # MeV/c
        charge = abs(self.beam_setup.get_reference()['charge'])  # eplus
        beta = (1 + _lc**
                2)**0.5 * 2. * charge * mom / common.constants['c_light'] / _b0
        self.window.set_text_entry('alpha_trans', 0.0)
        self.window.set_text_entry('beta_trans', beta)  # beta in mm
Пример #8
0
class BeamSetup:
    """GUI window to setup a reference particle and beam ellipse"""
    def __init__(self, main_window, parent):
        """Initialise the window"""
        self.window = Window(
            ROOT.gClient.GetRoot(),  # pylint: disable = E1101
            parent,
            SHARE_DIR + "beam_setup.json")
        self.main_window = main_window
        self.parent = parent
        self.matrix_select = None
        self.matrix = None

        self.window.set_button_action("&Penn", self.penn_action)
        self.window.set_button_action("&Twiss", self.twiss_action)
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)

    def twiss_action(self):
        """Handle Twiss button - open a window to set up Twiss parameters"""
        if self.matrix_select == None:
            self.matrix_select = TwissSetup(self, self.parent)

    def penn_action(self):
        """Handle Penn button - open a window to set up Penn parameters"""
        if self.matrix_select == None:
            self.matrix_select = PennSetup(self, self.parent)

    def okay_action(self):
        """
        Handle Okay button - set the beam in lattice and then close the window
        """
        if self.matrix_select != None:
            print "Close the matrix_select window"
            return
        self.main_window.lattice.set_beam(self.get_reference(),
                                          self.get_matrix())
        self.main_window.lattice.run_lattice()
        self.main_window.update_plot()
        self.window.close_window()
        self.main_window.beam_setup = None

    def cancel_action(self):
        """
        Handle Cancel button - just close the window
        """
        if self.matrix_select != None:
            print "Close the matrix_select window"
            return
        self.window.close_window()
        self.main_window.beam_setup = None

    def set_matrix(self, matrix):
        """
        Set the beam matrix (updating GUI elements)
        - matrix a maus_cpp.covariance_matrix.CovarianceMatrix object
        """
        self.matrix = matrix
        for i, var_i in enumerate(["t", "e", "x", "px", "y", "py"]):
            for j, var_j in enumerate(["t", "e", "x", "px", "y", "py"]):
                value = matrix.get_element(i + 1, j + 1)
                self.window.set_text_entry(var_i + var_j, value)

    def get_matrix(self):
        """
        Get the beam matrix (reading from GUI elements)

        Returns a maus_cpp.covariance_matrix.CovarianceMatrix object
        """
        np_matrix = numpy.zeros([6, 6])
        for i, var_i in enumerate(["t", "e", "x", "px", "y", "py"]):
            for j, var_j in enumerate(["t", "e", "x", "px", "y", "py"]):
                value1 = self.window.get_text_entry(var_i + var_j, type(1.))
                value2 = self.window.get_text_entry(var_j + var_i, type(1.))
                delta = abs(value1 - value2)
                if delta > 1e-6:
                    err = "Matrix has non-symmetric element Cov("+\
                            var_i+","+var_j+") with delta "+str(delta)+"."
                    err += "This element must be symmetric - upper "+\
                             "diagonal should be equal to lower diagonal."
                    raise GuiError(err)
                else:
                    np_matrix[i, j] = value1
        evals = numpy.linalg.eigvals(np_matrix)
        if not numpy.all(evals > 0):
            evals = [i for i in evals]
            raise GuiError("Beam ellipse should be positive definite but "+\
                           " eigenvalues were not all positive: "+str(evals))
        return maus_cpp.covariance_matrix.create_from_matrix(np_matrix)

    def set_reference(self, ref_hit):
        """
        Get the reference particle (writing to GUI elements)
        - ref_hit an xboa.Hit.Hit object
        """
        for var in ["x", "y", "z", "px", "py", "pz", "pid"]:
            self.window.set_text_entry(var, ref_hit[var])

    def get_reference(self):
        """
        Get the reference particle (reading from GUI elements)

        Returns an xboa.Hit.Hit object
        """
        ref_dict = {}
        for var in ["x", "y", "z", "px", "py", "pz"]:
            var_dict = self.window.get_frame_dict(var, "named_text_entry")
            ref_dict[var] = float(var_dict["text_entry"].text_entry.GetText())
        pid_dict = self.window.get_frame_dict("pid", "named_text_entry")
        try:
            ref_dict["pid"] = int(pid_dict["text_entry"].text_entry.GetText())
            ref_dict["mass"] = common.pdg_pid_to_mass[abs(ref_dict["pid"])]
            ref_dict["charge"] = common.pdg_pid_to_charge[abs(ref_dict["pid"])]
        except KeyError:
            raise GuiError("Did not recognise reference particle pid")
        try:
            hit = Hit.new_from_dict(ref_dict, "energy")
        except Exception:
            raise GuiError("Failed to generate a reference particle")
        return hit
Пример #9
0
class PlotSetup():
    """
    GUI handler for defining data to be plotted
    """
    def __init__(self, main_window, parent, selected):
        """
        Initialise the window
        """
        self.main_window = main_window
        self.parent = parent
        self.selected = selected
        self.window = Window(
            ROOT.gClient.GetRoot(),  # pylint: disable = E1101
            parent,
            self.share_dir + "plot_setup.json")
        for i, item in enumerate(self.selected):
            self.window.set_action("variable_type_" + str(i), "drop_down",
                                   "Selected(Int_t)", self.select_action)
            type_frame = self.window.get_frame("variable_type_" + str(i),
                                               "drop_down")
            type_frame.Select(item["variable_type"])
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)

    def okay_action(self):
        """
        Handle Okay button press; get the plot selection and push it to the
        main window
        """
        type_int = self.window.get_frame_dict(
            "variable_type_0", "drop_down")["frame"].GetSelected()
        first_int = self.window.get_frame_dict(
            "first_var_0", "drop_down")["frame"].GetSelected()
        plot_apertures = self.window.get_frame("plot_apertures",
                                               "check_button").IsOn()
        if type_int != 0 and first_int == 0:
            raise GuiError("Please select plot variable")
        self.main_window.plot_setup_options = [{
            "variable_type": type_int,
            "first_var": first_int,
            "plot_apertures": plot_apertures
        }]
        self.main_window.update_plot()
        self.window.close_window()
        self.main_window.plot_setup = None

    def cancel_action(self):
        """
        Handle Cancel button press; just close the window
        """
        self.window.close_window()
        self.main_window.plot_setup = None

    def select_action(self):
        """
        Dynamically change the list of items in first_var depending on
        variable_type
        """
        for i, item in enumerate(self.selected):
            type_select = self.window.get_frame_dict("variable_type_" + str(i),
                                                     "drop_down")["frame"]
            first_var_select = self.window.get_frame_dict(
                "first_var_" + str(i), "drop_down")["frame"]
            selected_type_int = type_select.GetSelected()
            selected_type_str = self.type_list[selected_type_int]
            var_list_name = self.type_variables_dict[selected_type_str][0]
            var_list = self.select_lists[var_list_name]
            first_var_select.RemoveAll()
            for i, entry in enumerate(var_list):
                first_var_select.AddEntry(entry, i)
            first_var_select.Select(item["first_var"])

    @staticmethod
    def get_variable_type(options):
        """
        Return a string corresponding to integer variable_type selected
        """
        var_type_int = options[0]["variable_type"]
        return PlotSetup.type_list[var_type_int]

    @staticmethod
    def get_first_var(options):
        """
        Return a string corresponding to integer first_var selected
        """
        var_type_string = PlotSetup.get_variable_type(options)
        select_list_key = PlotSetup.type_variables_dict[var_type_string][0]
        select_list = PlotSetup.select_lists[select_list_key]
        first_var_int = options[0]["first_var"]
        return select_list[first_var_int]

    my_hit_get_variables = [var for var in Hit.get_variables() if len(var)]

    type_list = [
        "<Select plot type>", "mean", "envelope", "RMS", "beta", "alpha",
        "gamma", "emittance", "dispersion", "dispersion_prime"
    ]

    select_lists = {
        "no_var": [
            "",
        ],
        "optics_var":
        ["<Select plot variable>", "x", "y", "transverse", "longitudinal"],
        "physics_var": ["<Select plot variable>"] + my_hit_get_variables,
        "kinematic_var":
        ["<Select plot variable>", "x", "y", "px", "py", "t", "energy"],
        "disp_var": ["<Select plot variable>", "x", "y"]
    }

    type_variables_dict = {
        "<Select plot type>": ["no_var"],
        "mean": ["physics_var"],
        "envelope": ["kinematic_var"],
        "RMS": ["kinematic_var"],
        "beta": ["optics_var"],
        "alpha": ["optics_var"],
        "gamma": ["optics_var"],
        "emittance": ["optics_var"],
        "dispersion": ["disp_var"],
        "dispersion_prime": ["disp_var"],
    }

    share_dir = ""
Пример #10
0
class PennSetup:
    """GUI window to setup a beam according to 4D Penn parameterisation"""
    def __init__(self, beam_setup, root_frame):
        """Initialise the window"""
        self.window = Window(ROOT.gClient.GetRoot(), # pylint: disable = E1101
                             root_frame,
                             SHARE_DIR+"penn_setup.json")
        self.beam_setup = beam_setup
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)
        self.window.set_button_action("Get &B0", self.get_b0_action)
        self.window.set_button_action("&Constant Beta",
                                      self.constant_beta_action)

    def okay_action(self):
        """Handle a click on the Okay button, building the beam matrix"""
        ref = self.beam_setup.get_reference()
        matrix = maus_cpp.covariance_matrix.create_from_penn_parameters(
          mass = ref['mass'],
          momentum = ref['p'],
          emittance_t = self.window.get_text_entry("emittance_trans", type(1.)),
          beta_t = self.window.get_text_entry("beta_trans", type(1.)),
          emittance_l = self.window.get_text_entry("emittance_l", type(1.)),
          beta_l = self.window.get_text_entry("beta_l", type(1.)),
          alpha_t = self.window.get_text_entry("alpha_trans", type(1.)),
          alpha_l = self.window.get_text_entry("alpha_l", type(1.)),
          charge = ref["charge"],
          # BUG - multiply by -1 to fix BUG in source code
          bz = -self.window.get_text_entry("B0", type(1.))*common.units['T'],
          ltwiddle = self.window.get_text_entry("Lc", type(1.)),
          dispersion_x = self.window.get_text_entry("disp_x", type(1.)),
          dispersion_prime_x = self.window.get_text_entry("disp_prime_x",
                                                          type(1.)),
          dispersion_y = self.window.get_text_entry("disp_y", type(1.)),
          dispersion_prime_y = self.window.get_text_entry("disp_prime_y",
                                                          type(1.))
        )
        self.beam_setup.set_matrix(matrix)
        self.beam_setup.get_matrix() # check for consistency
        self.window.close_window()
        self.beam_setup.matrix_select = None

    def cancel_action(self):
        """Handle a click on the Cancel button"""
        self.window.close_window()
        self.beam_setup.matrix_select = None

    def get_b0_action(self):
        """
        Handle a click on the Get B0 button, getting B magnitude at the location
        of the reference particle
        """
        ref = self.beam_setup.get_reference()
        b_vec = maus_cpp.field.get_field_value(ref['x'], ref['y'], ref['z'],
                                               ref['t'])
        _b0 = (b_vec[0]**2+b_vec[1]**2+b_vec[2]**2)**0.5
        self.window.set_text_entry('B0', _b0*common.units['T'])

    def constant_beta_action(self):
        """
        Handle a click on the Constant Beta button, getting beta function that
        would be required to get a constant beta function, given other settings
        (reference particle, B0, Lc)
        """
        _b0 = abs(self.window.get_text_entry("B0", type(1.))*1e-3) # kT
        if abs(_b0) < 1e-9:
            raise ValueError(
              "Constant beta requires B0 is not 0. "+\
              "Modify the B0 field and try again."
            )
        _lc = self.window.get_text_entry("Lc", type(1.)) # kT
        mom = self.beam_setup.get_reference()['p'] # MeV/c
        charge = abs(self.beam_setup.get_reference()['charge']) # eplus
        beta = (1+_lc**2)**0.5 * 2.*charge*mom/common.constants['c_light']/_b0
        self.window.set_text_entry('alpha_trans', 0.0)
        self.window.set_text_entry('beta_trans', beta) # beta in mm
Пример #11
0
class BeamSetup:
    """GUI window to setup a reference particle and beam ellipse"""
    def __init__(self, main_window, parent):
        """Initialise the window"""
        self.window = Window(ROOT.gClient.GetRoot(), # pylint: disable = E1101
                             parent,
                             SHARE_DIR+"beam_setup.json")
        self.main_window = main_window
        self.parent = parent
        self.matrix_select = None
        self.matrix = None

        self.window.set_button_action("&Penn", self.penn_action)
        self.window.set_button_action("&Twiss", self.twiss_action)
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)

    def twiss_action(self):
        """Handle Twiss button - open a window to set up Twiss parameters"""
        if self.matrix_select == None:
            self.matrix_select = TwissSetup(self, self.parent)

    def penn_action(self):
        """Handle Penn button - open a window to set up Penn parameters"""
        if self.matrix_select == None:
            self.matrix_select = PennSetup(self, self.parent)

    def okay_action(self):
        """
        Handle Okay button - set the beam in lattice and then close the window
        """
        if self.matrix_select != None:
            print "Close the matrix_select window"
            return
        self.main_window.lattice.set_beam(self.get_reference(),
                                          self.get_matrix())   
        self.main_window.lattice.run_lattice()
        self.main_window.update_plot()
        self.window.close_window()
        self.main_window.beam_setup = None
    
    def cancel_action(self):
        """
        Handle Cancel button - just close the window
        """
        if self.matrix_select != None:
            print "Close the matrix_select window"
            return
        self.window.close_window()
        self.main_window.beam_setup = None

    def set_matrix(self, matrix):
        """
        Set the beam matrix (updating GUI elements)
        - matrix a maus_cpp.covariance_matrix.CovarianceMatrix object
        """
        self.matrix = matrix
        for i, var_i in enumerate(["t", "e", "x", "px", "y", "py"]):
            for j, var_j in enumerate(["t", "e", "x", "px", "y", "py"]):
                value = matrix.get_element(i+1, j+1)
                self.window.set_text_entry(var_i+var_j, value)

    def get_matrix(self):
        """
        Get the beam matrix (reading from GUI elements)

        Returns a maus_cpp.covariance_matrix.CovarianceMatrix object
        """
        np_matrix = numpy.zeros([6, 6])
        for i, var_i in enumerate(["t", "e", "x", "px", "y", "py"]):
            for j, var_j in enumerate(["t", "e", "x", "px", "y", "py"]):
                value1 = self.window.get_text_entry(var_i+var_j, type(1.))
                value2 = self.window.get_text_entry(var_j+var_i, type(1.))
                delta = abs(value1 - value2)
                if delta > 1e-6:
                    err = "Matrix has non-symmetric element Cov("+\
                            var_i+","+var_j+") with delta "+str(delta)+"."
                    err += "This element must be symmetric - upper "+\
                             "diagonal should be equal to lower diagonal."
                    raise GuiError(err)
                else:
                    np_matrix[i, j] = value1
        evals = numpy.linalg.eigvals(np_matrix)
        if not numpy.all(evals > 0):
            evals = [i for i in evals]
            raise GuiError("Beam ellipse should be positive definite but "+\
                           " eigenvalues were not all positive: "+str(evals))
        return maus_cpp.covariance_matrix.create_from_matrix(np_matrix)

    def set_reference(self, ref_hit):
        """
        Get the reference particle (writing to GUI elements)
        - ref_hit an xboa.Hit.Hit object
        """
        for var in ["x", "y", "z", "px", "py", "pz", "pid"]:
            self.window.set_text_entry(var, ref_hit[var])

    def get_reference(self):
        """
        Get the reference particle (reading from GUI elements)

        Returns an xboa.Hit.Hit object
        """
        ref_dict = {}
        for var in ["x", "y", "z", "px", "py", "pz"]: 
            var_dict = self.window.get_frame_dict(var, "named_text_entry")
            ref_dict[var] = float(var_dict["text_entry"].text_entry.GetText())
        pid_dict = self.window.get_frame_dict("pid", "named_text_entry")
        try:
            ref_dict["pid"] = int(pid_dict["text_entry"].text_entry.GetText())
            ref_dict["mass"] = common.pdg_pid_to_mass[abs(ref_dict["pid"])]
            ref_dict["charge"] = common.pdg_pid_to_charge[abs(ref_dict["pid"])]
        except KeyError:
            raise GuiError("Did not recognise reference particle pid")
        try:
            hit = Hit.new_from_dict(ref_dict, "energy")
        except Exception:
            raise GuiError("Failed to generate a reference particle")
        return hit
Пример #12
0
class OnlineGui():  # pylint: disable = R0902
    """Main frame for the online gui window"""
    def __init__(self, datacards):
        """Initialise the main window"""
        self.last = datetime.datetime.now()
        self.window = Window(ROOT.gClient.GetRoot(), ROOT.gClient.GetRoot(),
                             SHARE_DIR + "online_gui.json")
        self._set_actions()
        self._docstore = None
        self._datacards = datacards
        self._draw_titles = Queue.Queue()
        for title in datacards["online_gui_default_canvases"]:
            self._draw_titles.put(title)
        self._canvases_read = Queue.Queue()
        self._canvases_draw = ThreadedValue([])
        self._redraw_target = None
        self._paused = ThreadedValue(False)
        rotate_tmp = self.window.get_text_entry("Rotate period", type(1.))
        self._rotate_period = ThreadedValue(rotate_tmp)
        reload_tmp = self.window.get_text_entry("Reload period", type(1.))
        self._reload_period = ThreadedValue(reload_tmp)
        self._poll_period = ThreadedValue(0.1)
        self._collections = ThreadedValue([])
        self._init_docstore()
        self._start_polling()

    def _set_actions(self):
        """
        Set up actions (responses to signals)
        """
        self.window.set_button_action("&Pause", self.pause_action)
        self.window.set_button_action("&Redraw", self.redraw_action)
        self.window.set_button_action("Re&load", self.reload_action)
        self.window.set_button_action("Re&connect", self.reconnect_action)
        self.window.set_button_action("E&xit", self.exit_action)
        self.window.set_button_action("&<", self.back_action)
        self.window.set_button_action("&>", self.forward_action)
        self.window.get_frame("canvas_options", "vertical_frame").\
                                                                Resize(200, 500)
        self.window.get_frame("canvas_select", "list_box").Resize(200, 200)
        self.window.set_action("canvas_select", "list_box",
                               "SelectionChanged()", self.canvas_select_action)
        self.window.set_action("Rotate period", "named_text_entry",
                               "ReturnPressed()", self.rotate_period_action)
        self.window.set_action("Reload period", "named_text_entry",
                               "ReturnPressed()", self.reload_period_action)

    def _init_docstore(self):
        """
        Initialise the docstore
        """
        datacards = self._datacards
        self._docstore = RootDocumentStore(
            datacards["root_document_store_timeout"],
            datacards["root_document_store_poll_time"], 10, 10)
        self._docstore.connect({
            "host": datacards["online_gui_host"],
            "port": datacards["online_gui_port"]
        })

    def _start_polling(self):
        """
        Start polling the database and rotating the canvases
        """
        my_thread = threading.Thread(target=self._poll_canvas_draw, args=())
        my_thread.daemon = True
        my_thread.start()
        my_thread = threading.Thread(target=self._poll_docstore, args=())
        my_thread.daemon = True
        my_thread.start()

    def pause_action(self):
        """Handle a Pause/Unpause button press"""
        pause_button = self.window.get_frame("&Pause", "button")
        if self._paused.get_value():
            self._paused.set_value(False)
            pause_button.SetText("&Pause")
        else:
            self._paused.set_value(True)
            pause_button.SetText("&Play")

    def redraw_action(self):
        """
        Handle a Redraw button press

        Redraws self._redraw_target
        """
        if self._redraw_target == None:
            raise GuiError("No canvas to draw")
        self._do_redraw(self._redraw_target)

    def reload_action(self):
        """Handle a Reload button press"""
        self._reload()
        if self._redraw_target != None:
            go_back_fifo_queue(self._draw_titles)
            self._draw_next()

    def reconnect_action(self):
        """Handle a Reconnect button press"""
        self._docstore.disconnect()
        self._init_docstore()

    def exit_action(self):
        """Handle a Exit button press"""
        self._paused.set_value(True)  # tell polling to wait
        time.sleep(self._poll_period.get_value())  # thinkapow
        self.window.close_window()

    def reload_period_action(self):
        """Handle a change to the reload period"""
        tmp_reload_period = self.window.get_text_entry("Reload period",
                                                       type(1.))
        tmp_reload_period = float(tmp_reload_period)
        self._reload_period.set_value(tmp_reload_period)

    def rotate_period_action(self):
        """Handle a change to the rotate period"""
        tmp_rotate_period = self.window.get_text_entry("Rotate period",
                                                       type(1.))
        tmp_rotate_period = float(tmp_rotate_period)
        self._rotate_period.set_value(tmp_rotate_period)

    def forward_action(self):
        """Handle a click on the > forward button"""
        self._draw_next()

    def back_action(self):
        """Handle a click on the < back button"""
        go_back_fifo_queue(self._draw_titles)
        go_back_fifo_queue(self._draw_titles)
        self._draw_next()

    def canvas_select_action(self):
        """Handle an update to the selected canvases"""
        selected = self._get_canvas_selected()
        new_draw_titles = Queue.Queue()
        for item in selected:
            new_draw_titles.put(item)
        self._draw_titles = new_draw_titles

    def _poll_canvas_draw(self):
        """
        Iterate over the canvas_draw queue, updating the GUI as appropriate

        * Update the queue from the canvas_read queue
        * Refresh the combo box with items in the canvases_draw queue
        """
        # poll_number = 0
        while True:
            # poll_number += 1
            self._draw_next()
            self._sleepish(self._rotate_period)
            # handle pause
            while self._paused.get_value():
                time.sleep(self._poll_period.get_value())

    def _draw_next(self):
        """
        Draw the next canvas in the selected queue
        """
        canvases_draw = self._filter_canvases(self._canvases_draw.get_value())
        if len(canvases_draw) > 0:
            try:
                title = self._draw_titles.get_nowait()
                new_title = copy.deepcopy(title)
                self._draw_titles.put(title)
                for wrap in canvases_draw:
                    if wrap.canvas_title() == new_title:
                        self._do_redraw(wrap)
                        break
            except Queue.Empty:
                pass  # no canvas selected
            except ValueError:
                pass
        self._canvases_draw.set_value(canvases_draw)

    def _sleepish(self, sleep_length):
        """
        Sleep - but keep polling for changes to the sleep_length
        * sleep_length - threaded value containing numeric sleep time
        """
        time_slept = 0.
        poll_period = self._poll_period.get_value()
        while sleep_length.get_value() - time_slept > poll_period:
            time_slept += self._poll_period.get_value()
            time.sleep(poll_period)

    def _update_canvases_draw(self):
        """update canvases"""
        draw = [wrap.deepcopy() \
                           for wrap in generate_fifo_queue(self._canvases_read)]
        draw.sort(key=lambda wrap: wrap.sort_key())
        return draw

    def _do_redraw(self, wrap):
        """
        Draw the canvas wrap in the main_canvas
        """
        # wrap_copy is for ROOT
        wrap_copy = wrap.deepcopy()
        embedded_canvas = self.window.get_frame("main_canvas", "canvas")
        # this eats the wrap_copy TCanvas
        wrap_copy.canvas_wrapper.EmbedCanvas(embedded_canvas)
        embedded_canvas.GetCanvas().Update()
        # this _redraw_target is for the main thread
        self._redraw_target = wrap.deepcopy()

    def _poll_docstore(self):
        """
        Poll the doc store for documents
        """
        while True:
            print "poll_docstore"
            self._reload()
            self._sleepish(self._reload_period)
            # handle pause
            while self._paused.get_value():
                time.sleep(self._poll_period.get_value())

    def _reload(self):
        """
        Reload canvases; if the list of canvases changed, update the multiselect
        """
        self._reload_canvases()
        # now update GUI elements (each call iterates over the whole
        # _canvases_draw queue)
        self._update_canvas_select()

    def _update_canvas_select(self):
        """
        Update the list of canvases in the canvas_select frame
        """
        all_canvas_titles = [wrap.canvas_wrapper.GetCanvas().GetTitle() \
                           for wrap in generate_fifo_queue(self._canvases_read)]

        select_box = self.window.get_frame("canvas_select", "list_box")
        select_box.RemoveAll()
        for i, title in enumerate(all_canvas_titles):
            select_box.AddEntry(title, i)
        for title in generate_fifo_queue(self._draw_titles):
            try:
                index = all_canvas_titles.index(title)
                select_box.Select(index)
            except ValueError:  # item was no longer in the select - ignore it
                pass
        go_back_fifo_queue(self._draw_titles)
        # ROOT doesn't like redrawing the histogram titles properly - need to
        # force it
        self.window.get_frame("canvas_select", "list_box").Resize(200, 10)
        self.window.get_frame("canvas_select", "list_box").Resize(200, 200)

    def _get_canvas_selected(self):
        """
        Get the list of canvases selected in the canvas_select frame
        """
        select_box = self.window.get_frame("canvas_select", "list_box")
        selected_root = ROOT.TList()
        selected = []
        select_box.GetSelectedEntries(selected_root)
        while selected_root.Last() != None:
            selected.insert(0, copy.deepcopy(selected_root.Last().GetText()))
            selected_root.RemoveLast()
        return selected

    def _reload_canvases(self):
        """
        Reload canvases from the docstore
        """
        # get all of the collections (one for each reducer)
        try:
            collections = self._docstore.collection_names()
        except (DocumentStoreException, SocketError):
            print "Failed to get collection names from the docstore"
            return
        self._collections.set_value(collections)
        temp_canvases = []
        # update "last accessed" timestamp
        self.last = datetime.datetime.now()
        # get new canvases
        for collection_name in collections:
            temp_canvases += self._get_new_canvases(collection_name)
        # purge any duplicate canvases
        filtered_canvases = self._filter_canvases(temp_canvases)
        # build a new queue and reassign _canvases_read
        canvases_read = Queue.Queue()
        for canvas in filtered_canvases:
            canvases_read.put_nowait(canvas)
        self._canvases_read = canvases_read
        self._collections.set_value(collections)

    def _get_new_canvases(self, collection_name):
        """
        Get a list of new canvases from the collection
        """
        try:
            doc_list = self._docstore.get_since(collection_name, self.last)
        except DocumentStoreException:
            print "Failed to get documents from the docstore"
            return []
        if len(doc_list) == 0:
            return []
        tree = doc_list[-1]["doc"]
        image_data = ROOT.MAUS.ImageData()
        tree.SetBranchAddress("data", image_data)
        if tree.GetEntries() == 0:
            return []
        tree.GetEntry()
        return CanvasRewrapped.new_list_from_image(collection_name,
                                                   image_data.GetImage())

    def _filter_canvases(self, new_canvases):
        """
        Enforce that we only have one copy with a given (collection, title)

        Note that we want to keep the item with the highest input_time if there
        is a choice
        """
        combined_list = [wrap.deepcopy() \
                           for wrap in generate_fifo_queue(self._canvases_read)]
        combined_list += new_canvases
        combined_list.sort(key=lambda wrap: wrap.sort_key())

        if len(combined_list) == 0:
            return []
        filtered_list = []
        for i, item_2 in enumerate(combined_list[1:]):
            item_1 = combined_list[i]
            if item_1.collection == item_2.collection and \
               item_1.canvas_title() == item_2.canvas_title():
                pass
            else:
                filtered_list.append(item_1)
        filtered_list.append(combined_list[-1])
        return filtered_list
Пример #13
0
class MagnetSetup:
    """
    GUI handler for field setup
    """
    def __init__(self, main_window, parent):
        """Initialise the GUI window"""
        self.main_window = main_window
        self.parent = parent
        manipulators = {"magnet_setup_action": self.magnet_setup_action}
        self.window = Window(
            ROOT.gClient.GetRoot(),  # pylint: disable = E1101
            parent,
            data_file=self.share_dir + "magnet_setup.json",
            manipulator_dict=manipulators)
        self.window.set_button_action("&Okay", self.okay_action)
        self.window.set_button_action("&Cancel", self.cancel_action)

    def okay_action(self):
        """Handle okay button press"""
        field_list_out = []
        field_frame = self.window.get_frame_dict("field_list",
                                                 "vertical_frame")
        for frame in field_frame["children"]:
            try:
                name = frame["name"]
                scale_factor = float(frame["text_entry"].text_entry.GetText())
                field_list_out.append({
                    "field_name": name,
                    "scale_factor": scale_factor
                })
            except KeyError:
                pass
        self.main_window.lattice.set_fields(field_list_out)
        self.main_window.lattice.run_lattice()
        self.main_window.update_plot()
        self.window.close_window()
        self.main_window.magnet_setup = None

    def cancel_action(self):
        """Handle cancel button press"""
        self.window.close_window()
        self.main_window.magnet_setup = None

    def magnet_setup_action(self, item):  # pylint: disable=W0613
        """Iterate over magnets and add a line per magnet (special)"""
        field_list = self.main_window.lattice.get_field_list()
        geometry_list = {
            "type": "vertical_frame",
            "layout": "close",
            "name": "field_list",
            "children": []
        }
        if len(field_list) == 0:
            geometry_list["children"].append({
                "type": "label",
                "name": "No fields in lattice",
                "label_length": 30
            })
            return geometry_list
        for field in field_list:
            geometry_list["children"].append({
                "type":
                "named_text_entry",
                "default_text":
                str(field["scale_factor"]),
                "label_length":
                10,
                "name":
                str(field["field_name"]),
                "tool_tip":
                "Longitudinal beta function"
            })
        return geometry_list

    share_dir = ""