def GP(data, init_haps, positions, N, r, times, numThreads=1, pool=None): if numThreads == 1: results = [ GP.singleLocus(data[:, (i,), :], init_haps[:, (i,)], N, r, times, pos) for i, pos in enumerate(positions) ] else: params = [ {"X": data[:, (i,), :], "initHaps": init_haps[:, (i,)], "N": N, "r": r, "times": times, "pos": pos} for i, pos in enumerate(positions) ] results = pool.map(gpHelper, params) df = pd.DataFrame(map(lambda x: x.values(), results), columns=results[0].keys()).astype(float) df.index = df.pos.astype(int) df.drop("pos", axis=1, inplace=True) return df
def gpHelper(args): return GP.singleLocus(**args)