refactor: standardise biallelic diplotypes to count ref allele and handle missing calls

Author: Tanisha127

malariagen_data/anoph/distance.py

@@ -217,6 +217,9 @@ def _biallelic_diplotype_pairwise_distances(
         n_snps = gn.shape[0]
 
         # Prepare data for pairwise distance calculation.
+        # Mask missing calls (-127) before computing distances.
+        gn = gn.astype(float)
+        gn[gn == -127] = np.nan
         X = np.ascontiguousarray(gn.T)
 
         # Look up distance function.

malariagen_data/anoph/pca.py

@@ -152,7 +152,7 @@ def pca(
         # df_pca.index = df_pca.index.astype(str)
 
         # Name the DataFrame's columns as PC1, PC2, etc.
-        df_pca.columns = pd.Index([f"PC{i+1}" for i in range(coords.shape[1])])
+        df_pca.columns = pd.Index([f"PC{i + 1}" for i in range(coords.shape[1])])
 
         # Load the sample metadata.
         df_samples = self.sample_metadata(
@@ -231,6 +231,23 @@ def _pca(
                 loc_keep_fit = np.ones(len(samples), dtype=bool)
                 gn_fit = gn
 
+            # Impute missing calls (-127) with the mean value at each site.
+            if max_missing_an is not None and max_missing_an > 0:
+                gn_fit = gn_fit.astype(float)
+                gn = gn.astype(float)
+                for arr in [gn_fit, gn]:
+                    missing_mask = arr == -127
+                    site_means = np.where(
+                        np.all(missing_mask, axis=1, keepdims=True),
+                        0,
+                        np.nanmean(
+                            np.where(missing_mask, np.nan, arr), axis=1, keepdims=True
+                        ),
+                    )
+                    arr[missing_mask] = np.take(
+                        site_means.flatten(), np.where(missing_mask)[0]
+                    )
+
             # Remove any sites where all genotypes are identical.
             loc_var = np.any(gn_fit != gn_fit[:, 0, np.newaxis], axis=1)
             gn_fit_var = np.compress(loc_var, gn_fit, axis=0)

malariagen_data/anoph/snp_data.py

@@ -2037,8 +2037,12 @@ def _biallelic_diplotypes(
         samples = ds["sample_id"].values.astype("U")
 
         # Compute diplotypes as the number of alt alleles per genotype call.
+        # with missing calls coded as -127.
         gt = allel.GenotypeDaskArray(ds["call_genotype"].data)
         with self._dask_progress(desc="Compute biallelic diplotypes"):
             gn = gt.to_n_alt().compute()
+        # Code missing calls as -127.
+        missing = np.all(ds["call_genotype"].values == -1, axis=2)
+        gn[missing] = -127
 
         return dict(samples=samples, gn=gn)