Apply suggestions from code review

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

Author: apragsdale

docs/stats.md

@@ -805,10 +805,10 @@ To compute a two-way two-locus statistic, the `index` 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 scalar
-indexes will produce a single two-dimensional matrix, while list of indexes
+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
-equal to the length of the list of indexes.
+equal to the length of the list.
 
 For concreteness, we would expect the following dimensions with the specified
 `sample_sets` and `indexes` arguments.
@@ -869,47 +869,46 @@ input. Each of our summary functions has the signature
   allele labelings, is zero. Uses the `total` normalisation method.
 
 `D_prime`
-: {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = \frac{D}{D_{\max}}`
+: {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = \frac{D}{D_{\max}}`,
 
-  Where {math}```
-        D_{\max} = \begin{cases}
+  where {math}
             \min\{p_A (1-p_B), p_B (1-p_B)\} & \textrm{if }D>=0 \\
             \min\{p_A p_B, (1-p_B) (1-p_B)\} & \textrm{otherwise}
         \end{cases}```
 
-  and {math}`D` is defined above.
+  and {math}`D` is defined above. Polarised, `total` weighted.
 
 `D2`
 : {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = (p_{ab} - p_{a} p_{b})^2`
 
-  Unpolarised, total weighted.
+  Unpolarised, `total` weighted.
 
 `Dz`
-:  {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = D (1 - 2 p_{a})(1 - 2p_{b})`
+:  {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = D (1 - 2 p_{a})(1 - 2p_{b})`,
 
-  Where {math}`D` is defined above. Unpolarised, total weighted.
+  where {math}`D` is defined above. Unpolarised, `total` weighted.
 
 `pi2`
 : {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = p_{a}p_{b}(1-p_{a})(1-p_{b})`
 
-  Unpolarised, total weighted.
+  Unpolarised, `total` weighted.
 
 `r`
-: {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = \frac{D}{\sqrt{p_{a}p_{b}(1-p_{a})(1-p_{b})}}`
+: {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = \frac{D}{\sqrt{p_{a}p_{b}(1-p_{a})(1-p_{b})}}`,
 
-  Where {math}`D` is defined above. Polarised, total weighted.
+  where {math}`D` is defined above. Polarised, `total` weighted.
 
 `r2`
-: {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = \frac{D^{2}}{p_{a}p_{b}(1-p_{a})(1-p_{b})}`
+: {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = \frac{D^{2}}{p_{a}p_{b}(1-p_{a})(1-p_{b})}`,
 
-  Where {math}`D` is defined above. Unpolarised, haplotype weighted.
+  where {math}`D` is defined above. Unpolarised, `haplotype` weighted.
 
 Unbiased two-locus statistics from the Hill-Robertson (1968) system are
-computed from haplotype counts. Derivations for these unbiased estimators can
+computed from haplotype counts. Definitions of these unbiased estimators can
 be found in [Ragsdale and Gravel
 (2020)](https://doi.org/10.1093/molbev/msz265). They require at least 4 samples
-to be valid and are specified as `stat=D2_unbiased`, `Dz_unbiased`, or
-`pi2_unbiased`.
+to be valid and are specified as `stat="D2_unbiased"`, `"Dz_unbiased"`, or
+`"pi2_unbiased"`.
 
 (sec_two_locus_summary_functions_two_way)=
 
@@ -921,9 +920,9 @@ Two-way statistics are indexed by sample sets {math}`i, j` and compute values
 using haplotype counts within pairs of sample sets.
 
 `D2`
-: {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = D_i * D_j`
+: {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = D_i * D_j`,
 
-Where {math}`D` is defined above.
+where {math}`D` is defined above.
 
 `r2`
 : {math}`f(w_{Ab}, w_{aB}, w_{AB}, n) = \frac{(p_{AB_i} - (p_{A_i}  p_{B_i})) (p_{AB_j} - (p_{A_j}  p_{B_j}))}{\sqrt{p_{A_i} (1 - p_{A_i}) p_{B_i} (1 - p_{B_i})}\sqrt{p_{A_j} (1 - p_{A_j}) p_{B_j} (1 - p_{B_j})}}`

python/tskit/trees.py

@@ -10967,7 +10967,8 @@ def ld_matrix(
         Some LD statistics are defined for two sample sets instead of within a
         single set of samples. If the ``indexes`` argument is specified, at
         least two sample sets must also be specified. ``indexes`` specifies the
-        sample set indexes between which to compute LD.
+        indexes of the sample sets in the ``sample_sets`` list
+        between which to compute LD.
 
         For more on how the ``indexes`` and ``sample_sets`` interact with the
         output dimensions, see the :ref:`sec_stats_two_locus_sample_sets`
@@ -10985,7 +10986,7 @@ def ld_matrix(
         :math:`D`                y          n               "D"
         :math:`D'`               y          n               "D_prime"
         :math:`D_z`              n          n               "Dz"
-        :math:`\pi2`             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"
@@ -11002,10 +11003,12 @@ def ld_matrix(
             specified as a list of lists to control the row and column sites.
             Only applicable in site mode. Specify as
             ``[[row_sites], [col_sites]]`` or ``[all_sites]``.
+            Defaults to all sites.
         :param list positions: A list of genomic positions where expected LD is
             computed. Only applicable in branch mode. Can be specified as a list
             of lists to control the row and column positions. Specify as
             ``[[row_positions], [col_positions]]`` or ``[all_positions]``.
+            Defaults to the leftmost coordinates of all trees.
         :param list indexes: A list of 2-tuples or a single 2-tuple, specifying
             the indexes of two sample sets over which to compute a two-way LD
             statistic. Only :math:`r^2`, :math:`D^2`, and :math:`\widehat{D^2}`