Simple demonstration of the use of arlexecute


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()