diff --git a/RMtools_1D/clean_model.py b/RMtools_1D/clean_model.py index 6759eca..a643874 100755 --- a/RMtools_1D/clean_model.py +++ b/RMtools_1D/clean_model.py @@ -10,8 +10,12 @@ import numpy as np from RMtools_1D.do_RMsynth_1D import readFile as read_freqFile +from RMutils.util_misc import create_frac_spectra +from RMutils.util_misc import calculate_StokesI_model +from RMutils.util_plotTk import plot_Ipqu_spectra_fig from RMutils.util_misc import calculate_StokesI_model import json +import matplotlib.pyplot as plt def calculate_QU_model(freqArr, phiArr, CCArr, lambdaSq_0, Iparms=None): """Compute the predicted Stokes Q and U values for each channel from a @@ -72,7 +76,6 @@ def read_files(freqfile,rmSynthfile, CCfile): Iparms: list of Stokes I model parameters. lambdaSq_0: scalar value of lambda^2_0, in m^2. """ - phiArr, CCreal, CCimag = np.loadtxt(CCfile, unpack=True, dtype='float') CCArr=CCreal + 1j * CCimag @@ -81,12 +84,66 @@ def read_files(freqfile,rmSynthfile, CCfile): Iparms=[float(x) for x in synth_mDict['polyCoeffs'].split(',')] lambdaSq_0=synth_mDict['lam0Sq_m2'] + # Parse the data array + # freq_Hz, I, Q, U, dI, dQ, dU data=read_freqFile(freqfile, 64, verbose=False, debug=False) - freqArr=data[0] - + freqArr_Hz=data[0] + + return phiArr, CCArr, Iparms, lambdaSq_0,freqArr_Hz - return phiArr, CCArr, Iparms, lambdaSq_0,freqArr +def plot_model(freqfile, QUarr,Imodel): + """Plot the Stokes Q and U model. + + Args: + freqfile (str): File to read original data from + QUarr (np.ndarray): CLEAN QU model + Imodel (np.ndarray): Stokes I model + """ + # Parse the data array + # freq_Hz, I, Q, U, dI, dQ, dU + data=read_freqFile(freqfile, 64, verbose=False, debug=False) + try: + (freqArr_Hz, IArr, QArr, UArr, + dIArr, dQArr, dUArr) = data + print("\nFormat [freq_Hz, I, Q, U, dI, dQ, dU]") + except Exception: + # freq_Hz, Q, U, dQ, dU + try: + (freqArr_Hz, QArr, UArr, dQArr, dUArr) = data + + print("\nFormat [freq_Hz, Q, U, dQ, dU]") + noStokesI = True + except Exception: + print("\nError: Failed to parse data file!") + # Compute fractional spectra + # TODO: Currently re-computes Stokes I model - should be passed in + dataArr = create_frac_spectra(freqArr=freqArr_Hz, + IArr=IArr, + QArr=QArr, + UArr=UArr, + dIArr=dIArr, + dQArr=dQArr, + dUArr=dUArr, + polyOrd=2, + verbose=True, + fit_function='log') + (IModArr, qArr, uArr, dqArr, duArr, IfitDict) = dataArr + + specFig = plt.figure(facecolor='w',figsize=(10, 6)) + plot_Ipqu_spectra_fig(freqArr_Hz = freqArr_Hz, + IArr = IArr, + qArr = qArr, + uArr = uArr, + dIArr = dIArr, + dqArr = dqArr, + duArr = duArr, + freqHirArr_Hz = freqArr_Hz, + IModArr = Imodel, + qModArr = QUarr.real, + uModArr = QUarr.imag, + fig = specFig) + plt.show() def main(): """Generate Stokes QU model based on clean components and (optional) @@ -104,13 +161,13 @@ def main(): # Parse the command line options parser = argparse.ArgumentParser(description=descStr, formatter_class=argparse.RawTextHelpFormatter) - parser.add_argument("freqfile", metavar="input.dat", nargs=1, + parser.add_argument("freqfile", metavar="input.dat", help="ASCII file containing original frequency spectra.") - parser.add_argument("rmSynthfile",metavar="_RMsynth.json",nargs=1, + parser.add_argument("rmSynthfile",metavar="_RMsynth.json", help="RMsynth1d output JSON file.") - parser.add_argument("CCfile",metavar="_FDFmodel.dat",nargs=1, + parser.add_argument("CCfile",metavar="_FDFmodel.dat", help="Clean component model file (_FDFmodel.dat)") - parser.add_argument("outfile",metavar="QUmodel.dat",nargs=1, + parser.add_argument("outfile",metavar="QUmodel.dat", help="Filename to save output model to.") args = parser.parse_args() @@ -119,6 +176,8 @@ def main(): args.rmSynthfile, args.CCfile) QUarr,Imodel=calculate_QU_model(freqArr, phiArr, CCArr, lambdaSq_0, Iparms) + + plot_model(args.freqfile, QUarr,Imodel) save_model(args.outfile,freqArr,Imodel,QUarr) diff --git a/RMutils/util_plotTk.py b/RMutils/util_plotTk.py index 79ccbb6..0e36975 100644 --- a/RMutils/util_plotTk.py +++ b/RMutils/util_plotTk.py @@ -301,12 +301,12 @@ def plot_PQU_vs_nu_ax(ax, freqArr_Hz, QArr, UArr, dQArr=None, # Plot the models if QmodArr is not None: - ax.plot(freqHirArr_Hz/1e9, QmodArr, color='tab:blue', lw=0.5, label='Model Q') + ax.plot(freqHirArr_Hz/1e9, QmodArr, color='tab:blue', lw=2.5, label='Model Q') if UmodArr is not None: - ax.plot(freqHirArr_Hz/1e9, UmodArr, color='tab:red', lw=0.5, label='Model U') + ax.plot(freqHirArr_Hz/1e9, UmodArr, color='tab:red', lw=2.5, label='Model U') if QmodArr is not None and UmodArr is not None: PmodArr = np.sqrt(QmodArr**2.0 + UmodArr**2.0 ) - ax.plot(freqHirArr_Hz/1e9, PmodArr, color='k', lw=0.5, label='Model P') + ax.plot(freqHirArr_Hz/1e9, PmodArr, color='k', lw=2.5, label='Model P') # Formatting ax.yaxis.set_major_locator(MaxNLocator(4)) @@ -415,12 +415,12 @@ def plot_pqu_vs_lamsq_ax(ax, lamSqArr_m2, qArr, uArr, pArr=None, dqArr=None, # Plot the models if qModArr is not None: - ax.plot(lamSqHirArr_m2, qModArr, color='tab:blue', alpha=1, lw=0.1, label='Model q') + ax.plot(lamSqHirArr_m2, qModArr, color='tab:blue', alpha=1, lw=2.5, label='Model q') if uModArr is not None: - ax.plot(lamSqHirArr_m2, uModArr, color='tab:red', alpha=1, lw=0.1, label='Model u') + ax.plot(lamSqHirArr_m2, uModArr, color='tab:red', alpha=1, lw=2.5, label='Model u') if qModArr is not None and uModArr is not None: pModArr = np.sqrt(qModArr**2.0 + uModArr**2.0 ) - ax.plot(lamSqHirArr_m2, pModArr, color='k', alpha=1, lw=0.1, label='Model p') + ax.plot(lamSqHirArr_m2, pModArr, color='k', alpha=1, lw=2.5, label='Model p') if model_dict is not None: errDict = {} # Sample the posterior randomly 100 times @@ -496,7 +496,7 @@ def plot_psi_vs_lamsq_ax(ax, lamSqArr_m2, qArr, uArr, dqArr=None, duArr=None, # Plot the model if qModArr is not None and uModArr is not None: psiHirArr_deg = np.degrees( np.arctan2(uModArr, qModArr) / 2.0 ) - ax.plot(lamSqHirArr_m2, psiHirArr_deg, color='tab:red', lw=0.1, + ax.plot(lamSqHirArr_m2, psiHirArr_deg, color='tab:red', lw=2.5, label='Model $\psi$') if model_dict is not None: errDict = {} @@ -552,7 +552,7 @@ def plot_q_vs_u_ax(ax, lamSqArr_m2, qArr, uArr, dqArr=None, duArr=None, # Plot the model if qModArr is not None and uModArr is not None: - ax.plot(qModArr, uModArr, color="k", lw=0.1, label='Model q & u', + ax.plot(qModArr, uModArr, color="k", lw=2.5, label='Model q & u', zorder=2) if model_dict is not None: errDict = {}