Understanding trait_linear_model

[quattro]

👋 Hi all, first off I'd like to state how fantastic this library is. It's a breeze to use such a highly engineered and robust package within simulations.

I've been looking into trait analysis using ARGs, and was curious as to how the trait_linear_model function works, as described in

Ralph, Peter, Kevin Thornton, and Jerome Kelleher. "Efficiently summarizing relationships in large samples: a general duality between statistics of genealogies and genomes." Genetics 215.3 (2020): 779-797.

I see there is a trait_linear_model method for a Tree object, but cannot seem to follow the logic to where self._ll_tree_sequence.trait_linear_model is defined separately. It appears this is called directly inside another function as method;

My understanding is that this is just a single instance of a more general pattern but would like to see where the final statistics are explicitly computed.

Any help would be greatly appreciated.

Cheers,
Nick

[petrelharp]

Gee, thanks for the nice words!

You're right, that trait_linear_model is a special case of the general stats framework. It uses the summary function listed in this table (the last one). The underlying math is written out in the appendix of the paper you cite:

The underlying C code where this is actually implemented is here.

Also, FYI, to my knowledge no-one has actually investigated how well this does in practice, so I'm curious to know what happens, if you do.

[quattro]

Phenomenal, thanks so much! I was a bit puzzled trying to stay in Python-land, but it makes sense that the computation is farmed out to C at some point. My initial reason for wanting to better understand things, is that the code

K = np.linalg.cholesky(np.matmul(Z.T, Z)).T
Z = np.matmul(Z, np.linalg.inv(K))

goes through the trouble of computing a cholesky factor, and then uses plain-ol inv. I couldn't trace exactly where these were used downstream.

I think a bit more robust approach that explicitly leverages the triangular structure of the cholesky factor may be,

K = np.linalg.cholesky(np.matmul(Z.T, Z))
Z = sp.linalg.solve_triangular(K, Z.T, lower=True).T

Under some super simple test data the results are the same, but this (in theory) should keep errors smaller than direct inverse and matmul, while giving a small improvement to runtime.

def rotate_old(Z):
  K = np.linalg.cholesky(np.matmul(Z.T, Z)).T
  return np.matmul(Z, np.linalg.inv(K))

def rotate(Z):
  K = np.linalg.cholesky(np.matmul(Z.T, Z))
  return sp.linalg.solve_triangular(K, Z.T, lower=True, check_finite=False).T

In [193]: N, P = 1000, 50

In [194]: np.random.seed(0)

In [195]: Z = np.random.normal(size=(N,P))

In [196]: Z
Out[196]:
array([[ 1.76405235,  0.40015721,  0.97873798, ...,  0.77749036,
        -1.61389785, -0.21274028],
       [-0.89546656,  0.3869025 , -0.51080514, ...,  1.78587049,
         0.12691209,  0.40198936],
       [ 1.8831507 , -1.34775906, -1.270485  , ...,  1.11701629,
        -1.31590741, -0.4615846 ],
       ...,
       [-1.60097004,  1.25286095, -0.57135417, ...,  1.02990903,
        -1.22198213, -0.92234101],
       [-0.04131705,  0.73273455, -0.41506565, ..., -0.24413417,
        -1.05984916,  0.06591413],
       [ 0.14517991, -1.36462743, -1.31893264, ..., -0.25034177,
         0.58472799, -1.25082696]])

In [197]: %timeit rotate_old(Z)
374 µs ± 18.2 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [198]: %timeit rotate(Z)
353 µs ± 3.68 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [199]: np.allclose(rotate_old(Z), rotate(Z))
Out[199]: True

[petrelharp]

Hey, thanks! Hm, let's see - the reason we do this is because the math above requires that Z is orthogonal (see note that says "although in the implementation we postprocess to make this the case").

And, I like the suggestion for the code... except that it would require pulling in a scipy dependency, which we've avoided so far? A good halfway would be that if this was a barrier for someone, we could have a flag to optionally skip the normalization, thus letting someone do it themselves. However, I'm not sure if this would ever be a bottleneck compared to the tree sequence computation in practice? If you run across a good use case, please let us know, however.