def main():
    """Runs test imaging pipeline using MPI."""
    # Check command line arguments.
    if len(sys.argv) < 2:
        raise RuntimeError(
            'Usage: mpiexec -n <np> '
            'python mpi_imager_test.py <settings_file> <dir>')

    # Get the MPI communicator and initialise broadcast variables.
    comm = MPI.COMM_WORLD
    settings = None
    inputs = None
    grid_weights = None

    # Create log.
    log = logging.getLogger()
    log.setLevel(logging.DEBUG)
    if len(log.handlers) == 0:
        log.addHandler(logging.StreamHandler(sys.stdout))

    if comm.Get_rank() == 0:
        # Load pipeline settings.
        with open(sys.argv[1]) as f:
            settings = json.load(f)

        # Get a list of input Measurement Sets to process.
        data_dir = str(sys.argv[2])
        inputs = glob(join(data_dir, '*.ms')) + glob(join(data_dir, '*.MS'))
        inputs = filter(None, inputs)
        log.info('Found input Measurement Sets: %s', ', '.join(inputs))

        # Distribute the list of Measurement Sets among processors.
        inputs = chunks(inputs, comm.Get_size())

    # Broadcast settings and scatter list of input files.
    comm.barrier()
    settings = comm.bcast(settings)
    inputs = comm.scatter(inputs)

    # Record which file(s) this node is working on.
    log.debug('Rank %d, processing [%s]', comm.Get_rank(), ', '.join(inputs))

    # Create an imager and configure it.
    precision = settings['precision']
    imager = oskar.Imager(precision)
    for key, value in settings['imager'].items():
        setattr(imager, key, value)

    # Allocate a local visibility grid.
    grid_norm = 0.
    grid_dim = [imager.plane_size, imager.plane_size]
    grid_data = numpy.zeros(grid_dim,
                            dtype='c8' if precision == 'single' else 'c16')

    # Process data according to mode.
    if settings['combine']:
        if imager.weighting == 'Uniform' or imager.algorithm == 'W-projection':
            # If necessary, generate a local weights grid.
            local_weights = None
            if imager.weighting == 'Uniform':
                grid_weights = numpy.zeros(grid_dim, dtype=precision)
                local_weights = numpy.zeros(grid_dim, dtype=precision)

            # Do a first pass for uniform weighting or W-projection.
            imager.coords_only = True
            for f in inputs:
                log.info('Reading coordinates from %s', f)
                process_input_data(f, imager, None, 0.0, local_weights)
            imager.coords_only = False

            # Get maximum number of W-projection planes required.
            num_w_planes = imager.num_w_planes
            num_w_planes = comm.allreduce(num_w_planes, op=MPI.MAX)
            imager.num_w_planes = num_w_planes

            # Combine (reduce) weights grids, and broadcast the result.
            if local_weights is not None:
                comm.Allreduce(local_weights, grid_weights, op=MPI.SUM)

        # Populate the local visibility grid.
        for f in inputs:
            log.info('Reading visibilities from %s', f)
            grid_norm = process_input_data(f, imager, grid_data, grid_norm,
                                           grid_weights)

        # Combine (reduce) visibility grids.
        grid = numpy.zeros_like(grid_data)
        comm.Reduce(grid_data, grid, op=MPI.SUM)
        grid_norm = comm.reduce(grid_norm, op=MPI.SUM)

        # Finalise grid and save image.
        if comm.Get_rank() == 0:
            save_image(imager, grid, grid_norm, settings['output_file'])
            log.info('Finished. Output file is %s', settings['output_file'])
    else:
        for f in inputs:
            # Clear the grid.
            grid_norm = 0.
            grid_data.fill(0)
            if imager.weighting == 'Uniform':
                grid_weights = numpy.zeros(grid_dim, dtype=precision)

            # Do a first pass for uniform weighting or W-projection.
            if imager.weighting == 'Uniform' or \
                    imager.algorithm == 'W-projection':
                imager.coords_only = True
                log.info('Reading coordinates from %s', f)
                process_input_data(f, imager, None, 0.0, grid_weights)
                imager.coords_only = False

            # Populate the local visibility grid.
            log.info('Reading visibilities from %s', f)
            grid_norm = process_input_data(f, imager, grid_data, grid_norm,
                                           grid_weights)

            # Save image by finalising grid.
            output_file = splitext(f)[0] + '.fits'
            save_image(imager, grid_data, grid_norm, output_file)
            log.info('Finished. Output file is %s', output_file)
Runs test imaging pipeline using MPI.