def sample_weighted_df()

in mmcif_utils.py [0:0]

• 29 lines of code
• 10 McCabe index (conditional complexity)

def sample_weighted_df(dfs, weights_array, uniform=False):
    """sample from rotamer library based off a weighted gaussian from nearby slots"""
    n = min(df["Probabil"].to_numpy().shape[0] for df in dfs)
    dfs = [df[:n].sort_values("chi1Val") for df in dfs]

    probs = []
    for weight, df in zip(weights_array, dfs):
        probs.append(df["Probabil"].to_numpy()[:n] * weight)

    probs = np.sum(np.array(probs), axis=0) / 100
    cum_prob = np.cumsum(probs)
    cutoff = np.random.uniform(0, cum_prob.max())
    ix = np.searchsorted(cum_prob, cutoff)

    if uniform:
        # ix = np.searchsorted(cum_prob, 0.99)
        if ix == 0:
            ix = 0
        else:
            ix = np.random.randint(ix)

    means = [
        weight * df[["chi{}Val".format(i) for i in range(1, 5)]].to_numpy()[:n]
        for weight, df in zip(weights_array, dfs)
    ]
    std = [
        weight * df[["chi{}Sig".format(i) for i in range(1, 5)]].to_numpy()[:n]
        for weight, df in zip(weights_array, dfs)
    ]

    means = np.sum(means, axis=0) / 100.0
    std = np.sum(std, axis=0) / 100

    chis = std[ix] * np.random.normal(0, 1, (4,)) + means[ix]

    # chis = (360 - chis) % 360

    chis[chis > 180] = chis[chis > 180] - 360
    chis[chis < -180] = chis[chis < -180] + 360
    return chis