def ks_distance(p_samples, q_samples): if isinstance(p_samples, tuple): idx, p_samples = p_samples return sc.ks_2samp(p_samples, q_samples)[0]