`diversity_stats` doesn't have a `max_cohort_size` parameter

[jonbrenas]

diversity_stats returns this kind of warning when a cohort is bigger than cohort_size: "Cohort downsampled from 16 to 10 samples. Set max_cohort_size=None to disable downsampling." However, it doesn't have a max_cohort_size, despite cohort_diversity_stats (which it calls) having it.

[adilraza99]

@jonbrenas diversity_stats is missing max_cohort_size (and min_cohort_size) in its parameters. It calls cohort_diversity_stats internally, but since these aren’t exposed, there’s no way to control the downsampling from the top-level function, even though the warning suggests setting max_cohort_size=None.

Other multi-cohort functions like pairwise_average_fst and h12_gwss already pass these through, so handling it the same way here would keep things consistent.

[jonbrenas]

Hi @adilraza99, looking more closely at the code, I think the problem is slightly more complex (and maybe dumber). cohort_size is a required parameter of diversity_stats (and cohort_diversity_stats). It trickles to _snp_calls which, if cohort_size is defined, overwrites min_cohort_size and max_cohort_size and raises the warning. Given that _snp_calls is used by other functions which could legitimately have cohort_size undefined, the warning is fine but it should be caught by cohort_diversity_stats and the advice to use max_cohort_size should be removed from a new warning.

[adilraza99]

Thanks for the guidance @jonbrenas.

I’ve updated the implementation so the warning is handled in cohort_diversity_stats and no longer suggests max_cohort_size, while keeping _snp_calls unchanged for other callers.

Everything looks consistent after testing. Happy to refine further if needed.