<jupyter_start><jupyter_text>Numerical accuracy of uvcontsub2021This notebook simulates MeasurementSets with known continuum shapes (for example, as polynomials of different order with known coefficients). It has functions to produce MSs including point sources, spectral lines, added noise, and polynomial continuum functions. These MSs can be used to test the numerical accuracy of the task uvcontsub2021, illustrate its usage and experiment with it under different simulated conditions.   Import required CASA tasks and tools, and other packages<jupyter_code>import os
import shutil
import numpy as np

import vizhelper

from casatools import componentlist, ctsys, measures, simulator, table
from casatasks import flagdata, listobs
from casatasks import uvcontsub2021
from casatasks.private import simutil

cl = componentlist()
sm = simulator()<jupyter_output><empty_output><jupyter_text>Definition of simulation building blocksThis section defines function to:* Create a MeasurementSet with basic structure* Plot visibilities* Populate and edit the visibilities of an MS (add sources, continuum, noise, etc.)Uses functions from https://github.com/urvashirau/Simulation-in-CASA/blob/master/Basic_Simulation_Demo/Simulation_Script_Demo.ipynb, adapted to this set of simulations for uvcontsub.<jupyter_code>def make_ms_frame(msname:str, ant_config:str, spw_params=None, verbose=False):
    """                                                                                                                                                       
    Construct an empty MeasurementSet that has the desired observation setup.                                                                                 
                                                                                                                                                              
    Args:                                                                                                                                                     
        msname: MS to create                                                                                                                                  
        ant_config: a telescope configuration from casadata (alma/simmos) (for                                                                                
                  example  alma/simmos/aca.cycle8.cfg)                                                                                                        
        swp_params: parameters such as SPW name, number of channels, frequency                                                                                
    """

    ## Open the simulatorFunctions                                                                                                                            
    sm.open(ms=msname);

    ## Read/create an antenna configuration.                                                                                                                  
    ## Canned antenna config text files are located here: <casadata>/alma/simmos/*cfg                                                                         

    antennalist = os.path.join(ctsys.resolve("alma/simmos"), ant_config)
    if verbose:
        print(f'Using antenna list file: {antennalist}')


    ## Fictitious telescopes can be simulated by specifying x, y, z, d, an, telname, antpos.                                                                  
    ##     x,y,z are locations in meters in ITRF (Earth centered) coordinates.                                                                                
    ##     d, an are lists of antenna diameter and name.                                                                                                      
    ##     telname and obspos are the name and coordinates of the observatory.                                                                                

    mysu = simutil.simutil()
    (x, y, z, d, an, an2, telname, obspos) = mysu.readantenna(antennalist)

    ## Set the antenna configuration                                                                                                                          
    metool = measures()
    sm.setconfig(
        telescopename=telname,
        x=x,
        y=y,
        z=z,
        dishdiameter=d,
        mount=['alt-az'],
        antname=an,
        coordsystem='local',
        referencelocation=metool.observatory(telname)
    )

    ## Set the polarization mode (this goes to the FEED subtable)                                                                                             
    sm.setfeed(mode='perfect X Y', pol=['']);    # X Y / R L                                                                                                  

    ## Set the spectral window and polarization (one data-description-id).                                                                                    
    ## Call multiple times with different names for multiple SPWs or pol setups.                                                                              
    sm.setspwindow(
        spwname=spw_params['name'],
        freq=spw_params['freq'],
        deltafreq='0.1GHz',
        freqresolution='0.2GHz',
        nchannels=spw_params['nchan'],
        stokes='XX YY'
    )

    ## Setup source/field information (i.e. where the observation phase center is)                                                                            
    ## Call multiple times for different pointings or source locations.                                                                                       
    source_name = 'simulated_source'
    sm.setfield(
        sourcename=source_name,
        sourcedirection=metool.direction(
            rf='J2000',
            v0='19h59m28.5s',
            v1='+40d44m01.5s'
        )
    )

    ## Set shadow/elevation limits (if you care). These set flags.                                                                                            
    sm.setlimits(shadowlimit=0.01, elevationlimit='1deg');

    ## Leave autocorrelations out of the MS.                                                                                                                  
    sm.setauto(autocorrwt=0.0);

    ## Set the integration time, and the convention to use for timerange specification                                                                        
    ## Note : It is convenient to pick the hourangle mode as all times specified in sm.observe()                                                              
    ##        will be relative to when the source transits.                                                                                                   
    sm.settimes(
        integrationtime='2s',
        usehourangle=True,
        referencetime=metool.epoch('UTC','2021/10/14/00:01:02')
    )

    ## Construct MS metadata and UVW values for one scan and did                                                                                              
    ## Call multiple times for multiple scans.                                                                                                                
    ## Call this with different sourcenames (fields) and spw/pol settings as defined above.                                                                   
    ## Timesteps will be defined in intervals of 'integrationtime', between starttime and stoptime.                                                           

    sm.observe(
        sourcename=source_name,
        spwname=spw_params['name'],
        starttime='0',
        stoptime='+8s'
    )

    ## Close the simulator                                                                                                                                    
    sm.close()

    ## Unflag everything (unless you care about elevation/shadow flags)                                                                                       
    flagdata(vis=msname, mode='unflag', flagbackup=False)
    
def add_gaussian_noise(msname:str, noise_level_sigma='0.1Jy'):
    """                                                                                                                                                       
    Adds Gaussian random noise using simulator / corrupt                                                                                                      
                                                                                                                                                              
    Args:                                                                                                                                                     
        ms_name: MS to modify                                                                                                                                 
        noise_level_sigma: noise sigma as used in the simulator corrupt function                                                                              
    """

    try:
        sm.openfromms(msname)
        sm.setseed(4321)
        sm.setnoise(mode='simplenoise', simplenoise=noise_level_sigma)
        sm.corrupt()
    finally:
        sm.close()

def add_point_source_component(msname, cl_name = 'sim_point_source.cl', freq=None,
                                flux=5.0, spectrumtype='constant', index=-1):
    """                                                                                                                                                       
    Adds a point source to the MS                                                                                                                             
                                                                                                                                                              
    Args:                                                                                                                                                     
        ms_name: MS to modify                                                                                                                                 
        freq:                                                                                                                                                 
        spectrumtype:                                                                                                                                         
        flux: In Jy, as used in componentlist.addcomponent                                                                                                    
    """

    make_point_source_comp_list(
        cl_name=cl_name,
        freq=freq,
        flux=flux,
        spectrumtype=spectrumtype,
        index=index
    )

    sim_from_comp_list(msname, cl_name)
    shutil.rmtree(cl_name)
        

def make_point_source_comp_list(cl_name:str, freq:str, flux:str, spectrumtype:str, index:int,
                                direction='J2000 19h59m28.5s +40d44m01.5s',
                                fluxunit='Jy', shape='point', label='sim_point_source'):
    """                                                                                                                                                       
    Makes a component list file with a point source                                                                                                           
                                                                                                                                                              
    Args:                                                                                                                                                     
        cl_name: name of the component list file to create                                                                                                    
        freq: freq quantity (with units) as used in the component list tool                                                                                   
        flux: flux, units are assumed in Jy                                                                                                                   
        spectrumtype: as used in component list tool: constant/spectral index                                                                                 
        index: spectral index                                                                                                                                 
    """

    try:
        cl.addcomponent(
            dir=direction,
            flux=flux,
            fluxunit=fluxunit,
            freq=freq,
            shape=shape,     # shape='gaussian',                                                                                                              
                             # majoraxis="5.0arcmin", minoraxis='2.0arcmin',                                                                                  
                             # polarization='linear', / 'Stokes'                                                                                              
                             # spectrumtype:'spectral index' / 'constant'                                                                                     
            spectrumtype=spectrumtype,
            index=index,
            label=label)
        cl.rename(filename=cl_name)
    finally:
        cl.close()

def sim_from_comp_list(msname:str, cl_name:str):
    """                                                                                                                                                       
    Updates the MS visibilities using simulator.predict to add                                                                                                
    components from the components list                                                                                                                       
                                                                                                                                                              
    Args:                                                                                                                                                     
        ms_name: MS to modify                                                                                                                                 
        cl_name: name of components list file to simulate                                                                                                     
    """
    try:
        sm.openfromms(msname)
        sm.predict(complist=cl_name, incremental=False)
    finally:
        sm.close()

def add_spectral_line(msname:str, line:'numpy.array', chan_range=[60, 86], amp_factor=None):
    """                                                                                                                                                       
    Adds a spectral line as a Gaussian function in the range of channels given                                                                                
                                                                                                                                                              
    Args:                                                                                                                                                     
        ms_name: MS to modify                                                                                                                                 
        line: function to produce spectral line. ex: fn(x, mu, sigma)                                                                                         
        chan_range: list with indices of the first and last channel                                                                                           
        amp_factor: factor to multiply the peak height / flux density                                                                                         
    """

    try:
        tbtool = table()
        tbtool.open(msname, nomodify=False)
        data = tbtool.getcol('DATA')



        if not amp_factor:
            amp_factor = 1.0/np.max(line)

        data[:,chan_range[0]:chan_range[1],:] += amp_factor * (1+0j) * line.reshape((1, len(line), 1))
        tbtool.putcol('DATA', data)

    finally:
        tbtool.done()

def add_polynomial_continuum(msname:str, pol_coeffs:list, nchan:int, amp_factor=1.0):
    """                                                                                                                                                       
    Update MS visibilities adding a polynomial evaluated for all channels.                                                                                    
                                                                                                                                                              
    Args:                                                                                                                                                     
        ms_name: MS to modify                                                                                                                                 
        pol_coeff: polynomial coefficients, evaluated [0.5, 0.25] => 0.5x + 0.25                                                                              
        nchan: number of channels in the SPW (x axis to eval polynomial)                                                                                      
    """
    try:
        tbtool = table()
        tbtool.open(msname, nomodify=False)
        data = tbtool.getcol('DATA')

        x = np.linspace(0, 1, nchan)
        polynomial = np.polyval(pol_coeffs, x)

        # Add same polynomial to real and imag part                                                                                                           
        data += amp_factor * (1+1j) * polynomial.reshape((1, len(polynomial),1))

        tbtool.putcol('DATA', data)
    finally:
        tbtool.done()<jupyter_output><empty_output><jupyter_text>Simple MS with continuum, line, and noise, plotted step by stepUsing order 0 continuum.The structure of this MS is the one used at the moment in the task test (test_task_req_uvcontsub2021). nchan=128. The spectral line is added to channels 60-85. The value of the fitspw parameter for this MS structure would be '0:0\~59;86\~127'.The following functions we produce MSs without adding polynomials to the real and imaginary part of the visibilities. These MSs can be used to test the behavior of the task with sources  Polynomial n=0 continuum<jupyter_code># Clean up and remove measurement files from previous runs

vizhelper.clean_measurement_files(prefix='sim_alma_noise')
ant_config = 'alma.cycle8.8.cfg'
spw_params = {
    'name': "Simulated_Band",
    'freq': '150GHz',
    'nchan': 128
}

msname='sim_alma_noise_cont_poly_order_0.ms' 

make_ms_frame(
    msname=msname, 
    ant_config=ant_config, 
    spw_params=spw_params, 
    verbose=True
)

# Add point source to measurement set
add_point_source_component(
    msname=msname,
    cl_name = 'sim_point_source.cl', 
    freq='100GHz', 
    flux=0.5, 
    spectrumtype='constant', 
    index=-1
)

# Define the spectral line shape
gaussian = lambda x, mu, sigma: np.exp(-np.power(x - mu, 2.) / (2 * np.power(sigma, 2.)))
x = np.linspace(-3, 3, 86-60)

# Add spectral line
add_spectral_line(
    msname=msname,
    line=gaussian(x, 0., 0.58), 
    chan_range=[60, 86]
)

# Add polynomial component to background
add_polynomial_continuum(
    msname=msname,
    pol_coeffs=[0.025], 
    nchan = spw_params['nchan'], 
    amp_factor=1.0)

# Add gaussian noise to all channels
add_gaussian_noise(
    msname=msname, 
    noise_level_sigma='0.1Jy')

# Make diagnostic plots
vizhelper.plot_ms_data(
    msname=msname,
    myplot='data_spectrum'
)<jupyter_output><empty_output><jupyter_text>Polynomial n=1 continuum<jupyter_code>ant_config = 'alma.cycle8.8.cfg'
spw_params = {
    'name': "Simulated_Band",
    'freq': '150GHz',
    'nchan': 128
}

msname='sim_alma_noise_cont_poly_order_1.ms' 

make_ms_frame(
    msname=msname, 
    ant_config=ant_config, 
    spw_params=spw_params, 
    verbose=True
)

# Add point source to measurement set
add_point_source_component(
    msname=msname,
    cl_name = 'sim_point_source.cl', 
    freq='100GHz', 
    flux=0.5, 
    spectrumtype='constant', 
    index=-1
)

# Define the spectral line shape
gaussian = lambda x, mu, sigma: np.exp(-np.power(x - mu, 2.) / (2 * np.power(sigma, 2.)))
x = np.linspace(-3, 3, 86-60)

# Add spectral line
add_spectral_line(
    msname=msname,
    line=gaussian(x, 0., 0.58), 
    chan_range=[60, 86]
)

# Add polynomial component to background
add_polynomial_continuum(
    msname=msname,
    pol_coeffs=[-1, 0.75], 
    nchan = spw_params['nchan'], 
    amp_factor=1.0)

# Add gaussian noise to all channels
add_gaussian_noise(
    msname=msname, 
    noise_level_sigma='0.1Jy')

# Make diagnostic plots
vizhelper.plot_ms_data(
    msname=msname,
    myplot='data_spectrum'
)<jupyter_output><empty_output><jupyter_text>Polynomial n=2 continuum<jupyter_code>ant_config = 'alma.cycle8.8.cfg'
spw_params = {
    'name': "Simulated_Band",
    'freq': '150GHz',
    'nchan': 128
}

msname='sim_alma_noise_cont_poly_order_2.ms' 

make_ms_frame(
    msname=msname, 
    ant_config=ant_config, 
    spw_params=spw_params, 
    verbose=True
)

# Add point source to measurement set
add_point_source_component(
    msname=msname,
    cl_name = 'sim_point_source.cl', 
    freq='100GHz', 
    flux=0.5, 
    spectrumtype='constant', 
    index=-1
)

# Define the spectral line shape
gaussian = lambda x, mu, sigma: np.exp(-np.power(x - mu, 2.) / (2 * np.power(sigma, 2.)))
x = np.linspace(-3, 3, 86-60)

# Add spectral line
add_spectral_line(
    msname=msname,
    line=gaussian(x, 0., 0.58), 
    chan_range=[60, 86]
)

# Add polynomial component to background
add_polynomial_continuum(
    msname=msname,
    pol_coeffs=[-1., 1, 0.15], 
    nchan = spw_params['nchan'], 
    amp_factor=1.0)

# Add gaussian noise to all channels
add_gaussian_noise(
    msname=msname, 
    noise_level_sigma='0.1Jy')

# Make diagnostic plots
vizhelper.plot_ms_data(
    msname=msname,
    myplot='data_spectrum'
)<jupyter_output><empty_output><jupyter_text>Polynomial n=3 continuum<jupyter_code>ant_config = 'alma.cycle8.8.cfg'
spw_params = {
    'name': "Simulated_Band",
    'freq': '150GHz',
    'nchan': 128
}

msname='sim_alma_noise_cont_poly_order_3.ms' 

make_ms_frame(
    msname=msname, 
    ant_config=ant_config, 
    spw_params=spw_params, 
    verbose=True
)

# Add point source to measurement set
add_point_source_component(
    msname=msname,
    cl_name = 'sim_point_source.cl', 
    freq='100GHz', 
    flux=0.5, 
    spectrumtype='constant', 
    index=-1
)

# Define the spectral line shape
gaussian = lambda x, mu, sigma: np.exp(-np.power(x - mu, 2.) / (2 * np.power(sigma, 2.)))
x = np.linspace(-3, 3, 86-60)

# Add spectral line
add_spectral_line(
    msname=msname,
    line=gaussian(x, 0., 0.58), 
    chan_range=[60, 86]
)

# Add polynomial component to background
add_polynomial_continuum(
    msname=msname,
    pol_coeffs=[1., -0.75, -0.25, 0.15], 
    nchan = spw_params['nchan'], 
    amp_factor=1.0)

# Add gaussian noise to all channels
add_gaussian_noise(
    msname=msname, 
    noise_level_sigma='0.1Jy')

# Make diagnostic plots
vizhelper.plot_ms_data(
    msname=msname,
    myplot='data_spectrum'
)<jupyter_output><empty_output><jupyter_text>uvcontsubEach of the simulated meaurement files are background subtracted using uvcontsub<jupyter_code># Clean up and remove measurement files from previous runs

vizhelper.clean_measurement_files(prefix='uvcont_noise')

from casatasks import uvcontsub2021
fitspw='0:0~59;86~127'
  
ms_uvcont = 'uvcont_noise_sub_0.ms'
res = uvcontsub2021(
    vis='sim_alma_noise_cont_poly_order_0.ms', 
    outputvis=ms_uvcont, 
    fitorder=0, 
    fitspw=fitspw, 
    writemodel=True
)

ms_uvcont = 'uvcont_noise_sub_1.ms'
res = uvcontsub2021(
    vis='sim_alma_noise_cont_poly_order_1.ms', 
    outputvis=ms_uvcont, 
    fitorder=1, 
    fitspw=fitspw, 
    writemodel=True
)

ms_uvcont = 'uvcont_noise_sub_2.ms'
res = uvcontsub2021(
    vis='sim_alma_noise_cont_poly_order_2.ms', 
    outputvis=ms_uvcont, 
    fitorder=2, 
    fitspw=fitspw, 
    writemodel=True
)

ms_uvcont = 'uvcont_noise_sub_3.ms'
res = uvcontsub2021(
    vis='sim_alma_noise_cont_poly_order_3.ms', 
    outputvis=ms_uvcont, 
    fitorder=3, 
    fitspw=fitspw, 
    writemodel=True
)<jupyter_output><empty_output><jupyter_text>Average PlotsAverage plots display the data spectrum vs channel averaged along the time/sample axis. `Fit line` show the fit line used by uvcontsub to subtract the background. Plots are shown for each polynomial order background. polynomial n=0 continuum<jupyter_code>vizhelper.make_spectrum_plot(
    msname=[
        'sim_alma_noise_cont_poly_order_0.ms', 
        'sim_alma_noise_cont_poly_order_1.ms',
        'sim_alma_noise_cont_poly_order_2.ms',
        'sim_alma_noise_cont_poly_order_3.ms'
    ], 
    myplot='data_spectrum', 
    fitline=[
        'uvcont_noise_sub_0.ms',
        'uvcont_noise_sub_1.ms',
        'uvcont_noise_sub_2.ms',
        'uvcont_noise_sub_3.ms',
    ]
)<jupyter_output><empty_output><jupyter_text>Synopsis UVCONTSUB Plots and Statistics FiguresSynopsis plots show the uncorrected and corrected data spectrum plots. Each plot, spectrum vs channel, is averaged along the time/sample axis. A histogram of the corrected continuum background is displayed with a gaussian fit and statistics detailing the mean, stddev and skew to determine normality of the distribution. Finally, an interactive table with a detailed statistical comparison, using `casatasks.visstat`, is displayed to compare the spectrum statistics before and after uvcontsub.<jupyter_code>vizhelper.make_synopsis_plotly(
    uncorr='sim_alma_noise_cont_poly_order_0.ms', 
    corr='uvcont_noise_sub_0.ms', 
    chan_list=[[60, 82]],
    title = 'Order n=0 Continuum Subtraction Overview'
)

vizhelper.make_synopsis_plotly(
    uncorr='sim_alma_noise_cont_poly_order_1.ms', 
    corr='uvcont_noise_sub_1.ms', 
    chan_list=[[60, 82]],
    title = 'Order n=1 Continuum Subtraction Overview'
)

vizhelper.make_synopsis_plotly(
    uncorr='sim_alma_noise_cont_poly_order_2.ms', 
    corr='uvcont_noise_sub_2.ms', 
    chan_list=[[60, 82]],
    title = 'Order n=2 Continuum Subtraction Overview'
)

vizhelper.make_synopsis_plotly(
    uncorr='sim_alma_noise_cont_poly_order_3.ms', 
    corr='uvcont_noise_sub_3.ms', 
    chan_list=[[60, 82]],
    title = 'Order n=3 Continuum Subtraction Overview'
)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.06526548672566372, 'fraction_numerical': 0.014454277286135693, 'mean_word_length': 1.0227476133651552, 'pattern_counts': {'":': 0, '<': 38, '<?xml version=': 0, '>': 39, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 388, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '591329', 'license': 'MIT', 'n_tokens_mistral': 6482, 'n_tokens_neox': 5888, 'n_words': 1487, 'path': 'simulations_uvcontsub_ALMA_plotly_restruct.ipynb', 'repo_name': 'jrhosk/uvcontsub', 'size': '27120'}