def sample_rotomor_angle_db()

in mmcif_utils.py [0:0]

• 37 lines of code
• 5 McCabe index (conditional complexity)

def sample_rotomor_angle_db(db, phi, psi, name, uniform=False, n=1):
    df = db[name]
    lower_phi = (phi // 10) * 10
    upper_phi = lower_phi + 10

    lower_psi = (psi // 10) * 10
    upper_psi = lower_psi + 10

    weights = [
        (10 - (phi - lower_phi)) * (10 - (psi - lower_psi)),
        (10 - (upper_phi - phi)) * (10 - (psi - lower_psi)),
        (10 - (phi - lower_phi)) * (10 - (upper_psi - psi)),
        (10 - (upper_phi - phi)) * (10 - (upper_psi - psi)),
    ]
    weights_array = weights

    weights = np.cumsum(weights)

    dfs = [
        df[(df.Phi == lower_phi) & (df.Psi == lower_psi)],
        df[(df.Phi == upper_phi) & (df.Psi == lower_psi)],
        df[(df.Phi == lower_phi) & (df.Psi == upper_psi)],
        df[(df.Phi == upper_phi) & (df.Psi == upper_psi)],
    ]

    calc = np.random.uniform(0, 100, (n,))

    if n == 1:
        idxs = np.searchsorted(weights, calc)
        chis = sample_df(dfs[idxs[0]], uniform=uniform)
        return chis
    else:
        idxs = np.searchsorted(weights, calc)
        chis = []

        for i in range(4):
            count = (idxs == i).sum()
            if count > 0:
                chi = sample_df(dfs[i], uniform=uniform, sample=count)

                if count > 1:
                    chis.extend(chi)
                else:
                    chis.append(chi)

        return chis