In [ ]:
import os
import sys
sys.path.append(os.path.join('..','..'))
import numpy
import astropy.units as u
from astropy.coordinates import SkyCoord
from wrappers.arlexecute.image.gather_scatter import image_scatter_facets, image_gather_facets
from data_models.polarisation import PolarisationFrame
from wrappers.serial.simulation.testing_support import create_test_image
from wrappers.arlexecute.execution_support.arlexecute import arlexecute
import logging
logging.basicConfig(filename='simple-dask.log',
filemode='a',
format='%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s',
datefmt='%H:%M:%S',
level=logging.DEBUG)
Set up arlexecute to use dask
In [ ]:
arlexecute.set_client(use_dask=True)
In [ ]:
frequency = numpy.array([1e8])
phasecentre = SkyCoord(ra=+15.0 * u.deg, dec=-35.0 * u.deg, frame='icrs', equinox='J2000')
model_graph = arlexecute.execute(create_test_image)(frequency=frequency, phasecentre=phasecentre, cellsize=0.001,
polarisation_frame=PolarisationFrame('stokesI'))
model = arlexecute.compute(model_graph, sync=True)
Define a simple function to take the square root of an image
In [ ]:
def imagerooter(im, **kwargs):
im.data = numpy.sqrt(numpy.abs(im.data))
return im
facets_graph = arlexecute.execute(image_scatter_facets, nout=16)(model_graph, facets=4)
root_graph = [arlexecute.execute(imagerooter)(facet) for facet in facets_graph]
gathered = arlexecute.execute(image_gather_facets)(root_graph, model_graph, facets=4)
Run the graph directly
In [ ]:
result = arlexecute.compute(gathered, sync=True)
numpy.testing.assert_array_almost_equal_nulp(result.data**2, numpy.abs(model.data), 7)
In [ ]:
arlexecute.close()