Saleh model GNU Radio python block

saleh.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):  # other base classes are basic_block, decim_block, interp_block
    """Saleh predistortion model"""

    def __init__(self, aa=1.0, ba=0.0, ap=0.0, bp=0.0):  # only default arguments here
        """arguments to this function show up as parameters in GRC"""
        gr.sync_block.__init__(
            self,
            name='Saleh Model',   # will show up in GRC
            in_sig=[np.complex64],
            out_sig=[np.complex64]
        )
        # if an attribute with the same name as a parameter is found,
        # a callback is registered (properties work, too).
        self.aa = aa
        self.ba = ba
        self.ap = ap
        self.bp = bp

    def work(self, input_items, output_items):
        """Apply Saleh model to complex baseband"""
        out_mag=(self.aa*np.absolute(input_items[0]))/(1.0+self.ba*np.absolute(input_items[0])*np.absolute(input_items[0]))
        out_phase=np.angle(input_items[0]) + (self.ap*np.absolute(input_items[0])*np.absolute(input_items[0]))/(1.0+self.bp*np.absolute(input_items[0])*np.absolute(input_items[0]))
        output_items[0][:] = out_mag * np.exp(1j*out_phase)
        return len(output_items[0])