sample_metadata.py

import difflib
import io
import json
import re
from itertools import cycle
from typing import (
    Any,
    Callable,
    DefaultDict,
    Dict,
    List,
    Mapping,
    Optional,
    Sequence,
    Tuple,
    Union,
)
from collections import defaultdict
import warnings

import ipyleaflet  # type: ignore
import numpy as np
import pandas as pd
import plotly.express as px  # type: ignore
from numpydoc_decorator import doc  # type: ignore

from ..util import _check_types
from . import base_params, map_params, plotly_params
from .base import AnophelesBase


class AnophelesSampleMetadata(AnophelesBase):
    def __init__(
        self,
        cohorts_analysis: Optional[str] = None,
        aim_analysis: Optional[str] = None,
        aim_metadata_dtype: Optional[Mapping[str, Any]] = None,
        taxon_colors: Optional[Mapping[str, str]] = None,
        aim_species_colors: Optional[Mapping[str, str]] = None,
        **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)

        # If provided, this analysis version will override the
        # default value provided in the release configuration.
        self._cohorts_analysis_override = cohorts_analysis

        # If provided, this analysis version will override the
        # default value provided in the release configuration.
        self._aim_analysis_override = aim_analysis

        # N.B., the expected AIM metadata columns may vary between
        # data resources, and so column names and dtype need to be
        # passed in as parameters.
        self._aim_metadata_columns: Optional[List[str]] = None
        self._aim_metadata_dtype: Optional[Mapping[str, Any]] = {}

        # Only apply the `aim_metadata_dtype` if it is a type of `Mapping`.
        if isinstance(aim_metadata_dtype, Mapping):
            # Convert all of the column names to lowercase.
            prepared_aim_metadata_dtype_dict = {
                k.lower(): v for k, v in aim_metadata_dtype.items()
            }

            # Get all the column names from the prepared dict.
            self._aim_metadata_columns = list(prepared_aim_metadata_dtype_dict.keys())

            # Update the _aim_metadata_dtype with the prepared dict.
            self._aim_metadata_dtype.update(prepared_aim_metadata_dtype_dict)

        # Add the sample_id to the _aim_metadata_dtype.
        self._aim_metadata_dtype["sample_id"] = "object"

        # Set up taxon colors.
        self._taxon_colors = taxon_colors

        self._aim_species_colors = aim_species_colors

        # Set up extra metadata.
        self._extra_metadata: List = []

        # Initialize cache attributes.
        self._cache_sample_metadata: Dict = dict()
        self._cache_cohorts: Dict = dict()
        self._cache_cohort_geometries: Dict = dict()

    def _metadata_paths(
        self,
        *,
        sample_sets: List[str],
        path_template: str,
        aim_analysis: Optional[str] = None,
        cohorts_analysis: Optional[str] = None,
    ) -> Dict[str, str]:
        paths = dict()
        for sample_set in sample_sets:
            release = self.lookup_release(sample_set=sample_set)
            release_path = self._release_to_path(release=release)
            if aim_analysis:
                path = path_template.format(
                    release_path=release_path,
                    sample_set=sample_set,
                    aim_analysis=aim_analysis,
                )
            elif cohorts_analysis:
                path = path_template.format(
                    release_path=release_path,
                    sample_set=sample_set,
                    cohorts_analysis=cohorts_analysis,
                )
            else:
                path = path_template.format(
                    release_path=release_path, sample_set=sample_set
                )
            paths[sample_set] = path
        return paths

    def _parse_metadata_paths(
        self,
        path_template: str,
        parse_metadata_func: Callable[[str, Union[bytes, Exception]], pd.DataFrame],
        sample_sets: List[str],
        aim_analysis: Optional[str] = None,
        cohorts_analysis: Optional[str] = None,
    ) -> pd.DataFrame:
        # Note: we don't use `_prep_sample_sets_param` in this function because that can cause a circular dependency, eventually raising a `RecursionError`.
        # For instance, `_prep_sample_sets_param` uses `_relevant_sample_sets`, which uses `_surveillance_flags`, which uses `_parse_metadata_paths`.
        # Instead, use `_prep_sample_sets_param` to prepare `sample_sets` as a `List[str]` before passing it to this function.

        # Obtain paths for all files we need to fetch.
        file_paths: Mapping[str, str] = self._metadata_paths(
            sample_sets=sample_sets,
            path_template=path_template,
            aim_analysis=aim_analysis,
            cohorts_analysis=cohorts_analysis,
        )

        # Fetch all files. N.B., here is an optimisation, this allows us to fetch
        # multiple files concurrently.
        files: Mapping[str, Union[bytes, Exception]] = self.read_files(
            paths=file_paths.values(), on_error="return"
        )

        # Parse files into DataFrames.
        dfs = []
        for sample_set in sample_sets:
            path = file_paths[sample_set]
            data = files[path]
            df = parse_metadata_func(sample_set, data)
            dfs.append(df)

        # Concatenate all DataFrames.
        df_ret = pd.concat(dfs, axis=0, ignore_index=True)

        return df_ret

    def _parse_general_metadata(
        self, sample_set: str, data: Union[bytes, Exception]
    ) -> pd.DataFrame:
        if isinstance(data, bytes):
            dtype_dict = {
                "sample_id": "object",
                "partner_sample_id": "object",
                "contributor": "object",
                "country": "object",
                "location": "object",
                "year": "int64",
                "month": "int64",
                "latitude": "float64",
                "longitude": "float64",
                "sex_call": "object",
            }
            # `dict[str, str]` is incompatible with the `dtype` of `pd.read_csv`
            dtype: DefaultDict[str, str] = defaultdict(lambda: "object", dtype_dict)
            df = pd.read_csv(io.BytesIO(data), dtype=dtype, na_values="")

            # Ensure all column names are lower case.
            df.columns = [c.lower() for c in df.columns]  # type: ignore

            # Add a couple of columns for convenience.
            df["sample_set"] = sample_set
            release = self.lookup_release(sample_set=sample_set)
            df["release"] = release

            # Derive a quarter column from month.
            # Vectorized operation: quarter = ((month - 1) // 3) + 1 if month > 0 else -1
            df["quarter"] = np.where(df["month"] > 0, ((df["month"] - 1) // 3) + 1, -1)

            # Add study columns.
            study_info = self.lookup_study_info(sample_set=sample_set)
            for column in study_info:
                df[column] = study_info[column]

            # Add terms-of-use columns.
            terms_of_use_info = self.lookup_terms_of_use_info(sample_set=sample_set)
            for column in terms_of_use_info:
                df[column] = terms_of_use_info[column]

            return df

        else:
            raise data

    @_check_types
    @doc(
        summary="""
            Read general sample metadata for one or more sample sets into a pandas
            DataFrame.
        """,
        returns="A pandas DataFrame, one row per sample.",
    )
    def general_metadata(
        self, sample_sets: Optional[base_params.sample_sets] = None
    ) -> pd.DataFrame:
        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
        del sample_sets

        return self._parse_metadata_paths(
            path_template="{release_path}/metadata/general/{sample_set}/samples.meta.csv",
            parse_metadata_func=self._parse_general_metadata,
            sample_sets=prepared_sample_sets,
        )

    @property
    def _sequence_qc_metadata_dtype(self):
        # Note: tests expect an ordered dictionary.
        # Note: insertion order in dictionary keys is guaranteed since Python 3.7
        # Note: using nullable dtypes (e.g. Int64 instead of int64) to allow missing data.

        dtype = {
            "sample_id": "object",
            "mean_cov": "Float64",
            "median_cov": "Int64",
            "modal_cov": "Int64",
        }

        for contig in sorted(self.config["CONTIGS"]):
            dtype[f"mean_cov_{contig}"] = "Float64"
            dtype[f"median_cov_{contig}"] = "Int64"
            dtype[f"mode_cov_{contig}"] = "Int64"

        dtype.update(
            {
                "frac_gen_cov": "Float64",
                "divergence": "Float64",
                "contam_pct": "Float64",
                "contam_LLR": "Float64",
            }
        )

        return dtype

    def _parse_sequence_qc_metadata(
        self, sample_set: str, data: Union[bytes, Exception]
    ) -> pd.DataFrame:
        if isinstance(data, bytes):
            # Get the dtype of the constant columns.
            dtype_dict = self._sequence_qc_metadata_dtype

            # `dict[str, str]` is incompatible with the `dtype` of `pd.read_csv`
            dtype: DefaultDict[str, str] = defaultdict(lambda: "object", dtype_dict)

            # Read the CSV using the dtype dict.
            df = pd.read_csv(io.BytesIO(data), dtype=dtype, na_values="")

            return df

        elif isinstance(data, FileNotFoundError):
            # Sequence QC metadata are missing for this sample set,
            # so return a blank DataFrame.

            # Copy the sample ids from the general metadata.
            df_general = self.general_metadata(sample_sets=sample_set)
            df = df_general[["sample_id"]].copy()

            # Add the sequence QC columns with appropriate missing values.
            # For each column, set the value to either NA or NaN.
            for c, datum_dtype in self._sequence_qc_metadata_dtype.items():
                if pd.api.types.is_integer_dtype(datum_dtype):
                    # Note: this creates a column with dtype int64.
                    df[c] = -1
                else:
                    # Note: this creates a column with dtype float64.
                    df[c] = np.nan

            # Set the column data types.
            df = df.astype(self._sequence_qc_metadata_dtype)

            return df

        else:
            raise data

    @_check_types
    @doc(
        summary="""
            Access sequence QC metadata for one or more sample sets.
        """,
        returns="""A pandas DataFrame, one row per sample. The columns are:
        `sample_id` is the identifier of the sample,
        `partner_sample_id` is the identifier of the sample used by the partners who contributed it,
        `contributor` is the partner who contributed the sample,
        `country` is the country the sample was collected in,
        `location` is the location the sample was collected in,
        `year` is the year the sample was collected,
        `month` is the month the sample was collected,
        `latitude` is the latitude of the location the sample was collected in,
        `longitude` is the longitude of the location the sample was collected in,
        `sex_call` is the sex of the sample,
        `sample_set` is the sample set containing the sample,
        `release` is the release containing the sample,
        `quarter` is the quarter of the year the sample was collected,
        `study_id* is the identifier of the study the sample set containing the sample came from,
        `study_url` is the URL of the study the sample set containing the sample came from,
        `terms_of_use_expiry_date` is the date the terms of use for the sample expire,
        `terms_of_use_url` is the URL of the terms of use for the sample,
        `unrestricted_use` indicates whether the sample can be used without restrictions (e.g., if the terms of use of expired),
        `mean_cov` is mean value of the coverage,
        `median_cov` is the median value of the coverage,
        `modal_cov` is the mode of the coverage,
        `mean_cov_2L` is mean value of the coverage on 2L,
        `median_cov_2L` is the median value of the coverage on 2L,
        `mode_cov_2L` is the mode of the coverage on 2L,
        `mean_cov_2R` is mean value of the coverage on 2R,
        `median_cov_2R` is the median value of the coverage on 2R,
        `mode_cov_2R` is the mode of the coverage on 2R,
        `mean_cov_3L` is mean value of the coverage on 3L,
        `median_cov_3L` is the median value of the coverage on 3L,
        `mode_cov_3L` is the mode of the coverage on 3L,
        `mean_cov_3R` is mean value of the coverage on 3R,
        `median_cov_3R` is the median value of the coverage on 3R,
        `mode_cov_3R` is the mode of the coverage on 3R,
        `mean_cov_X` is mean value of the coverage on X,
        `median_cov_X` is the median value of the coverage on X,
        `mode_cov_X` is the mode of the coverage on X,
        `frac_gen_cov` is the faction of the genome covered,
        `divergence` is the divergence,
        `contam_pct` is the percentage of contamination,
        `contam_LLR` is the log-likelihood ratio of contamination,
        `aim_species_fraction_arab` is the fraction of the gambcolu vs. arabiensis AIMs that indicated arabiensis (this column is only present for *Ag3*),
        `aim_species_fraction_colu` is the fraction of the gambiae vs. coluzzii AIMs that indicated coluzzii (this column is only present for *Ag3*),
        `aim_species_fraction_colu_no2l` is the fraction of the gambiae vs. coluzzii AIMs that indicated coluzzii, not including the chromosome arm 2L which contains an introgression (this column is only present for *Ag3*),
        `aim_species_gambcolu_arabiensis` is the taxonomic group assigned by the gambcolu vs. arabiensis AIMs (this column is only present for *Ag3*),
        `aim_species_gambiae_coluzzi` is the taxonomic group assigned by the gambiae vs. coluzzii AIMs (this column is only present for *Ag3*),
        `aim_species_gambcolu_arabiensis` is the taxonomic group assigned by the combination of both AIMs analyses (this column is only present for *Ag3*),
        `country_iso` is the ISO code of the country the sample was collected in,
        `admin1_name` is the name of the first administrative level the sample was collected in,
        `admin1_iso` is the ISO code of the first administrative level the sample was collected in,
        `admin2_name` is the name of the second administrative level the sample was collected in,
        `taxon` is the taxon assigned to the sample by the combination of the AIMs analysis and the cohort analysis,
        `cohort_admin1_year` is the cohort the sample belongs to when samples are grouped by first administrative level and year,
        `cohort_admin1_month` is the cohort the sample belongs to when samples are grouped by first administrative level and month,
        `cohort_admin1_quarter` is the cohort the sample belongs to when samples are grouped by first administrative level and quarter,
        `cohort_admin2_year` is the cohort the sample belongs to when samples are grouped by second administrative level and year,
        `cohort_admin2_month` is the cohort the sample belongs to when samples are grouped by second administrative level and month,
        `cohort_admin2_quarter` is the cohort the sample belong to when samples are grouped by second administrative level and quarter.
        """,
    )
    def sequence_qc_metadata(
        self, sample_sets: Optional[base_params.sample_sets] = None
    ) -> pd.DataFrame:
        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
        del sample_sets

        return self._parse_metadata_paths(
            path_template="{release_path}/metadata/curation/{sample_set}/sequence_qc_stats.csv",
            parse_metadata_func=self._parse_sequence_qc_metadata,
            sample_sets=prepared_sample_sets,
        )

    def _parse_surveillance_flags(
        self, sample_set: str, data: Union[bytes, Exception]
    ) -> pd.DataFrame:
        # Get the current warning filters.
        original_warning_filters = warnings.filters[:]

        # Specify the expected data type for each column.
        # Note: "bool" is not nullable and does not support `NaN`, which is required when missing data.
        # Otherwise `NaN` will be mis-translated to `True` when the dtype is applied to the DataFrame.
        dtype_dict = {
            "sample_id": "object",
            "is_surveillance": "boolean",
        }
        # `dict[str, str]` is incompatible with the `dtype` of `pd.read_csv`
        dtype: DefaultDict[str, str] = defaultdict(lambda: "object", dtype_dict)

        if isinstance(data, bytes):
            # Read the CSV data.
            df = pd.read_csv(io.BytesIO(data), dtype=dtype, na_values="")

            # If there are any nulls in these data, show a warning.
            if df.isnull().values.any():
                # Trigger the warning.
                warnings.simplefilter("default", UserWarning)
                warnings.warn(
                    f"WARNING: The surveillance flags data contains null values for sample set {sample_set}",
                    UserWarning,
                )

                # Restore the original warning filters.
                warnings.filters = original_warning_filters

            # Ensure all column names are lower case.
            df.columns = [c.lower() for c in df.columns]  # type: ignore

            return df

        elif isinstance(data, FileNotFoundError):
            # Surveillance flags are missing for this sample set.
            # Show a warning and return a blank DataFrame.

            # Trigger the warning.
            warnings.simplefilter("default", UserWarning)
            warnings.warn(
                f"WARNING: The surveillance flags data is missing for sample set {sample_set}",
                UserWarning,
            )

            # Restore the original warning filters.
            warnings.filters = original_warning_filters

            # Get a copy of the sample ids.
            df_general = self.general_metadata(sample_sets=sample_set)
            df = df_general[["sample_id"]].copy()

            # Set each column value to null.
            df["is_surveillance"] = np.nan

            # Set the data type.
            df = df.astype(dtype)

            return df

        else:
            raise data

    @_check_types
    @doc(
        summary="""
            Access surveillance flags for one or more sample sets.
        """,
        parameters=dict(
            sample_sets="List of sample sets.",
        ),
        returns="""A pandas DataFrame, one row per sample. The columns are:
        `sample_id` is the identifier of the sample,
        `is_surveillance` indicates whether the sample can be used for surveillance,
        """,
    )
    def _surveillance_flags(self, sample_sets: List[str]) -> pd.DataFrame:
        # Note: we don't use `_prep_sample_sets_param` in this function because that can cause a circular dependency, eventually raising a `RecursionError`.
        # For instance, `_prep_sample_sets_param` uses `_relevant_sample_sets`, which uses `_surveillance_flags`.
        # Instead, use `_prep_sample_sets_param` to prepare `sample_sets` as a `List[str]` before passing it to this function.

        return self._parse_metadata_paths(
            path_template="{release_path}/metadata/general/{sample_set}/surveillance.flags.csv",
            parse_metadata_func=self._parse_surveillance_flags,
            sample_sets=sample_sets,
        )

    @property
    def _cohorts_analysis(self):
        if self._cohorts_analysis_override:
            return self._cohorts_analysis_override
        else:
            # N.B., this will return None if the key is not present in the
            # config.
            return self.config.get("DEFAULT_COHORTS_ANALYSIS")

    @property
    def _cohorts_metadata_columns(self):
        # Handle changes to columns used in different analyses.
        cols = None
        if self._cohorts_analysis:
            if self._cohorts_analysis < "20230223":
                cols = (
                    "country_iso",
                    "admin1_name",
                    "admin1_iso",
                    "admin2_name",
                    "taxon",
                    "cohort_admin1_year",
                    "cohort_admin1_month",
                    "cohort_admin2_year",
                    "cohort_admin2_month",
                )
            # We assume that cohorts analyses from "20230223" onwards always include quarter
            # columns.
            else:
                cols = (
                    "country_iso",
                    "admin1_name",
                    "admin1_iso",
                    "admin2_name",
                    "taxon",
                    "cohort_admin1_year",
                    "cohort_admin1_month",
                    "cohort_admin1_quarter",
                    "cohort_admin2_year",
                    "cohort_admin2_month",
                    "cohort_admin2_quarter",
                )
        return cols

    @property
    def _cohorts_metadata_dtype(self):
        cols = self._cohorts_metadata_columns
        if cols:
            # All columns are string columns.
            dtype = {c: "object" for c in cols}
            dtype["sample_id"] = "object"
            return dtype

    def _parse_cohorts_metadata(
        self, sample_set: str, data: Union[bytes, Exception]
    ) -> pd.DataFrame:
        if isinstance(data, bytes):
            # Parse CSV data.
            dtype_dict = self._cohorts_metadata_dtype

            # `dict[str, str]` is incompatible with the `dtype` of `pd.read_csv`
            dtype: DefaultDict[str, str] = defaultdict(lambda: "object", dtype_dict)

            df = pd.read_csv(io.BytesIO(data), dtype=dtype, na_values="")

            # Ensure all column names are lower case.
            df.columns = [c.lower() for c in df.columns]  # type: ignore

            # Rename some columns for consistent naming.
            df.rename(
                columns={
                    "adm1_iso": "admin1_iso",
                    "adm1_name": "admin1_name",
                    "adm2_name": "admin2_name",
                },
                inplace=True,
            )

            return df

        elif isinstance(data, FileNotFoundError):
            # Cohorts metadata are missing for this sample set, fill with a blank
            # DataFrame.
            df_general = self.general_metadata(sample_sets=sample_set)
            df = df_general[["sample_id"]].copy()
            for c in self._cohorts_metadata_columns:
                df[c] = np.nan
            df = df.astype(self._cohorts_metadata_dtype)
            return df

        else:
            raise data

    def _require_cohorts_analysis(self):
        if not self._cohorts_analysis:
            raise NotImplementedError(
                "Cohorts data not available for this data resource."
            )

    @_check_types
    @doc(
        summary="""
            Access cohort membership metadata for one or more sample sets.
        """,
        returns="A pandas DataFrame, one row per sample.",
    )
    def cohorts_metadata(
        self, sample_sets: Optional[base_params.sample_sets] = None
    ) -> pd.DataFrame:
        self._require_cohorts_analysis()

        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
        del sample_sets

        return self._parse_metadata_paths(
            path_template="{release_path}/metadata/cohorts_{cohorts_analysis}/{sample_set}/samples.cohorts.csv",
            parse_metadata_func=self._parse_cohorts_metadata,
            sample_sets=prepared_sample_sets,
            cohorts_analysis=self._cohorts_analysis,
        )

    @property
    def _aim_analysis(self):
        if self._aim_analysis_override:
            return self._aim_analysis_override
        else:
            # N.B., this will return None if the key is not present in the
            # config.
            return self.config.get("DEFAULT_AIM_ANALYSIS")

    def _parse_aim_metadata(
        self, sample_set: str, data: Union[bytes, Exception]
    ) -> pd.DataFrame:
        assert self._aim_metadata_columns is not None
        assert self._aim_metadata_dtype is not None
        if isinstance(data, bytes):
            # Parse CSV data but don't apply the dtype yet.
            df = pd.read_csv(io.BytesIO(data), na_values="")

            # Convert all column names to lowercase.
            df.columns = [c.lower() for c in df.columns]  # type: ignore

            # For each column in the DataFrame...
            for c in df.columns:
                # Apply the corresponding dtype from `_aim_metadata_dtype`.
                # Convert the type to a NumPy dtype.
                col_dtype_as_np = np.dtype(self._aim_metadata_dtype[c])
                df[c] = df[c].astype(col_dtype_as_np)

            return df

        elif isinstance(data, FileNotFoundError):
            # AIM data are missing for this sample set, fill with a blank DataFrame.
            df_general = self.general_metadata(sample_sets=sample_set)
            df = df_general[["sample_id"]].copy()
            for c in self._aim_metadata_columns:
                df[c] = np.nan
            df = df.astype(self._aim_metadata_dtype)
            return df

        else:
            raise data

    def _require_aim_analysis(self):
        if not self._aim_analysis:
            raise NotImplementedError("AIM data not available for this data resource.")

    @_check_types
    @doc(
        summary="""
            Access ancestry-informative marker (AIM) metadata for one or more
            sample sets.
        """,
        returns="A pandas DataFrame, one row per sample.",
    )
    def aim_metadata(
        self, sample_sets: Optional[base_params.sample_sets] = None
    ) -> pd.DataFrame:
        self._require_aim_analysis()

        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
        del sample_sets

        return self._parse_metadata_paths(
            path_template="{release_path}/metadata/species_calls_aim_{aim_analysis}/{sample_set}/samples.species_aim.csv",
            parse_metadata_func=self._parse_aim_metadata,
            sample_sets=prepared_sample_sets,
            aim_analysis=self._aim_analysis,
        )

    @_check_types
    @doc(
        summary="""
            Add extra sample metadata, e.g., including additional columns
            which you would like to use to query and group samples.
        """,
        parameters=dict(
            data="""
                A data frame with one row per sample. Must include either a
                "sample_id" or "partner_sample_id" column.
            """,
            on="""
                Name of column to use when merging with sample metadata.
            """,
        ),
        notes="""
            The values in the column containing sample identifiers must be
            unique.
        """,
    )
    def add_extra_metadata(self, data: pd.DataFrame, on: str = "sample_id"):
        # Check parameters.
        if not isinstance(data, pd.DataFrame):
            raise TypeError("`data` parameter must be a pandas DataFrame")
        if on not in data.columns:
            raise ValueError(f"dataframe does not contain column {on!r}")
        if on not in {"sample_id", "partner_sample_id"}:
            raise ValueError(
                "`on` parameter must be either 'sample_id' or 'partner_sample_id'"
            )

        # Check for uniqueness.
        if not data[on].is_unique:
            raise ValueError(f"column {on!r} does not have unique values")

        # check there are matching samples.
        df_samples = self.sample_metadata()
        loc_isec = data[on].isin(df_samples[on])
        if not loc_isec.any():
            raise ValueError("no matching samples found")

        # store extra metadata
        self._extra_metadata.append((on, data.copy()))

    @doc(
        summary="Clear any extra metadata previously added",
    )
    def clear_extra_metadata(self):
        self._extra_metadata = []

    @_check_types
    @doc(
        summary="Access sample metadata for one or more sample sets.",
        returns="A dataframe of sample metadata, one row per sample.",
        notes="""
            Some samples in the dataset are lab crosses — mosquitoes bred in
            the laboratory that have no real collection date. These samples
            use ``year=-1`` and ``month=-1`` as sentinel values. They may
            cause unexpected results in date-based analyses (e.g.,
            ``pd.to_datetime`` will fail on negative year values).

            To exclude lab cross samples, use::

                df = api.sample_metadata(sample_query="year >= 0")
        """,
    )
    def sample_metadata(
        self,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
    ) -> pd.DataFrame:
        # Check that either sample_query xor sample_indices are provided.
        base_params._validate_sample_selection_params(
            sample_query=sample_query, sample_indices=sample_indices
        )

        # Prepare parameters.
        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
        prepared_sample_query = self._prep_sample_query_param(sample_query=sample_query)

        # Delete original parameters to prevent accidental use.
        del sample_sets
        del sample_query

        # Determine the cache key.
        cache_key = tuple(prepared_sample_sets)
        try:
            # Attempt to retrieve from the cache.
            df_samples = self._cache_sample_metadata[cache_key]

        except KeyError:
            with self._spinner(desc="Load sample metadata"):
                ## Build a single DataFrame using all available metadata.

                # Get the general sample metadata.
                # Note: this includes study and terms-of-use info.
                df_samples = self.general_metadata(sample_sets=prepared_sample_sets)

                # Merge with the sequence QC metadata.
                # Note: merging can change column dtypes, e.g. due to new NaNs.
                df_sequence_qc = self.sequence_qc_metadata(
                    sample_sets=prepared_sample_sets
                )
                df_samples = df_samples.merge(
                    df_sequence_qc, on="sample_id", sort=False, how="left"
                )

                # Merge with the surveillance flags.
                # Note: merging can change column dtypes, e.g. due to new NaNs.
                df_surveillance_flags = self._surveillance_flags(
                    sample_sets=prepared_sample_sets
                )
                df_samples = df_samples.merge(
                    df_surveillance_flags, on="sample_id", sort=False, how="left"
                )

                # If available, merge with the AIM metadata.
                if self._aim_analysis:
                    df_aim = self.aim_metadata(sample_sets=prepared_sample_sets)
                    df_samples = df_samples.merge(
                        df_aim, on="sample_id", sort=False, how="left"
                    )

                # If available, merge with the cohorts metadata.
                if self._cohorts_analysis:
                    df_cohorts = self.cohorts_metadata(sample_sets=prepared_sample_sets)
                    df_samples = df_samples.merge(
                        df_cohorts, on="sample_id", sort=False, how="left"
                    )

            # Store sample metadata in the cache.
            self._cache_sample_metadata[cache_key] = df_samples

        # Add extra metadata.
        for on, data in self._extra_metadata:
            df_samples = df_samples.merge(data, how="left", on=on)

        # Apply the sample_query, if there is one.
        # Note: this might have been internally modified, e.g. `is_surveillance == True`.
        if prepared_sample_query is not None:
            # Assume a pandas query string.
            sample_query_options = sample_query_options or {}

            # Save a reference to the pre-query DataFrame so we can detect
            # zero-result queries and provide a helpful warning.
            df_before_query = df_samples

            # Use the python engine in order to support extension array dtypes, e.g. Float64, Int64, boolean.
            df_samples = df_samples.query(
                prepared_sample_query, **sample_query_options, engine="python"
            )
            df_samples = df_samples.reset_index(drop=True)

            # Warn if query returned zero results on a non-empty dataset.
            # Provide fuzzy-match suggestions so users can spot typos,
            # case mismatches, or partial-value issues.
            if len(df_samples) == 0 and len(df_before_query) > 0:
                hint_lines = [
                    f"sample_metadata() returned 0 samples for query: {prepared_sample_query!r}.",
                ]

                # Extract column == 'value' pairs from the query.
                col_val_pairs = re.findall(
                    r"\b(\w+)\s*==\s*['\"]([^'\"]+)['\"]",
                    prepared_sample_query,
                )

                for col_name, queried_val in col_val_pairs:
                    # If the column name is not recognised, suggest
                    # close column names.
                    if col_name not in df_before_query.columns:
                        close_cols = difflib.get_close_matches(
                            col_name,
                            df_before_query.columns.tolist(),
                            n=3,
                            cutoff=0.6,
                        )
                        if close_cols:
                            hint_lines.append(
                                f"Column {col_name!r} not found. "
                                f"Did you mean: {close_cols}?"
                            )
                        continue

                    # For string columns, suggest close values.
                    if df_before_query[col_name].dtype == object:
                        valid_vals = (
                            df_before_query[col_name].dropna().unique().tolist()
                        )
                        close_vals = difflib.get_close_matches(
                            queried_val, valid_vals, n=5, cutoff=0.6
                        )
                        if close_vals:
                            hint_lines.append(
                                f"Value {queried_val!r} not found in "
                                f"column {col_name!r}. "
                                f"Did you mean: {close_vals}?"
                            )

                warnings.warn("\n".join(hint_lines), UserWarning, stacklevel=2)

        # Apply the sample_indices, if there are any.
        # Note: this might need to apply to the result of an internal sample_query, e.g. `is_surveillance == True`.
        if sample_indices is not None:
            # Assume it is an indexer.
            df_samples = df_samples.iloc[sample_indices]
            df_samples = df_samples.reset_index(drop=True)

        return df_samples.copy()

    @_check_types
    @doc(
        summary="""
            Create a pivot table showing numbers of samples available by space,
            time and taxon.
        """,
        parameters=dict(
            index="Sample metadata columns to use for the pivot table index.",
            columns="Sample metadata columns to use for the pivot table columns.",
        ),
        returns="Pivot table of sample counts. One row per admin2_year cohort. Unless otherwise specified using the `columns` parameters, the samples are grouped according to their taxon and then counted.",
    )
    def count_samples(
        self,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        index: Union[str, Sequence[str]] = (
            "country",
            "admin1_iso",
            "admin1_name",
            "admin2_name",
            "year",
        ),
        columns: Union[str, Sequence[str]] = "taxon",
    ) -> pd.DataFrame:
        # Load sample metadata.
        df_samples = self.sample_metadata(
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
        )

        # Create pivot table.
        df_pivot = df_samples.pivot_table(
            index=index,
            columns=columns,
            values="sample_id",
            aggfunc="count",
            fill_value=0,
        )

        return df_pivot

    @_check_types
    @doc(
        summary="""
            Plot an interactive map showing sampling locations using ipyleaflet.
        """,
        parameters=dict(
            min_samples="""
                Minimum number of samples required to show a marker for a given
                location.
            """,
            count_by="""
                Metadata column to report counts of samples by for each location.
            """,
        ),
        returns="Ipyleaflet map widget.",
    )
    def plot_samples_interactive_map(
        self,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        basemap: Optional[map_params.basemap] = map_params.basemap_default,
        center: map_params.center = map_params.center_default,
        zoom: map_params.zoom = map_params.zoom_default,
        height: map_params.height = map_params.height_default,
        width: map_params.width = map_params.width_default,
        min_samples: int = 1,
        count_by: str = "taxon",
    ) -> ipyleaflet.Map:
        # Normalise height and width to string
        if isinstance(height, int):
            height = f"{height}px"
        if isinstance(width, int):
            width = f"{width}px"

        # Load sample metadata.
        df_samples = self.sample_metadata(
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
        )

        # Pivot taxa by locations.
        location_composite_key = [
            "country",
            "admin1_iso",
            "admin1_name",
            "admin2_name",
            "location",
            "latitude",
            "longitude",
        ]
        df_pivot = df_samples.pivot_table(
            index=location_composite_key,
            columns=count_by,
            values="sample_id",
            aggfunc="count",
            fill_value=0,
        )

        # Append aggregations to pivot.
        df_location_aggs = df_samples.groupby(location_composite_key).agg(
            {
                "year": lambda x: ", ".join(str(y) for y in sorted(x.unique())),
                "sample_set": lambda x: ", ".join(str(y) for y in sorted(x.unique())),
                "contributor": lambda x: ", ".join(str(y) for y in sorted(x.unique())),
            }
        )
        df_pivot = df_pivot.merge(
            df_location_aggs, on=location_composite_key, validate="one_to_one"
        )

        # Handle basemap.
        basemap_abbrevs = map_params.basemap_abbrevs

        # Determine basemap_provider via basemap
        if isinstance(basemap, str):
            # Interpret string
            # Support case-insensitive basemap abbreviations
            basemap_str = basemap.lower()
            if basemap_str not in basemap_abbrevs:
                raise ValueError(
                    f"Basemap abbreviation not recognised: {basemap_str!r}; try one of {list(basemap_abbrevs.keys())}"
                )
            basemap_provider = basemap_abbrevs[basemap_str]
        elif basemap is None:
            # Default.
            basemap_provider = ipyleaflet.basemaps.Esri.WorldImagery
        else:
            # Expect dict or TileProvider or TileLayer.
            basemap_provider = basemap

        # Create a map.
        samples_map = ipyleaflet.Map(
            center=center,
            zoom=zoom,
            basemap=basemap_provider,
        )
        scale_control = ipyleaflet.ScaleControl(position="bottomleft")
        samples_map.add(scale_control)
        samples_map.layout.height = height
        samples_map.layout.width = width

        # Add markers.
        count_factors = df_samples[count_by].dropna().sort_values().unique()
        for _, row in df_pivot.reset_index().iterrows():
            title = (
                f"Location: {row.location} ({row.latitude:.3f}, {row.longitude:.3f})"
            )
            title += f"\nAdmin level 2: {row.admin2_name}"
            title += f"\nAdmin level 1: {row.admin1_name} ({row.admin1_iso})"
            title += f"\nCountry: {row.country}"
            title += f"\nYears: {row.year}"
            title += f"\nSample sets: {row.sample_set}"
            title += f"\nContributors: {row.contributor}"
            title += "\nNo. specimens: "
            all_n = 0
            for factor in count_factors:
                # Get the number of samples in this taxon
                n = row[factor]
                # Count the number of samples in all taxa
                all_n += n
                if n > 0:
                    title += f"{n} {factor}; "
            # Only show a marker when there are enough samples
            if all_n >= min_samples:
                marker = ipyleaflet.Marker(
                    location=(row.latitude, row.longitude),
                    draggable=False,
                    title=title,
                )
                samples_map.add(marker)

        return samples_map

    @_check_types
    @doc(
        summary="""
            Load a data catalog providing URLs for downloading BAM, VCF and Zarr
            files for samples in a given sample set.
        """,
        returns="""One row per sample, the columns are
        `sample_id`, the identifier of the sample,
        `alignments_bam`, the URL of the alignments BAM file,
        `snp_genotypes_vcf`, the URL of the SNP genotypes VCF file,
        `snp_genotypes_zarr`, the URL of the SNP genotypes Zarr file.""",
    )
    def wgs_data_catalog(self, sample_set: base_params.sample_set):
        # Look up release for sample set.
        release = self.lookup_release(sample_set=sample_set)
        release_path = self._release_to_path(release=release)

        # Load data catalog.
        path = f"{self._base_path}/{release_path}/metadata/general/{sample_set}/wgs_snp_data.csv"
        with self._fs.open(path) as f:
            df = pd.read_csv(
                f, na_values="", engine="python"
            )  # Force python engine to ensure thread-safety

        # Normalise columns.
        df = df[
            [
                "sample_id",
                "alignments_bam",
                "snp_genotypes_vcf",
                "snp_genotypes_zarr",
            ]
        ]

        # If `_unrestricted_use_only` is `True`, then only return data if this sample set has `unrestricted_use` set to `True`.
        if self._unrestricted_use_only and not self._sample_set_has_unrestricted_use(
            sample_set=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:
            surveillance_flags_df = self._surveillance_flags(sample_sets=[sample_set])
            df = df.merge(
                surveillance_flags_df[["sample_id", "is_surveillance"]],
                on="sample_id",
                how="left",
            )
            df = df[df["is_surveillance"]]
            df = df.drop(columns=["is_surveillance"])

        return df

    @_check_types
    @doc(
        summary="""
            Load a data accessions catalog providing ENA run accessions
            for samples in a given sample set.
        """,
        returns="One row per sample, columns provide run accessions.",
    )
    def wgs_run_accessions(self, sample_set: base_params.sample_set):
        # Look up release for sample set.
        release = self.lookup_release(sample_set=sample_set)
        release_path = self._release_to_path(release=release)

        # Load data catalog.
        path = f"{self._base_path}/{release_path}/metadata/general/{sample_set}/wgs_accession_data.csv"
        with self._fs.open(path) as f:
            df = pd.read_csv(f, na_values="")

        # Normalise columns.
        df = df[
            [
                "sample_id",
                "run_ena",
            ]
        ]

        return df

    def _prep_sample_selection_cache_params(
        self,
        *,
        sample_sets: Optional[base_params.sample_sets],
        sample_query: Optional[base_params.sample_query],
        sample_query_options: Optional[base_params.sample_query_options],
        sample_indices: Optional[base_params.sample_indices],
    ) -> Tuple[List[str], Optional[List[int]]]:
        # Check that either sample_query xor sample_indices are provided.
        base_params._validate_sample_selection_params(
            sample_query=sample_query, sample_indices=sample_indices
        )

        # Resolve query to a list of integers for more cache hits - we
        # do this because there are different ways to write the same pandas
        # query, and so it's better to evaluate the query and use a list of
        # integer indices instead.

        # Scenario 1: No `sample_query` nor `sample_indices` were given,
        #             and there is no internal `sample_query`,
        #             so no `sample_indices` will be returned.

        # Scenario 2: No `sample_query` nor `sample_indices` were given,
        #             but there is an internal `sample_query`,
        #             which will be converted into `sample_indices` and returned.

        # Scenario 3: Only `sample_query` has been provided,
        #             which will be converted into `sample_indices` and returned.
        #             This will be handled the same as Scenario 2.

        # Scenario 4: Only `sample_indices` has been provided,
        #             and there is no internal `sample_query`,
        #             simply return `sample_indices`.

        # Scenario 5: Only `sample_indices` has been provided,
        #             but there is also an internal `sample_query`, still return `sample_indices`,
        #             which ought to already align with `sample_metadata`.

        # Normalise sample sets.
        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
        prepared_sample_query = self._prep_sample_query_param(sample_query=sample_query)

        # Delete original parameters to prevent accidental use.
        del sample_sets
        del sample_query

        # Start with assuming there are no sample indices.
        # This can be returned if there is no `prepared_sample_query` nor `sample_indices`.
        prepared_sample_indices = None

        # If there is a `prepared_sample_query` but no `sample_indices`...
        if prepared_sample_query is not None and sample_indices is None:
            # Get the unfiltered sample metadata for the given sample sets.
            # Note: we don't want to pass the `sample_query` to `sample_metadata` here
            # because we want to get the sample indices that represent the `sample_query`.
            df_samples = self.sample_metadata(sample_sets=prepared_sample_sets)

            # Default the sample_query_options to an empty dict.
            sample_query_options = sample_query_options or {}

            # Use the python engine in order to support extension array dtypes, e.g. Float64, Int64, boolean.
            # Get the Pandas Series as a NumPy array of Boolean values.
            # Note: if `prepared_sample_query` is an internal query, this will select all samples,
            # since `sample_metadata` should have already applied the internal query.
            loc_samples = df_samples.eval(
                prepared_sample_query, **sample_query_options, engine="python"
            ).values

            # Convert the sample indices to a list.
            # Get the indices of the True values in the Boolean array and convert it to a list of integers.
            prepared_sample_indices = np.nonzero(loc_samples)[0].tolist()

        # If there is a `prepared_sample_query` and a `sample_indices`...
        elif prepared_sample_query is not None and sample_indices is not None:
            # Given that we don't allow both `sample_query` and `sample_indices` params in this function,
            # we can deduce that the `prepared_sample_query` has resulted from an internal query.
            # Given that `sample_indices` should be aligned with the results of `sample_metadata`,
            # which should already apply the internal query, simply return the given `sample_indices`.

            prepared_sample_indices = sample_indices

        # If there is no `prepared_sample_query` but there is a `sample_indices`...
        elif prepared_sample_query is None and sample_indices is not None:
            prepared_sample_indices = sample_indices

        return prepared_sample_sets, prepared_sample_indices

    def _results_cache_add_analysis_params(self, params: dict):
        super()._results_cache_add_analysis_params(params)
        params["cohorts_analysis"] = self._cohorts_analysis
        params["aim_analysis"] = self._aim_analysis

    @_check_types
    @doc(
        summary="Get the metadata for a specific sample and sample set.",
        returns="The metadata for the specified sample.",
    )
    def lookup_sample(
        self,
        sample: base_params.sample,
        sample_set: Optional[base_params.sample_set] = None,
    ) -> pd.Series:
        df_samples = self.sample_metadata(sample_sets=sample_set).set_index("sample_id")
        sample_rec = None
        if isinstance(sample, str):
            sample_rec = df_samples.loc[sample]
        else:
            assert isinstance(sample, int)
            sample_rec = df_samples.iloc[sample]
        return sample_rec

    @_check_types
    @doc(
        summary="""
            Plot a bar chart showing the number of samples available, grouped by
            some variable such as country or year.
        """,
        parameters=dict(
            x="Name of sample metadata column to plot on the X axis.",
            color="Name of the sample metadata column to color bars by.",
            sort="If True, sort the bars in size order.",
            kwargs="Passed through to px.bar().",
        ),
    )
    def plot_samples_bar(
        self,
        x: str,
        color: Optional[str] = None,
        sort: bool = True,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        template: plotly_params.template = "plotly_white",
        width: plotly_params.fig_width = 800,
        height: plotly_params.fig_height = 600,
        show: plotly_params.show = True,
        renderer: plotly_params.renderer = None,
        **kwargs,
    ) -> plotly_params.figure:
        # Load sample metadata.
        df_samples = self.sample_metadata(
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
        )

        # Special handling for plotting by year.
        if x == "year":
            # Remove samples with missing year.
            df_samples = df_samples.query("year > 0")

        # Construct a long-form dataframe to plot.
        if color:
            grouper: Union[str, List[str]] = [x, color]
        else:
            grouper = x
        df_plot = df_samples.groupby(grouper).agg({"sample_id": "count"}).reset_index()

        # Deal with request to sort by bar size.
        if sort:
            df_sort = (
                df_samples.groupby(x)
                .agg({"sample_id": "count"})
                .reset_index()
                .sort_values("sample_id")
            )
            x_order = df_sort[x].values
            category_orders = kwargs.get("category_orders", dict())
            category_orders.setdefault(x, x_order)
            kwargs["category_orders"] = category_orders

        # Make the plot.
        fig = px.bar(
            df_plot,
            x=x,
            y="sample_id",
            color=color,
            template=template,
            width=width,
            height=height,
            **kwargs,
        )

        # Visual styling.
        fig.update_layout(
            xaxis_title=x.capitalize(),
            yaxis_title="No. samples",
        )
        if color:
            fig.update_layout(legend_title=color.capitalize())

        if show:  # pragma: no cover
            fig.show(renderer=renderer)
            return None
        else:
            return fig

    def _setup_sample_symbol(
        self,
        *,
        data,
        symbol,
    ):
        if symbol is None:
            return None

        # Handle the symbol parameter.
        if isinstance(symbol, str):
            if "cohort_" + symbol in data.columns:
                # Convenience to allow things like "admin1_year" instead of "cohort_admin1_year".
                symbol_prepped = "cohort_" + symbol
            else:
                symbol_prepped = symbol
            if symbol_prepped not in data.columns:
                raise ValueError(
                    f"{symbol_prepped!r} is not a known column in the data."
                )

        else:
            # Custom grouping using queries.
            assert isinstance(symbol, Mapping)
            data["symbol"] = ""
            for key, value in symbol.items():
                data.loc[data.query(value).index, "symbol"] = key
            symbol_prepped = "symbol"

        # Handle missing data in a consistent way.
        data[symbol_prepped] = data[symbol_prepped].fillna("<NA>")

        return symbol_prepped

    def _setup_sample_colors_plotly(
        self,
        *,
        data,
        color,
        color_discrete_sequence,
        color_discrete_map,
        category_orders,
    ):
        # Check for no color.
        if color is None:
            # Bail out early.
            return None, None, None

        # Special handling for taxon colors.
        if color == "taxon" and color_discrete_map is None:
            # Special case, default taxon colors and order.
            color_discrete_map = self._taxon_colors

        # Special handling for aim_species colors.
        if color == "aim_species" and color_discrete_map is None:
            # Special case, default aim_species colors and order.
            color_discrete_map = self._aim_species_colors

        if isinstance(color, str):
            if "cohort_" + color in data.columns:
                # Convenience to allow things like "admin1_year" instead of "cohort_admin1_year".
                color_prepped = "cohort_" + color
            else:
                color_prepped = color

            if color_prepped not in data.columns:
                raise ValueError(
                    f"{color_prepped!r} is not a known column in the data."
                )

        else:
            # Custom grouping using queries.
            assert isinstance(color, Mapping)
            data["color"] = ""
            for key, value in color.items():
                data.loc[data.query(value).index, "color"] = key
            color_prepped = "color"

        # Finish handling of color parameter.
        del color

        # Handle missing data in a consistent way.
        data[color_prepped] = data[color_prepped].fillna("<NA>")

        # Obtain the values that we will be mapping to colors.
        color_data_unique_values = data[color_prepped].unique()

        # Now set up color choices.
        if color_discrete_map is None:
            # Choose a color palette.
            if color_discrete_sequence is None:
                if len(color_data_unique_values) <= 10:
                    color_discrete_sequence = px.colors.qualitative.Plotly
                else:
                    color_discrete_sequence = px.colors.qualitative.Alphabet

            # Map values to colors.
            color_discrete_map_prepped = {
                v: c
                for v, c in zip(
                    color_data_unique_values, cycle(color_discrete_sequence)
                )
            }

        else:
            color_discrete_map_prepped = color_discrete_map

        # Consistent color for missing data.
        color_discrete_map_prepped["<NA>"] = "black"

        # Finished handling of color map params.
        del color_discrete_map
        del color_discrete_sequence

        # Define category orders.
        if category_orders is None:
            # Default ordering.
            category_orders_prepped = {color_prepped: color_data_unique_values.tolist()}

        else:
            category_orders_prepped = category_orders

        # Finised handling of category orders.
        del category_orders

        return (
            color_prepped,
            color_discrete_map_prepped,
            category_orders_prepped,
        )

    def _setup_sample_hover_data_plotly(
        self,
        *,
        color,
        symbol,
    ):
        hover_data = [
            "sample_id",
            "partner_sample_id",
            "sample_set",
            "taxon",
            "country",
            "admin1_iso",
            "admin1_name",
            "admin2_name",
            "location",
            "year",
            "month",
        ]
        if color and color not in hover_data:
            hover_data.append(color)
        if symbol and symbol not in hover_data:
            hover_data.append(symbol)
        return hover_data

    def _setup_cohort_queries(
        self,
        cohorts: base_params.cohorts,
        sample_sets: Optional[base_params.sample_sets],
        sample_query: Optional[base_params.sample_query],
        sample_query_options: Optional[base_params.sample_query_options],
        cohort_size: Optional[base_params.cohort_size],
        min_cohort_size: Optional[base_params.min_cohort_size],
    ):
        """Convenience function to normalise the `cohorts` parameter to a
        dictionary mapping cohort labels to sample metadata queries."""

        if isinstance(cohorts, dict):
            # User has supplied a custom dictionary mapping cohort identifiers
            # to pandas queries.
            cohort_queries = cohorts

        else:
            assert isinstance(cohorts, str)
            # User has supplied a column in the sample metadata.
            df_samples = self.sample_metadata(
                sample_sets=sample_sets,
                sample_query=sample_query,
                sample_query_options=sample_query_options,
            )

            # Determine column in dataframe - allow abbreviation.
            if "cohort_" + cohorts in df_samples.columns:
                cohorts = "cohort_" + cohorts
            if cohorts not in df_samples.columns:
                raise ValueError(
                    f"{cohorts!r} is not a known column in the sample metadata"
                )

            # Find cohort labels and build queries dictionary.
            cohort_labels = sorted(df_samples[cohorts].dropna().unique())
            cohort_queries = {coh: f"{cohorts} == '{coh}'" for coh in cohort_labels}

        # Handle sample_query parameter.
        if sample_query is not None:
            cohort_queries = {
                cohort_label: f"({cohort_query}) and ({sample_query})"
                for cohort_label, cohort_query in cohort_queries.items()
            }

        # Check cohort sizes, drop any cohorts which are too small.
        cohort_queries_checked = dict()
        for cohort_label, cohort_query in cohort_queries.items():
            df_cohort_samples = self.sample_metadata(
                sample_sets=sample_sets, sample_query=cohort_query
            )
            n_samples = len(df_cohort_samples)
            if min_cohort_size is not None:
                cohort_size = min_cohort_size
            if cohort_size is not None and n_samples < cohort_size:
                warnings.warn(
                    f"Cohort ({cohort_label}) has insufficient samples ({n_samples}) for requested cohort size ({cohort_size}), dropping.",
                    stacklevel=2,
                )
            else:
                cohort_queries_checked[cohort_label] = cohort_query

        return cohort_queries_checked

    @_check_types
    @doc(
        summary="""
            Read data for a specific cohort set, including cohort size,
            country code, taxon, administrative units name, ISO code, geoBoundaries
            shape ID and representative latitude and longitude points.
        """,
        parameters=dict(
            cohort_set="""
                A cohort set name. Accepted values are:
                "admin1_month", "admin1_quarter", "admin1_year",
                "admin2_month", "admin2_quarter", "admin2_year".
            """,
            query="""
                An optional pandas query string to filter the resulting
                dataframe, e.g., "country == 'Burkina Faso'".
            """,
        ),
        returns="""A dataframe of cohort data, one row per cohort. There are up to 18 columns:
        `cohort_id` is the identifier of the cohort,
        `cohort_size` is the number of samples in the cohort,
        `country` is the country the cohort is from,
        `country_alpha2` is the ISO alpha-2 code for the country the cohort is from,
        `country_alpha3` is the ISO alpha-3 code for the country the cohort is from,
        `taxon` is the taxon of the samples in the cohort,
        `year` is the year the samples in the cohort were collected in,
        `quarter` is the quarter the samples in the cohort were collected in (this column is only present if the temporal dimension is *quarter* or *month*),
        `month` is the month the samples in the cohort were collected in (this column is only present if the temporal dimension is *month*),
        `admin1_name` is the name of the first administrative level the samples in the cohort were collected in,
        `admin1_iso` is the ISO code of the first administrative level the samples in the cohort were collected in,
        `admin1_geoboundaries_shape_id` is the identifier of the geoboundary shape corresponding to the first administrative level the samples in the cohort were collected in,
        `admin1_representative_longitude` is the representative longitude for the first administrative level the samples in the cohort were collected in,
        `admin1_representative_latitude` is the representative latitude for the first administrative level the samples in the cohort were collected in,
        `admin2_name` is the name of the second administrative level the samples in the cohort were collected in (this column is only present if the spatial dimension is *admin2*),
        `admin2_iso` is the ISO code of the second administrative level the samples in the cohort were collected in (this column is only present if the spatial dimension is *admin2*),
        `admin2_geoboundaries_shape_id` is the identifier of the geoboundary shape corresponding to the second administrative level the samples in the cohort were collected in (this column is only present if the spatial dimension is *admin2*),
        `admin2_representative_longitude` is the representative longitude for the second administrative level the samples in the cohort were collected in (this column is only present if the spatial dimension is *admin2*),
        `admin2_representative_latitude` is the representative latitude for the second administrative level the samples in the cohort were collected in (this column is only present if the spatial dimension is *admin2*).
        """,
    )
    def cohorts(
        self,
        cohort_set: base_params.cohorts,
        query: Optional[str] = None,
    ) -> pd.DataFrame:
        valid_cohort_sets = {
            "admin1_month",
            "admin1_quarter",
            "admin1_year",
            "admin2_month",
            "admin2_quarter",
            "admin2_year",
        }
        if cohort_set not in valid_cohort_sets:
            raise ValueError(
                f"{cohort_set!r} is not a valid cohort set. "
                f"Accepted values are: {sorted(valid_cohort_sets)}."
            )

        cohorts_analysis = self._cohorts_analysis

        # Cache to avoid repeated reads.
        cache_key = (cohorts_analysis, cohort_set)
        try:
            df_cohorts = self._cache_cohorts[cache_key]
        except KeyError:
            major_version_path = self._major_version_path
            path = f"{major_version_path[:2]}_cohorts/cohorts_{cohorts_analysis}/cohorts_{cohort_set}.csv"

            with self.open_file(path) as f:
                df_cohorts = pd.read_csv(f, sep=",", na_values="")

            # Ensure all column names are lower case.
            df_cohorts.columns = [c.lower() for c in df_cohorts.columns]  # type: ignore

            self._cache_cohorts[cache_key] = df_cohorts

        if query is not None:
            df_cohorts = df_cohorts.query(query)
            df_cohorts = df_cohorts.reset_index(drop=True)

        return df_cohorts.copy()

    @_check_types
    @doc(
        summary="""
            Read GeoJSON geometry data for a specific cohort set,
            providing boundary geometries for each cohort.
        """,
        parameters=dict(
            cohort_set="""
                A cohort set name. Accepted values are:
                "admin1_month", "admin1_quarter", "admin1_year",
                "admin2_month", "admin2_quarter", "admin2_year".
            """,
        ),
        returns="""
            A dict containing the parsed GeoJSON FeatureCollection,
            with boundary geometries for each cohort in the set.
        """,
    )
    def cohort_geometries(
        self,
        cohort_set: base_params.cohorts,
    ) -> dict:
        valid_cohort_sets = {
            "admin1_month",
            "admin1_quarter",
            "admin1_year",
            "admin2_month",
            "admin2_quarter",
            "admin2_year",
        }
        if cohort_set not in valid_cohort_sets:
            raise ValueError(
                f"{cohort_set!r} is not a valid cohort set. "
                f"Accepted values are: {sorted(valid_cohort_sets)}."
            )

        cohorts_analysis = self._cohorts_analysis

        # Cache to avoid repeated reads.
        cache_key = (cohorts_analysis, cohort_set)
        try:
            geojson_data = self._cache_cohort_geometries[cache_key]
        except KeyError:
            major_version_path = self._major_version_path
            path = f"{major_version_path[:2]}_cohorts/cohorts_{cohorts_analysis}/cohorts_{cohort_set}.geojson"

            with self.open_file(path) as f:
                geojson_data = json.load(f)

            self._cache_cohort_geometries[cache_key] = geojson_data

        return geojson_data

    @_check_types
    @doc(
        summary="""
            Plot markers on a map showing sample locations
            as a Mapbox scatter plot.
        """,
        parameters=dict(
            kwargs="Passed through to px.scatter_mapbox().",
        ),
    )
    def plot_sample_location_mapbox(
        self,
        *,
        sample_sets: Optional[base_params.sample_sets],
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        marker_size: plotly_params.marker_size = 10,
        color: plotly_params.color = "admin1_name",
        color_discrete_sequence: plotly_params.color_discrete_sequence = px.colors.qualitative.Prism,
        category_orders: plotly_params.category_order = None,
        hover_name: plotly_params.hover_name = "location",
        zoom: plotly_params.zoom = None,
        width: plotly_params.fig_width = 800,
        height: plotly_params.fig_height = 600,
        show: plotly_params.show = True,
        renderer: plotly_params.renderer = None,
        **kwargs,
    ) -> plotly_params.figure:
        # Get the sample metadata.
        df_samples = self.sample_metadata(
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
        )

        # Set the location columns to use from the sample metadata.
        location_columns = [
            "country",
            "admin1_iso",
            "admin1_name",
            "admin2_name",
            "location",
            "latitude",
            "longitude",
        ]

        # Trim and dedupe the sample locations.
        # Sort by `color` column by default, which can be overridden via category_orders.
        df_locations = df_samples[location_columns].drop_duplicates().sort_values(color)

        fig = px.scatter_map(
            df_locations,
            lat="latitude",
            lon="longitude",
            map_style="open-street-map",
            zoom=zoom,
            color=color,
            category_orders=category_orders,
            color_discrete_sequence=color_discrete_sequence,
            hover_name=hover_name,
            hover_data=location_columns,
            width=width,
            height=height,
            **kwargs,
        )

        # Set the size of the markers.
        fig.update_traces(marker=dict(size=marker_size))

        if show:  # pragma: no cover
            fig.show(renderer=renderer)
            return None
        else:
            return fig

    @_check_types
    @doc(
        summary="""
            Plot markers on a map showing sample locations
            as a geographic scatter plot.
        """,
        parameters=dict(
            kwargs="Passed through to px.scatter_mapbox().",
        ),
    )
    def plot_sample_location_geo(
        self,
        *,
        sample_sets: Optional[base_params.sample_sets],
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        marker_size: plotly_params.marker_size = 10,
        color: plotly_params.color = "admin1_name",
        color_discrete_sequence: plotly_params.color_discrete_sequence = px.colors.qualitative.Prism,
        category_orders: plotly_params.category_order = None,
        hover_name: plotly_params.hover_name = "location",
        fitbounds: plotly_params.fitbounds = "locations",
        scope: plotly_params.scope = "world",
        width: plotly_params.fig_width = 800,
        height: plotly_params.fig_height = 600,
        show: plotly_params.show = True,
        renderer: plotly_params.renderer = None,
        **kwargs,
    ) -> plotly_params.figure:
        # Get the sample metadata.
        df_samples = self.sample_metadata(
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
        )

        # Set the location columns to use from the sample metadata.
        location_columns = [
            "country",
            "admin1_iso",
            "admin1_name",
            "admin2_name",
            "location",
            "latitude",
            "longitude",
        ]

        # Trim and dedupe the sample locations.
        # Sort by `color` column by default, which can be overridden via category_orders.
        df_locations = df_samples[location_columns].drop_duplicates().sort_values(color)

        fig = px.scatter_geo(
            df_locations,
            lat="latitude",
            lon="longitude",
            scope=scope,
            height=height,
            width=width,
            color=color,
            hover_name=hover_name,
            hover_data=location_columns,
            category_orders=category_orders,
            color_discrete_sequence=color_discrete_sequence,
            fitbounds=fitbounds,
            **kwargs,
        )

        # Set the size of the markers.
        fig.update_traces(marker=dict(size=marker_size))

        if show:  # pragma: no cover
            fig.show(renderer=renderer)
            return None
        else:
            return fig


def _locate_cohorts(*, cohorts, data, min_cohort_size):
    # Build cohort dictionary where key=cohort_id, value=loc_coh.
    coh_dict = {}

    if isinstance(cohorts, Mapping):
        # User has supplied a custom dictionary mapping cohort identifiers
        # to pandas queries.

        for coh, query in cohorts.items():
            try:
                loc_coh = data.eval(query).values
            except Exception as e:
                raise ValueError(
                    f"Invalid query for cohort {coh!r}: {query!r}. Error: {e}"
                ) from e
            coh_dict[coh] = loc_coh

    else:
        assert isinstance(cohorts, str)
        # User has supplied the name of a sample metadata column.

        # Convenience to allow things like "admin1_year" instead of "cohort_admin1_year".
        if "cohort_" + cohorts in data.columns:
            cohorts = "cohort_" + cohorts

        # Check the given cohort set exists.
        if cohorts not in data.columns:
            raise ValueError(f"{cohorts!r} is not a known column in the data.")
        cohort_labels = data[cohorts].dropna().unique()

        # Remove the nans and sort.
        cohort_labels = sorted([c for c in cohort_labels if isinstance(c, str)])
        for coh in cohort_labels:
            loc_coh = data[cohorts] == coh
            coh_dict[coh] = loc_coh.values

    # Remove cohorts below minimum cohort size.
    coh_dict = {
        coh: loc_coh
        for coh, loc_coh in coh_dict.items()
        if np.count_nonzero(loc_coh) >= min_cohort_size
    }

    # Early check for no cohorts.
    if len(coh_dict) == 0:
        raise ValueError(
            "No cohorts available for the given sample selection parameters and minimum cohort size."
        )

    return coh_dict