Apply suggestions from code review

Co-authored-by: Peter Ralph <petrel.harp@gmail.com>

Author: apragsdale

docs/stats.md

@@ -705,9 +705,9 @@ Two-locus statistics can be computed using two modes, either `site` or
 single-site statistics. Within this framework, statistics may be either
 polarised or unpolarised. For statistics that are polarised, we compute
 statistic values for pairs of derived alleles. Unpolarised statistics compute
-statistics over all possible alleles, derived and ancestral, and the result in
-averaged over statistics computed for all pairs of alleles (see weighting
-schemes below). The option for polarisation is not exposed to the user.
+statistics over all pairs of alleles, derived and ancestral. In either case,
+the result is averaged over these values, using a weighting
+scheme descrbed below. The option for polarisation is not exposed to the user.
 Instead, we implement statistics that are polarised where appropriate.
 
 (sec_stats_two_locus_site)=
@@ -746,16 +746,17 @@ print(ld)
 
 Because we allow for two-locus statistics to be computed for multi-allelic
 data, we need to be able to combine statistical results from each pair of
-alleles into one summary for a pair of sites. We use two implementations for
+alleles into one summary for a pair of sites. This does not affect biallelic
+data (and so this section can be skipped on first reading).
+We use two implementations for
 combining results from multiple alleles: `hap_weighted` and `total_weighted`.
 These are statistic-specific and not chosen by the user.
 
 Briefly, consider a pair of sites with {math}`n` alleles at the first locus and
-{math}`m` alleles at the second. Write {math}`f_{i,j}` as the statistic
-computed for focal alleles {math}`A_i` and {math}`B_j`, with haplotype weights
-{math}`(A_i B_j, A_i b_j, a_i B_j)`, where {math}`a_i` and {math}`b_j` are the
-collection of alleles that are not the focal alleles {math}`A_i` or
-{math}`B_j`, respectively. Then the weighting schemes are defined as:
+{math}`m` alleles at the second. (Whether this includes the ancestral allele
+depends on whether the statistic is polarised.) Write {math}`f_{i,j}` as the statistic
+computed for focal alleles {math}`A_i` and {math}`B_j`.
+Then the weighting schemes are defined as:
 
 - `hap_weighted`: {math}`\sum_{i=1}^{n}\sum_{j=1}^{m}p(A_{i}B_{j})f_{ij}`,
   where {math}`p(A_{i}B_{j})` is the frequency of haplotype {math}`A_{i}B_{j}`.
@@ -780,7 +781,7 @@ The `branch` mode computes expected two-locus statistics between pairs of
 trees, conditioned on the marginal topologies and branch lengths of those
 trees. The trees for which we compute statistics are specified by positions,
 and for a pair of positions we consider all possible haplotypes that could be
-generated by a single mutation occurring at the two trees.
+generated by a single mutation occurring on each of the two trees.
 
 For two trees, one with {math}`n` branches and the other with {math}`m`
 branches, there are {math}`nm` possible pairs of branches that may carry the
@@ -827,13 +828,13 @@ in the same manner as the rest of the stats API (see
 the stats API in that we handle one-way and two-way statistics in the same
 function call.
 
-To compute a two-way two-locus statistic, the `index` argument must be
+To compute a two-way two-locus statistic, the `indexes` argument must be
 provided. The statistics are selected in the same way (with the `stat`
 argument), but we provide a restricted set of two-way statistics (see
 {ref}`sec_stats_two_locus_summary_functions_two_way`). The dimension-dropping
 rules for the result follow the rest of the tskit stats API in that a single list
 or tuple will produce a single two-dimensional matrix, while list of these
-will produce a three-dimensional array, with the outer dimension of length
+will produce a three-dimensional array, with the first dimension of length
 equal to the length of the list.
 
 For example, to compute the {math}`r^2` LD matrix over a subset of samples in
@@ -868,8 +869,8 @@ ts.ld_matrix(sample_sets=[[0, 1, 2, 3], [4, 5, 6, 7]], indexes=[(0, 1)]) -> 3 di
 
 #### Why are there nan values in the LD matrix?
 
-For some statistics, it is possible to observe nan entries in the LD matrix,
-which can be surprising or numerically impact downstream analyses. A nan entry
+For some statistics, it is possible to observe `nan` entries in the LD matrix,
+which can be surprising or numerically impact downstream analyses. A `nan` entry
 may occur when computing a ratio statistic (such as {math}`r` or {math}`r^2`)
 with a denominator of zero, indicating that one or both sites in the pair are
 monomorphic. This can happen for a number of reasons:

python/tskit/trees.py

@@ -10947,7 +10947,8 @@ def ld_matrix(
         between pairs of trees at all specified ``positions`` ("branch" mode,
         producing a num_positions-by-num_positions sized matrix).
 
-        In the site mode, the sites under consideration can be restricted using
+        The sites considered for "site" mode defaults to all sites (which may
+        result in a very large matrix!), but can be restricted using
         the ``sites`` argument. Sites can be passed as a list of lists,
         specifying the ``[[row_sites], [col_sites]]``, resulting in a
         rectangular matrix, or by specifying a single list of ``[sites]``, in
@@ -10975,7 +10976,9 @@ def ld_matrix(
         section.
 
         **Available Stats** (use ``Stat Name`` in the ``stat`` keyword
-        argument).
+        argument). Statistics marked as "multi sample set" allow
+        (but do not require) computation from two sample sets
+        via the ``indexes`` argument. 
 
         ======================= ========== ================ ==============
         Stat                     Polarised Multi Sample Set Stat Name
@@ -10986,7 +10989,7 @@ def ld_matrix(
         :math:`D`                y          n               "D"
         :math:`D'`               y          n               "D_prime"
         :math:`D_z`              n          n               "Dz"
-        :math:`\pi_2`             n          n               "pi2"
+        :math:`\pi_2`            n          n               "pi2"
         :math:`\widehat{D^2}`    n          y               "D2_unbiased"
         :math:`\widehat{D_z}`    n          n               "Dz_unbiased"
         :math:`\widehat{\pi_2}`  n          n               "pi2_unbiased"