def get_one_data_set(D, i_exp, get_po: bool = True):
    """ Get data for one experiment. Adapted from https://github.com/clinicalml/cfrnet"""
    D_exp = {}
    D_exp['X'] = D['X'][:, :, i_exp - 1]
    D_exp['w'] = D['w'][:, i_exp - 1:i_exp]
    D_exp['y'] = D['y'][:, i_exp - 1:i_exp]
    if D['HAVE_TRUTH']:
        D_exp['ycf'] = D['ycf'][:, i_exp - 1:i_exp]
    else:
        D_exp['ycf'] = None

    if get_po:
        D_exp['mu0'] = D['mu0'][:, i_exp - 1:i_exp]
        D_exp['mu1'] = D['mu1'][:, i_exp - 1:i_exp]

    return D_exp