AM/AM poly DPD python block for GNU Radio

dpd.py

"""
Embedded Python Blocks:

Each time this file is saved, GRC will instantiate the first class it finds
to get ports and parameters of your block. The arguments to __init__  will
be the parameters. All of them are required to have default values!
"""

import numpy as np
from gnuradio import gr

class blk(gr.sync_block):
    """Embedded Python Block - a simple quadratic polynomial DPD model.\nout = a1*in + a2*in*abs(in) + a3*in*abs(in^2)"""

    def __init__(self, a1=1.0, a2=0.0, a3=0.0):
        """arguments to this function show up as parameters in GRC"""
        gr.sync_block.__init__(
            self,
            name='AM/AM poly DPD',   # will show up in GRC
            in_sig=[np.complex64],
            out_sig=[np.complex64]
        )

        self.a1 = a1
        self.a2 = a2
        self.a3 = a3

    def work(self, input_items, output_items):
        """apply predistortion to samples"""
        output_items[0][:] = self.a1 * input_items[0] + self.a2 * input_items[0] * np.absolute(input_items[0]) + self.a3 * input_items[0] * np.absolute(np.power(input_items[0], 2.0))
        return len(output_items[0])