def load_im_data(results_dir):
"""Load the full-Stokes images into memory.
Args:
results_dir (str): directory to save results.
Returns:
image_temp: ARL image data.
frequency: array of observed frequencies in Hz.
weights: array of weights per channel (1/sigma**2).
"""
try:
if os.path.isdir(results_dir):
try:
# Load the channel 0 data as an image template:
image_temp = import_image_from_fits('%s/imaging_clean_WStack-%s.fits'
% (results_dir, 0))
# Fill image_temp with the multi-frequency data:
image_temp.data = np.concatenate(([import_image_from_fits(
'%s/imaging_clean_WStack-%s.fits'
% (results_dir, channel)).data
for channel in range(0, 40)]))
# Read the array of the channel frequencies:
frequency = image_temp.frequency
# Calculate the weights, in the form [channel, stokes, npix, npix]:
# Initially using std for weights, should consider more robust options.
weights = np.array([1.0/(np.std(image_temp.data[channel, 0, :, :])**2)
for channel in range(0, 40)])
except:
print("Unexpected error:", sys.exc_info()[0])
raise
except:
print("Input directory does not exist:", sys.exc_info()[0])
raise
return image_temp, frequency, weights
def test_create_gaintable_from_screen(self):
screen = import_image_from_fits(
arl_path('data/models/test_mpc_screen.fits'))
beam = create_test_image(cellsize=0.0015,
phasecentre=self.vis.phasecentre,
frequency=self.frequency)
beam = create_low_test_beam(beam)
gleam_components = create_low_test_skycomponents_from_gleam(
flux_limit=1.0,
phasecentre=self.phasecentre,
frequency=self.frequency,
polarisation_frame=PolarisationFrame('stokesI'),
radius=0.2)
pb_gleam_components = apply_beam_to_skycomponent(
gleam_components, beam)
actual_components = filter_skycomponents_by_flux(pb_gleam_components,
flux_min=1.0)
gaintables = create_gaintable_from_screen(self.vis, actual_components,
screen)
assert len(gaintables) == len(actual_components), len(gaintables)
assert gaintables[0].gain.shape == (3, 94, 1, 1,
1), gaintables[0].gain.shape
def buffer_data_model_to_memory(jbuff, dm):
"""Copy a buffer data model into memory data model
The file type is derived from the file extension. All are hdf only with the exception of Imaghe which can also be
fits.
:param jbuff: JSON describing buffer
:param dm: JSON describing data model
:return: data model
"""
import os
name = os.path.join(jbuff["directory"], dm["name"])
import os
_, file_extension = os.path.splitext(dm["name"])
if dm["data_model"] == "BlockVisibility":
return import_blockvisibility_from_hdf5(name)
elif dm["data_model"] == "Image":
if file_extension == ".fits":
return import_image_from_fits(name)
else:
return import_image_from_hdf5(name)
elif dm["data_model"] == "SkyModel":
return import_skymodel_from_hdf5(name)
elif dm["data_model"] == "GainTable":
return import_gaintable_from_hdf5(name)
else:
raise ValueError("Data model %s not supported" % dm["data_model"])
def moments_save_to_disk(moments_im, stokes_type='q', results_dir='./results_dir', outname='mean'):
"""Save the Faraday moments images to disk.
Args:
rmsynth (numpy array): array of complex numbers from RM Synthesis.
cellsize (float): advised cellsize in Faraday space.
maxrm_est (float): maximum observable RM (50 percent sensitivity).
rmtype (str): the component of the complex numbers to process and save.
results_dir (str): directory to save results.
outname (str): outname for saved file.
Returns:
None
"""
# Read in the first channel image, and appropriate it as the new moments image:
im_moments = import_image_from_fits('%s/imaging_clean_WStack-%s.fits' % (results_dir, 0))
# Place the data into the open image:
im_moments.data = moments_im
try:
if stokes_type == 'p':
stokes_val = 0.0
elif stokes_type == 'q':
stokes_val = 2.0
elif stokes_type == 'u':
stokes_val = 3.0
except:
print("Unknown value for stokes_type:", sys.exc_info()[0])
raise
# This line also adjusts the listed Stokes parameter (use 0.0=? for P):
im_moments.wcs.wcs.crval = [im_moments.wcs.wcs.crval[0], im_moments.wcs.wcs.crval[1],
stokes_val, im_moments.wcs.wcs.crval[3]]
# Output the file to disk:
export_image_to_fits(im_moments, '%s/%s_%s.fits' % (results_dir, outname, stokes_type))
return
def create_test_image(
canonical=True,
cellsize=None,
frequency=None,
channel_bandwidth=None,
phasecentre=None,
polarisation_frame=PolarisationFrame("stokesI")) -> Image:
"""Create a useful test image
This is the test image M31 widely used in ALMA and other simulations. It is actually part of an Halpha region in
M31.
:param canonical: Make the image into a 4 dimensional image
:param cellsize:
:param frequency: Frequency (array) in Hz
:param channel_bandwidth: Channel bandwidth (array) in Hz
:param phasecentre: Phase centre of image (SkyCoord)
:param polarisation_frame: Polarisation frame
:return: Image
"""
if frequency is None:
frequency = [1e8]
im = import_image_from_fits(arl_path("data/models/M31.MOD"))
if canonical:
if polarisation_frame is None:
im.polarisation_frame = PolarisationFrame("stokesI")
elif isinstance(polarisation_frame, PolarisationFrame):
im.polarisation_frame = polarisation_frame
else:
raise ValueError("polarisation_frame is not valid")
im = replicate_image(im,
frequency=frequency,
polarisation_frame=im.polarisation_frame)
if cellsize is not None:
im.wcs.wcs.cdelt[0] = -180.0 * cellsize / numpy.pi
im.wcs.wcs.cdelt[1] = +180.0 * cellsize / numpy.pi
if frequency is not None:
im.wcs.wcs.crval[3] = frequency[0]
if channel_bandwidth is not None:
im.wcs.wcs.cdelt[3] = channel_bandwidth[0]
else:
if len(frequency) > 1:
im.wcs.wcs.cdelt[3] = frequency[1] - frequency[0]
else:
im.wcs.wcs.cdelt[3] = 0.001 * frequency[0]
im.wcs.wcs.radesys = 'ICRS'
im.wcs.wcs.equinox = 2000.00
if phasecentre is not None:
im.wcs.wcs.crval[0] = phasecentre.ra.deg
im.wcs.wcs.crval[1] = phasecentre.dec.deg
# WCS is 1 relative
im.wcs.wcs.crpix[0] = im.data.shape[3] // 2 + 1
im.wcs.wcs.crpix[1] = im.data.shape[2] // 2 + 1
return im
def create_low_test_beam(model: Image) -> Image:
"""Create a test power beam for LOW using an image from OSKAR
:param model: Template image
:return: Image
"""
beam = import_image_from_fits(arl_path('data/models/SKA1_LOW_beam.fits'))
# Scale the image cellsize to account for the different in frequencies. Eventually we will want to
# use a frequency cube
log.info("create_low_test_beam: primary beam is defined at %.3f MHz" %
(beam.wcs.wcs.crval[2] * 1e-6))
nchan, npol, ny, nx = model.shape
# We need to interpolate each frequency channel separately. The beam is assumed to just scale with
# frequency.
reprojected_beam = create_empty_image_like(model)
for chan in range(nchan):
model2dwcs = model.wcs.sub(2).deepcopy()
model2dshape = [model.shape[2], model.shape[3]]
beam2dwcs = beam.wcs.sub(2).deepcopy()
# The frequency axis is the second to last in the beam
frequency = model.wcs.sub(['spectral']).wcs_pix2world([chan], 0)[0]
fscale = beam.wcs.wcs.crval[2] / frequency
beam2dwcs.wcs.cdelt = fscale * beam.wcs.sub(2).wcs.cdelt
beam2dwcs.wcs.crpix = beam.wcs.sub(2).wcs.crpix
beam2dwcs.wcs.crval = model.wcs.sub(2).wcs.crval
beam2dwcs.wcs.ctype = model.wcs.sub(2).wcs.ctype
model2dwcs.wcs.crpix = [
model.shape[2] // 2 + 1, model.shape[3] // 2 + 1
]
beam2d = create_image_from_array(beam.data[0, 0, :, :], beam2dwcs,
model.polarisation_frame)
reprojected_beam2d, footprint = reproject_image(beam2d,
model2dwcs,
shape=model2dshape)
assert numpy.max(
footprint.data) > 0.0, "No overlap between beam and model"
reprojected_beam2d.data *= reprojected_beam2d.data
reprojected_beam2d.data[footprint.data <= 0.0] = 0.0
for pol in range(npol):
reprojected_beam.data[chan,
pol, :, :] = reprojected_beam2d.data[:, :]
return reprojected_beam
def test_create_gradient(self):
real_vp = import_image_from_fits(
arl_path('data/models/MID_GRASP_VP_real.fits'))
gradx, grady = image_gradients(real_vp)
gradxx, gradxy = image_gradients(gradx)
gradyx, gradyy = image_gradients(grady)
gradx.data *= real_vp.data
grady.data *= real_vp.data
gradxx.data *= real_vp.data
gradxy.data *= real_vp.data
gradyx.data *= real_vp.data
gradyy.data *= real_vp.data
import matplotlib.pyplot as plt
plt.clf()
show_image(gradx, title='gradx')
plt.show()
plt.clf()
show_image(grady, title='grady')
plt.show()
export_image_to_fits(gradx,
"%s/test_image_gradients_gradx.fits" % (self.dir))
export_image_to_fits(grady,
"%s/test_image_gradients_grady.fits" % (self.dir))
plt.clf()
show_image(gradxx, title='gradxx')
plt.show()
plt.clf()
show_image(gradxy, title='gradxy')
plt.show()
plt.clf()
show_image(gradyx, title='gradyx')
plt.show()
plt.clf()
show_image(gradyy, title='gradyy')
plt.show()
export_image_to_fits(
gradxx, "%s/test_image_gradients_gradxx.fits" % (self.dir))
export_image_to_fits(
gradxy, "%s/test_image_gradients_gradxy.fits" % (self.dir))
export_image_to_fits(
gradyx, "%s/test_image_gradients_gradyx.fits" % (self.dir))
export_image_to_fits(
gradyy, "%s/test_image_gradients_gradyy.fits" % (self.dir))
def rmcube_save_to_disk(rmsynth, cellsize, maxrm_est, rmtype='abs', results_dir='./results_dir', outname='dirty'):
"""Save the RM cubes to disk.
Args:
rmsynth (numpy array): array of complex numbers from RM Synthesis.
cellsize (float): advised cellsize in Faraday space.
maxrm_est (float): maximum observable RM (50 percent sensitivity).
rmtype (str): the component of the complex numbers to process and save.
results_dir (str): directory to save results.
outname (str): outname for saved file.
"""
# Read in the first channel image, and appropriate it as the new RM cube:
im_rmsynth = import_image_from_fits('%s/imaging_clean_WStack-%s.fits' % (results_dir, 0))
# Output the polarised data:
try:
if rmtype == 'abs':
im_rmsynth.data = np.abs(rmsynth)
stokes_val = 0.0
elif rmtype == 'real':
im_rmsynth.data = np.real(rmsynth)
stokes_val = 2.0
elif rmtype == 'imag':
im_rmsynth.data = np.imag(rmsynth)
stokes_val = 3.0
except:
print("Unknown value for rmtype:", sys.exc_info()[0])
raise
# Adjust the various axes of the cube:
im_rmsynth.wcs.wcs.ctype = [im_rmsynth.wcs.wcs.ctype[0], im_rmsynth.wcs.wcs.ctype[1],
'FARDEPTH', im_rmsynth.wcs.wcs.ctype[2]]
im_rmsynth.wcs.wcs.cdelt = [im_rmsynth.wcs.wcs.cdelt[0], im_rmsynth.wcs.wcs.cdelt[1],
cellsize, im_rmsynth.wcs.wcs.cdelt[2]]
im_rmsynth.wcs.wcs.crpix = [im_rmsynth.wcs.wcs.crpix[0], im_rmsynth.wcs.wcs.crpix[1],
1.0, im_rmsynth.wcs.wcs.crpix[2]]
im_rmsynth.wcs.wcs.cunit = [im_rmsynth.wcs.wcs.cunit[0], im_rmsynth.wcs.wcs.cunit[1],
'rad / m^2', im_rmsynth.wcs.wcs.cunit[2]]
# This line also adjusts the listed Stokes parameter (use 0.0=? for P):
im_rmsynth.wcs.wcs.crval = [im_rmsynth.wcs.wcs.crval[0], im_rmsynth.wcs.wcs.crval[1],
-maxrm_est, stokes_val]
# Tweak the axes into a more sensible order:
im_rmsynth.data = np.rollaxis(im_rmsynth.data, 2, 0)
im_rmsynth.data = np.rollaxis(im_rmsynth.data, 2, 1)
# Output the file to disk:
export_image_to_fits(im_rmsynth, '%s/rmsynth-%s-%s.fits' % (results_dir, rmtype, outname))
return
def test_grid_gaintable_to_screen(self):
screen = import_image_from_fits(
arl_path('data/models/test_mpc_screen.fits'))
beam = create_test_image(cellsize=0.0015,
phasecentre=self.vis.phasecentre,
frequency=self.frequency)
beam = create_low_test_beam(beam, use_local=False)
gleam_components = create_low_test_skycomponents_from_gleam(
flux_limit=1.0,
phasecentre=self.phasecentre,
frequency=self.frequency,
polarisation_frame=PolarisationFrame('stokesI'),
radius=0.2)
pb_gleam_components = apply_beam_to_skycomponent(
gleam_components, beam)
actual_components = filter_skycomponents_by_flux(pb_gleam_components,
flux_min=1.0)
gaintables = create_gaintable_from_screen(self.vis, actual_components,
screen)
assert len(gaintables) == len(actual_components), len(gaintables)
assert gaintables[0].gain.shape == (3, 94, 3, 1,
1), gaintables[0].gain.shape
newscreen = create_empty_image_like(screen)
newscreen, weights = grid_gaintable_to_screen(self.vis, gaintables,
newscreen)
assert numpy.max(numpy.abs(screen.data)) > 0.0
if self.persist:
export_image_to_fits(
newscreen,
arl_path('test_results/test_mpc_screen_gridded.fits'))
if self.persist:
export_image_to_fits(
weights,
arl_path('test_results/test_mpc_screen_gridded_weights.fits'))
def test_read_screen(self):
screen = import_image_from_fits(
arl_path('data/models/test_mpc_screen.fits'))
assert screen.data.shape == (1, 3, 2000, 2000), screen.data.shape
def main(args):
"""
Initialising launch sequence.
"""
# ------------------------------------------------------
# Print some stuff to show that the code is running:
print("")
os.system(
"printf 'A demonstration of a \033[5mDPrepB/DPrepC\033[m SDP pipeline\n'"
)
print("")
# Set the directory for the moment images:
MOMENTS_DIR = args.outputs + '/MOMENTS'
# Check that the output directories exist, if not then create:
os.makedirs(args.outputs, exist_ok=True)
os.makedirs(MOMENTS_DIR, exist_ok=True)
# Set the polarisation definition of the instrument:
POLDEF = init_inst(args.inst)
# Setup Variables for SIP services
# ------------------------------------------------------
# Define the Queue Producer settings:
if args.queues:
queue_settings = {
'bootstrap.servers': 'scheduler:9092',
'message.max.bytes': 100000000
} #10.60.253.31:9092
# Setup the Confluent Kafka Queue
# ------------------------------------------------------
if args.queues:
from confluent_kafka import Producer
import pickle
# Create an SDP queue:
sip_queue = Producer(queue_settings)
# Define a Data Array Format
# ------------------------------------------------------
def gen_data(channel):
return np.array([
vis1[channel], vis2[channel], channel, None, None, False, False,
args.plots,
float(args.uvcut),
float(args.pixels), POLDEF, args.outputs,
float(args.angres), None, None, None, None, None, None, args.twod,
npixel_advice, cell_advice
])
# Setup the Dask Cluster
# ------------------------------------------------------
starttime = t.time()
dask.config.set(get=dask.distributed.Client.get)
client = Client(
args.daskaddress) # scheduler for Docker container, localhost for P3.
print("Dask Client details:")
print(client)
print("")
# Define channel range for 1 subband, each containing 40 channels:
channel_range = np.array(range(int(args.channels)))
# Load the data into memory:
"""
The input data should be interfaced with Buffer Management.
"""
print("Loading data:")
print("")
vis1 = [
load('%s/%s' % (args.inputs, args.ms1), range(channel, channel + 1),
POLDEF) for channel in range(0, int(args.channels))
]
vis2 = [
load('%s/%s' % (args.inputs, args.ms2), range(channel, channel + 1),
POLDEF) for channel in range(0, int(args.channels))
]
# Prepare Measurement Set
# ------------------------------------------------------
# Combine MSSS snapshots:
vis_advice = append_visibility(vis1[0], vis2[0])
# Apply a uv-distance cut to the data:
vis_advice = uv_cut(vis_advice, float(args.uvcut))
npixel_advice, cell_advice = uv_advice(vis_advice, float(args.uvcut),
float(args.pixels))
# Begin imaging via the Dask cluster
# ------------------------------------------------------
# Submit data for each channel to the client, and return an image:
# Scatter all the data in advance to all the workers:
"""
The data here could be passed via Data Queues.
Queues may not be ideal. Data throughput challenges.
Need to think more about the optimum approach.
"""
print("Scatter data to workers:")
print("")
big_job = [client.scatter(gen_data(channel)) for channel in channel_range]
# Submit jobs to the cluster and create a list of futures:
futures = [
client.submit(dprepb_imaging, big_job[channel], pure=False, retries=3)
for channel in channel_range
]
"""
The dprepb_imaging function could generate QA, logging, and pass this information via Data Queues.
Queues work well for this.
Python logging calls are preferable. Send them to a text file on the node.
Run another service that watches that file. Or just read from standard out.
The Dockerisation will assist with logs.
"""
print("Imaging on workers:")
# Watch progress:
progress(futures)
# Wait until all futures are complete:
wait(futures)
# Check that no futures have errors, if so resubmit:
for future in futures:
if future.status == 'error':
print("ERROR: Future", future, "has 'error' status, as:")
print(client.recreate_error_locally(future))
print("Rerunning...")
print("")
index = futures.index(future)
futures[index].cancel()
futures[index] = client.submit(dprepb_imaging,
big_job[index],
pure=False,
retries=3)
# Wait until all futures are complete:
wait(futures)
# Gather results from the futures:
results = client.gather(futures, errors='raise')
# Run QA on ARL objects and produce to queue:
if args.queues:
print("Adding QA to queue:")
for result in results:
sip_queue.produce('qa', pickle.dumps(qa_image(result), protocol=2))
sip_queue.flush()
# Return the data element of each ARL object, as a Dask future:
futures = [
client.submit(arl_data_future, result, pure=False, retries=3)
for result in results
]
progress(futures)
wait(futures)
# Calculate the Moment images
# ------------------------------------------------------
# Now use 'distributed Dask arrays' in order to parallelise the Moment image calculation:
# Construct a small Dask array for every future:
print("")
print("Calculating Moment images:")
print("")
arrays = [
da.from_delayed(future,
dtype=np.dtype('float64'),
shape=(1, 4, 512, 512)) for future in futures
]
# Stack all small Dask arrays into one:
stack = da.stack(arrays, axis=0)
# Combine chunks to reduce overhead - is initially (40, 1, 4, 512, 512):
stack = stack.rechunk((1, 1, 4, 64, 64))
# Spread the data around on the cluster:
stack = client.persist(stack)
# Data is now coordinated by the single logical Dask array, 'stack'.
# Save the Moment images:
"""
The output moment images should be interfaced with Buffer Management.
Need to know more about the Buffer specification.
Related to initial data distribution also/staging.
"""
print("Saving Moment images to disk:")
print("")
# First generate a template:
image_template = import_image_from_fits('%s/imaging_dirty_WStack-%s.fits' %
(args.outputs, 0))
# Output mean images:
# I:
image_template.data = stack[:, :, 0, :, :].mean(axis=0).compute()
# Run QA on ARL objects and produce to queue:
if args.queues:
sip_queue.produce('qa',
pickle.dumps(qa_image(image_template), protocol=2))
# Export the data to disk:
export_image_to_fits(image_template,
'%s/Mean-%s.fits' % (MOMENTS_DIR, 'I'))
# Q:
image_template.data = stack[:, :, 1, :, :].mean(axis=0).compute()
# Run QA on ARL objects and produce to queue:
if args.queues:
sip_queue.produce('qa',
pickle.dumps(qa_image(image_template), protocol=2))
# Export the data to disk:
export_image_to_fits(image_template,
'%s/Mean-%s.fits' % (MOMENTS_DIR, 'Q'))
# U:
image_template.data = stack[:, :, 2, :, :].mean(axis=0).compute()
# Run QA on ARL objects and produce to queue:
if args.queues:
sip_queue.produce('qa',
pickle.dumps(qa_image(image_template), protocol=2))
# Export the data to disk:
export_image_to_fits(image_template,
'%s/Mean-%s.fits' % (MOMENTS_DIR, 'U'))
# P:
image_template.data = da.sqrt(
(da.square(stack[:, :, 1, :, :]) +
da.square(stack[:, :, 2, :, :]))).mean(axis=0).compute()
# Run QA on ARL objects and produce to queue:
if args.queues:
sip_queue.produce('qa',
pickle.dumps(qa_image(image_template), protocol=2))
# Export the data to disk:
export_image_to_fits(image_template,
'%s/Mean-%s.fits' % (MOMENTS_DIR, 'P'))
# Output standard deviation images:
# I:
image_template.data = stack[:, :, 0, :, :].std(axis=0).compute()
# Run QA on ARL objects and produce to queue:
if args.queues:
sip_queue.produce('qa',
pickle.dumps(qa_image(image_template), protocol=2))
# Export the data to disk:
export_image_to_fits(image_template, '%s/Std-%s.fits' % (MOMENTS_DIR, 'I'))
# Q:
image_template.data = stack[:, :, 1, :, :].std(axis=0).compute()
# Run QA on ARL objects and produce to queue:
if args.queues:
sip_queue.produce('qa',
pickle.dumps(qa_image(image_template), protocol=2))
# Export the data to disk:
export_image_to_fits(image_template, '%s/Std-%s.fits' % (MOMENTS_DIR, 'Q'))
# U:
image_template.data = stack[:, :, 2, :, :].std(axis=0).compute()
# Run QA on ARL objects and produce to queue:
if args.queues:
sip_queue.produce('qa',
pickle.dumps(qa_image(image_template), protocol=2))
# Export the data to disk:
export_image_to_fits(image_template, '%s/Std-%s.fits' % (MOMENTS_DIR, 'U'))
# P:
image_template.data = da.sqrt(
(da.square(stack[:, :, 1, :, :]) +
da.square(stack[:, :, 2, :, :]))).std(axis=0).compute()
# Run QA on ARL objects and produce to queue:
if args.queues:
sip_queue.produce('qa',
pickle.dumps(qa_image(image_template), protocol=2))
# Export the data to disk:
export_image_to_fits(image_template, '%s/Std-%s.fits' % (MOMENTS_DIR, 'P'))
# Flush queue:
if args.queues:
sip_queue.flush()
# Make a tarball of moment images:
subprocess.call([
'tar', '-cvf',
'%s/moment.tar' % (MOMENTS_DIR),
'%s/' % (MOMENTS_DIR)
])
subprocess.call(['gzip', '-9f', '%s/moment.tar' % (MOMENTS_DIR)])
endtime = t.time()
print(endtime - starttime)
