Add unrestricted_use_only and surveillance_use_only constructor params

malariagen_data/af1.py

@@ -91,6 +91,8 @@ def __init__(
         discordant_read_calls_analysis=None,
         pre=False,
         tqdm_class=None,
+        unrestricted_use_only=False,
+        surveillance_use_only=False,
         **storage_options,  # used by fsspec via init_filesystem()
     ):
         super().__init__(
@@ -127,18 +129,23 @@ def __init__(
             virtual_contigs=None,
             gene_names=None,
             inversion_tag_path=None,
+            unrestricted_use_only=unrestricted_use_only,
+            surveillance_use_only=surveillance_use_only,
         )
 
     def __repr__(self):
         text = (
             f"<MalariaGEN Af1 API client>\n"
-            f"Storage URL             : {self._url}\n"
-            f"Data releases available : {', '.join(self.releases)}\n"
-            f"Results cache           : {self._results_cache}\n"
-            f"Cohorts analysis        : {self._cohorts_analysis}\n"
-            f"Site filters analysis   : {self._site_filters_analysis}\n"
-            f"Software version        : malariagen_data {malariagen_data.__version__}\n"
-            f"Client location         : {self.client_location}\n"
+            f"Storage URL                           : {self._url}\n"
+            f"Data releases available               : {', '.join(self._available_releases)}\n"
+            f"Results cache                         : {self._results_cache}\n"
+            f"Cohorts analysis                      : {self._cohorts_analysis}\n"
+            f"Site filters analysis                 : {self._site_filters_analysis}\n"
+            f"Software version                      : malariagen_data {malariagen_data.__version__}\n"
+            f"Client location                       : {self.client_location}\n"
+            f"Data filtered to unrestricted use only: {self._unrestricted_use_only}\n"
+            f"Data filtered to surveillance use only: {self._surveillance_use_only}\n"
+            f"Relevant data releases                : {', '.join(self.releases)}\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"
@@ -172,7 +179,7 @@ def _repr_html_(self):
                         <th style="text-align: left">
                             Data releases available
                         </th>
-                        <td>{', '.join(self.releases)}</td>
+                        <td>{', '.join(self._available_releases)}</td>
                     </tr>
                     <tr>
                         <th style="text-align: left">
@@ -204,6 +211,24 @@ def _repr_html_(self):
                         </th>
                         <td>{self.client_location}</td>
                     </tr>
+                    <tr>
+                        <th style="text-align: left">
+                            Data filtered for unrestricted use only
+                        </th>
+                        <td>{self._unrestricted_use_only}</td>
+                    </tr>
+                    <tr>
+                        <th style="text-align: left">
+                            Data filtered for surveillance use only
+                        </th>
+                        <td>{self._surveillance_use_only}</td>
+                    </tr>
+                    <tr>
+                        <th style="text-align: left">
+                            Relevant data releases
+                        </th>
+                        <td>{', '.join(self.releases)}</td>
+                    </tr>
                 </tbody>
             </table>
         """

malariagen_data/ag3.py

@@ -75,6 +75,18 @@ def _setup_aim_palettes():
     "unassigned": "black",
 }
 
+# Note: These column names will be treated as case-insensitive,
+# because these column names and the column names from the CSV
+# will be converted to lowercase before applying these dtypes.
+AIM_METADATA_DTYPE = {
+    "aim_species_fraction_arab": "float64",
+    "aim_species_fraction_colu": "float64",
+    "aim_species_fraction_colu_no2l": "float64",
+    "aim_species_gambcolu_arabiensis": "object",
+    "aim_species_gambiae_coluzzii": "object",
+    "aim_species": "object",
+}
+
 
 class Ag3(AnophelesDataResource):
     """Provides access to data from Ag3.x releases.
@@ -150,6 +162,8 @@ def __init__(
         discordant_read_calls_analysis=None,
         pre=False,
         tqdm_class=None,
+        unrestricted_use_only=False,
+        surveillance_use_only=False,
         **storage_options,  # used by fsspec via init_filesystem()
     ):
         super().__init__(
@@ -158,14 +172,7 @@ def __init__(
             config_path=CONFIG_PATH,
             cohorts_analysis=cohorts_analysis,
             aim_analysis=aim_analysis,
-            aim_metadata_dtype={
-                "aim_species_fraction_arab": "float64",
-                "aim_species_fraction_colu": "float64",
-                "aim_species_fraction_colu_no2l": "float64",
-                "aim_species_gambcolu_arabiensis": "object",
-                "aim_species_gambiae_coluzzii": "object",
-                "aim_species": "object",
-            },
+            aim_metadata_dtype=AIM_METADATA_DTYPE,
             aim_ids=("gambcolu_vs_arab", "gamb_vs_colu"),
             aim_palettes=AIM_PALETTES,
             site_filters_analysis=site_filters_analysis,
@@ -193,6 +200,8 @@ def __init__(
             virtual_contigs=VIRTUAL_CONTIGS,
             gene_names=GENE_NAMES,
             inversion_tag_path=INVERSION_TAG_PATH,
+            unrestricted_use_only=unrestricted_use_only,
+            surveillance_use_only=surveillance_use_only,
         )
 
         # set up caches
@@ -204,21 +213,24 @@ def v3_wild(self):
         3.0 release, excluding the lab crosses."""
         return [
             x
-            for x in self.sample_sets(release="3.0")["sample_set"].tolist()
+            for x in self._available_sample_sets(release="3.0")["sample_set"].tolist()
             if x != "AG1000G-X"
         ]
 
     def __repr__(self):
         text = (
             f"<MalariaGEN Ag3 API client>\n"
-            f"Storage URL             : {self._url}\n"
-            f"Data releases available : {', '.join(self.releases)}\n"
-            f"Results cache           : {self._results_cache}\n"
-            f"Cohorts analysis        : {self._cohorts_analysis}\n"
-            f"AIM analysis            : {self._aim_analysis}\n"
-            f"Site filters analysis   : {self._site_filters_analysis}\n"
-            f"Software version        : malariagen_data {malariagen_data.__version__}\n"
-            f"Client location         : {self.client_location}\n"
+            f"Storage URL                           : {self._url}\n"
+            f"Data releases available               : {', '.join(self._available_releases)}\n"
+            f"Results cache                         : {self._results_cache}\n"
+            f"Cohorts analysis                      : {self._cohorts_analysis}\n"
+            f"AIM analysis                          : {self._aim_analysis}\n"
+            f"Site filters analysis                 : {self._site_filters_analysis}\n"
+            f"Software version                      : malariagen_data {malariagen_data.__version__}\n"
+            f"Client location                       : {self.client_location}\n"
+            f"Data filtered to unrestricted use only: {self._unrestricted_use_only}\n"
+            f"Data filtered to surveillance use only: {self._surveillance_use_only}\n"
+            f"Relevant data releases                : {', '.join(self.releases)}\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"
@@ -252,7 +264,7 @@ def _repr_html_(self):
                         <th style="text-align: left">
                             Data releases available
                         </th>
-                        <td>{', '.join(self.releases)}</td>
+                        <td>{', '.join(self._available_releases)}</td>
                     </tr>
                     <tr>
                         <th style="text-align: left">
@@ -290,6 +302,24 @@ def _repr_html_(self):
                         </th>
                         <td>{self.client_location}</td>
                     </tr>
+                    <tr>
+                        <th style="text-align: left">
+                            Data filtered for unrestricted use only
+                        </th>
+                        <td>{self._unrestricted_use_only}</td>
+                    </tr>
+                    <tr>
+                        <th style="text-align: left">
+                            Data filtered for surveillance use only
+                        </th>
+                        <td>{self._surveillance_use_only}</td>
+                    </tr>
+                    <tr>
+                        <th style="text-align: left">
+                            Relevant data releases
+                        </th>
+                        <td>{', '.join(self.releases)}</td>
+                    </tr>
                 </tbody>
             </table>
         """
@@ -337,6 +367,34 @@ def cross_metadata(self):
             debug("drop 'phenotype' column, not used")
             df.drop("phenotype", axis="columns", inplace=True)
 
+            # Identify the crosses sample set.
+            # Note: this sample set identifier is also hard-coded in `v3_wild()`.
+            crosses_sample_set = "AG1000G-X"
+
+            # If `_unrestricted_use_only` is `True`, then only return data if the crosses sample set has `unrestricted_use` set to `True`.
+            if (
+                self._unrestricted_use_only
+                and not self._sample_set_has_unrestricted_use(
+                    sample_set=crosses_sample_set
+                )
+            ):
+                # Remove all the data from the DataFrame and reset its index.
+                df = df.iloc[0:0].reset_index(drop=True)
+
+            # If `_surveillance_use_only` is `True`, then only return samples that have `is_surveillance` set to `True`.
+            if self._surveillance_use_only:
+                crosses_surveillance_flags_df = self._surveillance_flags(
+                    sample_sets=[crosses_sample_set]
+                )
+                df = df.merge(
+                    crosses_surveillance_flags_df[["sample_id", "is_surveillance"]],
+                    on="sample_id",
+                    how="left",
+                )
+                df = df[df["is_surveillance"]]
+                df = df.drop(columns=["is_surveillance"])
+
+            # Cache the cross metadata.
             self._cache_cross_metadata = df
 
         return self._cache_cross_metadata.copy()

malariagen_data/anoph/aim_data.py

@@ -138,31 +138,45 @@ def aim_calls(
     ) -> xr.Dataset:
         self._require_aim_analysis()
 
-        # Normalise parameters.
-        aims = self._prep_aims_param(aims=aims)
-        sample_sets_prepped = self._prep_sample_sets_param(sample_sets=sample_sets)
+        # Prepare parameters.
+        prepared_aims = self._prep_aims_param(aims=aims)
+        del aims
+        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
         del sample_sets
-
-        # Access SNP calls and concatenate multiple sample sets and/or regions.
-        ly = []
-        for s in sample_sets_prepped:
-            y = self._aim_calls_dataset(
-                aims=aims,
-                sample_set=s,
+        prepared_sample_query = self._prep_sample_query_param(sample_query=sample_query)
+        del sample_query
+
+        # Start a list of AIM calls Datasets, one for each sample set.
+        aim_calls_datasets = []
+
+        # For each sample set...
+        for sample_set in prepared_sample_sets:
+            # Get the AIM calls for all samples in the set, as a Xarray Dataset.
+            aim_calls_dataset = self._aim_calls_dataset(
+                aims=prepared_aims,
+                sample_set=sample_set,
             )
-            ly.append(y)
+
+            # Add this Dataset to the list.
+            aim_calls_datasets.append(aim_calls_dataset)
 
         # Concatenate data from multiple sample sets.
-        ds = simple_xarray_concat(ly, dim=DIM_SAMPLE)
+        ds = simple_xarray_concat(aim_calls_datasets, dim=DIM_SAMPLE)
 
-        # Handle sample query.
-        if sample_query is not None:
-            df_samples = self.sample_metadata(sample_sets=sample_sets_prepped)
+        # If there's a sample query...
+        if prepared_sample_query is not None:
+            # Get the relevant sample metadata.
+            df_samples = self.sample_metadata(sample_sets=prepared_sample_sets)
+
+            # If there are no sample query options, then default to an empty dict.
             sample_query_options = sample_query_options or {}
-            loc_samples = df_samples.eval(sample_query, **sample_query_options).values
-            if np.count_nonzero(loc_samples) == 0:
-                raise ValueError(f"No samples found for query {sample_query!r}")
-            ds = ds.isel(samples=loc_samples)
+
+            ds = self._filter_sample_dataset(
+                ds=ds,
+                df_samples=df_samples,
+                sample_query=prepared_sample_query,
+                sample_query_options=sample_query_options,
+            )
 
         return ds