def _change_frame(ixds_block, observer, target, frequency, frame):
"""
This function will be called using xr.map_blocks
reference
http://xarray.pydata.org/en/stable/generated/xarray.map_blocks.html
and
https://xarray-contrib.github.io/xarray-tutorial/scipy-tutorial/06_xarray_and_dask.html#map_blocks
"""
if ("velocity"
in ixds_block.coords) and ("Velocity"
in ixds_block.attrs["units_map"]):
# update target to use velo coordinate value for this block
# since this is inside map_blocks, the only velo/chan available is the one we want to use
chan_vel = ixds_block.velocity.values * units.Unit(
ixds_block.attrs["units_map"]["Velocity"])
# assignment via `SkyCoord(target, radial_velocity=chan_vel)` doesn't seem to work
# see
# https://docs.astropy.org/en/stable/coordinates/velocities.html#adding-velocities-to-existing-frame-objects
# https://docs.astropy.org/en/stable/coordinates/representations.html#attaching-differential-objects-to-representation-objects
target = target.data.with_differentials(
{'s': RadialDifferential(chan_vel)})
for xda in ixds_block.data_vars:
# we don't want to transform boolean data arrays
# we also don't want to if there is no chan dimension
if ixds_block[xda].dtype != bool and "chan" in ixds_block[xda].dims:
# ixds_block[xda].values is the computed ndarray for this block
sc = SpectralCoord(
ixds_block[xda].values,
unit=ixds_block.attrs["units_map"]["Frequency"],
observer=observer,
target=target,
doppler_rest=frequency,
doppler_convention=ixds_block.attrs["velocity__type"],
)
# see https://docs.astropy.org/en/stable/coordinates/spectralcoord.html#specifying-an-observer-and-a-target-explicitly
new_sc = sc.with_observer_stationary_relative_to(frame)
ixds_block[xda] = new_sc.values
return ixds_block
def plot(obs_parameters='', n=0, m=0, f_rest=0, slope_correction=False, dB=False, vlsr=False, meta=False, avg_ylim=[0,0], cal_ylim=[0,0], rfi=[], xlim=[0,0], ylim=[0,0], dm=0,
obs_file='observation.dat', cal_file='', waterfall_fits='', spectra_csv='', power_csv='', plot_file='plot.png'):
'''
Process, analyze and plot data.
Args:
obs_parameters: dict. Observation parameters (identical to parameters used to acquire data)
dev_args: string. Device arguments (gr-osmosdr)
rf_gain: float. RF gain
if_gain: float. IF gain
bb_gain: float. Baseband gain
frequency: float. Center frequency [Hz]
bandwidth: float. Instantaneous bandwidth [Hz]
channels: int: Number of frequency channels (FFT size)
t_sample: float: Integration time per FFT sample
duration: float: Total observing duration [sec]
loc: string: latitude, longitude, and elevation of observation (float, separated by spaces)
ra_dec: string: right ascension and declination of observation target (float, separated by space)
az_alt: string: azimuth and altitude of observation target (float, separated by space; takes precedence over ra_dec)
n: int. Median filter factor (spectrum)
m: int. Median filter factor (time series)
f_rest: float. Spectral line reference frequency used for radial velocity (Doppler shift) calculations [Hz]
slope_correction: bool. Correct slope in poorly-calibrated spectra using linear regression
dB: bool. Display data in decibel scaling
vlsr: bool. Display graph in VLSR frame of reference
meta: bool. Display header with date, time, and target
rfi: list. Blank frequency channels contaminated with RFI ([low_frequency, high_frequency]) [Hz]
avg_ylim: list. Averaged plot y-axis limits ([low, high])
cal_ylim: list. Calibrated plot y-axis limits ([low, high])
xlim: list. x-axis limits ([low_frequency, high_frequency]) [Hz]
ylim: list. y-axis limits ([start_time, end_time]) [Hz]
dm: float. Dispersion measure for dedispersion [pc/cm^3]
obs_file: string. Input observation filename (generated with virgo.observe)
cal_file: string. Input calibration filename (generated with virgo.observe)
waterfall_fits: string. Output FITS filename
spectra_csv: string. Output CSV filename (spectra)
power_csv: string. Output CSV filename (time series)
plot_file: string. Output plot filename
'''
import matplotlib
matplotlib.use('Agg') # Try commenting this line if you run into display/rendering errors
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
plt.rcParams['legend.fontsize'] = 14
plt.rcParams['axes.labelsize'] = 14
plt.rcParams['axes.titlesize'] = 18
plt.rcParams['xtick.labelsize'] = 12
plt.rcParams['ytick.labelsize'] = 12
def decibel(x):
if dB: return 10.0*np.log10(x)
return x
def shift(phase_num, n_rows):
waterfall[:, phase_num] = np.roll(waterfall[:, phase_num], -n_rows)
def SNR(spectrum, mask=np.array([])):
'''Signal-to-Noise Ratio estimator, with optional masking.
If mask not given, then all channels will be used to estimate noise
(will drastically underestimate S:N - not robust to outliers!)'''
if mask.size == 0:
mask = np.zeros_like(spectrum)
noise = np.nanstd((spectrum[2:]-spectrum[:-2])[mask[1:-1] == 0])/np.sqrt(2)
background = np.nanmean(spectrum[mask == 0])
return (spectrum-background)/noise
def best_fit(power):
'''Compute best Gaussian fit'''
avg = np.nanmean(power)
var = np.var(power)
gaussian_fit_x = np.linspace(np.min(power),np.max(power),100)
gaussian_fit_y = 1.0/np.sqrt(2*np.pi*var)*np.exp(-0.5*(gaussian_fit_x-avg)**2/var)
return [gaussian_fit_x, gaussian_fit_y]
# Load observation parameters from dictionary argument/header file
if obs_parameters != '':
frequency = obs_parameters['frequency']
bandwidth = obs_parameters['bandwidth']
channels = obs_parameters['channels']
t_sample = obs_parameters['t_sample']
loc = obs_parameters['loc']
ra_dec = obs_parameters['ra_dec']
az_alt = obs_parameters['az_alt']
else:
header_file = '.'.join(obs_file.split('.')[:-1])+'.header'
warnings.warn('No observation parameters passed. Attempting to load from header file ('+header_file+')...')
with open(header_file, 'r') as f:
headers = [parameter.rstrip('\n') for parameter in f.readlines()]
for i in range(len(headers)):
if 'mjd' in headers[i]:
mjd = float(headers[i].strip().split('=')[1])
elif 'frequency' in headers[i]:
frequency = float(headers[i].strip().split('=')[1])
elif 'bandwidth' in headers[i]:
bandwidth = float(headers[i].strip().split('=')[1])
elif 'channels' in headers[i]:
channels = int(headers[i].strip().split('=')[1])
elif 't_sample' in headers[i]:
t_sample = float(headers[i].strip().split('=')[1])
elif 'loc' in headers[i]:
loc = tuple(map(float, headers[i].strip().split('=')[1].split(' ')))
elif 'ra_dec' in headers[i]:
ra_dec = tuple(map(str, headers[i].split('=')[1].split(' ')))
elif 'az_alt' in headers[i]:
az_alt = tuple(map(float, headers[i].split('=')[1].split(' ')))
# Transform frequency axis limits to MHz
xlim = [x / 1e6 for x in xlim]
# Transform to VLSR
if vlsr:
from astropy import units as u
from astropy.coordinates import SpectralCoord, EarthLocation, SkyCoord
from astropy.time import Time
obs_location = EarthLocation.from_geodetic(loc[0], loc[1], loc[2])
obs_time = obs_location.get_itrs(obstime=Time(str(mjd), format='mjd', scale='utc'))
if az_alt!='':
obs_coord = SkyCoord(az=az_alt[0]*u.degree, alt=az_alt[1]*u.degree, frame='altaz', location=obs_location, obstime=Time(str(mjd), format='mjd', scale='utc'))
obs_coord = obs_coord.icrs
print (obs_coord)
else:
obs_coord = SkyCoord(ra=ra_dec[0]*u.degree, dec=ra_dec[1]*u.degree, frame='icrs')
#Transform center frequency
frequency = SpectralCoord(frequency * u.MHz, observer=obs_time, target=obs_coord)
frequency = frequency.with_observer_stationary_relative_to('lsrk')
frequency = frequency.quantity.value
# Define Radial Velocity axis limits
left_velocity_edge = -299792.458*(bandwidth-2*frequency+2*f_rest)/(bandwidth-2*frequency)
right_velocity_edge = 299792.458*(-bandwidth-2*frequency+2*f_rest)/(bandwidth+2*frequency)
# Transform sampling time to number of bins
bins = int(t_sample*bandwidth/channels)
# Load observation & calibration data
offset = 1
waterfall = offset*np.fromfile(obs_file, dtype='float32').reshape(-1, channels)/bins
# Delete first 3 rows (potentially containing outlier samples)
waterfall = waterfall[3:, :]
# Mask RFI-contaminated channels
if rfi != []:
for j in range(len(rfi)):
# Frequency to channel transformation
current_rfi = rfi[j]
rfi_lo = channels*(current_rfi[0] - (frequency - bandwidth/2))/bandwidth
rfi_hi = channels*(current_rfi[1] - (frequency - bandwidth/2))/bandwidth
# Blank channels
for i in range(int(rfi_lo), int(rfi_hi)):
waterfall[:, i] = np.nan
if cal_file != '':
waterfall_cal = offset*np.fromfile(cal_file, dtype='float32').reshape(-1, channels)/bins
# Delete first 3 rows (potentially containing outlier samples)
waterfall_cal = waterfall_cal[3:, :]
# Mask RFI-contaminated channels
if rfi != []:
for j in range(len(rfi)):
# Frequency to channel transformation
current_rfi = rfi[j]
rfi_lo = channels*(current_rfi[0] - (frequency - bandwidth/2))/bandwidth
rfi_hi = channels*(current_rfi[1] - (frequency - bandwidth/2))/bandwidth
# Blank channels
for i in range(int(rfi_lo), int(rfi_hi)):
waterfall_cal[:, i] = np.nan
# Compute average spectra
with warnings.catch_warnings():
warnings.filterwarnings(action='ignore', message='Mean of empty slice')
avg_spectrum = decibel(np.nanmean(waterfall, axis=0))
if cal_file != '':
avg_spectrum_cal = decibel(np.nanmean(waterfall_cal, axis=0))
# Number of sub-integrations
subs = waterfall.shape[0]
# Compute Time axis
t = t_sample*np.arange(subs)
# Compute Frequency axis; convert Hz to MHz
frequency = np.linspace(frequency-0.5*bandwidth, frequency+0.5*bandwidth,
channels, endpoint=False)*1e-6
# Perform de-dispersion
if dm != 0:
deltaF = float(np.max(frequency)-np.min(frequency))/subs
f_start = np.min(frequency)
for t_bin in range(subs):
f_chan = f_start+t_bin*deltaF
deltaT = 4149*dm*((1/(f_chan**2))-(1/(np.max(frequency)**2)))
n = int((float(deltaT)/(float(1)/channels)))
shift(t_bin, n)
# Define array for Time Series plot
power = decibel(np.nanmean(waterfall, axis=1))
# Apply Mask
mask = np.zeros_like(avg_spectrum)
mask[np.logical_and(frequency > f_rest*1e-6-0.2, frequency < f_rest*1e-6+0.8)] = 1 # Margins OK for galactic HI
# Define text offset for axvline text label
text_offset = 0
# Calibrate Spectrum
if cal_file != '':
if dB:
spectrum = 10**((avg_spectrum-avg_spectrum_cal)/10)
else:
spectrum = avg_spectrum/avg_spectrum_cal
spectrum = SNR(spectrum, mask)
if slope_correction:
idx = np.isfinite(frequency) & np.isfinite(spectrum)
fit = np.polyfit(frequency[idx], spectrum[idx], 1)
ang_coeff = fit[0]
intercept = fit[1]
fit_eq = ang_coeff*frequency + intercept
spectrum = SNR(spectrum-fit_eq, mask)
# Mitigate RFI (Frequency Domain)
if n != 0:
spectrum_clean = SNR(spectrum.copy(), mask)
for i in range(0, int(channels)):
spectrum_clean[i] = np.nanmedian(spectrum_clean[i:i+n])
# Apply position offset for Spectral Line label
text_offset = 60
# Mitigate RFI (Time Domain)
if m != 0:
power_clean = power.copy()
for i in range(0, int(subs)):
power_clean[i] = np.nanmedian(power_clean[i:i+m])
# Write Waterfall to file (FITS)
if waterfall_fits != '':
from astropy.io import fits
# Load data
hdu = fits.PrimaryHDU(waterfall)
# Prepare FITS headers
hdu.header['NAXIS'] = 2
hdu.header['NAXIS1'] = channels
hdu.header['NAXIS2'] = subs
hdu.header['CRPIX1'] = channels/2
hdu.header['CRPIX2'] = subs/2
hdu.header['CRVAL1'] = frequency[int(channels/2)]
hdu.header['CRVAL2'] = t[int(subs/2)]
hdu.header['CDELT1'] = bandwidth*1e-6/channels
hdu.header['CDELT2'] = t_sample
hdu.header['CTYPE1'] = 'Frequency (MHz)'
hdu.header['CTYPE2'] = 'Relative Time (s)'
try:
hdu.header['MJD-OBS'] = mjd
except NameError:
warnings.warn('Observation MJD could not be found and will not be part of the FITS header.')
pass
# Delete pre-existing FITS file
try:
os.remove(waterfall_fits)
except OSError:
pass
# Write to file
hdu.writeto(waterfall_fits)
# Write Spectra to file (csv)
if spectra_csv != '':
if cal_file != '':
np.savetxt(spectra_csv, np.concatenate((frequency.reshape(channels, 1),
avg_spectrum.reshape(channels, 1), avg_spectrum_cal.reshape(channels, 1),
spectrum.reshape(channels, 1)), axis=1), delimiter=',', fmt='%1.6f')
else:
np.savetxt(spectra_csv, np.concatenate((frequency.reshape(channels, 1),
avg_spectrum.reshape(channels, 1)), axis=1), delimiter=',', fmt='%1.6f')
# Write Time Series to file (csv)
if power_csv != '':
np.savetxt(power_csv, np.concatenate((t.reshape(subs, 1), power.reshape(subs, 1)),
axis=1), delimiter=',', fmt='%1.6f')
# Initialize plot
if cal_file != '':
fig = plt.figure(figsize=(27, 15))
gs = GridSpec(2, 3)
else:
fig = plt.figure(figsize=(21, 15))
gs = GridSpec(2, 2)
if meta:
from astropy.coordinates import get_constellation
epoch = (mjd - 40587) * 86400.0
meta_title = 'Date and Time: ' + time.strftime('%Y-%m-%d %H:%M:%S %Z', time.localtime(epoch)) + ' '
meta_title += 'Target: ' + obs_coord.to_string('hmsdms', precision=0) + ' in ' + get_constellation(obs_coord) + '\n'
plt.suptitle(meta_title, fontsize=18)
# Plot Average Spectrum
ax1 = fig.add_subplot(gs[0, 0])
ax1.plot(frequency, avg_spectrum)
if xlim == [0,0]:
ax1.set_xlim(np.min(frequency), np.max(frequency))
else:
ax1.set_xlim(xlim[0], xlim[1])
ax1.ticklabel_format(useOffset=False)
ax1.set_xlabel('Frequency (MHz)')
if avg_ylim != [0,0]:
ax1.set_ylim(avg_ylim[0], avg_ylim[1])
if dB:
ax1.set_ylabel('Relative Power (dB)')
else:
ax1.set_ylabel('Relative Power')
if vlsr:
cal_title = r'$Average\ Spectrum\ (V_{LSR})$'
else:
cal_title = 'Average Spectrum'
if f_rest != 0:
cal_title += '\n'
ax1.set_title(cal_title)
ax1.grid()
if xlim == [0,0] and f_rest != 0:
# Add secondary axis for Radial Velocity
ax1_secondary = ax1.twiny()
ax1_secondary.set_xlabel('Radial Velocity (km/s)', labelpad=5)
ax1_secondary.axvline(x=0, color='brown', linestyle='--', linewidth=2, zorder=0)
ax1_secondary.annotate('Spectral Line\nRest Frequency', xy=(460-text_offset, 5),
xycoords='axes points', size=14, ha='left', va='bottom', color='brown')
ax1_secondary.set_xlim(left_velocity_edge, right_velocity_edge)
ax1_secondary.tick_params(axis='x', direction='in', pad=-22)
#Plot Calibrated Spectrum
if cal_file != '':
ax2 = fig.add_subplot(gs[0, 1])
ax2.plot(frequency, spectrum, label='Raw Spectrum')
if n != 0:
ax2.plot(frequency, spectrum_clean, color='orangered', label='Median (n = '+str(n)+')')
if cal_ylim !=[0,0]:
ax2.set_ylim(cal_ylim[0],cal_ylim[1])
else:
ax2.set_ylim()
if xlim == [0,0]:
ax2.set_xlim(np.min(frequency), np.max(frequency))
else:
ax2.set_xlim(xlim[0], xlim[1])
ax2.ticklabel_format(useOffset=False)
ax2.set_xlabel('Frequency (MHz)')
ax2.set_ylabel('Signal-to-Noise Ratio (S/N)')
if vlsr:
cal_title = r'$Calibrated\ Spectrum\ (V_{LSR})$' + '\n'
else:
cal_title = 'Calibrated Spectrum\n'
if f_rest != 0:
ax2.set_title(cal_title)
else:
ax2.set_title('Calibrated Spectrum')
if n != 0:
if f_rest != 0:
ax2.legend(bbox_to_anchor=(0.002, 0.96), loc='upper left')
else:
ax2.legend(loc='upper left')
if xlim == [0,0] and f_rest != 0:
# Add secondary axis for Radial Velocity
ax2_secondary = ax2.twiny()
ax2_secondary.set_xlabel('Radial Velocity (km/s)', labelpad=5)
ax2_secondary.axvline(x=0, color='brown', linestyle='--', linewidth=2, zorder=0)
ax2_secondary.annotate('Spectral Line\nRest Frequency', xy=(400, 5),
xycoords='axes points', size=14, ha='left', va='bottom', color='brown')
ax2_secondary.set_xlim(left_velocity_edge, right_velocity_edge)
ax2_secondary.tick_params(axis='x', direction='in', pad=-22)
ax2.grid()
# Plot Dynamic Spectrum
if cal_file != '':
ax3 = fig.add_subplot(gs[0, 2])
else:
ax3 = fig.add_subplot(gs[0, 1])
ax3.imshow(decibel(waterfall), origin='lower', interpolation='None', aspect='auto',
extent=[np.min(frequency), np.max(frequency), np.min(t), np.max(t)])
if xlim == [0,0] and ylim != [0,0]:
ax3.set_ylim(ylim[0], ylim[1])
elif xlim != [0,0] and ylim == [0,0]:
ax3.set_xlim(xlim[0], xlim[1])
elif xlim != [0,0] and ylim != [0,0]:
ax3.set_xlim(xlim[0], xlim[1])
ax3.set_ylim(ylim[0], ylim[1])
ax3.ticklabel_format(useOffset=False)
ax3.set_xlabel('Frequency (MHz)')
ax3.set_ylabel('Relative Time (s)')
ax3.set_title('Dynamic Spectrum (Waterfall)')
# Adjust Subplot Width Ratio
if cal_file != '':
gs = GridSpec(2, 3, width_ratios=[16.5, 1, 1])
else:
gs = GridSpec(2, 2, width_ratios=[7.6, 1])
# Plot Time Series (Power vs Time)
ax4 = fig.add_subplot(gs[1, 0])
ax4.plot(t, power, label='Raw Time Series')
if m != 0:
ax4.plot(t, power_clean, color='orangered', label='Median (n = '+str(m)+')')
ax4.set_ylim()
if ylim == [0,0]:
ax4.set_xlim(0, np.max(t))
else:
ax4.set_xlim(ylim[0], ylim[1])
ax4.set_xlabel('Relative Time (s)')
if dB:
ax4.set_ylabel('Relative Power (dB)')
else:
ax4.set_ylabel('Relative Power')
ax4.set_title('Average Power vs Time')
if m != 0:
ax4.legend(bbox_to_anchor=(1, 1), loc='upper right')
ax4.grid()
# Plot Total Power Distribution
if cal_file != '':
gs = GridSpec(2, 3, width_ratios=[7.83, 1.5, -0.325])
else:
gs = GridSpec(2, 2, width_ratios=[8.8, 1.5])
ax5 = fig.add_subplot(gs[1, 1])
ax5.hist(power, np.max([int(np.size(power)/50),10]), density=1, alpha=0.5, color='royalblue', orientation='horizontal', zorder=10)
ax5.plot(best_fit(power)[1], best_fit(power)[0], '--', color='blue', label='Best fit (Raw)', zorder=20)
if m != 0:
ax5.hist(power_clean, np.max([int(np.size(power_clean)/50),10]), density=1, alpha=0.5, color='orangered', orientation='horizontal', zorder=10)
ax5.plot(best_fit(power_clean)[1], best_fit(power_clean)[0], '--', color='red', label='Best fit (Median)', zorder=20)
ax5.set_xlim()
ax5.set_ylim()
ax5.get_shared_x_axes().join(ax5, ax4)
ax5.set_yticklabels([])
ax5.set_xlabel('Probability Density')
ax5.set_title('Total Power Distribution')
ax5.legend(bbox_to_anchor=(1, 1), loc='upper right')
ax5.grid()
# Save plots to file
plt.tight_layout()
plt.savefig(plot_file)
plt.clf()