malariagen_data 2.0.0__tar.gz → 2.2.0__tar.gz

{malariagen_data-2.0.0 → malariagen_data-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: malariagen-data
-Version: 2.0.0
+Version: 2.2.0
 Summary: A package for accessing MalariaGEN public data.
 License: MIT
 Author: Alistair Miles
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3.7
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Requires-Dist: BioPython
+Requires-Dist: bokeh
 Requires-Dist: dask[array]
 Requires-Dist: fsspec
 Requires-Dist: gcsfs

{malariagen_data-2.0.0 → malariagen_data-2.2.0}/malariagen_data/ag3.py

@@ -4,15 +4,13 @@ from collections import Counter
 
 import allel
 import dask.array as da
-import ipyleaflet
-import ipywidgets
 import numba
 import numpy as np
 import pandas as pd
-import plotly.express as px
 import xarray as xr
 import zarr
-from statsmodels.stats.proportion import proportion_confint
+
+import malariagen_data
 
 from . import veff
 from .util import (
@@ -145,15 +143,16 @@ class Ag3:
     def __repr__(self):
         return (
             f"<MalariaGEN Ag3 data resource API>\n"
-            f"Storage URL           : {self._url}\n"
-            f"Releases available    : {', '.join(self.releases)}\n"
-            f"Cohorts analysis      : {DEFAULT_COHORTS_ANALYSIS}\n"
-            f"Species analysis      : {DEFAULT_SPECIES_ANALYSIS}\n"
-            f"Site filters analysis : {DEFAULT_SITE_FILTERS_ANALYSIS}\n"
+            f"Storage URL             : {self._url}\n"
+            f"Data releases available : {', '.join(self.releases)}\n"
+            f"Cohorts analysis        : {DEFAULT_COHORTS_ANALYSIS}\n"
+            f"Species analysis        : {DEFAULT_SPECIES_ANALYSIS}\n"
+            f"Site filters analysis   : {DEFAULT_SITE_FILTERS_ANALYSIS}\n"
+            f"Software version        : {malariagen_data.__version__}\n"
             f"---\n"
             f"Please note that data are subject to terms of use,\n"
             f"for more information see https://www.malariagen.net/data\n"
-            f"or contact data@malariagen.net. For API documentation see: "
+            f"or contact data@malariagen.net. For API documentation see: \n"
             f"https://malariagen.github.io/vector-data/ag3/api.html"
         )
 
@@ -173,10 +172,11 @@ class Ag3:
                 </thead>
                 <tbody>
                     <tr><th style="text-align: left">Storage URL</th><td>{self._url}</td></tr>
-                    <tr><th style="text-align: left">Releases available</th><td>{', '.join(self.releases)}</td></tr>
+                    <tr><th style="text-align: left">Data releases available</th><td>{', '.join(self.releases)}</td></tr>
                     <tr><th style="text-align: left">Cohorts analysis</th><td>{DEFAULT_COHORTS_ANALYSIS}</td></tr>
                     <tr><th style="text-align: left">Species analysis</th><td>{DEFAULT_SPECIES_ANALYSIS}</td></tr>
                     <tr><th style="text-align: left">Site filters analysis</th><td>{DEFAULT_SITE_FILTERS_ANALYSIS}</td></tr>
+                    <tr><th style="text-align: left">Software version</th><td>{malariagen_data.__version__}</td></tr>
                 </tbody>
             </table>
         """
@@ -1287,7 +1287,11 @@ class Ag3:
             ann.get_effects(transcript=transcript, variants=df_snps)
 
             # add label
-            df_snps["label"] = df_snps.apply(_make_snp_label_effect, axis="columns")
+            df_snps["label"] = _pandas_apply(
+                _make_snp_label_effect,
+                df_snps,
+                columns=["contig", "position", "ref_allele", "alt_allele", "aa_change"],
+            )
 
             # set index
             df_snps.set_index(
@@ -1298,7 +1302,11 @@ class Ag3:
         else:
 
             # add label
-            df_snps["label"] = df_snps.apply(_make_snp_label, axis="columns")
+            df_snps["label"] = _pandas_apply(
+                _make_snp_label,
+                df_snps,
+                columns=["contig", "position", "ref_allele", "alt_allele"],
+            )
 
             # set index
             df_snps.set_index(
@@ -1709,7 +1717,7 @@ class Ag3:
 
         Parameters
         ----------
-        contig : str
+        contig : str or list of str
             Chromosome arm, e.g., "3R". Multiple values can be provided
             as a list, in which case data will be concatenated, e.g., ["2R", "3R"].
         sample_sets : str or list of str, optional
@@ -1884,7 +1892,7 @@ class Ag3:
 
         Parameters
         ----------
-        contig : str
+        contig : str or list of str
             Chromosome arm, e.g., "3R". Multiple values can be provided
             as a list, in which case data will be concatenated, e.g., ["2R", "3R"].
         sample_set : str
@@ -2037,7 +2045,7 @@ class Ag3:
 
         Parameters
         ----------
-        contig : str
+        contig : str or list of str
             Chromosome arm, e.g., "3R". Multiple values can be provided
             as a list, in which case data will be concatenated, e.g., ["2R", "3R"].
         sample_sets : str or list of str, optional
@@ -2088,8 +2096,9 @@ class Ag3:
 
         Parameters
         ----------
-        contig : str
-            Chromosome arm, e.g., "3R".
+        contig : str or list of str
+            Chromosome arm, e.g., "3R". Multiple values can be provided
+            as a list, in which case data will be concatenated, e.g., ["2R", "3R"].
         sample_sets : str or list of str
             Can be a sample set identifier (e.g., "AG1000G-AO") or a list of sample set
             identifiers (e.g., ["AG1000G-BF-A", "AG1000G-BF-B"]) or a release identifier (e.g.,
@@ -2102,6 +2111,22 @@ class Ag3:
 
         """
 
+        # handle multiple contigs
+        if isinstance(contig, str):
+            contig = [contig]
+
+        ds = xarray_concat(
+            [self._gene_cnv(contig=c, sample_sets=sample_sets) for c in contig],
+            dim="genes",
+        )
+
+        return ds
+
+    def _gene_cnv(self, *, contig, sample_sets):
+
+        # sanity check
+        assert isinstance(contig, str)
+
         # access HMM data
         ds_hmm = self.cnv_hmm(contig=contig, sample_sets=sample_sets)
         pos = ds_hmm["variant_position"].values
@@ -2111,6 +2136,7 @@ class Ag3:
         # access genes
         df_geneset = self.geneset()
         df_genes = df_geneset.query(f"type == 'gene' and contig == '{contig}'")
+        n_genes = len(df_genes)
 
         # setup intermediates
         windows = []
@@ -2143,14 +2169,16 @@ class Ag3:
         ds_out = xr.Dataset(
             coords={
                 "gene_id": (["genes"], df_genes["ID"].values),
-                "gene_start": (["genes"], df_genes["start"].values),
-                "gene_end": (["genes"], df_genes["end"].values),
                 "sample_id": (["samples"], ds_hmm["sample_id"].values),
             },
             data_vars={
+                "gene_contig": (["genes"], np.array([contig] * n_genes, dtype=object)),
+                "gene_start": (["genes"], df_genes["start"].values),
+                "gene_end": (["genes"], df_genes["end"].values),
                 "gene_windows": (["genes"], windows),
                 "gene_name": (["genes"], df_genes["Name"].values),
                 "gene_strand": (["genes"], df_genes["strand"].values),
+                "gene_description": (["genes"], df_genes["description"].values),
                 "CN_mode": (["genes", "samples"], modes),
                 "CN_mode_count": (["genes", "samples"], counts),
                 "sample_coverage_variance": (
@@ -2183,8 +2211,9 @@ class Ag3:
 
         Parameters
         ----------
-        contig : str
-            Chromosome arm, e.g., "3R".
+        contig : str or list of str
+            Chromosome arm, e.g., "3R". Multiple values can be provided
+            as a list, in which case data will be concatenated, e.g., ["2R", "3R"].
         cohorts : str or dict
             If a string, gives the name of a predefined cohort set, e.g., one of
             {"admin1_month", "admin1_year", "admin2_month", "admin2_year"}.
@@ -2221,6 +2250,46 @@ class Ag3:
         # check parameters
         _check_param_min_cohort_size(min_cohort_size)
 
+        if isinstance(contig, str):
+            contig = [contig]
+
+        df = pd.concat(
+            [
+                self._gene_cnv_frequencies(
+                    contig=c,
+                    cohorts=cohorts,
+                    sample_query=sample_query,
+                    cohorts_analysis=cohorts_analysis,
+                    min_cohort_size=min_cohort_size,
+                    species_analysis=species_analysis,
+                    sample_sets=sample_sets,
+                    drop_invariant=drop_invariant,
+                    max_coverage_variance=max_coverage_variance,
+                )
+                for c in contig
+            ],
+            axis=0,
+        )
+
+        return df
+
+    def _gene_cnv_frequencies(
+        self,
+        *,
+        contig,
+        cohorts,
+        sample_query,
+        cohorts_analysis,
+        min_cohort_size,
+        species_analysis,
+        sample_sets,
+        drop_invariant,
+        max_coverage_variance,
+    ):
+
+        # sanity check - this function is one contig at a time
+        assert isinstance(contig, str)
+
         # load sample metadata
         df_samples = self.sample_metadata(
             sample_sets=sample_sets,
@@ -2245,9 +2314,6 @@ class Ag3:
             else:
                 loc_samples = loc_pass_samples
 
-        # get genes
-        df_genes = self.geneset().query(f"type == 'gene' and contig == '{contig}'")
-
         # figure out expected copy number
         if contig == "X":
             is_male = (df_samples["sex_call"] == "M").values
@@ -2255,20 +2321,25 @@ class Ag3:
         else:
             expected_cn = 2
 
-        # setup output dataframe
-        n_genes = len(df_genes)
+        # setup output dataframe - two rows for each gene, one for amplification and one for deletion
+        n_genes = ds_cnv.dims["genes"]
+        df_genes = ds_cnv[
+            [
+                "gene_id",
+                "gene_name",
+                "gene_strand",
+                "gene_description",
+                "gene_contig",
+                "gene_start",
+                "gene_end",
+            ]
+        ].to_dataframe()
         df = pd.concat([df_genes, df_genes], axis=0).reset_index(drop=True)
-
-        # drop columns we don't need
-        df.drop(columns=["source", "type", "score", "phase", "Parent"], inplace=True)
-
-        # rename some columns
         df.rename(
             columns={
-                "ID": "gene_id",
-                "Name": "gene_name",
-                "strand": "gene_strand",
-                "description": "gene_description",
+                "gene_contig": "contig",
+                "gene_start": "start",
+                "gene_end": "end",
             },
             inplace=True,
         )
@@ -2329,9 +2400,12 @@ class Ag3:
         df.reset_index(drop=True, inplace=True)
         df = pd.concat([df, df_freqs, df_extras], axis=1)
         df.sort_values(["contig", "start", "cnv_type"], inplace=True)
+        df.reset_index(drop=True, inplace=True)
 
         # add label
-        df["label"] = df.apply(_make_gene_cnv_label, axis="columns")
+        df["label"] = _pandas_apply(
+            _make_gene_cnv_label, df, columns=["gene_id", "gene_name", "cnv_type"]
+        )
 
         # deal with invariants
         if drop_invariant:
@@ -2709,7 +2783,11 @@ class Ag3:
         df_aaf["max_af"] = df_aaf[freq_cols].max(axis=1)
 
         # add label
-        df_aaf["label"] = df_aaf.apply(_make_snp_label_aa, axis="columns")
+        df_aaf["label"] = _pandas_apply(
+            _make_snp_label_aa,
+            df_aaf,
+            columns=["aa_change", "contig", "position", "ref_allele", "alt_allele"],
+        )
 
         # sort by genomic position
         df_aaf = df_aaf.sort_values(["position", "aa_change"])
@@ -2743,20 +2821,23 @@ class Ag3:
         title=True,
         **kwargs,
     ):
-
-        """Plot a heatmap from a pandas DataFrame of frequencies, e.g., output from
-            `Ag3.snp_allele_frequencies()` or `Ag3.gene_cnv_frequencies()`. It's recommended to
-            filter the input DataFrame to just rows of interest, i.e., fewer rows than `max_len`.
+        """Plot a heatmap from a pandas DataFrame of frequencies, e.g., output
+        from `Ag3.snp_allele_frequencies()` or `Ag3.gene_cnv_frequencies()`.
+        It's recommended to filter the input DataFrame to just rows of interest,
+        i.e., fewer rows than `max_len`.
 
         Parameters
         ----------
         df : pandas DataFrame
-           A DataFrame of frequencies, e.g., output from `snp_allele_frequencies()` or `gene_cnv_frequencies()`.
+           A DataFrame of frequencies, e.g., output from
+           `snp_allele_frequencies()` or `gene_cnv_frequencies()`.
         index : str or list of str
-            One or more column headers that are present in the input dataframe. This becomes the heatmap y-axis
-            row labels. The column/s must produce a unique index.
+            One or more column headers that are present in the input dataframe.
+            This becomes the heatmap y-axis row labels. The column/s must
+            produce a unique index.
         max_len : int, optional
-            Displaying large styled dataframes may cause ipython notebooks to crash.
+            Displaying large styled dataframes may cause ipython notebooks to
+            crash.
         x_label : str, optional
             This is the x-axis label that will be displayed on the heatmap.
         y_label : str, optional
@@ -2781,6 +2862,8 @@ class Ag3:
 
         """
 
+        import plotly.express as px
+
         # check len of input
         if len(df) > max_len:
             raise ValueError(f"Input DataFrame is longer than {max_len}")
@@ -2871,39 +2954,40 @@ class Ag3:
         species_analysis=DEFAULT_SPECIES_ANALYSIS,
         site_filters_analysis=DEFAULT_SITE_FILTERS_ANALYSIS,
     ):
-        """Group samples by taxon, area (space) and period (time), then compute SNP allele counts
-        and frequencies.
+        """Group samples by taxon, area (space) and period (time), then compute
+        SNP allele counts and frequencies.
 
         Parameters
         ----------
         transcript : str
             Gene transcript ID (AgamP4.12), e.g., "AGAP004707-RD".
         area_by : str
-            Column name in the sample metadata to use to group samples spatially. E.g.,
-            use "adm1_ISO" or "adm1_name" to group by level 1 administrative divisions,
-            or use "adm2_name" to group by level 2 administrative divisions.
+            Column name in the sample metadata to use to group samples
+            spatially. E.g., use "admin1_iso" or "admin1_name" to group by level 1
+            administrative divisions, or use "admin2_name" to group by level 2
+            administrative divisions.
         period_by : {"year", "quarter", "month"}
             Length of time to group samples temporally.
         sample_sets : str or list of str, optional
-            Can be a sample set identifier (e.g., "AG1000G-AO") or a list of sample set
-            identifiers (e.g., ["AG1000G-BF-A", "AG1000G-BF-B"]) or a release identifier (e.g.,
-            "3.0") or a list of release identifiers.
+            Can be a sample set identifier (e.g., "AG1000G-AO") or a list of
+            sample set identifiers (e.g., ["AG1000G-BF-A", "AG1000G-BF-B"]) or a
+            release identifier (e.g., "3.0") or a list of release identifiers.
         sample_query : str, optional
-            A pandas query string which will be evaluated against the sample metadata e.g.,
-            "taxon == 'coluzzii' and country == 'Burkina Faso'".
+            A pandas query string which will be evaluated against the sample
+            metadata e.g., "taxon == 'coluzzii' and country == 'Burkina Faso'".
         min_cohort_size : int, optional
             Minimum cohort size. Any cohorts below this size are omitted.
         drop_invariant : bool, optional
-            If True, variants with no alternate allele calls in any cohorts are dropped from
-            the result.
+            If True, variants with no alternate allele calls in any cohorts are
+            dropped from the result.
         variant_query : str, optional
         site_mask : str, optional
             Site filters mask to apply.
         nobs_mode : {"called", "fixed"}
-            Method for calculating the denominator when computing frequencies. If "called"
-            then use the number of called alleles, i.e., number of samples with non-missing
-            genotype calls multiplied by 2. If "fixed" then use the number of samples
-            multiplied by 2.
+            Method for calculating the denominator when computing frequencies.
+            If "called" then use the number of called alleles, i.e., number of
+            samples with non-missing genotype calls multiplied by 2. If "fixed"
+            then use the number of samples multiplied by 2.
         ci_method : {"normal", "agresti_coull", "beta", "wilson", "binom_test"}, optional
             Method to use for computing confidence intervals, passed through to
             `statsmodels.stats.proportion.proportion_confint`.
@@ -2918,11 +3002,11 @@ class Ag3:
         -------
         ds : xarray.Dataset
             The resulting dataset contains data has dimensions "cohorts" and "variants".
-            Variables prefixed with "cohort_" are 1-dimensional arrays with data about
+            Variables prefixed with "cohort" are 1-dimensional arrays with data about
             the cohorts, such as the area, period, taxon and cohort size. Variables
-            prefixed with "variant_" are 1-dimensional arrays with data about the variants,
+            prefixed with "variant" are 1-dimensional arrays with data about the variants,
             such as the contig, position, reference and alternate alleles. Variables prefixed
-            with "event_" are 2-dimensional arrays with the allele counts and frequency
+            with "event" are 2-dimensional arrays with the allele counts and frequency
             calculations.
 
         """
@@ -2976,8 +3060,8 @@ class Ag3:
         # bring genotypes into memory
         gt = gt.compute()
 
-        # convert genotypes to more convenient representation
-        gac, gan = _genotypes_to_alt_allele_counts_melt(gt, max_allele=3)
+        # # convert genotypes to more convenient representation
+        # gac, gan = _genotypes_to_alt_allele_counts_melt(gt, max_allele=3)
 
         # set up variant variables
         contigs = ds_snps.attrs["contigs"]
@@ -3013,13 +3097,18 @@ class Ag3:
             # compute cohort allele counts
             # cohort_gac = np.take(gac, sample_indices, axis=1)
             # np.sum(cohort_gac, axis=1, out=count[:, cohort_index])
-            count[:, cohort_index] = _take_sum_cols(gac, sample_indices)
+            # count[:, cohort_index] = _take_sum_cols(gac, sample_indices)
+            cohort_ac, cohort_an = _cohort_alt_allele_counts_melt(
+                gt, sample_indices, max_allele=3
+            )
+            count[:, cohort_index] = cohort_ac
 
             # compute cohort allele numbers
             if nobs_mode == "called":
                 # cohort_gan = np.take(gan, sample_indices, axis=1)
                 # np.sum(cohort_gan, axis=1, out=nobs[:, cohort_index])
-                nobs[:, cohort_index] = _take_sum_cols(gan, sample_indices)
+                # nobs[:, cohort_index] = _take_sum_cols(gan, sample_indices)
+                nobs[:, cohort_index] = cohort_an
             elif nobs_mode == "fixed":
                 nobs[:, cohort_index] = cohort.size * 2
             else:
@@ -3053,7 +3142,7 @@ class Ag3:
         # deal with SNP alleles not observed
         if drop_invariant:
             loc_variant = max_af > 0
-            df_variants = df_variants.loc[loc_variant]
+            df_variants = df_variants.loc[loc_variant].reset_index(drop=True)
             count = np.compress(loc_variant, count, axis=0)
             nobs = np.compress(loc_variant, nobs, axis=0)
             frequency = np.compress(loc_variant, frequency, axis=0)
@@ -3065,7 +3154,11 @@ class Ag3:
         ann.get_effects(transcript=transcript, variants=df_variants)
 
         # add variant labels
-        df_variants["label"] = df_variants.apply(_make_snp_label_effect, axis="columns")
+        df_variants["label"] = _pandas_apply(
+            _make_snp_label_effect,
+            df_variants,
+            columns=["contig", "position", "ref_allele", "alt_allele", "aa_change"],
+        )
 
         # build the output dataset
         ds_out = xr.Dataset()
@@ -3129,8 +3222,8 @@ class Ag3:
             Gene transcript ID (AgamP4.12), e.g., "AGAP004707-RD".
         area_by : str
             Column name in the sample metadata to use to group samples spatially. E.g.,
-            use "adm1_ISO" or "adm1_name" to group by level 1 administrative divisions,
-            or use "adm2_name" to group by level 2 administrative divisions.
+            use "admin1_iso" or "admin1_name" to group by level 1 administrative divisions,
+            or use "admin2_name" to group by level 2 administrative divisions.
         period_by : {"year", "quarter", "month"}
             Length of time to group samples temporally.
         sample_sets : str or list of str, optional
@@ -3164,11 +3257,11 @@ class Ag3:
         -------
         ds : xarray.Dataset
             The resulting dataset contains data has dimensions "cohorts" and "variants".
-            Variables prefixed with "cohort_" are 1-dimensional arrays with data about
+            Variables prefixed with "cohort" are 1-dimensional arrays with data about
             the cohorts, such as the area, period, taxon and cohort size. Variables
-            prefixed with "variant_" are 1-dimensional arrays with data about the variants,
+            prefixed with "variant" are 1-dimensional arrays with data about the variants,
             such as the contig, position, reference and alternate alleles. Variables prefixed
-            with "event_" are 2-dimensional arrays with the allele counts and frequency
+            with "event" are 2-dimensional arrays with the allele counts and frequency
             calculations.
 
         """
@@ -3241,7 +3334,11 @@ class Ag3:
         df_variants.columns = [c.split("variant_")[1] for c in df_variants.columns]
 
         # assign new variant label
-        label = df_variants.apply(_make_snp_label_aa, axis=1)
+        label = _pandas_apply(
+            _make_snp_label_aa,
+            df_variants,
+            columns=["aa_change", "contig", "position", "ref_allele", "alt_allele"],
+        )
         ds_aa_frq["variant_label"] = "variants", label
 
         # apply variant query if given
@@ -3285,12 +3382,13 @@ class Ag3:
 
         Parameters
         ----------
-        contig : str
-            Chromosome arm, e.g., "3R".
+        contig : str or list of str
+            Chromosome arm, e.g., "3R". Multiple values can be provided
+            as a list, in which case data will be concatenated, e.g., ["2R", "3R"].
         area_by : str
             Column name in the sample metadata to use to group samples spatially. E.g.,
-            use "adm1_ISO" or "adm1_name" to group by level 1 administrative divisions,
-            or use "adm2_name" to group by level 2 administrative divisions.
+            use "admin1_iso" or "admin1_name" to group by level 1 administrative divisions,
+            or use "admin2_name" to group by level 2 administrative divisions.
         period_by : {"year", "quarter", "month"}
             Length of time to group samples temporally.
         sample_sets : str or list of str, optional
@@ -3319,11 +3417,11 @@ class Ag3:
         -------
         ds : xarray.Dataset
             The resulting dataset contains data has dimensions "cohorts" and "variants".
-            Variables prefixed with "cohort_" are 1-dimensional arrays with data about
+            Variables prefixed with "cohort" are 1-dimensional arrays with data about
             the cohorts, such as the area, period, taxon and cohort size. Variables
-            prefixed with "variant_" are 1-dimensional arrays with data about the variants,
+            prefixed with "variant" are 1-dimensional arrays with data about the variants,
             such as the contig, position, reference and alternate alleles. Variables prefixed
-            with "event_" are 2-dimensional arrays with the allele counts and frequency
+            with "event" are 2-dimensional arrays with the allele counts and frequency
             calculations.
 
         """
@@ -3331,6 +3429,52 @@ class Ag3:
         # check parameters
         _check_param_min_cohort_size(min_cohort_size)
 
+        if isinstance(contig, str):
+            contig = [contig]
+
+        ds = xarray_concat(
+            [
+                self._gene_cnv_frequencies_advanced(
+                    contig=c,
+                    area_by=area_by,
+                    period_by=period_by,
+                    sample_sets=sample_sets,
+                    sample_query=sample_query,
+                    min_cohort_size=min_cohort_size,
+                    variant_query=variant_query,
+                    drop_invariant=drop_invariant,
+                    max_coverage_variance=max_coverage_variance,
+                    ci_method=ci_method,
+                    cohorts_analysis=cohorts_analysis,
+                    species_analysis=species_analysis,
+                )
+                for c in contig
+            ],
+            dim="variants",
+        )
+
+        return ds
+
+    def _gene_cnv_frequencies_advanced(
+        self,
+        *,
+        contig,
+        area_by,
+        period_by,
+        sample_sets,
+        sample_query,
+        min_cohort_size,
+        variant_query,
+        drop_invariant,
+        max_coverage_variance,
+        ci_method,
+        cohorts_analysis,
+        species_analysis,
+    ):
+
+        # sanity check - here we deal with one contig only
+        assert isinstance(contig, str)
+
         # load sample metadata
         df_samples = self.sample_metadata(
             sample_sets=sample_sets,
@@ -3439,7 +3583,11 @@ class Ag3:
         )
 
         # add variant label
-        df_variants["label"] = df_variants.apply(_make_gene_cnv_label, axis="columns")
+        df_variants["label"] = _pandas_apply(
+            _make_gene_cnv_label,
+            df_variants,
+            columns=["gene_id", "gene_name", "cnv_type"],
+        )
 
         # build the output dataset
         ds_out = xr.Dataset()
@@ -3461,7 +3609,7 @@ class Ag3:
         if drop_invariant:
             loc_variant = df_variants["max_af"].values > 0
             ds_out = ds_out.isel(variants=loc_variant)
-            df_variants = df_variants.loc[loc_variant]
+            df_variants = df_variants.loc[loc_variant].reset_index(drop=True)
 
         # apply variant query
         if variant_query is not None:
@@ -3508,6 +3656,8 @@ class Ag3:
 
         """
 
+        import plotly.express as px
+
         # handle title
         if title is True:
             title = ds.attrs.get("title", None)
@@ -3610,6 +3760,9 @@ class Ag3:
 
         """
 
+        import ipyleaflet
+        import ipywidgets
+
         # slice dataset to variant of interest
         if isinstance(variant, int):
             ds_variant = ds.isel(variants=variant)
@@ -3710,6 +3863,9 @@ class Ag3:
 
         """
 
+        import ipyleaflet
+        import ipywidgets
+
         # handle title
         if title is True:
             title = ds.attrs.get("title", None)
@@ -3744,6 +3900,248 @@ class Ag3:
 
         return out
 
+    def resolve_region(self, region):
+        """Convert a genome region into a standard data structure.
+
+        Parameters
+        ----------
+        region: str
+            Chromosome arm (e.g., "2L"), gene name (e.g., "AGAP007280") or genomic
+            region defined with coordinates (e.g., "2L:44989425-44998059").
+
+        Returns
+        -------
+        out : Region
+            A named tuple with attributes contig, start and end.
+
+        """
+
+        return resolve_region(self, region)
+
+    def plot_genes(
+        self,
+        region,
+        width=700,
+        height=120,
+        show=True,
+        toolbar_location="above",
+        x_range=None,
+        title="Genes",
+    ):
+        """@@TODO"""
+
+        import bokeh.models as bkmod
+        import bokeh.plotting as bkplt
+
+        # handle region parameter - this determines the genome region to plot
+        region = self.resolve_region(region)
+        contig = region.contig
+        start = region.start
+        end = region.end
+        if start is None:
+            start = 0
+        if end is None:
+            end = len(self.genome_sequence(contig))
+
+        # define x axis range
+        if x_range is None:
+            x_range = bkmod.Range1d(start, end, bounds="auto")
+
+        # select the genes overlapping the requested region
+        df_geneset = self.geneset(attributes=["ID", "Name", "Parent", "description"])
+        # df_geneset = df_geneset.set_index("ID")
+        data = df_geneset.query(
+            f"type == 'gene' and contig == '{contig}' and start < {end} and end > {start}"
+        ).copy()
+
+        # we're going to plot each gene as a rectangle, so add some additional
+        # columns
+        data["bottom"] = np.where(data["strand"] == "+", 1, 0)
+        data["top"] = data["bottom"] + 0.8
+
+        # tidy up some columns for presentation
+        data["Name"].fillna("", inplace=True)
+        data["description"].fillna("", inplace=True)
+
+        # define tooltips for hover
+        tooltips = [
+            ("ID", "@ID"),
+            ("Name", "@Name"),
+            ("Description", "@description"),
+            ("Location", "@contig:@start{,}-@end{,}"),
+        ]
+
+        # make a figure
+        fig = bkplt.figure(
+            title=title,
+            plot_width=width,
+            plot_height=height,
+            tools="xpan,xzoom_in,xzoom_out,xwheel_zoom,reset,tap,hover",
+            toolbar_location=toolbar_location,
+            active_scroll="xwheel_zoom",
+            active_drag="xpan",
+            tooltips=tooltips,
+            x_range=x_range,
+        )
+        fig.toolbar.logo = None
+
+        # add functionality to click through to vectorbase
+        url = "https://vectorbase.org/vectorbase/app/record/gene/@ID"
+        taptool = fig.select(type=bkmod.TapTool)
+        taptool.callback = bkmod.OpenURL(url=url)
+
+        # now plot the genes as rectangles
+        fig.quad(
+            bottom="bottom",
+            top="top",
+            left="start",
+            right="end",
+            source=data,
+            line_width=0.5,
+            fill_alpha=0.5,
+        )
+
+        # tidy up the plot
+        fig.xaxis.axis_label = f"Contig {contig} position (bp)"
+        fig.y_range = bkmod.Range1d(-0.4, 2.2)
+        fig.ygrid.visible = False
+        yticks = [0.4, 1.4]
+        yticklabels = ["-", "+"]
+        fig.yaxis.ticker = yticks
+        fig.yaxis.major_label_overrides = {k: v for k, v in zip(yticks, yticklabels)}
+        fig.xaxis[0].formatter = bkmod.NumeralTickFormatter(format="0,0")
+
+        if show:
+            bkplt.show(fig)
+
+        return fig
+
+    def plot_transcript(
+        self,
+        transcript,
+        width=700,
+        height=120,
+        show=True,
+        x_range=None,
+        toolbar_location="above",
+        title=True,
+    ):
+
+        import bokeh.models as bkmod
+        import bokeh.plotting as bkplt
+
+        # find the transcript annotation
+        df_geneset = self.geneset().set_index("ID")
+        parent = df_geneset.loc[transcript]
+
+        if title is True:
+            title = f"{transcript} ({parent.strand})"
+
+        # define tooltips for hover
+        tooltips = [
+            ("Type", "@type"),
+            ("Location", "@contig:@start{,}-@end{,}"),
+        ]
+
+        # make a figure
+        fig = bkplt.figure(
+            title=title,
+            plot_width=width,
+            plot_height=height,
+            tools="xpan,xzoom_in,xzoom_out,xwheel_zoom,reset,hover",
+            toolbar_location=toolbar_location,
+            active_scroll="xwheel_zoom",
+            active_drag="xpan",
+            tooltips=tooltips,
+            x_range=x_range,
+        )
+        fig.toolbar.logo = None
+
+        # find child components of the transcript
+        data = df_geneset.set_index("Parent").loc[transcript].copy()
+        data["bottom"] = -0.4
+        data["top"] = 0.4
+
+        # plot exons
+        exons = data.query("type == 'exon'")
+        fig.quad(
+            bottom="bottom",
+            top="top",
+            left="start",
+            right="end",
+            source=exons,
+            fill_color=None,
+            line_color="black",
+            line_width=0.5,
+            fill_alpha=0,
+        )
+
+        # plot introns
+        for intron_start, intron_end in zip(exons[:-1]["end"], exons[1:]["start"]):
+            intron_midpoint = (intron_start + intron_end) / 2
+            line_data = pd.DataFrame(
+                {
+                    "x": [intron_start, intron_midpoint, intron_end],
+                    "y": [0, 0.1, 0],
+                    "type": "intron",
+                    "contig": parent.contig,
+                    "start": intron_start,
+                    "end": intron_end,
+                }
+            )
+            fig.line(
+                x="x",
+                y="y",
+                source=line_data,
+                line_width=1,
+                line_color="black",
+            )
+
+        # plot UTRs
+        fig.quad(
+            bottom="bottom",
+            top="top",
+            left="start",
+            right="end",
+            source=data.query("type == 'five_prime_UTR'"),
+            fill_color="green",
+            line_width=0,
+            fill_alpha=0.5,
+        )
+        fig.quad(
+            bottom="bottom",
+            top="top",
+            left="start",
+            right="end",
+            source=data.query("type == 'three_prime_UTR'"),
+            fill_color="red",
+            line_width=0,
+            fill_alpha=0.5,
+        )
+
+        # plot CDSs
+        fig.quad(
+            bottom="bottom",
+            top="top",
+            left="start",
+            right="end",
+            source=data.query("type == 'CDS'"),
+            fill_color="blue",
+            line_width=0,
+            fill_alpha=0.5,
+        )
+
+        # tidy up the figure
+        fig.yaxis.ticker = []
+        fig.y_range = bkmod.Range1d(-0.6, 0.6)
+        fig.xaxis.axis_label = f"Contig {parent.contig} position (bp)"
+        fig.xaxis[0].formatter = bkmod.NumeralTickFormatter(format="0,0")
+
+        if show:
+            bkplt.show(fig)
+
+        return fig
+
 
 @numba.njit("Tuple((int8, int64))(int8[:], int8)")
 def _cn_mode_1d(a, vmax):
@@ -3818,109 +4216,59 @@ def _make_sample_period_year(row):
         return pd.NaT
 
 
-# def _genotypes_to_alt_allele_counts_melt(gt, max_allele):
-#     """Convert a genotype array to an array of alt allele counts, melted to
-#     store one row per alt allele."""
-#
-#     n_variants = gt.shape[0]
-#     n_samples = gt.shape[1]
-#
-#     # convert to genotype allele counts
-#     gac = allel.GenotypeArray(gt).to_allele_counts(max_allele=max_allele)
-#     assert gac.shape == (n_variants, n_samples, max_allele + 1)
-#
-#     # sum total observations over alleles
-#     gan = gac.sum(axis=2)
-#
-#     # keep only alt allele counts
-#     gac_alt = gac[:, :, 1:]
-#     assert gac_alt.shape == (n_variants, n_samples, max_allele)
-#
-#     # use some numpy tricks to melt alleles into rows
-#     gac_alt_melt = gac_alt.swapaxes(2, 1).reshape(-1, n_samples)
-#     assert gac_alt_melt.shape == (n_variants * max_allele, n_samples)
-#     gan_melt = np.repeat(gan, max_allele, axis=0)
-#     assert gan_melt.shape == (n_variants * max_allele, n_samples)
-#
-#     return gac_alt_melt, gan_melt
-
-
 @numba.njit
-def _genotypes_to_alt_allele_counts_melt_kernel(gt, max_allele):
-    """Convert a genotype array to an array of alt allele counts, melted to
-    store one row per alt allele."""
+def _cohort_alt_allele_counts_melt_kernel(gt, indices, max_allele):
 
     n_variants = gt.shape[0]
-    n_samples = gt.shape[1]
+    n_indices = indices.shape[0]
     ploidy = gt.shape[2]
 
-    gac_alt_melt = np.zeros((n_variants * max_allele, n_samples), dtype=np.uint8)
-    gan = np.zeros((n_variants, n_samples), dtype=np.uint8)
+    ac_alt_melt = np.zeros(n_variants * max_allele, dtype=np.int64)
+    an = np.zeros(n_variants, dtype=np.int64)
 
     for i in range(n_variants):
         out_i_offset = (i * max_allele) - 1
-        for j in range(n_samples):
+        for j in range(n_indices):
+            ix = indices[j]
             for k in range(ploidy):
-                allele = gt[i, j, k]
+                allele = gt[i, ix, k]
                 if allele > 0:
                     out_i = out_i_offset + allele
-                    gac_alt_melt[out_i, j] += 1
-                    gan[i, j] += 1
+                    ac_alt_melt[out_i] += 1
+                    an[i] += 1
                 elif allele == 0:
-                    gan[i, j] += 1
+                    an[i] += 1
 
-    return gac_alt_melt, gan
+    return ac_alt_melt, an
 
 
-def _genotypes_to_alt_allele_counts_melt(gt, max_allele):
-    gac_alt_melt, gan = _genotypes_to_alt_allele_counts_melt_kernel(gt, max_allele)
-    gan_melt = np.repeat(gan, max_allele, axis=0)
-    return gac_alt_melt, gan_melt
+def _cohort_alt_allele_counts_melt(gt, indices, max_allele):
+    ac_alt_melt, an = _cohort_alt_allele_counts_melt_kernel(gt, indices, max_allele)
+    an_melt = np.repeat(an, max_allele, axis=0)
+    return ac_alt_melt, an_melt
 
 
-@numba.njit
-def _take_sum_cols(a, indices):
-    n_variants = a.shape[0]
-    n_indices = indices.shape[0]
-    out = np.zeros(n_variants, dtype=np.int64)
-    for i in range(n_variants):
-        v_sum = 0
-        for j in range(n_indices):
-            ix = indices[j]
-            v = a[i, ix]
-            v_sum += v
-        out[i] = v_sum
-    return out
-
-
-def _make_snp_label(row):
-    label = (
-        f"{row['contig']}:{row['position']:,} {row['ref_allele']}>{row['alt_allele']}"
-    )
-    return label
+def _make_snp_label(contig, position, ref_allele, alt_allele):
+    return f"{contig}:{position:,} {ref_allele}>{alt_allele}"
 
 
-def _make_snp_label_effect(row):
-    label = (
-        f"{row['contig']}:{row['position']:,} {row['ref_allele']}>{row['alt_allele']}"
-    )
-    aa_change = row["aa_change"]
+def _make_snp_label_effect(contig, position, ref_allele, alt_allele, aa_change):
+    label = f"{contig}:{position:,} {ref_allele}>{alt_allele}"
     if isinstance(aa_change, str):
         label += f" ({aa_change})"
     return label
 
 
-def _make_snp_label_aa(row):
-    label = f"{row['aa_change']} ({row['contig']}:{row['position']:,} {row['ref_allele']}>{row['alt_allele']})"
+def _make_snp_label_aa(aa_change, contig, position, ref_allele, alt_allele):
+    label = f"{aa_change} ({contig}:{position:,} {ref_allele}>{alt_allele})"
     return label
 
 
-def _make_gene_cnv_label(row):
-    label = row["gene_id"]
-    gene_name = row["gene_name"]
+def _make_gene_cnv_label(gene_id, gene_name, cnv_type):
+    label = gene_id
     if isinstance(gene_name, str):
         label += f" ({gene_name})"
-    label += f" {row['cnv_type']}"
+    label += f" {cnv_type}"
     return label
 
 
@@ -3963,6 +4311,8 @@ def _map_snp_to_aa_change_frq_ds(ds):
 
 
 def _add_frequency_ci(ds, ci_method):
+    from statsmodels.stats.proportion import proportion_confint
+
     if ci_method is not None:
         count = ds["event_count"].values
         nobs = ds["event_nobs"].values
@@ -4047,3 +4397,11 @@ def _check_param_min_cohort_size(min_cohort_size):
         raise ValueError(
             f"Value of parameter min_cohort_size must be at least 1; found {min_cohort_size}."
         )
+
+
+def _pandas_apply(f, df, columns):
+    """Optimised alternative to pandas apply."""
+    df = df.reset_index(drop=True)
+    iterator = zip(*[df[c].values for c in columns])
+    ret = pd.Series((f(*vals) for vals in iterator))
+    return ret

{malariagen_data-2.0.0 → malariagen_data-2.2.0}/malariagen_data/util.py

@@ -404,3 +404,35 @@ def xarray_concat(
             join=join,
             **kwargs,
         )
+
+
+def type_error(
+    name,
+    value,
+    expectation,
+):
+    message = (
+        f"Bad type for parameter {name}; expected {expectation}, "
+        f"found {type(value)}"
+    )
+    raise TypeError(message)
+
+
+def value_error(
+    name,
+    value,
+    expectation,
+):
+    message = (
+        f"Bad value for parameter {name}; expected {expectation}, " f"found {value!r}"
+    )
+    raise ValueError(message)
+
+
+def check_type(
+    name,
+    value,
+    expectation,
+):
+    if not isinstance(value, expectation):
+        type_error(name, value, expectation)

{malariagen_data-2.0.0 → malariagen_data-2.2.0}/pyproject.toml

@@ -1,8 +1,8 @@
 [tool.poetry]
 name = "malariagen_data"
-version = "2.0.0"
+version = "2.2.0"
 description = "A package for accessing MalariaGEN public data."
-authors = ["Alistair Miles <alistair.miles@sanger.ac.uk>", "Chris Clarkson <cc28@sanger.ac.uk>", "Kathryn Murie <km22@sanger.ac.uk>", "Nace Kranjc <n.kranjc@imperial.ac.uk>"]
+authors = ["Alistair Miles <alistair.miles@sanger.ac.uk>", "Chris Clarkson <cc28@sanger.ac.uk>", "Kathryn Murie <km22@sanger.ac.uk>", "Nace Kranjc <n.kranjc@imperial.ac.uk>", "Kelly Bennett <kb25@sanger.ac.uk>"]
 license = "MIT"
 
 [tool.poetry.dependencies]
@@ -19,6 +19,7 @@ scikit-allel = "*"
 xarray = "*"
 numba = "*"
 plotly = "*"
+bokeh = "*"
 statsmodels = "*"
 ipyleaflet = "*"
 ipywidgets = "*"

{malariagen_data-2.0.0 → malariagen_data-2.2.0}/setup.py

@@ -9,6 +9,7 @@ package_data = \
 
 install_requires = \
 ['BioPython',
+ 'bokeh',
  'dask[array]',
  'fsspec',
  'gcsfs',
@@ -27,7 +28,7 @@ install_requires = \
 
 setup_kwargs = {
     'name': 'malariagen-data',
-    'version': '2.0.0',
+    'version': '2.2.0',
     'description': 'A package for accessing MalariaGEN public data.',
     'long_description': None,
     'author': 'Alistair Miles',