From 156f6a2f4fb220c14b4da49e85e7913bb01e58b9 Mon Sep 17 00:00:00 2001 From: Samuel Rowlinson Date: Fri, 7 Dec 2018 14:56:26 +0000 Subject: [PATCH] adding functions for generating amp detectors and their output up to specified maxtem, in pykat.tools.detecting --- pykat/tools/detecting.py | 197 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 197 insertions(+) create mode 100644 pykat/tools/detecting.py diff --git a/pykat/tools/detecting.py b/pykat/tools/detecting.py new file mode 100644 index 0000000..bc38f55 --- /dev/null +++ b/pykat/tools/detecting.py @@ -0,0 +1,197 @@ +import numpy as np + +def all_bp_detectors(_base, param='q', direction='x'): + """Generates a string of kat code consisting of + beam parameter detectors for every node in `_base`. + + Parameters + ---------- + _base : :class:`.kat` + Base object containing the optical configuration. + + param : str, optional + Parameter of the `bp` detector to detect (defaults to 'q'). + + direction : str, optional + Plane of detection ('x' for tangential, 'y' for sagittal, + "both" for generating `bp` detectors for both planes). + + Returns + ------- + out : str + A string of the kat code for the beam parameter detectors. + """ + code = "" + if direction == "both": direction = ['x', 'y'] + else: direction = [direction] + for node in _base.nodes.getNodes().keys(): + for d in direction: + code += "bp bp{n} {d} {p} {n}\n".format( + n=node, d=d, p=param + ) + return code + +def generate_amp_detectors(node, maxtem, f=0, ad_prefix="ad"): + """Generates a string of kat code consisting of amplitude detectors + for each TEM up to a given maxtem. + + Parameters + ---------- + node : str + Name of the node at which to place the amplitude detectors. + + maxtem : int + Maximum mode order up to which to generate amplitude detectors. + + f : double + Frequency offset to carrier [Hz]. + + ad_prefix : str + Prefix of amplitude detector names. These will be named + as ``[nm, ...]``. + + Returns + ------- + out : str + A string of the kat code for the amplitude detectors. + + author: George Smetana + date: 04/12/2018 + """ + temsize = maxtem + 1 + code = "" + for n in range(temsize): + for m in range(temsize - n): + code += "ad {prefix}{n}{m} {n} {m} {f} {node} \n".format( + prefix=ad_prefix, n=n, m=m, f=f, node=node + ) + return code + +def get_tem_array(out, maxtem, ad_prefix='ad', power=False, order_only=True, table=False, plot=False): + """ + Return array of TEM mode amplitudes/powers obtained from amplitude detectors at a node. + + Parameters + ---------- + out : :class:`.KatRun` + Output from a Finesse simulation. + + maxtem : int + Maximum TEM order to which detectors should be generated. + + ad_prefix : str, optional + Prefix of amplitude detector names. + + power : bool, optional + If `True` returns the TEM powers, else returns the TEM amplitudes. + + order_only : bool, optional + If `True`, returns quadrature sum of modes of the same TEM order, + else returns the full 2D array. + + table : bool, optional + If `True`, print the output using the ``tabulate`` module. + + plot : bool, optional + If `True`, generates a plot of the output. + + Returns + ------- + If `order_only == True` and noaxis set: + 1D array where array[k] is amp/power in kth TEM order + If `order_only == True` and xaxis set: + 2D array where array[k, i] is amp/power in kth TEM order of ith xaxis measurement. + If `order_only == False` and noaxis set: + 2D array where array[n, m] is amp/power in TEM_nm mode + If `order_only == False` and xaxis set: + 3D array where array[n, m, i] is amp/power in TEM_nm mode of ith xaxis measurement. + + author: George Smetana + date: 04/12/2018 + """ + from tabulate import tabulate + if plot: + import matplotlib as mpl + import matplotlib.pyplot as plt + + temsize = maxtem + 1 + # Depth is how many measurements each detector makes, this must be set manually to 1 for special + # case of a single measurement with noxaxis set. + depth = 1 if out.x.shape is () else out.x.size + # Initialise the amplitude array with all nans + tem_amps = np.empty((temsize, temsize, depth)) + tem_amps[:] = np.nan + for n in range(temsize): + for m in range(temsize - n): + tem_amps[n, m, :] = out["{prefix}{n}{m}".format( + prefix=ad_prefix, n=n, m=m + )] + + if depth > 1 and (table or plot): + # If multiple measurements are made, printing a table/plotting doesn't make sense. + raise Exception("Table printing and plotting is not supported " + "for multiple outputs per detector.") + + if order_only: + orders = list(range(temsize)) + tem_powers = tem_amps*np.conj(tem_amps) + + x = np.array(orders) + y = x[::-1] + # Rotate array so that diagonal elements (elements of the same TEM order) are stacked in columns. + rotated_powers = tem_powers[x, np.add.outer(x, y) - (temsize - 1), :] + order_powers = np.nansum(rotated_powers, axis=1) + order_vals = np.sqrt(order_powers) if not power else order_powers + + if depth == 1: + # If only one measurement is made per detector, keeping a 2D array is redundant. + order_vals = order_vals[:, 0] + + if table and depth == 1: + value_name = ['Power'] if power else ['Amplitude'] + headers = ['Order'] + orders + row = value_name + list(order_vals) + print(tabulate([row], headers=headers, tablefmt='grid')) + + if plot and depth == 1: + plt.figure() + ax = plt.axes() + ax.scatter(orders, order_vals) + ax.set_xticks(orders) + ax.grid(True, 'both', 'both') + ax.set_yscale('log') + ax.set_ylim((min(order_vals[order_vals > 0])*0.8, max(order_vals)*1.2)) + ax.set_xlabel('TEM modes order') + ax.set_ylabel('Power' if power else 'Amplitude') + ax.set_title('{} of Each TEM Order'.format('Power' if power else 'Amplitude')) + plt.show() + + return order_vals + else: + if depth == 1: + # If only one measurement is made per detector, keeping a 3D array is redundant. + tem_amps = tem_amps[:, :, 0] + tem_vals = tem_amps*np.conj(tem_amps) if power else tem_amps + orders = list(range(temsize)) + + if table and depth == 1: + value_name = 'Power' if power else 'Amplitude' + print('{} of All Modes - Location (n, m) for TEM_nm'.format(value_name)) + print(tabulate(tem_vals, tablefmt='grid')) + + if plot and depth == 1: + m_list = orders * temsize + n_list = sorted(m_list) + plt.figure() + ax = plt.axes() + grid = ax.scatter(n_list, m_list, c=tem_vals.flatten(), norm=mpl.colors.LogNorm()) + plt.colorbar(grid) + ax.set_xticks(orders) + ax.set_yticks(orders) + ax.grid(True, 'major', 'both') + ax.set_xlabel('n index in $TEM_{nm}$') + ax.set_ylabel('m index in $TEM_{nm}$') + ax.set_title('{} of Each TEM Mode'.format('Power' if power else 'Amplitude')) + plt.show() + + return tem_vals \ No newline at end of file -- GitLab