ttcreator.py 文件源码

def min_traveltimes(modelname, source_depths, receiver_depths, distances, phases, callback=None):
    '''computes the minimum travel times using multiprocessing to speed up
    calculations
    min_traveltimes(modelname, ARRAY[N], ARRAY[N], SCALAR, ...) -> [N-length array]
    min_traveltimes(modelname, ARRAY[N], ARRAY[N], ARRAY[M]) -> [NxM] matrix
    min_traveltimes(modelname, SCALAR, SCALAR, ARRAY[M]) -> [M-length array]
    '''
    model = taumodel(modelname)
    source_depths, receiver_depths = np.broadcast_arrays(source_depths, receiver_depths)
    # assert we passed arrays, not matrices:
    if len(source_depths.shape) > 1 or len(distances.shape) > 1:
        raise ValueError("Need to have arrays, not matrices")
    norowdim = source_depths.ndim == 0
    if norowdim:
        source_depths = np.array([source_depths])
        receiver_depths = np.array([receiver_depths])
    nocoldim = distances.ndim == 0
    if nocoldim:
        distances = np.array([distances])
    ttimes = np.full(shape=(source_depths.shape[0], distances.shape[0]), fill_value=np.nan)

    def mp_callback(index, array, _callback=None):
        def _(result):
            array[index] = result
            if _callback is not None:
                _callback()
        return _

    pool = Pool()
    for idx, sd, rd in zip(count(), source_depths, receiver_depths):
        tmp_ttimes = ttimes[idx]
        for i, d in enumerate(distances):
            pool.apply_async(min_traveltime, (model, sd, rd, d, phases),
                             callback=mp_callback(i, tmp_ttimes, callback))
    pool.close()
    pool.join()

    if norowdim or nocoldim:
        ttimes = ttimes.flatten()

    return ttimes