From 54f1fd48a5ab9c2fad98cb7c4bbd326e2d490f07 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Fri, 1 Dec 2023 10:26:03 +0000
Subject: [PATCH] numpify normalisation

---
 python/tskit/trees.py | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/python/tskit/trees.py b/python/tskit/trees.py
index 80ce009be1..28b30da186 100644
--- a/python/tskit/trees.py
+++ b/python/tskit/trees.py
@@ -8142,16 +8142,11 @@ def genetic_relatedness_matrix(
         def _normalise(B):
             if len(B) == 0:
                 return B
-            K = np.zeros_like(B)
-            N = K.shape[0]
-            B_mean = np.mean(B)
-            Bi_mean = np.mean(B, axis=0)
-            # TODO numpify - should be easy enough by creating full matrices
-            # for the row and column means.
-            for i in range(N):
-                for j in range(N):
-                    K[i, j] = B[i, j] - Bi_mean[i] - Bi_mean[j] + B_mean
-            # FIXME I don't know what this factor -2 is about
+            K = B + np.mean(B)
+            y = np.mean(B, axis=0)
+            X = y[:, np.newaxis] + y[np.newaxis, :]
+            K -= X
+            # FIXME I don't know what this factor of -2 is about
             return K / -2
 
         if windows is None: