m-thesis-documentation/figures/beacon/src/beacon_field.py

#!/usr/bin/env python3

__doc__ = \
"""
Generate a figure showing a value combining
the delays between a transmitter and a set of antennas
"""


import matplotlib.pyplot as plt
import numpy as np
from itertools import chain, combinations, product

c_light = 3e8
f_beacon = 50e6

class Antenna:
    """
    Simple Antenna class
    """
    def __init__(self,x=0,y=0,z=0,t0=0,name=""):
        self.x = x
        self.y = y
        self.z = z
        self.t = t0
        self.name = name

    def __repr__(self):
        cls = self.__class__.__name__

        return f'{cls}(x={self.x!r},y={self.y!r},z={self.z!r},t0={self.t!r},name={self.name!r})'

def phase_mod(phase, low=np.pi):
    """
    Modulo phase such that it falls within the
    interval $[-low, 2\pi - low)$.
    """
    return (phase + low) % (2*np.pi) - low

def dist(a,b):
    return ((a.x-b.x)**2+(a.y-b.y)**2+(a.z-b.z)**2)**0.5

def phase(a,b,f=f_beacon,wrap=False):
    t = dist(a,b)/c_light
    phase = t*f*2*np.pi

    if wrap:
        phase = phase_mod(phase)

    return phase

def grid_plot(grid, ax=None, **plot_kwargs):
    if ax is None:
        ax = plt.gca()

    if not grid:
        return

    x = [a.x for a in grid]
    y = [a.y for a in grid]
    l = [a.name for a in grid]
    ax.plot(x,y,'kx', **plot_kwargs)
    for x_,y_,l_ in zip(x,y,l):
        ax.annotate(l_,(x_,y_))

def antenna_combinations(ants, ref_ant=None):
    if ref_ant is not None: # use only one reference antenna for the baselines
        ref = ref_ant
        ant_combi = [ (ref, j) for j in ants if j is not ref ]
    else: # use all baselines
        ant_combi = combinations(ants, 2)

    return list(ant_combi)

def calculate_field(field, tx, ant_combi, ref_ant=None, calculate_phase=True):
    if plot_phase:
        p_ij = [ (phase(tx,i) - phase(tx,j)) for i,j in ant_combi ]
    else:
        t_ij = [ (dist(tx,i) - dist(tx, j))/c_light for i,j in ant_combi ]

    val = []
    xx = []
    yy = []

    xs = field[0]
    ys = field[1]

    for x in xs:
        for y in ys:
            tx_t = Antenna(x=x,y=y)

            if plot_phase: # phase
                dp_ij = np.array([ (phase(tx_t,i) - phase(tx_t,j)) - p_ij[k] for k,(i,j) in enumerate(ant_combi) ])

                if True: # phase wrap
                    dp_ij = phase_mod(dp_ij)

                val.append(np.sum(dp_ij**2)**0.5)
            else: # time
                dt_ij = np.array([ (dist(tx_t,i) - dist(tx_t, j))/c_light - t_ij[k] for k,(i,j) in enumerate(ant_combi) ])

                val.append(np.sum(dt_ij**2)**0.5)

            xx.append(x)
            yy.append(y)

    return np.array(xx), np.array(yy), np.array(val)

def plot_field(
        tx, ants, xx, yy, val,
        ref_ant=None, plot_phase=None, mask=None,
        color_label='$\\left( t - \\tau  \\right)^2$',
        ax=None, bin_type='square', colorbar=True,
        **scatter_kwargs
        ):
    if ax is None:
        ax = plt.gca()

    default_scatter_kwargs = {}
    default_scatter_kwargs['cmap'] = 'Spectral_r'

    if mask is not None:
        val[mask] = np.nan

    if plot_phase is None:
        pass
    elif plot_phase:
        default_scatter_kwargs['vmax'] = len(ants)*np.pi
        pass

    scatter_kwargs = {**default_scatter_kwargs, **scatter_kwargs}

    if tx or ants:
        grid_plot([tx] + ants, ax=ax)

    if ref_ant is not None:
        ax.set_title("Single baseline\n reference antenna={}, f={}MHz".format(ref_ant.name, f_beacon/1e6))
    else:
        ax.set_title("All baselines\n f={} MHz".format(f_beacon/1e6))

    if bin_type == 'hex': # hexbin
        sc = ax.hexbin(xx, yy, C=val, **scatter_kwargs)
    elif bin_type == 'square': # squarebinning
        _, _, _, sc = ax.hist2d(xx, yy, weights=val, bins=int(len(xx)**0.5), **scatter_kwargs)
    else: # overlayed markers
        sc = ax.scatter(xx,yy,c=val, **scatter_kwargs)

    if colorbar:
        fig = ax.get_figure()
        fig.colorbar(sc, ax=ax, label=color_label)

    return ax

def square_grid(dx=1, N_x=10, dy=None, N_y=None, x_start=0, y_start=0):
    N_y = N_x if N_y is None else N_y
    dy = dx if dy is None else dy

    return product(
                (x_start + n*dx for n in range(N_x)),
                (y_start + n*dy for n in range(N_y)),
            )

def triangular_grid(dx=1, N_x=10, dy=None, N_y=None, x_start=0, y_start=0):

    return chain(
                square_grid(dx=dx,dy=dy,N_x=N_x,N_y=N_y,x_start=x_start,y_start=y_start),
                square_grid(dx=dx,dy=dy,N_x=N_x,N_y=N_y,x_start=x_start + dx/2,y_start=y_start + dy/2),
            )



if __name__ == "__main__":
    from argparse import ArgumentParser
    import os.path as path

    parser = ArgumentParser(description=__doc__)
    parser.add_argument("fname", metavar="path/to/figure[/]", nargs="?", help="Location for generated figure, will append __file__ if a directory. If not supplied, figure is shown.")


    parser.add_argument("type", choices=['single-left', 'single-center', 'square', 'tri', 'preset'])


    command_group = parser.add_mutually_exclusive_group(required=True)
    command_group.add_argument('--time',  help='Use the time difference for the field', action='store_false')
    command_group.add_argument('--phase', help='Use wrapped phase for the field', action='store_true')

    parser.add_argument('--ref', dest='ref_ant', metavar='ref_antenna', type=int, help='Number of antenna to use as reference')

    args = parser.parse_args()

    if args.fname is not None and path.isdir(args.fname):
        args.fname = path.join(args.fname, path.splitext(path.basename(__file__))[0] + ".pdf")

    args.plot_phase = args.phase or args.time
    del args.time, args.phase


    if 'single' in args.type: # single baseline
        ### Field
        x_low, x_high, N_x = -300, 300, 151
        y_low, y_high, N_y = -300, 300, 151

        ### Geometry
        ants = [
                Antenna(x=-50,y=0,z=0,name="a"),
                Antenna(x=50,y=0,z=0,name="b"),
            ]

        if args.type == 'single-center':
            tx = Antenna(x=-000,y=200,z=0,name="tx")
        else:
            tx = Antenna(x=-200,y=200,z=0,name="tx")

    elif args.type == 'square' or args.type == 'tri': # from grid definition
        ### Field
        x_low, x_high, N_x = -1800, 1800, 261
        y_low, y_high, N_y = -x_low, -x_high, N_x

        ### Geometry
        tx = Antenna(x=-800,y=300,z=0,name="tx")
        x_start, dx, ant_N_x = 0, 50, 2
        y_start, dy, ant_N_y = 0, dx, ant_N_x

        if args.type == 'square': # square grid
            grid_func = square_grid
        elif args.type == 'tri': # triangular
            grid_func = triangular_grid

        grid = grid_func(dx=dx, dy=dy, N_x=ant_N_x, N_y=ant_N_y, x_start=x_start, y_start=y_start)


        ants = [ Antenna(x=x,y=y,z=0,name=i) for i, (x,y) in enumerate(grid) ]
    else:
        ### Field
        x_low, x_high, N_x = -400, 400, 161
        y_low, y_high, N_y = -x_low, -x_high, N_x

        ### Geometry
        tx = Antenna(x=-300,y=300,z=0,name="tx")
        ants = [
                Antenna(x=100,y=0,z=0,name="a"),
                Antenna(x=0,y=-50,z=0,name="b"),
                Antenna(x=+50,y=-180,z=0,name="c"),
                Antenna(x=125,y=180,z=0,name="d"),
            ]

    ###
    ### Options
    ###
    plot_phase = args.plot_phase
    if args.ref_ant:
        ref_ant = ants[args.ref_ant]
    else:
        ref_ant = None


    ant_combi = antenna_combinations(ants, ref_ant=ref_ant)

    xs = np.linspace(x_low, x_high, N_x)
    ys = np.linspace(y_low, y_high, N_y)

    xx, yy, val = calculate_field((xs, ys), tx, ant_combi, calculate_phase=plot_phase)


    kwargs = {}
    mask = None
    if plot_phase:
        color_label='$\\sqrt{ \\sum \\left(\\varphi(x) - \\Delta \\varphi\\right)^2}$'
        mask = abs(val) > np.pi
    else:
        color_label='$\\sqrt{ \\sum \\left(t(x) - \\Delta t\\right)^2}$ [ns]'
        val *= 1e9
        kwargs['vmax'] = 100
        if not True:
            mask = abs(val) > 1.1*kwargs['vmax']


    if mask is not None:
        ax = plot_field([], [], xx, yy, val, cmap='Greys', colorbar=False)

    ax = plot_field(tx, ants, xx, yy, val, ref_ant=ref_ant, mask=mask, color_label=color_label, **kwargs)
    
#    if plot_phase:
#        N_lowest = np.min(len(ant_combi)-1, 10)
#        lowest_idx = np.argpartition(val, N_lowest)[:N_lowest]
#
#        print(lowest_idx)
#        print(val[lowest_idx])
#        print( list(zip(np.array(xx)[lowest_idx], np.array(yy)[lowest_idx])) )

    if args.fname is not None:
        plt.savefig(args.fname)
    else:
        plt.show()