def diffract(parameters):
    """
    Calculate all the diffraction patterns based on the parameters provided as a dictionary.
    Save all results in one single file. Not used in MPI.

    :param parameters: dictionary of command line arguments
    """

    pmi_start_id = int(parameters['pmiStartID'])
    pmi_end_id = int(parameters['pmiEndID'])
    num_dp = int(parameters['numDP'])
    ntasks = (pmi_end_id - pmi_start_id + 1) * num_dp
    rotation_axis = parameters['rotationAxis']
    uniform_rotation = parameters['uniformRotation']
    myquaternions = generate_rotations(uniform_rotation, rotation_axis, ntasks)
    output_name = parameters['outputDir'] + '/diffr_out_0000001.h5'
    if os.path.exists(output_name):
        os.remove(output_name)
    pu.prep_h5(output_name)
    for ntask in range(ntasks):
        make_one_diffr(myquaternions, ntask, parameters, output_name)
def master_diffract(comm, parameters):
    """
    Master node. Get the diffraction patterns with mpi

    :param comm: MPI comm
    :param parameters: dictionary of command line arguments
    :return:
    """
    pmi_start_id = int(parameters['pmiStartID'])
    pmi_end_id = int(parameters['pmiEndID'])
    num_dp = int(parameters['numDP'])

    # Number of processes
    num_process = comm.Get_size()
    ntasks = (pmi_end_id - pmi_start_id + 1) * num_dp

    if num_process == 1:
        rotation_axis = parameters['rotationAxis']
        uniform_rotation = parameters['uniformRotation']
        my_quaternion = generate_rotations(uniform_rotation, rotation_axis,
                                           ntasks)
        output_name = parameters['outputDir'] + '/diffr_out_0000001.h5'
        if os.path.exists(output_name):
            os.remove(output_name)
        pu.prep_h5(output_name)
        for ntask in range(ntasks):
            make_one_diffr(my_quaternion, ntask, parameters, output_name)
    else:
        for ntask in range(ntasks):
            status = MPI.Status()
            # Waiting for messages from slave
            # Successful message reciving means slave is ready for simulation
            comm.recv(source=MPI.ANY_SOURCE, tag=1, status=status)
            rnk = status.Get_source()
            # Trigger calculation on slave.
            comm.send(ntask, dest=rnk)

    # Final send: stop all processes from waiting for tasks
    for process in range(1, num_process):
        comm.send(-1, dest=process)
def slave_diffract(comm, parameters):
    """
    Slave node. Get the diffraction patterns with mpi

    :param comm: MPI comm
    :param parameters: dictionary of command line arguments
    :return:
    """
    pmi_start_id = int(parameters['pmiStartID'])
    pmi_end_id = int(parameters['pmiEndID'])
    num_dp = int(parameters['numDP'])
    ntasks = (pmi_end_id - pmi_start_id + 1) * num_dp
    rotation_axis = parameters['rotationAxis']
    uniform_rotation = parameters['uniformRotation']
    my_quaternion = generate_rotations(uniform_rotation, rotation_axis, ntasks)

    # Setup output file
    output_name = parameters['outputDir'] + '/diffr_out_' + '{0:07}'.format(
        comm.Get_rank()) + '.h5'
    if os.path.exists(output_name):
        os.remove(output_name)
    pu.prep_h5(output_name)

    # Init a local counter
    counter = 0
    # Wave to master, we're good to go.
    comm.send(counter, dest=0, tag=1)
    # Start event loop and generate the diffraction images.
    while True:
        counter = comm.recv(source=0)
        if counter < 0:
            # end of simulation
            return None
        make_one_diffr(my_quaternion, counter, parameters, output_name)
        # Show master we're ready for another task
        comm.send(counter, dest=0, tag=1)