GnuRadio/gr-jay/python/peak_manager.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 
# Copyright 2019 Jay Arrowfield.
# 
# This is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3, or (at your option)
# any later version.
# 
# This software is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this software; see the file COPYING.  If not, write to
# the Free Software Foundation, Inc., 51 Franklin Street,
# Boston, MA 02110-1301, USA.
# 

import numpy
import pmt
from gnuradio import gr

class peak_manager(gr.basic_block):
	"""
	docstring for block peak_manager
	"""
	def __init__(self, numPeaksMax):
		gr.basic_block.__init__(self,
								name="peak_manager",
								in_sig=None,
								out_sig=None)

		self.numPeaksMax = numPeaksMax
		self.peak_dict = {}
		self.ports_avail = []
		self.ports_used = []

		# Register in / out message ports
		self.message_port_register_in(pmt.intern("control"))

		if self.numPeaksMax < 2:
			out_port_name = "status"
			self.ports_avail.append(out_port_name)
			self.message_port_register_out(pmt.intern(out_port_name))
		else:
			for n in range(0, self.numPeaksMax):
				out_port_name = "status" + str(n)
				self.ports_avail.append(out_port_name)
				self.message_port_register_out(pmt.intern(out_port_name))

		# Register in message handler
		self.set_msg_handler(pmt.intern("control"), self.msg_handler)


	def peak_add(self, key, peak):
		if len(self.ports_avail) > 0:
			port_name = self.ports_avail.pop()
			peak['port_name'] = port_name
			self.peak_dict[key] = peak
			return peak
		return None

	def peak_del(self, key):
		if key in self.peak_dict:
			port_name = self.peak_dict[key]['port_name']
			self.ports_avail.append(port_name)
			self.peak_dict.pop(key, None)
			return True
		return False

	def peak_get(self, key):
		result = None
		if key in self.peak_dict:
			result = self.peak_dict[key]
		return result

	def msg_handler(self, msg):
		if pmt.is_dict(msg):
			items = pmt.dict_items(msg)

			peakMsg = {}
			for n in range(0, pmt.length(items)):
				pair = pmt.nth(n, items)
				key = pmt.car(pair)
				val = pmt.cdr(pair)

				key_str = pmt.symbol_to_string(key)

				if pmt.is_integer(val):
					value = pmt.to_long(val)
					peakMsg[key_str] = value
				elif pmt.is_real(val):
					value = pmt.to_float(val)
					peakMsg[key_str] = value
				elif pmt.is_bool(val):
					value = pmt.to_bool(val)
					peakMsg[key_str] = value

			id = peakMsg['id']
			freq_MHz = peakMsg['freq_abs_Hz']*1e-6

			peak_key = 'peak' + str(id)
			peakEntry = self.peak_get(peak_key)

			if peakMsg['state'] == 1:
				if peakEntry is not None:
					msg_dict = pmt.make_dict()
					msg_dict = pmt.dict_add(msg_dict, pmt.intern("freq_rel_Hz"), pmt.from_float(peakMsg['freq_rel_Hz']))
					msg_dict = pmt.dict_add(msg_dict, pmt.intern("freq_abs_Hz"), pmt.from_float(peakMsg['freq_abs_Hz']))
					msg_dict = pmt.dict_add(msg_dict, pmt.intern("gain"), pmt.from_float(1.0))
					portname = peakEntry['port_name']
					self.message_port_pub(pmt.intern(portname), msg_dict)
					print ("Updated peak #{} at {:0.6f} MHz on port {}".format(id, freq_MHz, portname))
				else:
					peakEntry = self.peak_add(peak_key, peakMsg)
					if peakEntry is not None:
						msg_dict = pmt.make_dict()
						msg_dict = pmt.dict_add(msg_dict, pmt.intern("freq_rel_Hz"), pmt.from_float(peakMsg['freq_rel_Hz']))
						msg_dict = pmt.dict_add(msg_dict, pmt.intern("freq_abs_Hz"), pmt.from_float(peakMsg['freq_abs_Hz']))
						msg_dict = pmt.dict_add(msg_dict, pmt.intern("gain"), pmt.from_float(1.0))
						portname = peakEntry['port_name']
						self.message_port_pub(pmt.intern(portname), msg_dict)
						print ("New peak #{} at {:0.6f} MHz on port {}".format(id, freq_MHz, portname))
			else:
				if peakEntry is not None:
					pair = pmt.cons(pmt.intern("gain"), pmt.from_float(0.0))
					portname = peakEntry['port_name']
					self.message_port_pub(pmt.intern(portname), pair)
					if self.peak_del(peak_key):
						print ("Lost peak #{} at {:0.2f} MHz on port {}".format(id, freq_MHz, portname))


		else:
			print(pmt.symbol_to_string(msg))

	def forecast(self, noutput_items, ninput_items_required):
		#setup size of input_items[i] for work call
		for i in range(len(ninput_items_required)):
			ninput_items_required[i] = noutput_items

	def general_work(self, input_items, output_items):
		output_items[0][:] = input_items[0]
		consume(0, len(input_items[0]))
		#self.consume_each(len(input_items[0]))
		return len(output_items[0])