Пример #1
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
Пример #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 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
Пример #4
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
Пример #5
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
Пример #6
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
Пример #7
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
Пример #8
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