Merge branch 'master' into GH407_allow_gene_labels

Author: leehart

docs/source/Af1.rst

@@ -71,6 +71,7 @@ SNP data access
     is_accessible
     biallelic_snp_calls
     biallelic_diplotypes
+    biallelic_snps_to_plink
 
 Haplotype data access
 ---------------------

docs/source/Ag3.rst

@@ -72,6 +72,7 @@ SNP data access
     is_accessible
     biallelic_snp_calls
     biallelic_diplotypes
+    biallelic_snps_to_plink
 
 Haplotype data access
 ---------------------

malariagen_data/__init__.py

@@ -4,6 +4,7 @@
 from .amin1 import Amin1
 from .anopheles import AnophelesDataResource, Region
 from .pf7 import Pf7
+from .pf8 import Pf8
 from .pv4 import Pv4
 from .util import SiteClass
 

malariagen_data/anoph/distance.py

@@ -6,7 +6,6 @@
 import numba  # type: ignore
 import numpy as np
 from numpydoc_decorator import doc  # type: ignore
-import anjl.params  # type: ignore
 
 # Internal imports.
 from .snp_data import AnophelesSnpData
@@ -410,10 +409,10 @@ def plot_njt(
         metric: distance_params.distance_metric = distance_params.default_distance_metric,
         distance_sort: Optional[tree_params.distance_sort] = None,
         count_sort: Optional[tree_params.count_sort] = None,
-        center_x: anjl.params.center_x = 0,
-        center_y: anjl.params.center_y = 0,
-        arc_start: anjl.params.arc_start = 0,
-        arc_stop: anjl.params.arc_stop = 2 * math.pi,
+        center_x: distance_params.center_x = 0,
+        center_y: distance_params.center_y = 0,
+        arc_start: distance_params.arc_start = 0,
+        arc_stop: distance_params.arc_stop = 2 * math.pi,
         width: plotly_params.fig_width = 800,
         height: plotly_params.fig_height = 600,
         show: plotly_params.show = True,
@@ -426,8 +425,8 @@ def plot_njt(
         color_discrete_sequence: plotly_params.color_discrete_sequence = None,
         color_discrete_map: plotly_params.color_discrete_map = None,
         category_orders: plotly_params.category_order = None,
-        edge_legend: anjl.params.edge_legend = False,
-        leaf_legend: anjl.params.leaf_legend = True,
+        edge_legend: distance_params.edge_legend = False,
+        leaf_legend: distance_params.leaf_legend = True,
         legend_sizing: plotly_params.legend_sizing = "constant",
         thin_offset: base_params.thin_offset = 0,
         sample_sets: Optional[base_params.sample_sets] = None,
@@ -449,6 +448,10 @@ def plot_njt(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
     ) -> plotly_params.figure:
+        # Only import anjl if needed, as it requires a couple of seconds to compile
+        # functions.
+        import anjl  # type: ignore
+
         # Normalise params.
         if count_sort is None and distance_sort is None:
             count_sort = True

malariagen_data/anoph/distance_params.py

@@ -19,3 +19,20 @@
 ]
 
 default_nj_algorithm: nj_algorithm = "dynamic"
+
+center_x: TypeAlias = Annotated[int | float, "X coordinate where plotting is centered."]
+
+center_y: TypeAlias = Annotated[int | float, "Y coordinate where plotting is centered."]
+
+arc_start: TypeAlias = Annotated[int | float, "Angle where tree layout begins."]
+
+arc_stop: TypeAlias = Annotated[int | float, "Angle where tree layout ends."]
+
+edge_legend: TypeAlias = Annotated[
+    bool, "Show legend entries for the different edge (line) colors."
+]
+
+leaf_legend: TypeAlias = Annotated[
+    bool,
+    "Show legend entries for the different leaf node (scatter) colors and symbols.",
+]

malariagen_data/anoph/plink_params.py

@@ -0,0 +1,18 @@
+"""Parameters for Plink converter functions."""
+
+from typing_extensions import Annotated, TypeAlias
+
+overwrite: TypeAlias = Annotated[
+    bool,
+    """
+    A boolean indicating whether a previously written file with the same name ought
+    to be overwritten. Default is False.
+    """,
+]
+
+output_dir: TypeAlias = Annotated[
+    str,
+    """
+    A string indicating the desired output file location.
+    """,
+]

malariagen_data/anoph/to_plink.py

@@ -0,0 +1,132 @@
+from typing import Optional
+
+import allel  # type: ignore
+import numpy as np
+import os
+import bed_reader
+
+from ..util import dask_compress_dataset
+from .snp_data import AnophelesSnpData
+from . import base_params
+from . import plink_params
+from . import pca_params
+from numpydoc_decorator import doc  # type: ignore
+
+
+class PlinkConverter(
+    AnophelesSnpData,
+):
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        # N.B., this class is designed to work cooperatively, and
+        # so it's important that any remaining parameters are passed
+        # to the superclass constructor.
+        super().__init__(**kwargs)
+
+    @doc(
+        summary="""
+            Write Anopheles biallelic SNP data to the Plink binary file format.
+        """,
+        extended_summary="""
+            This function writes biallelic SNPs to the Plink binary file format. It enables
+            subsetting to specific regions (`region`), selecting specific sample sets, or lists of
+            samples, randomly downsampling sites, and specifying filters based on missing data and
+            minimum minor allele count (see the docs for `biallelic_snp_calls` for more information).
+            The `overwrite` parameter, set to true, will enable overwrite of data with the same
+            SNP selection parameter values.
+        """,
+        returns="""
+        Base path to files containing binary Plink output files. Append .bed,
+        .bim or .fam to obtain paths for the binary genotype table file, variant
+        information file and sample information file respectively.
+        """,
+        notes="""
+            This computation may take some time to run, depending on your computing
+            environment. Unless the `overwrite` parameter is set to `True`, results will be returned
+            from a previous computation, if available.
+        """,
+    )
+    def biallelic_snps_to_plink(
+        self,
+        output_dir: plink_params.output_dir,
+        region: base_params.regions,
+        n_snps: base_params.n_snps,
+        overwrite: plink_params.overwrite = False,
+        thin_offset: base_params.thin_offset = 0,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        sample_query: Optional[base_params.sample_query] = None,
+        sample_indices: Optional[base_params.sample_indices] = None,
+        site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
+        min_minor_ac: Optional[
+            base_params.min_minor_ac
+        ] = pca_params.min_minor_ac_default,
+        max_missing_an: Optional[
+            base_params.max_missing_an
+        ] = pca_params.max_missing_an_default,
+        random_seed: base_params.random_seed = 42,
+        inline_array: base_params.inline_array = base_params.inline_array_default,
+        chunks: base_params.chunks = base_params.native_chunks,
+    ):
+        # Define output files
+        plink_file_path = f"{output_dir}/{region}.{n_snps}.{min_minor_ac}.{max_missing_an}.{thin_offset}"
+
+        bed_file_path = f"{plink_file_path}.bed"
+
+        # Check to see if file exists and if overwrite is set to false, return existing file
+        if os.path.exists(bed_file_path):
+            if not overwrite:
+                return plink_file_path
+
+        # Get snps
+        ds_snps = self.biallelic_snp_calls(
+            region=region,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_indices=sample_indices,
+            site_mask=site_mask,
+            min_minor_ac=min_minor_ac,
+            max_missing_an=max_missing_an,
+            n_snps=n_snps,
+            thin_offset=thin_offset,
+            random_seed=random_seed,
+            inline_array=inline_array,
+            chunks=chunks,
+        )
+
+        # Set up dataset with required vars for plink conversion
+
+        # Compute gt ref counts
+        with self._dask_progress("Computing genotype ref counts"):
+            gt_asc = ds_snps["call_genotype"].data  # dask array
+            gn_ref = allel.GenotypeDaskArray(gt_asc).to_n_ref(fill=-127)
+            gn_ref = gn_ref.compute()
+
+        # Ensure genotypes vary
+        loc_var = np.any(gn_ref != gn_ref[:, 0, np.newaxis], axis=1)
+
+        # Load final data
+        ds_snps_final = dask_compress_dataset(ds_snps, loc_var, dim="variants")
+
+        # Init vars for input to bed reader
+        gn_ref_final = gn_ref[loc_var]
+        val = gn_ref_final.T
+        with self._spinner("Prepare output data"):
+            alleles = ds_snps_final["variant_allele"].values
+            properties = {
+                "iid": ds_snps_final["sample_id"].values,
+                "chromosome": ds_snps_final["variant_contig"].values,
+                "bp_position": ds_snps_final["variant_position"].values,
+                "allele_1": alleles[:, 0],
+                "allele_2": alleles[:, 1],
+            }
+
+        bed_reader.to_bed(
+            filepath=bed_file_path,
+            val=val,
+            properties=properties,
+            count_A1=True,
+        )
+
+        return plink_file_path

malariagen_data/anopheles.py

@@ -46,6 +46,7 @@
 from .anoph.distance import AnophelesDistanceAnalysis
 from .anoph.sample_metadata import AnophelesSampleMetadata, locate_cohorts
 from .anoph.snp_data import AnophelesSnpData
+from .anoph.to_plink import PlinkConverter
 from .anoph.g123 import AnophelesG123Analysis
 from .anoph.fst import AnophelesFstAnalysis
 from .anoph.h12 import AnophelesH12Analysis
@@ -100,6 +101,7 @@ class AnophelesDataResource(
     AnophelesSnpFrequencyAnalysis,
     AnophelesDistanceAnalysis,
     AnophelesPca,
+    PlinkConverter,
     AnophelesIgv,
     AnophelesAimData,
     AnophelesHapData,

malariagen_data/pf8.py

@@ -0,0 +1,43 @@
+import os
+
+from .plasmodium import PlasmodiumDataResource
+
+
+class Pf8(PlasmodiumDataResource):
+    """Provides access to data from the Pf8 release.
+
+    Parameters
+    ----------
+    url : str, optional
+        Base path to data. Default uses Google Cloud Storage "gs://pf8-release/",
+        or specify a local path on your file system if data have been downloaded.
+    data_config : str, optional
+        Path to config for structure of Pf8 data resource. Defaults to config included
+        with the malariagen_data package.
+    **kwargs
+        Passed through to fsspec when setting up file system access.
+
+    Examples
+    --------
+    Access data from Google Cloud Storage (default):
+
+        >>> import malariagen_data
+        >>> pf8 = malariagen_data.Pf8()
+
+    Access data downloaded to a local file system:
+
+        >>> pf8 = malariagen_data.Pf8("/local/path/to/pf8-release/")
+
+    """
+
+    def __init__(
+        self,
+        url=None,
+        data_config=None,
+        **kwargs,
+    ):
+        # setup filesystem
+        if not data_config:
+            working_dir = os.path.dirname(os.path.abspath(__file__))
+            data_config = os.path.join(working_dir, "pf8_config.json")
+        super().__init__(data_config=data_config, url=url)