InteractiveStreamingExample/python/streamDet_JanaSub.py

0001 #!/usr/bin/python3.6
0002
0003 # Copyright 2020, Jefferson Science Associates, LLC.
0004 # Subject to the terms in the LICENSE file found in the top-level directory.
0005
0006 import zmq, struct, pickle, random
0007 import numpy as np
0008 import matplotlib.pyplot as plt
0009
0010 # define the subscriber and publisher ports
0011 subPort = 5557
0012 pubPort = 5558
0013 # define global variables
0014 numChans = 80
0015
0016 # configure the subscriber
0017 subContext = zmq.Context()
0018 subscriber = subContext.socket(zmq.SUB)
0019 subscriber.setsockopt(zmq.SUBSCRIBE, b'')
0020 subscriber.connect('tcp://127.0.0.1:%d' % subPort)
0021 print('\nSubscribing to JANA ZMQ Messages on socket tcp://127.0.0.1:%d\n' % subPort)
0022
0023 # # configure the publisher
0024 # pubContext = zmq.Context()
0025 # publisher  = pubContext.socket(zmq.PUB)
0026 # publisher.bind('tcp://127.0.0.1:%d' % pubPort)
0027 # print('\nPublishing on tcp://127.0.0.1:%d\n'  % pubPort)
0028
0029
0030 class IndraMessage:
0031     """Class to handle decoding of ZMQ INDRA Messages"""
0032
0033     def __init__(self, zmqMessage):
0034         # def __init__(self, zmqMessage, zmqPublisher):
0035         # define the data dictionary
0036         self.dataDict = {}
0037         # size up the zmq message
0038         self.messageSize = len(zmqMessage)
0039         # indra message header is a fixed 56 bytes
0040         self.payloadBytes = self.messageSize - 56
0041         # unpack the message according to the indra message protocol
0042         self.message = struct.unpack('IIIIIIIQqq%ds' % self.payloadBytes, zmqMessage)
0043         # decode the indra message members
0044         self.sourceId = self.message[0]
0045         self.totalBytes = self.message[1]
0046         self.payloadBytes = self.message[2]
0047         self.compressedBytes = self.message[3]
0048         self.magic = self.message[4]
0049         self.formatVersion = self.message[5]
0050         self.flags = self.message[6]
0051         self.recordCounter = self.message[7]
0052         self.timeStampSec = self.message[8]
0053         self.timeStampNanoSec = self.message[9]
0054         self.payload = self.message[10]
0055         # print message info
0056         print('INDRA Message recieved -> event = %d, size = %d bytes' % (self.recordCounter, self.messageSize))
0057         # adc samples are uints of length 4 plus 1 space delimiter, convert to array of ints
0058         self.adcSamplesStr = np.frombuffer(self.payload, dtype='S5')
0059         self.adcSamples = np.reshape(self.adcSamplesStr.astype(np.int), (1024, 80))
0060         # define data dictionary keys and data types
0061         for chan in range(1, numChans + 1):
0062             self.dataDict.update({'adcSamplesChan_%s' % chan: np.array([])})
0063             self.dataDict.update(
0064                 {'tdcSamplesChan_%s' % chan: np.arange((self.recordCounter - 1) * 1024, self.recordCounter * 1024)})
0065         # enumerate the samples object and populate the data dictionary
0066         for index, sample in np.ndenumerate(self.adcSamples):
0067             self.dataDict['adcSamplesChan_%s' % str(index[1] + 1)] = \
0068                 np.append(self.dataDict['adcSamplesChan_%s' % str(index[1] + 1)], sample)
0069         # serialize the data dictionary via pickle and publish it
0070         # self.eventDataDict = pickle.dumps(self.dataDict)
0071         # zmqPublisher.send_pyobj(self.eventDataDict)
0072         # remove event data after being published
0073         # for chan in range(1, numChans + 1) :
0074         #     self.dataDict.pop('adcSamplesChan_%s' % str(chan), None)
0075         #     self.dataDict.pop('tdcSamplesChan_%s' % str(chan), None)
0076
0077
0078 class MonitoringFigure:
0079     """Class to create figure and subplots based on user input"""
0080
0081     # define figure and plot attributes
0082     def __init__(self, rows, cols):
0083         self.numRows = rows
0084         self.numCols = cols
0085         self.fig, self.axs = plt.subplots(self.numRows, self.numCols)
0086         self.fig.set_size_inches(18.5, 10.5, forward=True)
0087
0088     def get_num_rows(self):
0089         return self.numRows
0090
0091     def get_num_cols(self):
0092         return self.numCols
0093
0094     def get_fig_obj(self):
0095         return self.fig
0096
0097     def get_axs_obj(self):
0098         return self.axs
0099
0100
0101 class MonitoringPlots:
0102     """Class to plot streaming ADC vs. TDC data"""
0103
0104     # instance variables unique to each instance
0105     def __init__(self, dataDict, thresh, rows, cols, fig, axs):
0106         # figure parameters
0107         self.numRows = rows
0108         self.numCols = cols
0109         self.fig = fig
0110         self.axs = axs
0111         # lists for colors and markers
0112         self.cl = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple',
0113                    'tab:brown', 'tab:pink', 'tab:gray', 'tab:olive', 'tab:cyan']
0114         self.ml = ['o', '^', 's', 'p', 'P', '*', 'X', 'd']
0115         # random channel list, ascending and non-repeating
0116         self.rcl = random.sample(range(1, numChans + 1), self.numRows * self.numCols)
0117         self.rcl.sort()
0118         # event number list, hit threshold (adc channels)
0119         self.enl = []
0120         self.hitThresh = thresh
0121         self.ic = 0
0122         for row in range(self.numRows):
0123             for column in range(self.numCols):
0124                 self.axs[row, column].cla()
0125                 self.axs[row, column].plot(dataDict['tdcSamplesChan_%d' % self.rcl[self.ic]],
0126                                            dataDict['adcSamplesChan_%d' % self.rcl[self.ic]],
0127                                            color=self.cl[self.ic % len(self.cl)],
0128                                            marker=self.ml[self.ic % len(self.ml)],
0129                                            ls='', label='Channel %d' % self.rcl[self.ic])
0130                 hitLoc = np.where(dataDict['adcSamplesChan_%d' % self.rcl[self.ic]] > 100)[0] + \
0131                          np.min(dataDict['tdcSamplesChan_%d' % self.rcl[self.ic]])
0132                 if len(hitLoc) != 0: self.axs[row, column].set_xlim(np.min(hitLoc) - 10, np.max(hitLoc) + 20)
0133                 self.axs[row, column].set_ylim(0, 1024)
0134                 if column == 0:
0135                     self.axs[row, column].set_ylabel('ADC Value')
0136                 if row == self.numRows - 1:
0137                     self.axs[row, column].set_xlabel('TDC Sample Number')
0138                 self.axs[row, column].legend(loc='best', markerscale=0, handletextpad=0, handlelength=0)
0139                 self.ic += 1
0140         plt.tight_layout()
0141         # plt.savefig('plots/event_%d.png' % ec)
0142         plt.pause(0.05)
0143         # remove event data after being published
0144         for chan in range(1, numChans + 1):
0145             dataDict.pop('adcSamplesChan_%s' % str(chan), None)
0146             dataDict.pop('tdcSamplesChan_%s' % str(chan), None)
0147
0148
0149 # instantiate the monitoring figure class
0150 monFig = MonitoringFigure(2, 2)
0151 # receive zmq messages from jana publisher
0152 while True:
0153     # receive zmq packets from jana publisher
0154     janaMessage = subscriber.recv()
0155     # instantiate the indra message class
0156     jevent = IndraMessage(janaMessage)
0157     # jevent = IndraMessage(janaMessage, publisher)
0158     # instantiate the monitoring plot class
0159     monPlots = MonitoringPlots(jevent.dataDict, 100,
0160                                monFig.get_num_rows(), monFig.get_num_cols(),
0161                                monFig.get_fig_obj(), monFig.get_axs_obj())
0162
0163 # We never get here but clean up anyhow
0164 # subscriber.close()
0165 # subContext.term()
0166 # publisher.close()
0167 # pubContext.term()